import numpy as np
import pyqtgraph as pg
from pyqtgraph.flowchart import Flowchart, Node
from pyqtgraph.flowchart.NodeLibrary import NodeLibrary
from pyqtgraph.flowchart.library.common import CtrlNode
from PyQt5 import QtWidgets, QtCore, QtGui

class Beam:
    def __init__(self, wavelengths=None):
        if wavelengths is None:
            self.wavelengths = np.arange(300, 900, 1) # 1nm resolution
        else:
            self.wavelengths = wavelengths
        
        # Components: {source_name: intensity_array}
        self.components = {}

    def add_component(self, name, intensity):
        if len(intensity) != len(self.wavelengths):
            # Intepolate to match beam wavelengths
            # Assuming intensity is on same range if not specified, but here we require matching length for simplicity
            # or we could pass x,y and interpolate.
            # For now, assume aligned.
            pass
        self.components[name] = intensity

    def apply_transmission(self, transmission_spectrum):
        """
        Apply a transmission spectrum (0-1) to all components.
        transmission_spectrum: array matching wavelengths
        """
        new_beam = Beam(self.wavelengths)
        for name, intensity in self.components.items():
            new_beam.components[name] = intensity * transmission_spectrum
        return new_beam
    
    def copy(self):
        new_beam = Beam(self.wavelengths)
        for name, intensity in self.components.items():
            new_beam.components[name] = intensity.copy()
        return new_beam

    def __repr__(self):
        return f"<Beam with {len(self.components)} components>"

class LaserNode(CtrlNode):
    """
    Generates a beam with a specific wavelength and power.
    """
    nodeName = 'Laser'

    def __init__(self, name):
        ui = [
            {'name': 'name', 'type': 'text', 'value': 'Laser 1'},
            {'name': 'wavelength', 'type': 'spin', 'value': 488, 'min': 300, 'max': 2000},
            {'name': 'power', 'type': 'spin', 'value': 100, 'suffix': 'mW', 'min': 0},
            {'name': 'bandwidth', 'type': 'spin', 'value': 2.0, 'suffix': 'nm', 'min': 0.1}, # FWHM
        ]
        super().__init__(name, ui=ui, terminals={'Out': {'io': 'out'}})
        
    def process(self, **kwargs):
        wl_center = self.ctrls['wavelength'].value()
        power = self.ctrls['power'].value()
        width = self.ctrls['bandwidth'].value()
        name = self.ctrls['name'].text()

        beam = Beam()
        x = beam.wavelengths
        
        # Gaussian profile
        sigma = width / 2.355
        profile = np.exp(-(x - wl_center)**2 / (2 * sigma**2))
        # Normalize so area/sum represents power (simple approximation)
        if np.sum(profile) > 0:
            profile = profile / np.sum(profile) * power
        
        beam.add_component(name, profile)
        return {'Out': beam}

class FilterNode(CtrlNode):
    """
    Applies a transmission filter to the beam.
    """
    nodeName = 'Filter'

    def __init__(self, name):
        ui = [
            {'name': 'cutoff', 'type': 'spin', 'value': 500, 'min': 300, 'max': 900},
            {'name': 'type', 'type': 'combo', 'values': ['LongPass', 'ShortPass', 'BandPass']},
            {'name': 'width', 'type': 'spin', 'value': 20, 'min': 1, 'max': 500}, # For bandpass
        ]
        super().__init__(name, ui=ui, terminals={'In': {'io': 'in'}, 'Out': {'io': 'out'}})

    def process(self, In, **kwargs):
        if In is None:
            return {'Out': None}
        
        cutoff = self.ctrls['cutoff'].value()
        ftype = self.ctrls['type'].currentText()
        width = self.ctrls['width'].value()
        
        x = In.wavelengths
        transmission = np.ones_like(x, dtype=float)
        
        # Simple step functions for simulation (ideal filters)
        # Could be improved by loading real spectra
        if ftype == 'LongPass':
            transmission[x < cutoff] = 0.0
        elif ftype == 'ShortPass':
            transmission[x > cutoff] = 0.0
        elif ftype == 'BandPass':
            lower = cutoff - width/2
            upper = cutoff + width/2
            transmission[(x < lower) | (x > upper)] = 0.0
            
        return {'Out': In.apply_transmission(transmission)}

class SampleNode(CtrlNode):
    """
    Absorbs excitation and emits fluorescence.
    """
    nodeName = 'Sample'

    def __init__(self, name):
        ui = [
            {'name': 'name', 'type': 'text', 'value': 'Dye A'},
            {'name': 'ext_coef', 'type': 'spin', 'value': 100000, 'min': 0},
            {'name': 'qy', 'type': 'spin', 'value': 0.8, 'min': 0, 'max': 1.0},
            {'name': 'abs_max', 'type': 'spin', 'value': 490, 'min': 300, 'max': 900},
            {'name': 'em_max', 'type': 'spin', 'value': 520, 'min': 300, 'max': 900},
            {'name': 'conc', 'type': 'spin', 'value': 1.0, 'suffix': 'uM'},
            {'name': 'path_len', 'type': 'spin', 'value': 1.0, 'suffix': 'cm'},
        ]
        super().__init__(name, ui=ui, terminals={'In': {'io': 'in'}, 'Out': {'io': 'out'}})
        
    def process(self, In, **kwargs):
        if In is None:
            return {'Out': None}
        
        name = self.ctrls['name'].text()
        ext_coef_max = self.ctrls['ext_coef'].value()
        qy = self.ctrls['qy'].value()
        abs_max = self.ctrls['abs_max'].value()
        em_max = self.ctrls['em_max'].value()
        conc = self.ctrls['conc'].value() * 1e-6 # convert uM to M
        path_len = self.ctrls['path_len'].value()
        
        x = In.wavelengths
        
        # 1. Create Absorption Spectrum (Gaussian approximation)
        sigma_abs = 30 / 2.355 # guess width
        abs_spectrum = np.exp(-(x - abs_max)**2 / (2 * sigma_abs**2))
        # Normalize peak to 1 (shape factor)
        
        # Extinction coefficient at each wavelength
        epsilon = abs_spectrum * ext_coef_max
        
        # Absorbance A = eps * c * l
        A = epsilon * conc * path_len
        
        # Transmission T = 10^(-A)
        transmission = 10**(-A)
        
        # Absorbed fraction = 1 - T
        absorbed_fraction = 1.0 - transmission
        
        # 2. Create Emission Spectrum (Gaussian approximation)
        sigma_em = 40 / 2.355
        em_shape = np.exp(-(x - em_max)**2 / (2 * sigma_em**2))
        if np.sum(em_shape) > 0:
            em_shape = em_shape / np.sum(em_shape) # Normalized to area 1 (probability distribution)
            
        # 3. Calculate Emission for each component in input beam
        out_beam = In.apply_transmission(transmission) # The transmitted excitation
        
        for comp_name, intensity in In.components.items():
            # Calculate total absorbed power for this component
            # We approximate by summing (intensity * absorbed_fraction)
            # In reality, should integrate. intensity is "power per bin".
            absorbed_power = np.sum(intensity * absorbed_fraction)
            
            if absorbed_power > 0:
                # Emitted power = Absorbed * QY
                emitted_power = absorbed_power * qy
                
                # Add emission component
                # We name it "Emission_{DyeName}_from_{SourceName}"
                # But to keep it simple for crosstalk matrix, maybe just "Emission_{DyeName}"
                # If multiple sources excite it, they sum up?
                # Let's separate by Dye Name.
                
                emission_comp_name = f"{name} (em)"
                
                current_emission = out_beam.components.get(emission_comp_name, np.zeros_like(x))
                out_beam.components[emission_comp_name] = current_emission + (em_shape * emitted_power)
                
        return {'Out': out_beam}

class DichroicNode(CtrlNode):
    """
    Splits beam: Reflects one range, Transmits another.
    """
    nodeName = 'Dichroic'

    def __init__(self, name):
        ui = [
            {'name': 'cutoff', 'type': 'spin', 'value': 500, 'min': 300, 'max': 900},
            {'name': 'type', 'type': 'combo', 'values': ['LongPass (Reflect Short)', 'ShortPass (Reflect Long)']},
        ]
        super().__init__(name, ui=ui, terminals={
            'In': {'io': 'in'},
            'Transmitted': {'io': 'out'},
            'Reflected': {'io': 'out'}
        })

    def process(self, In, **kwargs):
        if In is None:
            return {'Transmitted': None, 'Reflected': None}
            
        cutoff = self.ctrls['cutoff'].value()
        dtype = self.ctrls['type'].currentText()
        
        x = In.wavelengths
        transmission = np.zeros_like(x, dtype=float)
        
        if dtype.startswith('LongPass'):
            transmission[x > cutoff] = 1.0
        else:
            transmission[x < cutoff] = 1.0
            
        reflection = 1.0 - transmission
        
        t_beam = In.apply_transmission(transmission)
        r_beam = In.apply_transmission(reflection)
        
        return {'Transmitted': t_beam, 'Reflected': r_beam}

class DetectorNode(CtrlNode):
    """
    Integrates signal.
    """
    nodeName = 'Detector'

    def __init__(self, name):
        ui = [
            {'name': 'name', 'type': 'text', 'value': 'Det 1'},
            {'name': 'min_wl', 'type': 'spin', 'value': 300, 'min': 300, 'max': 2000},
            {'name': 'max_wl', 'type': 'spin', 'value': 900, 'min': 300, 'max': 2000},
        ]
        super().__init__(name, ui=ui, terminals={'In': {'io': 'in'}, 'Signal': {'io': 'out'}})
        
    def process(self, In, **kwargs):
        if In is None:
            return {'Signal': {}}
        
        name = self.ctrls['name'].text()
        min_wl = self.ctrls['min_wl'].value()
        max_wl = self.ctrls['max_wl'].value()
        
        x = In.wavelengths
        mask = (x >= min_wl) & (x <= max_wl)
        
        signals = {}
        for comp_name, intensity in In.components.items():
            # Integrate intensity over bandwidth
            total_signal = np.sum(intensity[mask])
            signals[comp_name] = total_signal
            
        # Return a dict of signals {component_name: value}
        # We wrap it in a dict with detector name for the aggregator
        return {'Signal': {name: signals}}


class CrosstalkCalculatorWidget(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.layout = QtWidgets.QHBoxLayout(self)
        
        # Create library
        self.library = NodeLibrary()
        for node_class in [LaserNode, FilterNode, SampleNode, DichroicNode, DetectorNode]:
            self.library.addNodeType(node_class, [('Optical',)])
        
        # Flowchart
        self.flowchart = Flowchart(terminals={}, library=self.library)
        self.fc_widget = self.flowchart.widget()
        
        # Splitter
        splitter = QtWidgets.QSplitter(QtCore.Qt.Horizontal)
        self.layout.addWidget(splitter)
        
        # Left: Flowchart List / Controls
        left_widget = QtWidgets.QWidget()
        left_layout = QtWidgets.QVBoxLayout(left_widget)
        left_layout.addWidget(self.fc_widget)
        splitter.addWidget(left_widget)
        
        # Center: Plot of Flowchart (The visual graph)
        # Note: Flowchart.widget() shows the tree. Flowchart.viewBox (in a PlotWidget or similar) shows the nodes.
        # Actually, fc.widget() usually contains the Ctrl panel.
        # We need to add the chart graphics view.
        # Wait, Flowchart.widget() is a container that has 'chartWidget' (graphics view) and 'ctrlWidget'.
        # But by default it might be just the ctrl.
        # Let's constructing it manually if needed, but typically flowchart.widget() gives a splitter with both.
        # Let's check pyqtgraph source if I can't remember.
        # Usually:
        # w = fc.widget()
        # self.layout.addWidget(w)
        # w contains the node graph view.
        
        # Right: Results
        right_widget = QtWidgets.QWidget()
        right_layout = QtWidgets.QVBoxLayout(right_widget)
        
        self.calc_button = QtWidgets.QPushButton("Calculate Matrix")
        self.calc_button.clicked.connect(self.calculate_matrix)
        right_layout.addWidget(self.calc_button)
        
        self.results_table = QtWidgets.QTableWidget()
        right_layout.addWidget(self.results_table)
        
        splitter.addWidget(right_widget)
        
        # Add default nodes
        self.init_default_graph()

    def init_default_graph(self):
        # Add a simple Laser -> Sample -> Detector
        laser = self.flowchart.createNode('Laser', pos=(0, 0))
        sample = self.flowchart.createNode('Sample', pos=(150, 0))
        detector = self.flowchart.createNode('Detector', pos=(300, 0))
        
        self.flowchart.connectTerminals(laser['Out'], sample['In'])
        self.flowchart.connectTerminals(sample['Out'], detector['In'])

    def calculate_matrix(self):
        # Trigger processing
        # output = self.flowchart.output() # This gets output of terminal nodes if they are connected to output
        # But our detectors are nodes. We need to harvest data from all DetectorNodes.
        
        nodes = self.flowchart.nodes()
        detectors = [n for n in nodes.values() if isinstance(n, DetectorNode)]
        
        if not detectors:
            return

        # We need to force update?
        # Flowchart updates automatically on change usually.
        # But we can get the output of each detector node directly.
        
        # Structure of results:
        # Rows: Components (Laser 1, Dye A (em), ...)
        # Cols: Detectors
        
        all_signals = {} # {detector_name: {component_name: value}}
        all_components = set()
        
        for det in detectors:
            # Get output value
            # We can access .outputValues() of the node?
            # Or we can process().
            # Ideally, we look at the cached output.
            # 'Signal' is the output terminal name.
            val = det.terminals['Signal'].value() 
            # val should be {'DetName': {comp: val}}
            if val:
                for det_name, sigs in val.items():
                    all_signals[det_name] = sigs
                    all_components.update(sigs.keys())
                    
        # Update Table
        comps = sorted(list(all_components))
        dets = sorted(list(all_signals.keys()))
        
        self.results_table.setRowCount(len(comps))
        self.results_table.setColumnCount(len(dets))
        self.results_table.setHorizontalHeaderLabels(dets)
        self.results_table.setVerticalHeaderLabels(comps)
        
        for r, comp in enumerate(comps):
            for c, det in enumerate(dets):
                val = all_signals[det].get(comp, 0.0)
                self.results_table.setItem(r, c, QtWidgets.QTableWidgetItem(f"{val:.2e}"))

