"""
Python rewrite of NEURON extracellular stimulation code
Replaces the HOC workflow for loading SWC, applying potentials, and detecting firing
"""

import numpy as np
from neuron import h
from neuron.units import ms, mV
import sys

# Load standard files
h.load_file("stdrun.hoc")
h.load_file("import3d.hoc")

class NeuronModel:
    def __init__(self, swc_file="pruned.swc"):
        """Initialize neuron model from SWC file"""
        self.swc_file = swc_file
        self.sections = []
        self.fired = False
        self.detector = None
        
        # Import morphology
        self._import_swc()
        
        # Set parameters
        self._set_parameters()
        
        # Calculate segment coordinates
        self._calculate_segment_coords()
        
    def _import_swc(self):
        """Load topology from SWC file"""
        cell = h.Import3d_SWC_read()
        cell.input(self.swc_file)
        i3d = h.Import3d_GUI(cell, 0)
        i3d.instantiate(None)
        
        # Store sections
        self.sections = [sec for sec in h.allsec()]
        print(f"Loaded {len(self.sections)} sections from {self.swc_file}")
        
    def _set_parameters(self):
        """Set up geometry and membrane parameters"""
        # Global parameters for all sections
        for sec in h.allsec():
            sec.Ra = 100
            sec.insert('pas')
            sec.g_pas = 3e-5
            sec.e_pas = -75
            sec.insert('extracellular')
            sec.insert('xtra')
        
        # Soma-specific
        for sec in h.allsec():
            if 'soma' in sec.name():
                sec.cm = 1
                sec.insert('CaDynamics_E2_soma')
                sec.insert('Ca_HVA_soma')
                sec.insert('Ca_LVAst_soma')
                sec.insert('Ih')
                sec.insert('NaTs2_t_soma')
                sec.insert('SK_E2_soma')
                sec.insert('SKv3_1_soma')
        
        # Axon-specific
        for sec in h.allsec():
            if 'axon' in sec.name():
                sec.cm = 1
                sec.insert('CaDynamics_E2')
                sec.insert('Ca_HVA')
                sec.insert('Ca_LVAst')
                sec.insert('K_Pst')
                sec.insert('K_Tst')
                sec.insert('NaTa_t')
                sec.insert('Nap_Et2')
                sec.insert('SK_E2')
                sec.insert('SKv3_1')
        
        # Dendrite-specific
        for sec in h.allsec():
            if 'dend' in sec.name():
                sec.cm = 2
                sec.insert('Ih')
        
        # Apical dendrite-specific
        for sec in h.allsec():
            if 'apic' in sec.name():
                sec.cm = 2
                sec.insert('Ih')
                sec.insert('Im')
                sec.insert('NaTs2_t_apic')
                sec.insert('SKv3_1_apic')
        
        # Set nseg based on AC length constant
        freq = 100  # Hz
        d_lambda = 0.01
        
        for sec in h.allsec():
            sec.nseg = int((sec.L / (d_lambda * self._lambda_f(sec, freq)) + 0.999) / 2) * 2 + 1
    
    def _lambda_f(self, sec, freq):
        """Calculate AC length constant for a section"""
        # Simplified lambda_f calculation
        return 1e5 * np.sqrt(sec.diam / (4 * np.pi * freq * sec.Ra * sec.cm))
    
    def _calculate_segment_coords(self):
        """Compute xyz coords of segments (interpxyz.hoc equivalent)"""
        for sec in h.allsec():
            if not hasattr(sec, 'n3d') or sec.n3d() == 0:
                continue
                
            # Get 3D point data
            n = sec.n3d()
            xx = np.array([sec.x3d(i) for i in range(n)])
            yy = np.array([sec.y3d(i) for i in range(n)])
            zz = np.array([sec.z3d(i) for i in range(n)])
            length = np.array([sec.arc3d(i) for i in range(n)])
            
            # Normalize length
            if length[-1] > 0:
                length_norm = length / length[-1]
            else:
                continue
            
            # Create range for segments
            nseg = sec.nseg
            range_vals = np.linspace(0, 1, nseg + 2)
            range_vals = range_vals - 1/(2*nseg)
            range_vals[0] = 0
            range_vals[-1] = 1
            
            # Interpolate coordinates
            x_interp = np.interp(range_vals, length_norm, xx)
            y_interp = np.interp(range_vals, length_norm, yy)
            z_interp = np.interp(range_vals, length_norm, zz)
            
            # Assign to segments (skip endpoints 0 and 1)
            for i, seg in enumerate(sec, start=1):
                seg.x_xtra = x_interp[i]
                seg.y_xtra = y_interp[i]
                seg.z_xtra = z_interp[i]
    
    def _set_pointers(self):
        """Link extracellular potential to xtra mechanism"""
        for sec in h.allsec():
            if hasattr(sec, 'ismembrane'):
                for seg in sec:
                    h.setpointer(seg._ref_e_extracellular, 'ex', seg.xtra)
    
    def load_potentials_from_file(self, potential_file):
        """
        Load extracellular potentials from file.
        File format: x y z potential (one per line)
        """
        print(f"Loading potentials from {potential_file}")
        
        try:
            data = np.loadtxt(potential_file)
            coords = data[:, :3]
            potentials = data[:, 3]
            
            # For each point in the file, find closest segment and assign potential
            for x, y, z, v in data:
                min_dist = float('inf')
                closest_sec = None
                closest_x = 0.5
                
                for sec in h.allsec():
                    if sec.n3d() == 0:
                        continue
                    
                    for i in range(sec.n3d()):
                        dx = x - sec.x3d(i)
                        dy = y - sec.y3d(i)
                        dz = z - sec.z3d(i)
                        dist = np.sqrt(dx*dx + dy*dy + dz*dz)
                        
                        if dist < min_dist:
                            min_dist = dist
                            closest_sec = sec
                            if sec.L > 0:
                                closest_x = sec.arc3d(i) / sec.L
                            else:
                                closest_x = 0.5
                
                if closest_sec is not None:
                    closest_sec(closest_x).e_extracellular = v
            
            print(f"Loaded {len(data)} potential values")
            
        except Exception as e:
            print(f"Error loading potential file: {e}")
            sys.exit(1)
    
    def load_potentials_ordered(self, potential_file):
        """
        Load potentials from file where order matches segment order.
        File format: one potential value per line
        """
        print(f"Loading ordered potentials from {potential_file}")
        
        try:
            potentials = np.loadtxt(potential_file)
            
            idx = 0
            for sec in h.allsec():
                for seg in sec:
                    if idx < len(potentials):
                        seg.es_xtra = potentials[idx]
                        idx += 1
            
            print(f"Loaded {idx} potential values")
            
        except Exception as e:
            print(f"Error loading potential file: {e}")
            sys.exit(1)
    
    def setup_stimulus(self, delay=0, duration=0.1, scale=-70):
        """
        Set up extracellular stimulus waveform.
        
        Parameters:
        -----------
        delay : float
            Stimulus delay in ms
        duration : float
            Stimulus duration in ms
        scale : float
            Stimulus scaling factor (unitless)
        """
        # Create stimulus vectors
        stim_scale = h.Vector([0, 0, 1, 1, 0, 0])
        stim_scale.mul(scale)
        
        stim_time = h.Vector([0, delay, delay, delay+duration, 
                              delay+duration, delay+duration+1])
        
        # Attach to xtra mechanism (only need to do once since stim_xtra is GLOBAL)
        for sec in h.allsec():
            if hasattr(sec(0.5), 'xtra'):
                stim_scale.play(sec(0.5).xtra._ref_stim_xtra, stim_time, 1)
                break
        
        print(f"Stimulus: delay={delay}ms, duration={duration}ms, scale={scale}")
    
    def setup_detector(self, section_name='axon', threshold=0):
        """Set up spike detector"""
        self.fired = False
        
        # Find the target section
        target_sec = None
        for sec in h.allsec():
            if section_name in sec.name():
                target_sec = sec
                break
        
        if target_sec is None:
            print(f"Warning: Could not find section matching '{section_name}'")
            target_sec = self.sections[0]
        
        # Create detector
        self.detector = h.NetCon(target_sec(0.5)._ref_v, None, sec=target_sec)
        self.detector.threshold = threshold
        self.detector.record(self._handle_spike)
    
    def _handle_spike(self):
        """Callback for spike detection"""
        self.fired = True
        h.stoprun = 1
    
    def run_simulation(self, tstop=10):
        """Run simulation"""
        self.fired = False
        h.finitialize(-65 * mV)
        h.continuerun(tstop * ms)
        return self.fired
    
    def find_threshold(self, delay=0, duration=0.1, output_file="threshold_results.txt"):
        """
        Binary search to find threshold stimulus amplitude.
        
        Parameters:
        -----------
        delay : float
            Stimulus delay in ms
        duration : float  
            Stimulus duration in ms
        output_file : str
            File to save results
        """
        print("Finding threshold...")
        
        # Initialize detector
        self.setup_detector('axon')
        
        # Binary search parameters
        sca = -20
        sca_top = sca
        sca_bottom = sca
        first_fired = False
        first_unfired = False
        max_iter = 100
        tolerance = 1e-2
        
        results = []
        
        for iteration in range(1, max_iter + 1):
            if abs(sca_top - sca_bottom) < tolerance and iteration > 1:
                break
            
            # Set stimulus and run
            self.setup_stimulus(delay, duration, sca)
            fired = self.run_simulation()
            
            result = f"Iteration {iteration}: SCA = {sca:.6f}, fired = {int(fired)}"
            print(result)
            results.append(result)
            
            # Update search bounds
            if fired and not first_unfired:
                sca_top = sca
                sca = sca / 2
                sca_bottom = sca
                first_fired = True
            elif fired and first_unfired:
                sca_top = sca
                sca = (sca_top + sca_bottom) / 2
                first_fired = True
            elif not fired and not first_fired:
                sca_bottom = sca
                sca = sca * 2
                sca_top = sca
                first_unfired = True
            elif not fired and first_fired:
                sca_bottom = sca
                sca = (sca_top + sca_bottom) / 2
                first_unfired = True
        
        # Save results
        with open(output_file, 'w') as f:
            for result in results:
                f.write(result + '\n')
        
        print(f"Threshold: {sca:.6f}")
        print(f"Results saved to {output_file}")
        return sca
    
    def export_segment_locations(self, output_file="locs_all_seg.txt"):
        """Export segment positions and parent connections"""
        with open(output_file, 'w') as f:
            for sec in h.allsec():
                parent_sec = sec.parentseg()
                has_parent = parent_sec is not None
                
                for seg in sec:
                    x_coord = seg.x_xtra
                    y_coord = seg.y_xtra
                    z_coord = seg.z_xtra
                    
                    # Determine parent coordinates
                    if seg == list(sec)[0] and has_parent:
                        # First segment - parent is last segment of parent section
                        px = parent_sec.x_xtra
                        py = parent_sec.y_xtra
                        pz = parent_sec.z_xtra
                    elif seg == list(sec)[0] and not has_parent:
                        # Root section first segment
                        px = py = pz = float('nan')
                    else:
                        # Parent is previous segment in same section
                        prev_seg = list(sec)[list(sec).index(seg) - 1]
                        px = prev_seg.x_xtra
                        py = prev_seg.y_xtra
                        pz = prev_seg.z_xtra
                    
                    f.write(f"{x_coord:.4f} {y_coord:.4f} {z_coord:.4f} "
                           f"{px:.4f} {py:.4f} {pz:.4f}\n")
        
        print(f"Segment locations exported to {output_file}")


def main():
    """Main execution function"""
    # Configuration
    SWC_FILE = "pruned.swc"
    #POTENTIAL_FILE = "pruned.ptntl"  # x y z potential format
    POTENTIAL_TXT = "pruned.txt"  # ordered potential values
    
    DELAY = 0      # ms
    DURATION = 0.1 # ms
    SCALE = -70    # unitless
    
    OUTPUT_FILE = "threshold_results.txt"
    
    # Create model
    model = NeuronModel(swc_file=SWC_FILE)
    
    # Load potentials (choose one method)
    #model.load_potentials_from_file(POTENTIAL_FILE)
    model.load_potentials_ordered(POTENTIAL_TXT)
    
    # Export segment locations
    model.export_segment_locations()
    
    # Find threshold
    threshold = model.find_threshold(DELAY, DURATION, OUTPUT_FILE)
    
    print(f"\nFinal threshold: {threshold:.6f}")


if __name__ == "__main__":
    main()