"""
Filename: updateoptodes.py
Description: Methods to update optode locations for FLARES
Note: Compliant with pylance strict type checking

Author: Tyler de Zeeuw
License: GPL-3.0
"""

# Built-in imports
import os
from pathlib import Path
from typing import Dict, Optional, Union

# External library imports
import pandas as pd
import numpy as np
import numpy.typing as npt

from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QHBoxLayout, QMessageBox, QLineEdit, QPushButton, QFileDialog
from PySide6.QtCore import Qt

from mne.io import read_raw_snirf  #type: ignore
from mne_nirs.io import write_raw_snirf  #type: ignore
from mne.channels import make_dig_montage  #type: ignore

from src.shared.shareddata import APP_NAME


class UpdateOptodesWindow(QWidget):

    def __init__(self, parent: Optional[QWidget] = None) -> None:
        super().__init__(parent, Qt.WindowType.Window)
        self.setWindowTitle(f"Update optode positions - {APP_NAME.upper()}")
        self.resize(760, 200)

        self.selected_snirf_files: list[str] = []

        self.label_file_a = QLabel("SNIRF files:")
        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: self.show_help_popup(help_text_a))
        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: self.show_help_popup(help_text_b))
        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: self.show_help_popup(help_text_suffix))
        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: str) -> None:
        msg = QMessageBox(self)
        msg.setWindowTitle(f"Parameter Info - {APP_NAME.upper()}")
        msg.setText(text)
        msg.exec()

    def handle_link_click(self, link: str) -> None:
        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) -> None:
        file_paths, _ = QFileDialog.getOpenFileNames(
            self,
            "Select SNIRF Files",
            "",
            "SNIRF Files (*.snirf)"
        )

        if file_paths:
            self.selected_snirf_files = file_paths
            self.line_edit_file_a.setText("; ".join(Path(p).name for p in file_paths))

    def browse_file_b(self) -> None:
        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) -> None:
        self.selected_snirf_files.clear()
        self.line_edit_file_a.clear()
        self.line_edit_file_b.clear()

    def go_action(self) -> None:
        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 self.selected_snirf_files:
            QMessageBox.critical(self, "Missing File", "Please select at least one SNIRF file.")
            return
        if not file_b:
            QMessageBox.critical(self, "Missing File", "Please select a TXT or XLSX digitization file.")
            return

        output_dir = QFileDialog.getExistingDirectory(
            self,
            "Select Output Directory"
        )

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

        output_path = Path(output_dir)

        successful_files: list[str] = []
        failed_files: list[str] = []

        for file_a in self.selected_snirf_files:
            input_path = Path(file_a)

            # Keep original filename and independently add suffix
            save_path = output_path / f"{input_path.stem}_{suffix}.snirf"

            try:
                self.update_optode_positions(
                    file_a=file_a,
                    file_b=file_b,
                    save_path=save_path
                )

                successful_files.append(save_path.name)

            except Exception as e:
                failed_files.append(
                    f"{input_path.name}: {e}"
                )

        # Build summary
        message_parts: list[str] = []

        if successful_files:
            message_parts.append(
                f"Successfully processed {len(successful_files)} "
                f"SNIRF file(s):\n\n"
                + "\n".join(successful_files)
            )

        if failed_files:
            message_parts.append(
                f"Failed to process {len(failed_files)} "
                f"SNIRF file(s):\n\n"
                + "\n".join(failed_files)
            )

        if failed_files:
            QMessageBox.warning(
                self,
                "Processing Complete",
                "\n\n".join(message_parts)
            )
        else:
            QMessageBox.information(
                self,
                "Files Saved",
                "\n\n".join(message_parts)
            )

    def update_optode_positions(
        self, 
        file_a: Union[str, Path], 
        file_b: Union[str, Path], 
        save_path: Union[str, Path]
    ) -> None:
                
        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') # type: ignore

            def _get_block_data(
                target_df: pd.DataFrame, 
                block_id: int, 
                row_mapping: Union[Dict[int, str], str], 
                scale: float = 0.001
            ) -> Dict[str, npt.NDArray[np.float64]]:
                
                """Isolates a block, cleans numeric data, and returns a scaled dictionary."""
                # 1. Isolate and clean
                block = target_df[target_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: Dict[str, npt.NDArray[np.float64]] = {}
                
                # 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
                else:
                    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(str(file_a), preload=True)
        raw.set_montage(initial_montage) # type: ignore
        write_raw_snirf(raw, save_path)