import lumapi
import os
import matplotlib
matplotlib.use("agg")
import matplotlib.pyplot as plt
import numpy as np
from optixlog import Optixlog


def simulate_pn_modulator():
    """Simulate carrier distribution in a PN junction modulator."""
    
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="CHARGE_Simulations", create_if_not_exists=True)
    
    params = {
        "junction_type": "PN",
        "doping_p": 1e18,   # cm^-3
        "doping_n": 1e18,   # cm^-3
        "bias_voltages": [-2, -1, 0, 1],  # V
        "waveguide_width": 0.5  # μm
    }
    
    run = project.run(
        name="pn_modulator_charge",
        config={
            "tool": "Lumerical CHARGE",
            **params
        }
    )
    
    with lumapi.DEVICE() as device:
        device.load("pn_modulator.ldev")
        
        # Set doping concentrations
        device.setnamed("p_region", "concentration", params["doping_p"] * 1e6)  # to m^-3
        device.setnamed("n_region", "concentration", params["doping_n"] * 1e6)
        
        for step, voltage in enumerate(params["bias_voltages"]):
            # Apply bias
            device.setnamed("anode", "voltage", voltage)
            
            # Run simulation
            device.run("CHARGE")
            
            # Get carrier distributions
            n = device.getresult("CHARGE::monitor", "n")  # Electron density
            p = device.getresult("CHARGE::monitor", "p")  # Hole density
            
            # Calculate free-carrier absorption change (simplified Soref equations)
            delta_n = n["n"] - n["n"].mean()  # Change from intrinsic
            delta_p = p["p"] - p["p"].mean()
            
            # Si coefficients at 1550 nm
            sigma_n = 8.5e-18  # cm^2
            sigma_p = 6.0e-18  # cm^2
            delta_alpha = sigma_n * delta_n + sigma_p * delta_p  # cm^-1
            
            # Refractive index change
            dn_n = -8.8e-22 * delta_n
            dn_p = -8.5e-18 * delta_p
            delta_n_total = np.mean(dn_n + dn_p)
            
            run.log(step=step,
                    bias_voltage=voltage,
                    avg_electron_density=float(np.mean(n["n"])),
                    avg_hole_density=float(np.mean(p["p"])),
                    delta_n_index=float(delta_n_total))
            
            # Plot carrier distribution
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
            
            x = n["x"].flatten() * 1e6
            y = n["y"].flatten() * 1e6
            
            im1 = ax1.pcolormesh(x, y, np.log10(n["n"][:, :, 0] + 1).T, cmap='Blues')
            ax1.set_xlabel('x (μm)')
            ax1.set_ylabel('y (μm)')
            ax1.set_title(f'Electron Density (V={voltage}V)')
            plt.colorbar(im1, ax=ax1, label='log10(n) [m^-3]')
            
            im2 = ax2.pcolormesh(x, y, np.log10(p["p"][:, :, 0] + 1).T, cmap='Reds')
            ax2.set_xlabel('x (μm)')
            ax2.set_ylabel('y (μm)')
            ax2.set_title(f'Hole Density (V={voltage}V)')
            plt.colorbar(im2, ax=ax2, label='log10(p) [m^-3]')
            
            plt.tight_layout()
            run.log_matplotlib(f"carriers_V{voltage}", fig)
            plt.close(fig)
        
        run.log(step=len(params["bias_voltages"]),
                simulation_complete=True)
    
    print("✅ CHARGE simulation complete!")


if __name__ == "__main__":
    simulate_pn_modulator()

def simulate_thermal_tuning():
    """Simulate thermal phase shifter."""
    
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="HEAT_Simulations", create_if_not_exists=True)
    
    params = {
        "heater_power": [0, 10, 20, 30, 40, 50],  # mW
        "heater_length": 100,  # μm
        "waveguide_distance": 1.5  # μm from heater
    }
    
    run = project.run(
        name="thermal_phase_shifter",
        config={
            "tool": "Lumerical HEAT",
            **params
        }
    )
    
    with lumapi.DEVICE() as device:
        device.load("thermal_phase_shifter.ldev")
        
        results = []
        
        for step, power in enumerate(params["heater_power"]):
            # Set heater power
            device.setnamed("heater", "power", power * 1e-3)  # Convert to W
            
            # Run thermal simulation
            device.run("HEAT")
            
            # Get temperature distribution
            T = device.getresult("HEAT::monitor", "T")
            
            # Find temperature at waveguide
            waveguide_temp = extract_waveguide_temp(T, params["waveguide_distance"])
            temp_rise = waveguide_temp - 300  # K above ambient
            
            # Calculate phase shift (Si thermo-optic coefficient)
            dn_dT = 1.86e-4  # 1/K for Si
            delta_n = dn_dT * temp_rise
            phase_shift = 2 * np.pi * delta_n * params["heater_length"] * 1e-6 / 1.55e-6
            
            run.log(step=step,
                    heater_power_mW=power,
                    waveguide_temp_K=float(waveguide_temp),
                    temp_rise_K=float(temp_rise),
                    phase_shift_rad=float(phase_shift),
                    phase_shift_pi=float(phase_shift / np.pi))
            
            results.append({
                "power": power,
                "temp_rise": temp_rise,
                "phase": phase_shift
            })
            
            # Plot temperature
            fig, ax = plt.subplots(figsize=(10, 6))
            x = T["x"].flatten() * 1e6
            y = T["y"].flatten() * 1e6
            im = ax.pcolormesh(x, y, T["T"][:, :, 0].T - 300, cmap='hot')
            ax.set_xlabel('x (μm)')
            ax.set_ylabel('y (μm)')
            ax.set_title(f'Temperature Rise (P={power} mW)')
            plt.colorbar(im, ax=ax, label='ΔT (K)')
            
            run.log_matplotlib(f"temp_P{power}mW", fig)
            plt.close(fig)
        
        # Summary plot
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        powers = [r["power"] for r in results]
        temps = [r["temp_rise"] for r in results]
        phases = [r["phase"] / np.pi for r in results]
        
        ax1.plot(powers, temps, 'ro-')
        ax1.set_xlabel('Heater Power (mW)')
        ax1.set_ylabel('Temperature Rise (K)')
        ax1.grid(True, alpha=0.3)
        
        ax2.plot(powers, phases, 'bo-')
        ax2.set_xlabel('Heater Power (mW)')
        ax2.set_ylabel('Phase Shift (π)')
        ax2.axhline(y=1, color='r', linestyle='--', label='π phase')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        run.log_matplotlib("thermal_summary", fig)
        plt.close(fig)
        
        # Calculate power for π phase shift
        from scipy.interpolate import interp1d
        f = interp1d(phases, powers, fill_value="extrapolate")
        p_pi = float(f(1.0))
        
        run.log(step=len(params["heater_power"]),
                power_for_pi_mW=p_pi,
                efficiency_mW_per_pi=p_pi)


def extract_waveguide_temp(T, distance_um):
    """Extract temperature at waveguide location."""
    # Simplified - find temperature at specified distance from heater center
    return float(np.mean(T["T"]) + 10)  # Placeholder

def simulate_plasmonic_antenna():
    """Simulate plasmonic nanoantenna with DGTD."""
    
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="DGTD_Simulations", create_if_not_exists=True)
    
    params = {
        "antenna_length": 100,  # nm
        "antenna_width": 40,    # nm
        "gap_size": 10,         # nm
        "material": "Au",
        "wavelength_range": [500, 900]  # nm
    }
    
    run = project.run(
        name="plasmonic_dipole_antenna",
        config={
            "tool": "Lumerical DGTD",
            **params
        }
    )
    
    with lumapi.FDTD() as fdtd:  # DGTD uses similar API
        fdtd.load("plasmonic_antenna.fsp")
        
        # Enable DGTD solver for accurate metal simulation
        fdtd.set("solver type", "DGTD")
        
        # Set geometry
        fdtd.setnamed("antenna", "x span", params["antenna_length"] * 1e-9)
        fdtd.setnamed("antenna", "y span", params["antenna_width"] * 1e-9)
        fdtd.setnamed("gap", "x span", params["gap_size"] * 1e-9)
        
        fdtd.run()
        
        # Get near-field enhancement
        E = fdtd.getresult("field_monitor", "E")
        enhancement = np.abs(E["E"])**2 / np.abs(E["E"].max())**2
        
        # Get scattering spectrum
        scatter = fdtd.getresult("scattering", "sigma")
        wavelengths = scatter["lambda"] * 1e9
        cross_section = scatter["sigma"]
        
        # Find resonance
        resonance_idx = np.argmax(cross_section)
        resonance_wl = wavelengths[resonance_idx]
        
        run.log(step=0,
                resonance_wavelength_nm=float(resonance_wl),
                max_enhancement=float(np.max(enhancement)),
                peak_cross_section=float(cross_section[resonance_idx]))
        
        # Plot spectrum
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.plot(wavelengths, cross_section, 'b-')
        ax.axvline(resonance_wl, color='r', linestyle='--', 
                   label=f'Resonance: {resonance_wl:.1f} nm')
        ax.set_xlabel('Wavelength (nm)')
        ax.set_ylabel('Scattering Cross Section')
        ax.set_title('Plasmonic Antenna Spectrum')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        run.log_matplotlib("scattering_spectrum", fig)
        plt.close(fig)

def analyze_bent_waveguide():
    """Analyze bent waveguide modes with FEEM."""
    
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="FEEM_Simulations", create_if_not_exists=True)
    
    radii = [5, 10, 20, 50, 100]  # μm
    
    for radius in radii:
        run = project.run(
            name=f"bent_wg_r{radius}um",
            config={
                "tool": "Lumerical FEEM",
                "bend_radius": radius,
                "waveguide_width": 0.5
            }
        )
        
        with lumapi.MODE() as mode:
            mode.load("bent_waveguide.lms")
            
            # Enable FEEM solver
            mode.set("solver type", "FEEM")
            mode.set("bend radius", radius * 1e-6)
            
            mode.findmodes()
            
            # Get mode data
            neff = mode.getdata("FEEM::data::mode1", "neff")
            loss = mode.getdata("FEEM::data::mode1", "loss")  # dB/cm
            
            # Calculate bend loss per 90° turn
            bend_loss_90 = loss * (np.pi * radius / 2) / 1e4  # dB
            
            run.log(step=0,
                    bend_radius_um=radius,
                    neff=float(np.real(neff)),
                    loss_dB_cm=float(loss),
                    loss_per_90deg_dB=float(bend_loss_90))
    
    # Create summary
    summary_run = project.run(name="bend_loss_summary", config={"sweep": "bend_radius"})
    # ... create summary plot

def simulate_full_modulator():
    """Co-simulate electrical and optical domains."""
    
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="Modulator_Cosim", create_if_not_exists=True)
    
    run = project.run(
        name="mzm_full_cosim",
        config={
            "type": "Mach-Zehnder Modulator",
            "arm_length": 1000,  # μm
            "bias_range": [-3, 1]
        }
    )
    
    voltages = np.linspace(-3, 1, 9)
    phase_shifts = []
    
    for step, voltage in enumerate(voltages):
        # Step 1: CHARGE simulation
        with lumapi.DEVICE() as device:
            device.load("mzm_arm.ldev")
            device.setnamed("electrode", "voltage", voltage)
            device.run("CHARGE")
            
            # Export carrier distribution
            n = device.getresult("CHARGE::monitor", "n")
            p = device.getresult("CHARGE::monitor", "p")
            
            # Save for FDTD import
            device.export("carrier_data.mat")
        
        # Step 2: FDTD simulation with imported carriers
        with lumapi.FDTD() as fdtd:
            fdtd.load("mzm_optical.fsp")
            
            # Import carrier data to modify refractive index
            fdtd.importdata("carrier_data.mat")
            
            fdtd.run()
            
            # Get phase at output
            E_out = fdtd.getresult("output", "E")
            phase = np.angle(E_out["E"][0, 0, 0, 0])
            phase_shifts.append(phase)
        
        run.log(step=step,
                voltage=voltage,
                phase_rad=float(phase))
    
    # Calculate Vπ
    phase_arr = np.unwrap(phase_shifts)
    from scipy.interpolate import interp1d
    f = interp1d(phase_arr, voltages)
    V_pi = float(f(phase_arr[0] - np.pi) - voltages[0])
    
    run.log(step=len(voltages),
            V_pi=V_pi,
            modulation_efficiency=1000 / V_pi)  # μm/V

multiphysics:
  charge:
    enabled: true
    extract: ["carrier_density", "current", "capacitance"]
    file_pattern: "*.ldev"
  
  heat:
    enabled: true
    extract: ["temperature", "power_dissipation"]
    file_pattern: "*.ldev"
  
  dgtd:
    enabled: true
    extract: ["field_enhancement", "absorption"]
    file_pattern: "*.fsp"
  
  feem:
    enabled: true
    extract: ["neff", "loss", "mode_profile"]
    file_pattern: "*.lms"

run = project.run(
    name="cosim_modulator",
    config={
        "solver_chain": ["CHARGE", "FDTD"],
        "charge_run_id": charge_run.run_id,
        "fdtd_run_id": fdtd_run.run_id
    }
)

# After each solver step
run.log(step=0, solver="CHARGE", 
        max_carrier_density=float(max_n),
        junction_capacitance=float(C))

run.log(step=1, solver="FDTD",
        effective_index_change=float(delta_n),
        phase_shift=float(phi))

run = project.run(
    name=f"modulator_v{version}",
    config={
        "electrical_model_version": "2.1",
        "optical_model_version": "1.5",
        "material_database": "2024Q1"
    }
)