"""
MPB Band Structure Tutorial with OptixLog Integration

Introduction to photonic band structure calculations.
"""

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


def main():
    api_key = os.getenv("OPTIX_API_KEY", "your_api_key_here")
    
    client = Optixlog(api_key=api_key)
    project = client.project(name="MeepExamples", create_if_not_exists=True)
    
    # Parameters
    num_bands = 8
    resolution = 32
    
    run = project.run(
        name="mpb_band_structure_tutorial",
        config={
            "simulation_type": "mpb_band_structure",
            "resolution": resolution,
            "num_bands": num_bands,
            "lattice_type": "square",
            "rod_radius": 0.2,
            "rod_epsilon": 12
        }
    )
    
    run.log(step=0, resolution=resolution, num_bands=num_bands)
    
    # Define square lattice
    geometry_lattice = mp.Lattice(size=mp.Vector3(1, 1))
    
    # Dielectric rod at center
    r = 0.2  # rod radius
    eps = 12  # rod permittivity (silicon)
    
    geometry = [
        mp.Cylinder(
            radius=r,
            material=mp.Medium(epsilon=eps)
        )
    ]
    
    # High-symmetry k-points for square lattice
    # Γ (0,0) → X (0.5,0) → M (0.5,0.5) → Γ (0,0)
    k_points = [
        mp.Vector3(),           # Γ
        mp.Vector3(0.5, 0),     # X
        mp.Vector3(0.5, 0.5),   # M
        mp.Vector3(),           # Γ
    ]
    
    k_points_interp = mp.interpolate(16, k_points)
    
    # Create ModeSolver
    ms = mpb.ModeSolver(
        geometry_lattice=geometry_lattice,
        geometry=geometry,
        resolution=resolution,
        num_bands=num_bands,
        k_points=k_points_interp,
    )
    
    print("⚡ Computing TM band structure...")
    ms.run_tm()
    tm_freqs = ms.all_freqs
    
    run.log(step=1, phase="tm_bands_computed", num_k_points=len(k_points_interp))
    
    print("⚡ Computing TE band structure...")
    ms.run_te()
    te_freqs = ms.all_freqs
    
    run.log(step=2, phase="te_bands_computed")
    
    # Process results
    k_dist = []
    current_dist = 0
    prev_k = k_points_interp[0]
    
    for k in k_points_interp:
        dk = mp.Vector3(k.x - prev_k.x, k.y - prev_k.y)
        current_dist += np.sqrt(dk.x**2 + dk.y**2)
        k_dist.append(current_dist)
        prev_k = k
    
    # High-symmetry point positions
    hs_points = [0]
    for i, k in enumerate(k_points_interp):
        if any(k.x == kp.x and k.y == kp.y for kp in k_points[1:-1]):
            if i > 0:
                hs_points.append(k_dist[i])
    hs_points.append(k_dist[-1])
    hs_labels = ['Γ', 'X', 'M', 'Γ']
    
    # Create band structure plot
    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
    
    # TM bands
    for band_idx in range(num_bands):
        band = [tm_freqs[k_idx][band_idx] for k_idx in range(len(k_points_interp))]
        axes[0].plot(k_dist, band, 'b-', linewidth=1.5)
    
    axes[0].set_xlabel('Wave Vector')
    axes[0].set_ylabel('Frequency (c/a)')
    axes[0].set_title('TM Band Structure')
    axes[0].set_xlim(0, k_dist[-1])
    axes[0].set_ylim(0, 0.8)
    
    # Add high-symmetry labels
    for pos in hs_points:
        axes[0].axvline(pos, color='gray', linestyle='--', alpha=0.5)
    axes[0].set_xticks(hs_points)
    axes[0].set_xticklabels(hs_labels)
    axes[0].grid(True, alpha=0.3, axis='y')
    
    # TE bands
    for band_idx in range(num_bands):
        band = [te_freqs[k_idx][band_idx] for k_idx in range(len(k_points_interp))]
        axes[1].plot(k_dist, band, 'r-', linewidth=1.5)
    
    axes[1].set_xlabel('Wave Vector')
    axes[1].set_ylabel('Frequency (c/a)')
    axes[1].set_title('TE Band Structure')
    axes[1].set_xlim(0, k_dist[-1])
    axes[1].set_ylim(0, 0.8)
    
    for pos in hs_points:
        axes[1].axvline(pos, color='gray', linestyle='--', alpha=0.5)
    axes[1].set_xticks(hs_points)
    axes[1].set_xticklabels(hs_labels)
    axes[1].grid(True, alpha=0.3, axis='y')
    
    plt.tight_layout()
    run.log_matplotlib("band_structure", fig)
    plt.close(fig)
    
    # Find band gaps
    tm_gap = find_gap(tm_freqs)
    te_gap = find_gap(te_freqs)
    
    run.log(step=3,
            tm_gap_size=float(tm_gap['size']) if tm_gap else 0,
            te_gap_size=float(te_gap['size']) if te_gap else 0,
            simulation_completed=True)
    
    if tm_gap:
        print(f"📊 TM gap: {tm_gap['size']:.4f} at ω = {tm_gap['mid']:.4f}")
    if te_gap:
        print(f"📊 TE gap: {te_gap['size']:.4f} at ω = {te_gap['mid']:.4f}")
    
    print(f"\n✅ Band structure calculation complete!")


def find_gap(all_freqs):
    """Find the largest complete band gap"""
    num_bands = len(all_freqs[0])
    
    for band in range(num_bands - 1):
        # Get max of lower band and min of upper band
        lower_max = max(freqs[band] for freqs in all_freqs)
        upper_min = min(freqs[band + 1] for freqs in all_freqs)
        
        if upper_min > lower_max:
            gap_size = (upper_min - lower_max) / ((upper_min + lower_max) / 2)
            return {
                'size': gap_size,
                'bottom': lower_max,
                'top': upper_min,
                'mid': (upper_min + lower_max) / 2,
                'bands': (band, band + 1)
            }
    
    return None


if __name__ == "__main__":
    main()

geometry_lattice = mp.Lattice(size=mp.Vector3(1, 1))

ms = mpb.ModeSolver(
    geometry_lattice=geometry_lattice,
    geometry=geometry,
    resolution=resolution,
    num_bands=num_bands,
    k_points=k_points,
)

ms.run_tm()  # TM polarization
ms.run_te()  # TE polarization