"""
Filename: updateoptodes.py
Description: Methods to update optode locations for FLARES

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

import os
from pathlib import Path

import pandas as pd
import numpy as np

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

from mne.io import read_raw_snirf
from mne_nirs.io import write_raw_snirf
from mne.channels import make_dig_montage

from src.shared.shareddata import APP_NAME


class UpdateOptodesWindow(QWidget):

    def __init__(self, parent=None):
        super().__init__(parent, Qt.WindowType.Window)
        self.setWindowTitle(f"Update optode positions - {APP_NAME.upper()}")
        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(f"Parameter Info - {APP_NAME.upper()}")
        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)