# =============================================================================
# QUANTUM BATTERY – FINAL PUBLICATION SCRIPT (QuTiP 5+ COMPATIBLE)
# Includes: extended disorder, stronger coupling/detuning disorder,
#           reduced pump, extended kappa_ext, 20‑seed averaging,
#           ODE solver fix (dict options with BDF), CSV export, α fitting.
# =============================================================================

import sys
import time
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from scipy.constants import c, epsilon_0, hbar
from scipy.signal import savgol_filter

# QuTiP imports
try:
    from qutip import basis, destroy, mesolve, qeye, sigmam, sigmap, sigmaz, tensor
except ImportError:
    raise SystemExit("Install QuTiP: python -m pip install qutip")

# Plotly for 3D (optional)
try:
    import plotly.graph_objects as go
except ImportError:
    go = None

# Pandas for CSV (optional but recommended)
try:
    import pandas as pd
except ImportError:
    pd = None

# ----------------------------------------------------------------------------
# 1. PHYSICAL PARAMETERS
# ----------------------------------------------------------------------------
lambda_c = 550e-9
omega_c = 2 * np.pi * c / lambda_c
L_cav = lambda_c / 2
finesse = 10000                     # High finesse → lower loss
kappa = omega_c / finesse
mu_0 = 6.0e-29
waist_radius = 0.5e-6
T_room = 300
gamma_sp = 1e9
gamma_phi_0 = 2 * np.pi * 1e11

PULSE_T0 = 5e-12
PULSE_TAU = 1e-12
PULSE_AMP = 4e11                  # Reduced for linear regime

def gamma_phi(T):
    return gamma_phi_0 * (T / 300.0)

def safe_savgol(vals, window=11, poly=3):
    if len(vals) < 5:
        return np.asarray(vals)
    w = min(window, len(vals) if len(vals) % 2 else len(vals) - 1)
    if w < poly + 2:
        return np.asarray(vals)
    return savgol_filter(vals, window_length=w, polyorder=poly)

# ----------------------------------------------------------------------------
# 2. DISORDERED ENSEMBLE (coupling + detuning disorder)
# ----------------------------------------------------------------------------
def generate_disordered_ensemble(N, sigma, waist=waist_radius, L=L_cav):
    r_max = 2.0 * waist
    r = r_max * np.sqrt(np.random.rand(N))
    phi = np.random.rand(N) * 2 * np.pi
    z = (np.random.rand(N) - 0.5) * L
    x = r * np.cos(phi)
    y = r * np.sin(phi)
    pos = np.column_stack((x, y, z))

    theta = np.arccos(2 * np.random.rand(N) - 1)
    phi_d = np.random.rand(N) * 2 * np.pi
    mu_x = np.sin(theta) * np.cos(phi_d)
    mu_y = np.sin(theta) * np.sin(phi_d)
    mu_z = np.cos(theta)
    dipoles = np.column_stack((mu_x, mu_y, mu_z))

    k = 2 * np.pi / lambda_c
    e0 = np.sqrt(2 * hbar * omega_c / (epsilon_0 * np.pi * waist**2 * L_cav))
    envelope = np.exp(-(y**2 + z**2) / waist**2) * np.cos(k * x)
    g_clean = np.abs(mu_0 * dipoles[:, 0] * envelope * e0 / hbar)

    g_disordered = g_clean * np.exp(3.0 * sigma * np.random.randn(N))
    delta_disorder = sigma * 3e12 * np.random.randn(N)   # stronger detuning disorder

    return pos, g_disordered, delta_disorder

# ----------------------------------------------------------------------------
# 3. PULSED PUMP
# ----------------------------------------------------------------------------
def pulsed_pump(t, args):
    return args['amp'] * np.exp(-((t - args['t0']) / args['tau'])**2)

# ----------------------------------------------------------------------------
# 4. HAMILTONIAN AND OPERATORS
# ----------------------------------------------------------------------------
def build_operators(N, N_cav=3, N_load=3):
    a = tensor(destroy(N_cav), qeye(N_load), *[qeye(2) for _ in range(N)])
    b = tensor(qeye(N_cav), destroy(N_load), *[qeye(2) for _ in range(N)])
    return a, b

def build_hamiltonian(N, g_list, delta_list, kappa_ext, N_cav=3, N_load=3):
    a, b = build_operators(N, N_cav, N_load)
    H = kappa_ext * (a.dag() * b + a * b.dag())
    for j in range(N):
        sz_list = [qeye(2)] * N
        sz_list[j] = sigmaz()
        sz = tensor(qeye(N_cav), qeye(N_load), *sz_list)
        H += 0.5 * delta_list[j] * sz

        sp_list = [qeye(2)] * N
        sp_list[j] = sigmap()
        sp = tensor(qeye(N_cav), qeye(N_load), *sp_list)

        sm_list = [qeye(2)] * N
        sm_list[j] = sigmam()
        sm = tensor(qeye(N_cav), qeye(N_load), *sm_list)

        H += g_list[j] * (a * sp + a.dag() * sm)
    return H, a, b

def get_collapse_operators(N, k_main, k_load, g_sp, g_phi, N_cav=3, N_load=3):
    a, b = build_operators(N, N_cav, N_load)
    c_ops = [np.sqrt(k_main) * a, np.sqrt(k_load) * b]
    for j in range(N):
        sm_list = [qeye(2)] * N
        sm_list[j] = sigmam()
        sm = tensor(qeye(N_cav), qeye(N_load), *sm_list)
        c_ops.append(np.sqrt(g_sp) * sm)

        sz_list = [qeye(2)] * N
        sz_list[j] = sigmaz()
        sz = tensor(qeye(N_cav), qeye(N_load), *sz_list)
        c_ops.append(np.sqrt(g_phi / 2.0) * sz)
    return c_ops

def excitation_operator(N, N_cav=3, N_load=3):
    n_exc = 0 * tensor(qeye(N_cav), qeye(N_load), *[qeye(2) for _ in range(N)])
    proj_e = (qeye(2) + sigmaz()) / 2
    for j in range(N):
        pl = [qeye(2)] * N
        pl[j] = proj_e
        n_exc += tensor(qeye(N_cav), qeye(N_load), *pl)
    return n_exc

# ----------------------------------------------------------------------------
# 5. SINGLE SIMULATION (with ODE solver fix using dictionary options)
# ----------------------------------------------------------------------------
def simulate_once(N, disorder, kappa_ext, kappa_load=2e11, t_max=20e-12, n_times=400):
    # Use smaller Fock space for larger N to avoid QuTiP memory bug
    N_cav = 2 if N >= 5 else 3
    N_load = 2 if N >= 5 else 3

    pos, g_list, delta_list = generate_disordered_ensemble(N, disorder)
    H0, a, b = build_hamiltonian(N, g_list, delta_list, kappa_ext, N_cav=N_cav, N_load=N_load)
    pump_args = {'t0': PULSE_T0, 'tau': PULSE_TAU, 'amp': PULSE_AMP}
    H = [H0, [a + a.dag(), pulsed_pump]]
    c_ops = get_collapse_operators(N, kappa, kappa_load, gamma_sp, gamma_phi(T_room), N_cav=N_cav, N_load=N_load)

    psi0 = tensor(basis(N_cav, 0), basis(N_load, 0), *[basis(2, 0) for _ in range(N)])
    tlist = np.linspace(0.0, t_max, n_times)
    n_main = a.dag() * a
    n_load = b.dag() * b
    n_exc = excitation_operator(N, N_cav=N_cav, N_load=N_load)

    opts = {
        'nsteps': 50000,
        'method': 'bdf',
        'rtol': 1e-6,
        'atol': 1e-8
    }

    result = mesolve(H, psi0, tlist, c_ops, [n_main, n_load, n_exc], args=pump_args, options=opts)
    nm = np.asarray(result.expect[0])
    nl = np.asarray(result.expect[1])
    ne = np.asarray(result.expect[2])

    e_sys = hbar * omega_c * (nm + nl + ne)
    pulse_env = pulsed_pump(tlist, pump_args)
    idx_end = np.where(pulse_env < 0.01 * PULSE_AMP)[0][0]
    e_stored = e_sys[idx_end]
    p_ext = kappa_load * hbar * omega_c * nl
    mask = tlist >= tlist[idx_end]
    e_extracted = np.trapezoid(p_ext[mask], tlist[mask])
    eta = e_extracted / e_stored if e_stored > 0 else 0.0

    # Charging power (for Figure 3 only)
    e_ch = e_sys[:idx_end+1]
    if len(e_ch) > 5:
        e_sm = safe_savgol(e_ch)
        p_ch = np.gradient(e_sm, tlist[1] - tlist[0])
        p_max_charge = np.max(p_ch)
    else:
        p_max_charge = 0.0

    return {
        'eta': eta,
        'e_stored': e_stored,
        'e_extracted': e_extracted,
        'P_max_charge': p_max_charge,
        'positions': pos,
        'g_list': g_list
    }
# ----------------------------------------------------------------------------
# 6. AVERAGING WRAPPER
# ----------------------------------------------------------------------------
def simulate_averaged(N, disorder, kappa_ext, n_seeds=20):
    etas, e_storeds, p_charges = [], [], []
    for seed in range(n_seeds):
        np.random.seed(seed)
        res = simulate_once(N, disorder, kappa_ext)
        etas.append(res['eta'])
        e_storeds.append(res['e_stored'])
        p_charges.append(res['P_max_charge'])
    return {
        'eta_mean': np.mean(etas),
        'eta_std': np.std(etas),
        'e_stored_mean': np.mean(e_storeds),
        'e_stored_std': np.std(e_storeds),
        'P_max_charge_mean': np.mean(p_charges),
        'positions': res['positions'],   # from last seed
        'g_list': res['g_list']
    }

# ----------------------------------------------------------------------------
# 7. RUN FULL SCAN AND GENERATE ALL FIGURES
# ----------------------------------------------------------------------------
def main():
    out_dir = Path("publication_figures")
    out_dir.mkdir(exist_ok=True)

    # ------------------------------------------------------------------------
    # Figure 1: Efficiency vs Disorder (20‑seed averaged)
    # ------------------------------------------------------------------------
    print("=== Figure 1: Efficiency vs Disorder (20‑seed average) ===")
    N_vals = [2, 3, 4, 5, 6]
    sig_vals = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 1.2]
    kappa_fixed = 1e11

    fig1_data = {}
    for N in N_vals:
        for sig in sig_vals:
            print(f"  N={N}, σ={sig:.1f}")
            avg = simulate_averaged(N, sig, kappa_fixed, n_seeds=20)
            fig1_data[(N, sig)] = avg

    fig1, ax1 = plt.subplots(figsize=(10,6))
    for N in N_vals:
        x, y, yerr = [], [], []
        for sig in sig_vals:
            d = fig1_data[(N, sig)]
            x.append(sig)
            y.append(d['eta_mean'])
            yerr.append(d['eta_std'])
        ax1.errorbar(x, y, yerr=yerr, fmt='o-', capsize=4, label=f'N={N}')
    ax1.set_xlabel('Disorder Strength σ')
    ax1.set_ylabel('Extraction Efficiency η')
    ax1.set_title('Disorder Reduces Extractable Energy (20‑seed average)')
    ax1.legend()
    ax1.grid(True)
    fig1.savefig(out_dir / 'fig1_efficiency_vs_disorder.png', dpi=300, bbox_inches='tight')
    plt.close(fig1)

    # ------------------------------------------------------------------------
    # Figure 2: Optimal κ_ext shift (20‑seed averaged, N=4)
    # ------------------------------------------------------------------------
    print("\n=== Figure 2: Optimal κ_ext (20‑seed average, N=4) ===")
    N4 = 4
    sig2 = [0.0, 0.4]
    kappa_ext_vals = [0.5e11, 1e11, 2e11, 5e11, 10e11, 15e11, 20e11, 30e11, 50e11]

    fig2_data = {}
    for sig in sig2:
        for kex in kappa_ext_vals:
            print(f"  σ={sig:.1f}, κ_ext={kex/1e9:.1f} GHz")
            avg = simulate_averaged(N4, sig, kex, n_seeds=20)
            fig2_data[(sig, kex)] = avg

    fig2, ax2 = plt.subplots(figsize=(10,6))
    for sig in sig2:
        x, y, yerr = [], [], []
        for kex in kappa_ext_vals:
            d = fig2_data[(sig, kex)]
            x.append(kex/1e9)
            y.append(d['eta_mean'])
            yerr.append(d['eta_std'])
        ax2.errorbar(x, y, yerr=yerr, fmt='o-', capsize=4, label=f'σ={sig}')
    ax2.set_xlabel('Extraction Coupling κ_ext (GHz)')
    ax2.set_ylabel('Extraction Efficiency η')
    ax2.set_title('Optimal Extraction Rate Shifts with Disorder (N=4)')
    ax2.legend()
    ax2.grid(True)
    fig2.savefig(out_dir / 'fig2_optimal_kappa.png', dpi=300, bbox_inches='tight')
    plt.close(fig2)

    # ------------------------------------------------------------------------
    # CSV EXPORT
    # ------------------------------------------------------------------------
    if pd is not None:
        rows = []
        for (N, sig), d in fig1_data.items():
            rows.append({'Figure': 1, 'N': N, 'sigma': sig, 'kappa_ext': kappa_fixed/1e9,
                         'eta_mean': d['eta_mean'], 'eta_std': d['eta_std']})
        for (sig, kex), d in fig2_data.items():
            rows.append({'Figure': 2, 'N': 4, 'sigma': sig, 'kappa_ext': kex/1e9,
                         'eta_mean': d['eta_mean'], 'eta_std': d['eta_std']})
        df = pd.DataFrame(rows)
        df.to_csv(out_dir / 'simulation_data.csv', index=False)
        print(f"\nExported simulation_data.csv to {out_dir}")
    else:
        print("\nPandas not installed, skipping CSV export.")

    # ------------------------------------------------------------------------
    # Figure 3: Stored Energy Scaling (single‑seed, full N sweep)
    # ------------------------------------------------------------------------
    print("\n=== Figure 3: Stored Energy Scaling ===")
    N_all = [2, 3, 4, 5, 6]
    sig3 = [0.0, 0.2, 0.4, 0.6]

    fig3_data = {}
    for N in N_all:
        for sig in sig3:
            print(f"  N={N}, σ={sig:.1f}")
            res = simulate_once(N, sig, kappa_fixed)
            fig3_data[(N, sig)] = res

    fig3, ax3 = plt.subplots(figsize=(10,6))
    for sig in sig3:
        x, y = [], []
        for N in N_all:
            d = fig3_data[(N, sig)]
            x.append(N)
            y.append(d['e_stored'])
        ax3.loglog(x, y, 'o-', label=f'σ={sig}')
    # Reference lines
    N_ref = np.array(N_all)
    ref_val = fig3_data[(2, 0.0)]['e_stored']
    ax3.loglog(N_ref, ref_val * (N_ref/2)**2, 'k--', label='N²')
    ax3.loglog(N_ref, ref_val * (N_ref/2), 'k:', label='N¹')
    ax3.set_xlabel('Number of Molecules N')
    ax3.set_ylabel('Stored Energy E_stored (J)')
    ax3.set_title('Superextensive Charging Degrades with Disorder')
    ax3.legend()
    ax3.grid(True)
    fig3.savefig(out_dir / 'fig3_power_scaling.png', dpi=300, bbox_inches='tight')
    plt.close(fig3)

    # Optional: fit scaling exponent α and print
    from scipy.optimize import curve_fit
    def power_law(n, a, alpha):
        return a * n**alpha
    print("\nFitted scaling exponents α:")
    for sig in sig3:
        n_vals = []
        e_vals = []
        for N in N_all:
            d = fig3_data[(N, sig)]
            n_vals.append(N)
            e_vals.append(d['e_stored'])
        if len(n_vals) > 2:
            popt, _ = curve_fit(power_law, n_vals, e_vals)
            print(f"  σ={sig:.1f}: α = {popt[1]:.2f}")

    # ------------------------------------------------------------------------
    # Figure 4: 3D Visualization (one example: N=4, σ=0.4)
    # ------------------------------------------------------------------------
    if go is not None:
        print("\n=== Figure 4: 3D Visualization ===")
        np.random.seed(0)
        pos, g_vals, _ = generate_disordered_ensemble(4, 0.4)
        norm = (g_vals - g_vals.min()) / (g_vals.max() - g_vals.min() + 1e-12)
        fig3d = go.Figure(data=[
            go.Scatter3d(
                x=pos[:,0]*1e6, y=pos[:,1]*1e6, z=pos[:,2]*1e6,
                mode='markers',
                marker=dict(size=8, color=norm, colorscale='Viridis',
                            colorbar=dict(title='Coupling Strength')),
                text=[f"g = {g/1e9:.2f} GHz" for g in g_vals]
            ),
            go.Surface(
                x=[[2,2],[2,2]], y=[[-2,2],[-2,2]], z=[[-0.5,-0.5],[0.5,0.5]],
                colorscale=[[0,'rgba(200,200,255,0.3)'],[1,'rgba(200,200,255,0.3)']],
                showscale=False, name='Extraction Port'
            )
        ])
        fig3d.update_layout(
            title='3D Disordered Ensemble with Extraction Port',
            scene=dict(xaxis_title='X (µm)', yaxis_title='Y (µm)', zaxis_title='Z (µm)'),
            width=900, height=700
        )
        fig3d.write_html(out_dir / 'fig4_disordered_ensemble.html')
    else:
        print("\nPlotly not installed, skipping 3D HTML.")

    print(f"\nAll figures saved to: {out_dir.resolve()}")

if __name__ == "__main__":
    start = time.time()
    main()
    print(f"Total runtime: {(time.time()-start)/60:.1f} minutes")