"""nsram.export — Export NS-RAM models to external frameworks.

Bridges to:
  - Brian2 (neuroscience SNN simulator)
  - NESTML (NEST neuron description language)
  - SPICE netlist (for ngspice / LTSpice / Xyce)
  - Hdl21 (Python-native circuit description)

Each export produces a string or file that researchers can
directly use in their own tools.
"""

from nsram.physics import DeviceParams


def to_brian2(params: DeviceParams = None, with_stp: bool = True) -> str:
    """Export NS-RAM neuron model as Brian2 equations string.

    Returns a string usable as:
        from brian2 import *
        eqs = nsram.export.to_brian2()
        G = NeuronGroup(128, eqs, threshold='Vm > V_thresh',
                        reset='Vm = V_reset', refractory='t_ref')

    The equations model:
      - Leaky integration with exponential near-threshold (AdEx)
      - Body charge dynamics (floating bulk potential)
      - Optional: SRH charge trapping → threshold modulation
    """
    if params is None:
        params = DeviceParams()

    eqs = f'''
    dVm/dt = (I_leak + I_aval + I_syn + I_ext + sigma*xi*sqrt(Hz)) / Cm : volt
    I_leak = g_leak * (E_L - Vm) : amp
    I_aval = I0 * exp(clip((Vm - BVpar) / (Ne * Vt), -20, 20)) : amp
    BVpar = BV0 - k_vg * Vg1 : volt
    I_syn : amp  # from Synapses object
    I_ext : amp  # external input

    # Device parameters (from Pazos et al., Nature 640, 2025)
    Cm = {params.C_mem} * farad : farad (constant)
    g_leak = {params.g_leak:.4e} * siemens : siemens (constant)
    E_L = {params.V_leak_rest} * volt : volt (constant)
    I0 = {params.Is} * amp : amp (constant)
    BV0 = {params.BV0} * volt : volt (constant)
    k_vg = {params.k_vg} : 1 (constant)
    Vg1 : volt  # gate voltage (input-modulated)
    Ne = {params.Ne} : 1 (constant)
    Vt = 0.02585 * volt : volt (constant)
    sigma = 1e-12 * amp : amp (constant)
    '''

    if with_stp:
        eqs += f'''
    # Charge trapping (SRH kinetics = Tsodyks-Markram STP)
    dQ/dt = (k_cap * (1 - Q) * rate_est - k_em * Q) : 1
    drate_est/dt = -rate_est / tau_rate : Hz
    delta_Vth = -alpha_Q * Q : volt
    V_thresh_eff = V_thresh + delta_Vth : volt

    k_cap : Hz  # capture rate (depends on VG2)
    k_em = 200 * Hz : Hz (constant)
    tau_rate = 1 * ms : second (constant)
    alpha_Q = 0.5 * volt : volt (constant)
    V_thresh = {params.V_thresh} * volt : volt (constant)
    '''
    else:
        eqs += f'''
    V_thresh = {params.V_thresh} * volt : volt (constant)
    '''

    return eqs.strip()


def to_brian2_example() -> str:
    """Complete Brian2 example script using NS-RAM neuron model."""
    return '''#!/usr/bin/env python3
"""NS-RAM neuron model in Brian2.

Based on: Pazos et al., Nature 640, 69-76 (2025)
Generated by: nsram.export.to_brian2_example()
"""
from brian2 import *
from nsram.export import to_brian2

# Get NS-RAM equations
eqs = to_brian2(with_stp=True)

# Create neuron group
N = 128
G = NeuronGroup(N, eqs,
                threshold='Vm > V_thresh',
                reset='Vm = 0*volt; rate_est += 1*Hz',
                refractory=1.6*us,
                method='euler')

# Set per-neuron gate voltages (heterogeneous)
G.Vg1 = '(0.25 + 0.20 * rand()) * volt'
G.k_cap = '1000 / (1 + exp((0.40 - 0.35 + 0.12*rand()) / 0.05)) * Hz'

# Input current
G.I_ext = '0.5 * nA * sin(2*pi*1*kHz*t)'

# Sparse recurrent synapses
S = Synapses(G, G, 'w : amp', on_pre='I_syn_post += w')
S.connect(p=0.15)
S.w = '(0.1 * randn() + 0.05) * nA'

# Monitors
spikemon = SpikeMonitor(G)
statemon = StateMonitor(G, ['Vm', 'Q'], record=[0, 1, 2, 3])

# Run
run(10*ms)

# Plot
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
ax1.plot(spikemon.t/ms, spikemon.i, '.k', markersize=1)
ax1.set_ylabel('Neuron #'); ax1.set_title('NS-RAM Spike Raster (Brian2)')
for i in range(4):
    ax2.plot(statemon.t/ms, statemon.Vm[i]/mV, label=f'N{i}')
ax2.set_xlabel('Time (ms)'); ax2.set_ylabel('Vm (mV)'); ax2.legend()
plt.tight_layout()
plt.savefig('nsram_brian2_example.png', dpi=150)
print(f"Spikes: {spikemon.num_spikes}")
'''


def to_nestml() -> str:
    """Export NS-RAM neuron model as NESTML description.

    The generated .nestml file can be compiled by NESTML into
    a NEST C++ module for large-scale simulation.

    Usage:
        with open('nsram_neuron.nestml', 'w') as f:
            f.write(nsram.export.to_nestml())
        # Then: nestml --input nsram_neuron.nestml --target NEST
    """
    return '''
/**
 * NS-RAM Neuron Model — Floating-Body Transistor with Charge Trapping
 *
 * Based on: Pazos et al., Nature 640, 69-76 (2025)
 * Generated by: nsram.export.to_nestml()
 *
 * Implements AdEx-LIF with:
 *   - Impact ionization avalanche (exponential near threshold)
 *   - SRH charge trapping (= Tsodyks-Markram STP on threshold)
 *   - Die-to-die variability via parameter heterogeneity
 */
model nsram_neuron:

    state:
        V_m mV = 0.0 mV           # Membrane potential (body voltage)
        Q real = 0.0               # Trapped charge fraction [0, 1]
        rate_est Hz = 0.0 Hz       # Spike rate estimator

    equations:
        # Membrane dynamics (AdEx-LIF with avalanche exponential)
        I_leak pA = g_L * (E_L - V_m)
        I_aval pA = I_0 * exp(min((V_m - BVpar) / (Ne * V_T), 20.0))
        BVpar mV = BV_0 - k_vg * Vg1

        V_m' = (I_leak + I_aval + I_syn + I_ext + I_noise) / C_m

        # Charge trapping (SRH kinetics)
        Q' = k_cap * (1.0 - Q) * rate_est - k_em * Q
        rate_est' = -rate_est / tau_rate

        # Threshold modulation
        V_thresh_eff mV = V_thresh - alpha_Q * Q

    parameters:
        C_m pF = 0.050            # Membrane capacitance (50 fF)
        g_L nS = 50.0             # Leak conductance
        E_L mV = 0.0 mV          # Resting potential
        V_thresh mV = 1364.0 mV   # Base threshold (1.364 V)
        V_reset mV = 0.0 mV      # Reset voltage
        t_ref ms = 0.0016         # Refractory period (1.6 us)
        Delta_T mV = 100.0 mV     # Exponential slope (0.1 V)

        # Avalanche parameters (from Pazos SPICE model)
        I_0 pA = 0.0001           # Saturation current (1e-16 A = 0.0001 pA)
        BV_0 mV = 3500.0 mV       # Base breakdown voltage (3.5 V)
        k_vg real = 1.5           # Gate voltage sensitivity
        Ne real = 1.5             # Avalanche relaxation factor
        V_T mV = 25.85 mV         # Thermal voltage at 300K
        Vg1 mV = 350.0 mV         # Gate voltage (input-modulated)

        # Charge trapping (SRH = Tsodyks-Markram)
        k_cap Hz = 500.0 Hz       # Capture rate (depends on VG2)
        k_em Hz = 200.0 Hz        # Emission rate
        tau_rate ms = 1.0 ms      # Rate estimator time constant
        alpha_Q mV = 500.0 mV     # Max threshold shift from trapping (0.5 V)

    input:
        I_syn pA <- spike
        I_ext pA <- continuous
        I_noise pA <- continuous

    output:
        spike

    update:
        integrate_odes()

    onReceive(spike):
        V_m += 0.1 mV  # Synaptic weight (scaled by Synapse)

    onThreshold:
        V_m = V_reset
        rate_est += 1000.0 Hz  # Spike impulse
        emit_spike()

    @threshold_condition
        V_m >= V_thresh_eff
'''.strip()


def to_spice_subcircuit(params: DeviceParams = None) -> str:
    """Export NS-RAM cell as SPICE subcircuit (.subckt).

    Can be included in any SPICE simulator (ngspice, LTSpice, Xyce).
    """
    if params is None:
        params = DeviceParams()

    return f"""* NS-RAM 2T Cell Subcircuit
* Pazos et al., Nature 640, 69-76 (2025)
* Generated by nsram.export.to_spice_subcircuit()
*
* Ports: drain gate1 gate2 source
* Internal: body (floating bulk)

.subckt NSRAM_2T drain gate1 gate2 source
+ params: BV0={params.BV0} k_vg={params.k_vg} Cb={params.Cb}
+         Rb_default={params.Rb_default} Is={params.Is} Bf={params.Bf}

* Main NMOS (M1) with floating body
M1 drain gate1 source body NMOS W=10u L=0.25u

* Parasitic BJT (avalanche action)
.model avalBJT NPN(Is={{Is}} Bf={{Bf}} Nf=0.9 Ne=1.5
+ Vaf=40 BVcbo='{{BV0}}-{{k_vg}}*V(gate1)' nBVcbo=0)
Q1 drain body source avalBJT

* Avalanche Zener diode
.model avalD D(Is=1e-21 bv='0.9*({{BV0}}-{{k_vg}}*V(gate1))'
+ nbv=7 Tbv1=-21.3u type=Zener)
D1 source body avalD

* Floating body capacitance
Cb body 0 {{Cb}}

* Bulk resistance (controlled by M2/gate2)
* Simplified: Rb = Rb_default * exp(-(V(gate2)-0.4)/0.05)
Rb body 0 {{Rb_default}}

.ends NSRAM_2T
"""
