import os
from dataclasses import dataclass
from typing import Dict, List

import numpy as np
from matplotlib import pyplot as plt

AIR_INDEX = 1.0  # normal incidence, non-absorbing ambient
WAVELENGTH_MIN = 0.4  # µm
WAVELENGTH_MAX = 25.0  # µm


def cauchy_index(wavelength_um: np.ndarray) -> np.ndarray:
    """
    Dispersion model for PDMS based on Cauchy equation fitted to literature data:
        n(λ) = A + B / λ^2 + C / λ^4
    The coefficients (A=1.385, B=0.0125, C=0.00045) match n≈1.43 @ 0.55 µm
    and n≈1.40 beyond 2 µm.
    """
    A = 1.385
    B = 0.0125
    C = 0.00045
    lam2 = wavelength_um ** 2
    return A + B / lam2 + C / (lam2 * wavelength_um ** 2)


def extinction_coeff(wavelength_um: np.ndarray) -> np.ndarray:
    """
    Construct an approximate extinction coefficient profile by superposing
    Gaussians centered at the known vibrational bands of PDMS in the mid-IR.
    Peaks taken from FTIR measurements (Si-O-Si at ~9.2 µm, Si-CH3 rocking at
    12.7 µm, CH stretches at 3.4 µm, etc.).
    """
    peaks = [
        # (center µm, amplitude, width µm)
        (3.4, 0.06, 0.25),
        (8.0, 0.30, 0.50),
        (9.2, 0.65, 0.60),
        (10.6, 0.35, 0.45),
        (12.7, 0.45, 0.35),
        (14.5, 0.25, 0.60),
    ]
    k = np.full_like(wavelength_um, 5e-4)
    for center, amp, width in peaks:
        k += amp * np.exp(-0.5 * ((wavelength_um - center) / width) ** 2)
    # enforce plausible upper bound
    return np.clip(k, 0, 2.0)


def thin_film_emissivity(
    wavelength_um: np.ndarray, thickness_um: float, n_complex: np.ndarray
) -> np.ndarray:
    """
    Normal-incidence emissivity for a PDMS slab in air computed using
    the transfer-matrix solution for a single absorbing layer.
    """
    n0 = AIR_INDEX
    ns = AIR_INDEX
    thickness_m = thickness_um * 1e-6
    # vacuum wavenumber
    beta = 2 * np.pi / (wavelength_um * 1e-6)
    delta = beta * n_complex * thickness_m

    r01 = (n0 - n_complex) / (n0 + n_complex)
    r12 = (n_complex - ns) / (n_complex + ns)
    exp_term = np.exp(2j * delta)
    numerator_r = r01 + r12 * exp_term
    denominator = 1 + r01 * r12 * exp_term
    r = numerator_r / denominator

    t01 = 2 * n0 / (n0 + n_complex)
    t12 = 2 * n_complex / (n_complex + ns)
    exp_half = np.exp(1j * delta)
    t = (t01 * t12 * exp_half) / denominator

    R = np.abs(r) ** 2
    T = (ns.real / n0) * np.abs(t) ** 2
    A = 1 - R - T
    return np.clip(A.real, 0, 1)


@dataclass
class EmissivityResult:
    wavelengths: np.ndarray
    emissivity_map: Dict[float, np.ndarray]
    window_avg: Dict[float, float]


def compute_emissivity(thicknesses_um: List[float]) -> EmissivityResult:
    wavelengths = np.linspace(WAVELENGTH_MIN, WAVELENGTH_MAX, 1200)
    n = cauchy_index(wavelengths)
    k = extinction_coeff(wavelengths)
    n_complex = n + 1j * k

    emissivity_map = {}
    window_avg = {}
    window_mask = (wavelengths >= 8) & (wavelengths <= 13)

    for d in thicknesses_um:
        eps = thin_film_emissivity(wavelengths, d, n_complex)
        emissivity_map[d] = eps
        window_avg[d] = float(np.trapz(eps[window_mask], wavelengths[window_mask]) /
                              np.trapz(np.ones_like(wavelengths[window_mask]),
                                       wavelengths[window_mask]))

    return EmissivityResult(wavelengths, emissivity_map, window_avg)


def plot_emissivity(result: EmissivityResult, outdir: str) -> str:
    plt.figure(figsize=(9, 5))
    for d, eps in result.emissivity_map.items():
        label = f"{d:.0f} µm (ε̄₈₋₁₃={result.window_avg[d]:.2f})"
        plt.plot(result.wavelengths, eps, label=label)

    plt.axvspan(8, 13, color="tab:gray", alpha=0.15, label="Atmospheric window")
    plt.xlabel("Wavelength (µm)")
    plt.ylabel("Spectral emissivity")
    plt.title("PDMS Thin-Film Emissivity vs. Wavelength")
    plt.ylim(0, 1.05)
    plt.xlim(WAVELENGTH_MIN, WAVELENGTH_MAX)
    plt.legend()
    plt.grid(alpha=0.3)

    os.makedirs(outdir, exist_ok=True)
    output_path = os.path.join(outdir, "question1_emissivity.png")
    plt.tight_layout()
    plt.savefig(output_path, dpi=300)
    plt.close()
    return output_path


def main():
    thicknesses = [1, 5, 10, 25, 50, 100]
    result = compute_emissivity(thicknesses)
    outdir = os.path.join(os.path.dirname(__file__), "outputs")
    figure_path = plot_emissivity(result, outdir)

    summary_lines = ["thickness_um,avg_emissivity_8_13um"]
    for d in thicknesses:
        summary_lines.append(f"{d},{result.window_avg[d]:.4f}")
    csv_path = os.path.join(outdir, "question1_emissivity_summary.csv")
    with open(csv_path, "w", encoding="utf-8") as f:
        f.write("\n".join(summary_lines))

    print(f"Figure saved to: {figure_path}")
    print(f"Summary saved to: {csv_path}")


if __name__ == "__main__":
    main()