#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# pylint: disable=no-name-in-module, import-error
import csv
import os

import numpy as np
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtGui import QBrush, QColor
from PyQt5.QtWidgets import (QButtonGroup, QDoubleSpinBox, QFileDialog,
                             QFormLayout, QGroupBox, QHBoxLayout, QHeaderView,
                             QLabel, QMessageBox, QPushButton, QRadioButton,
                             QSpinBox, QTableWidget, QTableWidgetItem,
                             QVBoxLayout, QWidget)


class RangeInputWidget(QWidget):
    """Widget for input range (start, end, num_points)."""

    def __init__(self, label, parent=None):
        super().__init__(parent)

        # Main layout
        layout = QFormLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)

        # Create group with label but remove border styling
        group = QGroupBox(label)
        group.setStyleSheet("QGroupBox { border: none; font-weight: bold; }")
        group_layout = QHBoxLayout(group)
        group_layout.setContentsMargins(10, 10, 10, 10)

        # Start value input with label
        start_widget = QWidget()
        start_layout = QFormLayout(start_widget)
        start_layout.setContentsMargins(0, 0, 0, 0)

        self.start_input = QDoubleSpinBox()
        self.start_input.setRange(-100.0, 100)
        self.start_input.setDecimals(3)
        self.start_input.setValue(0.0)
        start_layout.addRow("Start:", self.start_input)

        group_layout.addWidget(start_widget)

        # End value input with label
        end_widget = QWidget()
        end_layout = QFormLayout(end_widget)
        end_layout.setContentsMargins(0, 0, 0, 0)

        self.end_input = QDoubleSpinBox()
        self.end_input.setRange(-100.0, 100)
        self.end_input.setDecimals(3)
        self.end_input.setValue(-0.3)
        end_layout.addRow("End:", self.end_input)

        group_layout.addWidget(end_widget)

        layout.addRow(group)

    def get_range(self):
        """Get start and end values."""
        return (
            self.start_input.value(),
            self.end_input.value(),
        )

    def set_range(self, start, end):
        """Set start and end values."""
        self.start_input.setValue(start)
        self.end_input.setValue(end)

    def set_enabled(self, enabled):
        """Enable or disable widget."""
        self.start_input.setEnabled(enabled)
        self.end_input.setEnabled(enabled)

    def set_visible(self, visible):
        """Show or hide widget."""
        self.setVisible(visible)


class HKLScanControls(QWidget):
    """Widget for HKL scan controls."""

    # Signal emitted when calculate button is clicked
    calculateClicked = pyqtSignal()

    def __init__(self, parent=None):
        super().__init__(parent)

        # Main layout
        main_layout = QVBoxLayout(self)

        # Unified Fixed Angles panel
        fixed_angles_group = QGroupBox("Fixed Angles")
        fixed_angles_layout = QVBoxLayout(fixed_angles_group)

        # Top row: Fix χ/Fix φ buttons
        angle_selection = QWidget()
        angle_selection_layout = QHBoxLayout(angle_selection)
        angle_selection_layout.setContentsMargins(0, 0, 0, 0)

        self.fix_chi_btn = QPushButton("Fix χ")
        self.fix_phi_btn = QPushButton("Fix φ")
        
        # Make buttons checkable for toggle behavior
        for btn in (self.fix_chi_btn, self.fix_phi_btn):
            btn.setCheckable(True)
        
        self.fix_chi_btn.setChecked(True)  # Default to fixed chi

        # Create a button group for mutual exclusion
        self.angle_button_group = QButtonGroup(self)
        self.angle_button_group.addButton(self.fix_chi_btn)
        self.angle_button_group.addButton(self.fix_phi_btn)

        angle_selection_layout.addWidget(self.fix_chi_btn)
        angle_selection_layout.addWidget(self.fix_phi_btn)

        fixed_angles_layout.addWidget(angle_selection)

        # Bottom row: tth on left, chi/phi angles on right
        angle_values = QWidget()
        angle_values_layout = QHBoxLayout(angle_values)
        angle_values_layout.setContentsMargins(0, 0, 0, 0)

        # tth input (left side)
        self.tth_widget = QWidget()
        tth_layout = QFormLayout(self.tth_widget)
        tth_layout.setContentsMargins(0, 0, 0, 0)
        self.tth_input = QDoubleSpinBox()
        self.tth_input.setRange(0.0, 180.0)
        self.tth_input.setValue(150.0)
        self.tth_input.setSuffix(" °")
        tth_layout.addRow("tth:", self.tth_input)
        angle_values_layout.addWidget(self.tth_widget)
        
        # Chi input
        self.chi_widget = QWidget()
        chi_layout = QFormLayout(self.chi_widget)
        chi_layout.setContentsMargins(0, 0, 0, 0)
        self.chi_input = QDoubleSpinBox()
        self.chi_input.setRange(-180.0, 180.0)
        self.chi_input.setValue(0.0)
        self.chi_input.setSuffix(" °")
        chi_layout.addRow("χ:", self.chi_input)
        angle_values_layout.addWidget(self.chi_widget)

        # Phi input
        self.phi_widget = QWidget()
        phi_layout = QFormLayout(self.phi_widget)
        phi_layout.setContentsMargins(0, 0, 0, 0)
        self.phi_input = QDoubleSpinBox()
        self.phi_input.setRange(-180.0, 180.0)
        self.phi_input.setValue(0.0)
        self.phi_input.setSuffix(" °")
        phi_layout.addRow("φ:", self.phi_input)
        angle_values_layout.addWidget(self.phi_widget)

        fixed_angles_layout.addWidget(angle_values)
        main_layout.addWidget(fixed_angles_group)

        # HKL index selection
        hkl_group = QGroupBox("Choose a plane")
        hkl_layout = QVBoxLayout(hkl_group)

        # Plane selection - using plane-based naming like Structure Factor Calculator
        index_selection = QWidget()
        index_layout = QHBoxLayout(index_selection)
        index_layout.setContentsMargins(0, 0, 0, 0)

        self.hk_plane_toggle = QPushButton("HK plane")
        self.hl_plane_toggle = QPushButton("HL plane")
        self.kl_plane_toggle = QPushButton("KL plane")
        
        # Make buttons checkable for toggle behavior
        for btn in (self.hk_plane_toggle, self.hl_plane_toggle, self.kl_plane_toggle):
            btn.setCheckable(True)
        
        self.hk_plane_toggle.setChecked(True)  # Default to HK plane (L deactivated)

        # Create a button group for mutual exclusion
        index_button_group = QButtonGroup(self)
        index_button_group.addButton(self.hk_plane_toggle)
        index_button_group.addButton(self.hl_plane_toggle)
        index_button_group.addButton(self.kl_plane_toggle)

        index_layout.addWidget(self.hk_plane_toggle)
        index_layout.addWidget(self.hl_plane_toggle)
        index_layout.addWidget(self.kl_plane_toggle)

        hkl_layout.addWidget(index_selection)

        # Create range widgets for H, K, L
        ranges_widget = QWidget()
        ranges_layout = QVBoxLayout(ranges_widget)
        ranges_layout.setContentsMargins(0, 0, 0, 0)

        self.h_range = RangeInputWidget("H Range")
        self.k_range = RangeInputWidget("K Range")
        self.l_range = RangeInputWidget("L Range")

        ranges_layout.addWidget(self.h_range)
        ranges_layout.addWidget(self.k_range)
        ranges_layout.addWidget(self.l_range)

        hkl_layout.addWidget(ranges_widget)

        # Number of points
        points_widget = QWidget()
        points_layout = QFormLayout(points_widget)
        points_layout.setContentsMargins(0, 0, 0, 0)

        self.num_points = QSpinBox()
        self.num_points.setRange(2, 100)
        self.num_points.setValue(10)
        points_layout.addRow("Number of points:", self.num_points)

        hkl_layout.addWidget(points_widget)

        main_layout.addWidget(hkl_group)

        # Calculate button
        self.calculate_button = QPushButton("Calculate")
        self.calculate_button.clicked.connect(self.calculateClicked.emit)
        self.calculate_button.setObjectName("calculateButton")
        main_layout.addWidget(self.calculate_button)

        # Connect signals
        self.hk_plane_toggle.clicked.connect(lambda: self._set_active_plane("HK"))
        self.hl_plane_toggle.clicked.connect(lambda: self._set_active_plane("HL"))
        self.kl_plane_toggle.clicked.connect(lambda: self._set_active_plane("KL"))
        self.fix_chi_btn.clicked.connect(lambda: self._set_active_fixed_angle("chi"))
        self.fix_phi_btn.clicked.connect(lambda: self._set_active_fixed_angle("phi"))

        # Initialize widget states
        self._set_active_plane("HK")  # Set initial plane and apply styling
        self._set_active_fixed_angle("chi")  # Set initial fixed angle and apply styling

    def _update_widget_states(self):
        """Update visibility of widgets based on current plane selection.
        
        HK plane: H and K ranges visible, L range hidden
        HL plane: H and L ranges visible, K range hidden  
        KL plane: K and L ranges visible, H range hidden
        """
        if self.hk_plane_toggle.isChecked():
            # HK plane: H and K ranges visible, L range hidden
            self.h_range.set_visible(True)
            self.k_range.set_visible(True)
            self.l_range.set_visible(False)
        elif self.hl_plane_toggle.isChecked():
            # HL plane: H and L ranges visible, K range hidden
            self.h_range.set_visible(True)
            self.k_range.set_visible(False)
            self.l_range.set_visible(True)
        elif self.kl_plane_toggle.isChecked():
            # KL plane: K and L ranges visible, H range hidden
            self.h_range.set_visible(False)
            self.k_range.set_visible(True)
            self.l_range.set_visible(True)

    def _update_fixed_angle_ui(self):
        """Update UI based on which angle is fixed.
        
        If chi is fixed: Show chi input, hide phi input
        If phi is fixed: Show phi input, hide chi input
        """
        is_chi_fixed = self.fix_chi_btn.isChecked()
        # set both invisible first
        self.chi_widget.setVisible(False)
        self.phi_widget.setVisible(False)
        self.chi_widget.setVisible(is_chi_fixed)
        self.phi_widget.setVisible(not is_chi_fixed)
        self.phi_input.setEnabled(not is_chi_fixed)
        self.chi_input.setEnabled(is_chi_fixed)

    def _update_fixed_angle_styles(self, active: str):
        """Update fixed angle button colors based on active selection."""
        mapping = {
            "chi": self.fix_chi_btn,
            "phi": self.fix_phi_btn,
        }
        for name, btn in mapping.items():
            if name == active:
                btn.setChecked(True)
                btn.setProperty("class", "activeToggle")
            else:
                btn.setChecked(False)
                btn.setProperty("class", "inactiveToggle")
            # Force style refresh
            btn.style().unpolish(btn)
            btn.style().polish(btn)

    def _set_active_fixed_angle(self, angle: str):
        """Set the active fixed angle and update widget states and styling."""
        angle = angle.lower()
        self._update_fixed_angle_styles(angle)
        self._update_fixed_angle_ui()

    def _update_toggle_styles(self, active: str):
        """Update toggle button colors based on active plane."""
        mapping = {
            "HK": self.hk_plane_toggle,
            "HL": self.hl_plane_toggle,
            "KL": self.kl_plane_toggle,
        }
        for name, btn in mapping.items():
            if name == active:
                btn.setChecked(True)
                btn.setProperty("class", "activeToggle")
            else:
                btn.setChecked(False)
                btn.setProperty("class", "inactiveToggle")
            # Force style refresh
            btn.style().unpolish(btn)
            btn.style().polish(btn)

    def _set_active_plane(self, plane: str):
        """Set the active plane and update widget states and styling."""
        plane = plane.upper()
        self._update_toggle_styles(plane)
        self._update_widget_states()

    def get_scan_parameters(self):
        """Get parameters for scan."""
        # Get deactivated index based on plane selection
        # HK plane means L is fixed (deactivated)
        # HL plane means K is fixed (deactivated)
        # KL plane means H is fixed (deactivated)
        deactivated_index = None
        if self.hk_plane_toggle.isChecked():
            deactivated_index = "L"  # HK plane: L is fixed
        elif self.hl_plane_toggle.isChecked():
            deactivated_index = "K"  # HL plane: K is fixed
        elif self.kl_plane_toggle.isChecked():
            deactivated_index = "H"  # KL plane: H is fixed

        # Get fixed angle
        fixed_angle_name = "chi" if self.fix_chi_btn.isChecked() else "phi"
        fixed_angle_value = (
            self.chi_input.value()
            if self.fix_chi_btn.isChecked()
            else self.phi_input.value()
        )

        # Get ranges
        h_start, h_end = self.h_range.get_range()
        k_start, k_end = self.k_range.get_range()
        l_start, l_end = self.l_range.get_range()

        return {
            "tth": self.tth_input.value(),
            "deactivated_index": deactivated_index,
            "fixed_angle_name": fixed_angle_name,
            "fixed_angle": fixed_angle_value,
            "start_points": (h_start, k_start, l_start),
            "end_points": (h_end, k_end, l_end),
            "num_points": self.num_points.value(),
        }


class HKLScanResultsTable(QTableWidget):
    """Table to display HKL scan results with multiple solutions per HKL point."""

    def __init__(self, parent=None):
        super().__init__(parent)

        # Set up table - added Sol# column to show solution 1 or 2
        self.setColumnCount(8)
        self.setHorizontalHeaderLabels(["H", "K", "L", "Sol#", "θ (°)", "φ (°)", "χ (°)", "β (°)"])
        self.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)

        # Hide vertical header (row numbers)
        self.verticalHeader().setVisible(False)

        # Define colors for alternating HKL groups and solutions
        self.group_colors = [
            QColor(255, 255, 255),  # White - group 0
            QColor(230, 230, 240),  # Light gray - group 1
        ]
        self.solution_colors = [
            QColor(240, 248, 255),  # Alice blue - solution 1
            QColor(255, 250, 240),  # Floral white - solution 2
        ]
        self.infeasible_color = QColor(255, 199, 206)  # Light red

        # Enable sorting
        self.setSortingEnabled(True)

        # Add export button
        self.layout_wrapper = QVBoxLayout()
        self.layout_wrapper.setContentsMargins(0, 0, 0, 0)

        self.export_button = QPushButton("Export to CSV")
        self.export_button.clicked.connect(self.export_to_csv)
        self.export_button.setEnabled(False)  # Initially disabled until we have results

        self.layout_wrapper.addWidget(self)
        self.layout_wrapper.addWidget(self.export_button)

        # Store the last results for export
        self.last_results = None

    def display_results(self, results):
        """Display results in the table with visual separation for solution groups."""
        self.setSortingEnabled(False)  # Temporarily disable sorting
        self.setRowCount(0)  # Clear table

        # Check if we have results
        if not results or not results.get("success", False):
            self.export_button.setEnabled(False)
            self.last_results = None
            return

        # Store results for later export
        self.last_results = results
        
        # Get data from results
        h_values = results["H"]
        k_values = results["K"]
        l_values = results["L"]
        tth_values = results["tth"]
        theta_values = results["theta"]
        chi_values = results["chi"]
        phi_values = results["phi"]
        feasible_values = results.get("feasible", [True] * len(h_values))
        solution_groups = results.get("solution_group", list(range(len(h_values))))
        solution_indices = results.get("solution_index", [1] * len(h_values))

        # Add a row for each result
        for i in range(len(h_values)):
            row_position = self.rowCount()
            self.insertRow(row_position)

            # Add HKL values
            self.setItem(row_position, 0, QTableWidgetItem(f"{h_values[i]:.4f}"))
            self.setItem(row_position, 1, QTableWidgetItem(f"{k_values[i]:.4f}"))
            self.setItem(row_position, 2, QTableWidgetItem(f"{l_values[i]:.4f}"))
            
            # Add solution number
            sol_idx = solution_indices[i] if i < len(solution_indices) else 1
            self.setItem(row_position, 3, QTableWidgetItem(f"{sol_idx}"))

            # Add angle values
            self.setItem(row_position, 4, QTableWidgetItem(f"{theta_values[i]:.1f}"))
            self.setItem(row_position, 5, QTableWidgetItem(f"{chi_values[i]:.1f}"))
            self.setItem(row_position, 6, QTableWidgetItem(f"{phi_values[i]:.1f}"))
            self.setItem(row_position, 7, QTableWidgetItem(f"{tth_values[0]-theta_values[i]:.1f}"))

            # Determine row color based on feasibility and group
            if not feasible_values[i]:
                row_color = self.infeasible_color
            else:
                # Alternate colors by HKL group
                group_idx = solution_groups[i] if i < len(solution_groups) else i
                row_color = self.group_colors[group_idx % 2]
            
            for col in range(self.columnCount()):
                item = self.item(row_position, col)
                if item:
                    item.setBackground(QBrush(row_color))

        # Re-enable sorting and export button
        self.setSortingEnabled(True)
        self.export_button.setEnabled(True)



    def export_to_csv(self):
        """Export results to a CSV file."""
        if not self.last_results:
            QMessageBox.warning(self, "Export Error", "No results to export.")
            return

        # Open file dialog to get save location
        file_path, _ = QFileDialog.getSaveFileName(
            self, "Save Results", "", "CSV Files (*.csv);;All Files (*)"
        )

        if not file_path:
            return  # User cancelled

        # Add .csv extension if not present
        if not file_path.endswith(".csv"):
            file_path += ".csv"

        try:
            with open(file_path, "w", newline="") as csvfile:
                writer = csv.writer(csvfile)

                # Write header
                writer.writerow(
                    [
                        "H",
                        "K",
                        "L",
                        "Solution#",
                        "tth (deg)",
                        "theta (deg)",
                        "chi (deg)",
                        "phi (deg)",
                        "feasible",
                    ]
                )

                # Write data
                solution_indices = self.last_results.get("solution_index", [1] * len(self.last_results["tth"]))
                feasible_values = self.last_results.get("feasible", [True] * len(self.last_results["tth"]))
                
                for i in range(len(self.last_results["tth"])):
                    writer.writerow(
                        [
                            f"{self.last_results['H'][i]:.6f}",
                            f"{self.last_results['K'][i]:.6f}",
                            f"{self.last_results['L'][i]:.6f}",
                            f"{solution_indices[i]}",
                            f"{self.last_results['tth'][i]:.6f}",
                            f"{self.last_results['theta'][i]:.6f}",
                            f"{self.last_results['chi'][i]:.6f}",
                            f"{self.last_results['phi'][i]:.6f}",
                            f"{feasible_values[i]}",
                        ]
                    )

            QMessageBox.information(
                self, "Export Success", f"Results exported to {file_path}"
            )
        except Exception as e:
            QMessageBox.critical(
                self, "Export Error", f"Error exporting results: {str(e)}"
            )

    def get_widget(self):
        """Return the widget containing the table and export button."""
        container = QWidget()
        container.setLayout(self.layout_wrapper)
        return container

