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


def analyze_waveguide_modes():
    # Initialize OptixLog
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="MODE_Analysis", create_if_not_exists=True)
    
    # Parameters
    params = {
        "waveguide_width": 0.5,    # μm
        "waveguide_height": 0.22,  # μm
        "wavelength": 1.55,        # μm
        "num_modes": 4,
        "core_material": "Si (Silicon) - Palik",
        "clad_material": "SiO2 (Glass) - Palik"
    }
    
    run = project.run(
        name=f"fde_w{params['waveguide_width']}",
        config={
            "tool": "Lumerical MODE",
            "solver": "FDE",
            **params
        }
    )
    
    with lumapi.MODE() as mode:
        # Create simulation
        mode.addfde()
        mode.set("solver type", "2D X normal")
        mode.set("x", 0)
        mode.set("y", 0)
        mode.set("y span", 3e-6)
        mode.set("z", 0)
        mode.set("z span", 2e-6)
        mode.set("wavelength", params["wavelength"] * 1e-6)
        mode.set("number of trial modes", params["num_modes"])
        
        # Add waveguide
        mode.addrect()
        mode.set("name", "waveguide")
        mode.set("y", 0)
        mode.set("y span", params["waveguide_width"] * 1e-6)
        mode.set("z", 0)
        mode.set("z span", params["waveguide_height"] * 1e-6)
        mode.set("material", params["core_material"])
        
        # Add cladding (substrate)
        mode.addrect()
        mode.set("name", "substrate")
        mode.set("y", 0)
        mode.set("y span", 3e-6)
        mode.set("z", -1e-6)
        mode.set("z span", 2e-6)
        mode.set("material", params["clad_material"])
        
        # Set mesh
        mode.set("define y mesh by", "maximum mesh step")
        mode.set("dy", 0.02e-6)
        mode.set("define z mesh by", "maximum mesh step")
        mode.set("dz", 0.02e-6)
        
        # Find modes
        mode.findmodes()
        
        # Extract mode data
        neff_data = mode.getresult("FDE::data::mode1", "neff")
        
        # Log each mode
        for i in range(params["num_modes"]):
            try:
                mode_name = f"FDE::data::mode{i+1}"
                neff = mode.getdata(mode_name, "neff")
                loss = mode.getdata(mode_name, "loss")  # dB/cm
                
                run.log(step=i,
                        mode_number=i+1,
                        neff_real=float(np.real(neff)),
                        neff_imag=float(np.imag(neff)),
                        loss_dB_cm=float(loss) if loss else 0)
                
                # Get and plot mode profile
                E = mode.getresult(mode_name, "E")
                Ey = E["Ey"]
                y = E["y"].flatten() * 1e6
                z = E["z"].flatten() * 1e6
                
                fig, ax = plt.subplots(figsize=(8, 6))
                im = ax.pcolormesh(y, z, np.abs(Ey[:, :, 0, 0].T)**2, cmap='hot')
                ax.set_xlabel('y (μm)')
                ax.set_ylabel('z (μm)')
                ax.set_title(f'Mode {i+1}: neff = {np.real(neff):.4f}')
                ax.set_aspect('equal')
                plt.colorbar(im, ax=ax, label='|Ey|²')
                
                run.log_matplotlib(f"mode_{i+1}_profile", fig)
                plt.close(fig)
                
            except Exception as e:
                print(f"Could not extract mode {i+1}: {e}")
        
        run.log(step=params["num_modes"], analysis_complete=True)
        
    print("✅ MODE analysis complete!")


if __name__ == "__main__":
    analyze_waveguide_modes()

def sweep_waveguide_width():
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="WaveguideDesign", create_if_not_exists=True)
    
    widths = np.linspace(0.3, 0.7, 9)  # μm
    
    results = []
    
    for width in widths:
        run = project.run(
            name=f"width_{width:.2f}um",
            config={"waveguide_width": width, "wavelength": 1.55}
        )
        
        with lumapi.MODE() as mode:
            mode.load("waveguide_template.lms")
            mode.setnamed("waveguide", "y span", width * 1e-6)
            mode.findmodes()
            
            neff = mode.getdata("FDE::data::mode1", "neff")
            ng = mode.getdata("FDE::data::mode1", "ng")  # Group index
            
            run.log(step=0,
                    width_um=width,
                    neff=float(np.real(neff)),
                    ng=float(np.real(ng)))
            
            results.append({
                "width": width,
                "neff": float(np.real(neff)),
                "ng": float(np.real(ng))
            })
    
    # Summary plot
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    widths_arr = [r["width"] for r in results]
    neff_arr = [r["neff"] for r in results]
    ng_arr = [r["ng"] for r in results]
    
    ax1.plot(widths_arr, neff_arr, 'bo-')
    ax1.set_xlabel('Waveguide Width (μm)')
    ax1.set_ylabel('Effective Index (neff)')
    ax1.grid(True, alpha=0.3)
    
    ax2.plot(widths_arr, ng_arr, 'ro-')
    ax2.set_xlabel('Waveguide Width (μm)')
    ax2.set_ylabel('Group Index (ng)')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    # Log to a summary run
    summary_run = project.run(name="width_sweep_summary", config={"sweep_type": "width"})
    summary_run.log_matplotlib("neff_vs_width", fig)
    plt.close(fig)

def analyze_dispersion():
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="DispersionAnalysis", create_if_not_exists=True)
    
    run = project.run(
        name="waveguide_dispersion",
        config={"waveguide_width": 0.5, "wavelength_range": [1.5, 1.6]}
    )
    
    wavelengths = np.linspace(1.5, 1.6, 21)  # μm
    neff_vs_lambda = []
    
    with lumapi.MODE() as mode:
        mode.load("waveguide.lms")
        
        for wl in wavelengths:
            mode.set("wavelength", wl * 1e-6)
            mode.findmodes()
            
            neff = mode.getdata("FDE::data::mode1", "neff")
            neff_vs_lambda.append(float(np.real(neff)))
            
            run.log(step=len(neff_vs_lambda)-1,
                    wavelength_um=wl,
                    neff=float(np.real(neff)))
    
    # Calculate group index and dispersion
    neff_arr = np.array(neff_vs_lambda)
    wl_arr = np.array(wavelengths)
    
    # ng = neff - λ * dneff/dλ
    dneff_dlambda = np.gradient(neff_arr, wl_arr)
    ng = neff_arr - wl_arr * dneff_dlambda
    
    # D = -(λ/c) * d²neff/dλ²
    c = 3e8  # m/s
    d2neff_dlambda2 = np.gradient(dneff_dlambda, wl_arr)
    D = -(wl_arr * 1e-6 / c) * d2neff_dlambda2 * 1e6  # ps/(nm·km)
    
    # Log dispersion at 1.55 μm
    idx_1550 = np.argmin(np.abs(wl_arr - 1.55))
    run.log(step=len(wavelengths),
            ng_at_1550=float(ng[idx_1550]),
            D_at_1550=float(D[idx_1550]))
    
    # Plot
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    
    axes[0, 0].plot(wl_arr * 1000, neff_arr, 'b-')
    axes[0, 0].set_xlabel('Wavelength (nm)')
    axes[0, 0].set_ylabel('neff')
    axes[0, 0].set_title('Effective Index')
    axes[0, 0].grid(True, alpha=0.3)
    
    axes[0, 1].plot(wl_arr * 1000, ng, 'r-')
    axes[0, 1].set_xlabel('Wavelength (nm)')
    axes[0, 1].set_ylabel('ng')
    axes[0, 1].set_title('Group Index')
    axes[0, 1].grid(True, alpha=0.3)
    
    axes[1, 0].plot(wl_arr * 1000, D, 'g-')
    axes[1, 0].set_xlabel('Wavelength (nm)')
    axes[1, 0].set_ylabel('D (ps/nm/km)')
    axes[1, 0].set_title('Dispersion Parameter')
    axes[1, 0].grid(True, alpha=0.3)
    
    axes[1, 1].axis('off')
    axes[1, 1].text(0.5, 0.5, 
                    f"At λ = 1550 nm:\n\nneff = {neff_arr[idx_1550]:.4f}\nng = {ng[idx_1550]:.4f}\nD = {D[idx_1550]:.2f} ps/nm/km",
                    ha='center', va='center', fontsize=14,
                    transform=axes[1, 1].transAxes)
    
    plt.tight_layout()
    run.log_matplotlib("dispersion_analysis", fig)
    plt.close(fig)

def eme_taper_analysis():
    client = Optixlog(api_key=os.getenv("OPTIX_API_KEY"))
    project = client.project(name="EME_Analysis", create_if_not_exists=True)
    
    run = project.run(
        name="linear_taper",
        config={
            "solver": "EME",
            "width_in": 0.5,
            "width_out": 3.0,
            "taper_length": 20.0
        }
    )
    
    with lumapi.MODE() as mode:
        mode.load("taper_eme.lms")
        
        # Run EME
        mode.run()
        
        # Get S-parameters
        S = mode.getresult("EME", "S")
        S21 = S["S21"]
        
        # Log transmission
        transmission = np.abs(S21)**2
        run.log(step=0,
                transmission=float(transmission.flatten()[0]),
                insertion_loss_dB=float(-10 * np.log10(transmission.flatten()[0])))
        
        # Get field profile along taper
        mode.setemeanalysis("propagate", True)
        mode.emepropagation()
        
        E = mode.getresult("EME::field profile", "E")
        
        # Plot
        fig, ax = plt.subplots(figsize=(12, 4))
        extent = [E["x"].min()*1e6, E["x"].max()*1e6, 
                  E["y"].min()*1e6, E["y"].max()*1e6]
        im = ax.imshow(np.abs(E["Ey"][:, :, 0, 0]).T**2, 
                       extent=extent, aspect='auto', cmap='hot', origin='lower')
        ax.set_xlabel('Propagation (μm)')
        ax.set_ylabel('y (μm)')
        ax.set_title('EME Field Propagation')
        plt.colorbar(im, ax=ax, label='|Ey|²')
        
        run.log_matplotlib("eme_propagation", fig)
        plt.close(fig)

mode:
  enabled: true
  extract:
    - fde_modes: true
    - neff: true
    - ng: true
    - loss: true
    - mode_profiles: true  # Optional, can be large
  file_pattern: "*.lms"

# Classify TE vs TM modes
TE_fraction = mode.getdata(mode_name, "TE polarization fraction")
run.log(step=i,
        mode_number=i+1,
        neff=float(np.real(neff)),
        polarization="TE" if TE_fraction > 0.5 else "TM",
        te_fraction=float(TE_fraction))

# Check mode quality
mode_quality = mode.getdata(mode_name, "mode profile quality")
run.log(step=i, mode_quality=float(mode_quality))

# Sweep process corners
for width_var in [-0.02, 0, 0.02]:  # ±20 nm variation
    for height_var in [-0.01, 0, 0.01]:  # ±10 nm variation
        actual_width = nominal_width + width_var
        actual_height = nominal_height + height_var
        
        run = project.run(
            name=f"corner_w{width_var:+.3f}_h{height_var:+.3f}",
            config={"width_variation": width_var, "height_variation": height_var}
        )
        # ... run simulation