import os
from pathlib import Path
import pandas as pd
import numpy as np
from mne.io import read_raw_snirf
from mne_nirs.io import write_raw_snirf
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QHBoxLayout, QMessageBox, QLineEdit, QPushButton, QFileDialog
from PySide6.QtCore import Qt

class UpdateOptodesWindow(QWidget):

    def __init__(self, parent=None):
        super().__init__(parent, Qt.WindowType.Window)
        self.setWindowTitle("Update optode positions")
        self.resize(760, 200)

        self.label_file_a = QLabel("SNIRF file:")
        self.line_edit_file_a = QLineEdit()
        self.line_edit_file_a.setReadOnly(True)
        self.btn_browse_a = QPushButton("Browse .snirf")
        self.btn_browse_a.clicked.connect(self.browse_file_a)

        self.label_file_b = QLabel("Text file:")
        self.line_edit_file_b = QLineEdit()
        self.line_edit_file_b.setReadOnly(True)
        self.btn_browse_b = QPushButton("Browse .txt/.xlsx")
        self.btn_browse_b.clicked.connect(self.browse_file_b)

        self.label_suffix = QLabel("Suffix to append to filename:")
        self.line_edit_suffix = QLineEdit()
        self.line_edit_suffix.setText("flare")

        self.btn_clear = QPushButton("Clear")
        self.btn_go = QPushButton("Go")
        self.btn_clear.clicked.connect(self.clear_files)
        self.btn_go.clicked.connect(self.go_action)

        # --- 
        layout = QVBoxLayout()
        self.description = QLabel()
        self.description.setTextFormat(Qt.TextFormat.RichText)
        self.description.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
        self.description.setOpenExternalLinks(False) # Handle the click internally
            
        self.description.setText("Some software when creating snirf files will insert a template of optode positions as the correct position of the optodes for the participant.<br>"
        "This is rarely correct as each head differs slightly in shape or size, and a lot of calculations require the optodes to be in the correct location.<br>"
        "Using a .txt or .xlsx file, we can update the positions in the snirf file to match those of a digitization system such as one from Polhemus or elsewhere.<br>"
        "The .txt file should have the fiducials, detectors, and sources clearly labeled, followed by the x, y, and z coordinates seperated by a space.<br>"
        "An example format of what a digitization text file should look like can be found <a href='custom_link'>by clicking here</a>. Currently only .xlsx files directly exported from a<br>"
        "Polhemus system are supported.")

        self.description.linkActivated.connect(self.handle_link_click)
        layout.addWidget(self.description)

        help_text_a = "Select the SNIRF (.snirf) file to update with new optode positions."

        file_a_layout = QHBoxLayout()

        # Help button on the left
        help_btn_a = QPushButton("?")
        help_btn_a.setFixedWidth(25)
        help_btn_a.setToolTip(help_text_a)
        help_btn_a.clicked.connect(lambda _, text=help_text_a: self.show_help_popup(text))
        file_a_layout.addWidget(help_btn_a)

        # Container for label + line_edit + browse button with tooltip
        file_a_container = QWidget()
        file_a_container_layout = QHBoxLayout()
        file_a_container_layout.setContentsMargins(0, 0, 0, 0)
        file_a_container_layout.addWidget(self.label_file_a)
        file_a_container_layout.addWidget(self.line_edit_file_a)
        file_a_container_layout.addWidget(self.btn_browse_a)
        file_a_container.setLayout(file_a_container_layout)
        file_a_container.setToolTip(help_text_a)

        file_a_layout.addWidget(file_a_container)
        layout.addLayout(file_a_layout)

        help_text_b = "Provide a .txt file with labeled optodes (e.g., nz, rpa, lpa, d1, s1) and their x, y, z coordinates, or a .xlsx file from a Polhemius system."

        file_b_layout = QHBoxLayout()

        help_btn_b = QPushButton("?")
        help_btn_b.setFixedWidth(25)
        help_btn_b.setToolTip(help_text_b)
        help_btn_b.clicked.connect(lambda _, text=help_text_b: self.show_help_popup(text))
        file_b_layout.addWidget(help_btn_b)

        file_b_container = QWidget()
        file_b_container_layout = QHBoxLayout()
        file_b_container_layout.setContentsMargins(0, 0, 0, 0)
        file_b_container_layout.addWidget(self.label_file_b)
        file_b_container_layout.addWidget(self.line_edit_file_b)
        file_b_container_layout.addWidget(self.btn_browse_b)
        file_b_container.setLayout(file_b_container_layout)
        file_b_container.setToolTip(help_text_b)

        file_b_layout.addWidget(file_b_container)
        layout.addLayout(file_b_layout)


        help_text_suffix = "This text will be appended to the original filename when saving. Default is 'flare'."

        suffix_layout = QHBoxLayout()

        help_btn_suffix = QPushButton("?")
        help_btn_suffix.setFixedWidth(25)
        help_btn_suffix.setToolTip(help_text_suffix)
        help_btn_suffix.clicked.connect(lambda _, text=help_text_suffix: self.show_help_popup(text))
        suffix_layout.addWidget(help_btn_suffix)

        suffix_container = QWidget()
        suffix_container_layout = QHBoxLayout()
        suffix_container_layout.setContentsMargins(0, 0, 0, 0)
        suffix_container_layout.addWidget(self.label_suffix)
        suffix_container_layout.addWidget(self.line_edit_suffix)
        suffix_container.setLayout(suffix_container_layout)
        suffix_container.setToolTip(help_text_suffix)

        suffix_layout.addWidget(suffix_container)
        layout.addLayout(suffix_layout)

        buttons_layout = QHBoxLayout()
        buttons_layout.addStretch()
        buttons_layout.addWidget(self.btn_clear)
        buttons_layout.addWidget(self.btn_go)
        layout.addLayout(buttons_layout)

        self.setLayout(layout)

    def show_help_popup(self, text):
        msg = QMessageBox(self)
        msg.setWindowTitle("Parameter Info - FLARES")
        msg.setText(text)
        msg.exec()

    def handle_link_click(self, link):
        if link == "custom_link":
            msg = QMessageBox(self)
            msg.setWindowTitle("Example Digitization File")

            text = "nz: -1.91 85.175 -31.1525\n" \
                "rpa: 80.3825 -17.1925 -57.2775\n" \
                "lpa: -81.815 -17.1925 -57.965\n" \
                "d1: 0.01 -97.5175 62.5875\n" \
                "d2: 25.125 -103.415 45.045\n" \
                "d3: 49.095 -97.9025 30.2075\n" \
                "s1: 0.01 -112.43 32.595\n" \
                "s2: 30.325 -84.3125 71.8975\n" \
                "s3: 0.01 -70.6875 89.0925\n"
            msg.setText(text)
            msg.exec()

    def browse_file_a(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select SNIRF File", "", "SNIRF Files (*.snirf)")
        if file_path:
            self.line_edit_file_a.setText(file_path)

    def browse_file_b(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "Select File", "", "Supported Files (*.txt *.xlsx)")
        if file_path:
            self.line_edit_file_b.setText(file_path)

    def clear_files(self):
        self.line_edit_file_a.clear()
        self.line_edit_file_b.clear()

    def go_action(self):
        file_a = self.line_edit_file_a.text()
        file_b = self.line_edit_file_b.text()
        suffix = self.line_edit_suffix.text().strip() or "flare"

        if not file_a:
            QMessageBox.critical(self, "Missing File", "Please select a SNIRF file.")
            return
        if not file_b:
            QMessageBox.critical(self, "Missing File", "Please select a TXT file.")
            return

        # Get original filename without extension
        base_name = os.path.splitext(os.path.basename(file_a))[0]
        suggested_name = f"{base_name}_{suffix}.snirf"

        # Open save dialog with default name
        save_path, _ = QFileDialog.getSaveFileName(
            self,
            "Save SNIRF File As",
            suggested_name,
            "SNIRF Files (*.snirf)"
        )

        if not save_path:
            print("Save cancelled.")
            return

        # Ensure .snirf extension
        if not save_path.lower().endswith(".snirf"):
            save_path += ".snirf"

        try:
            self.update_optode_positions(file_a=file_a, file_b=file_b, save_path=save_path)
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to write file:\n{e}")
            return
        
        QMessageBox.information(self, "File Saved", f"File was saved to:\n{save_path}")

    def update_optode_positions(self, file_a, file_b, save_path):
                
        fiducials = {}
        ch_positions = {}

        extension = Path(file_b).suffix

        # Read the lines from the optode file
        if extension == '.txt':
            with open(file_b, 'r') as f:
                for line in f:
                    if line.strip():
                        # Split by the semicolon and convert to meters
                        ch_name, coords_str = line.split(":")
                        coords = np.array(list(map(float, coords_str.strip().split()))) * 0.001

                        # The key we have is a fiducial
                        if ch_name.lower() in ['lpa', 'nz', 'rpa']:
                            fiducials[ch_name.lower()] = coords

                        # The key we have is a source or detector
                        else:
                            ch_positions[ch_name.upper()] = coords

        elif extension == '.xlsx':

            # TODO: Bad! Why assume sheet1 has the contents?
            df = pd.read_excel(file_b, sheet_name='Sheet1')

            def _get_block_data(df, block_id, row_mapping, scale=0.001):
                """Isolates a block, cleans numeric data, and returns a scaled dictionary."""
                # 1. Isolate and clean
                block = df[df['block_id'] == block_id].iloc[:, [1, 2, 3]].copy()
                block = block.apply(pd.to_numeric, errors='coerce')
                
                # 2. Extract into dictionary based on mapping
                result = {}
                
                # If row_mapping is a dict (like {0: 'nz'}), use it directly
                if isinstance(row_mapping, dict):
                    for row_idx, key in row_mapping.items():
                        if row_idx < len(block):
                            result[key] = block.iloc[row_idx].to_numpy(dtype=float) * scale
                
                # If row_mapping is a string prefix (like 'D' or 'S'), auto-generate keys
                elif isinstance(row_mapping, str):
                    for i in range(len(block)):
                        result[f"{row_mapping}{i+1}"] = block.iloc[i].to_numpy(dtype=float) * scale
                        
                return result
                    
            # Identify blocks
            is_empty = df.isnull().all(axis=1)
            df['block_id'] = is_empty.cumsum()
            clean_df = df[~is_empty].copy()

            # Process Block 2: Landmarks
            fiducials = _get_block_data(clean_df, 2, {0: 'nz', 2: 'rpa', 3: 'lpa'})

            # Process Block 3: D-Points
            d_points = _get_block_data(clean_df, 3, 'D')

            # Process Block 4: S-Points
            s_points = _get_block_data(clean_df, 4, 'S')

            ch_positions = {**d_points, **s_points}

        # Create montage with updated coords in head space
        initial_montage = make_dig_montage(ch_pos=ch_positions, nasion=fiducials.get('nz'), lpa=fiducials.get('lpa'), rpa=fiducials.get('rpa'), coord_frame='head') # type: ignore

        # Read the SNIRF file, set the montage, and write it back
        raw = read_raw_snirf(file_a, preload=True)
        raw.set_montage(initial_montage)
        write_raw_snirf(raw, save_path)