oxiphoton 0.1.1

use num_complex::Complex64;
use std::f64::consts::PI;

/// 1D finite-difference BPM using the Crank-Nicolson scheme.
///
/// Solves the scalar paraxial wave equation:
///   ∂E/∂z = i/(2k₀n₀) · ∂²E/∂x² + ik₀ Δn · E
///
/// The Crank-Nicolson scheme yields an unconditionally stable solver:
///   (I - α H) E^{n+1} = (I + α H) E^n  with α = i dz / (4 k₀ n₀)
///
/// where H is the finite-difference Laplacian + material phase operator.
pub struct FdBpm1d {
    pub nx: usize,
    pub dx: f64,
    pub n_ref: f64,
    pub wavelength: f64,
    pub n_profile: Option<Vec<f64>>,
    pub field: Vec<Complex64>,
}

impl FdBpm1d {
    pub fn new(nx: usize, dx: f64, n_ref: f64, wavelength: f64) -> Self {
        Self {
            nx,
            dx,
            n_ref,
            wavelength,
            n_profile: None,
            field: vec![Complex64::new(0.0, 0.0); nx],
        }
    }

    pub fn set_index_profile(&mut self, n_profile: Vec<f64>) {
        assert_eq!(n_profile.len(), self.nx);
        self.n_profile = Some(n_profile);
    }

    pub fn set_gaussian_input(&mut self, amplitude: f64, x_center: f64, w0: f64) {
        for i in 0..self.nx {
            let x = i as f64 * self.dx - x_center;
            let env = (-x * x / (w0 * w0)).exp();
            self.field[i] = Complex64::new(amplitude * env, 0.0);
        }
    }

    pub fn set_field(&mut self, field: Vec<Complex64>) {
        self.field = field;
    }

    /// Propagate one step dz using Crank-Nicolson.
    ///
    /// Solves: (I - α H) E^{n+1} = (I + α H) E^n
    /// where H_ii = -2/dx² + k₀Δn·2 (material)
    ///       H_{i,i±1} = 1/dx²
    ///       α = i·dz / (4 k₀ n₀)
    pub fn step(&mut self, dz: f64) {
        let k0 = 2.0 * PI / self.wavelength;
        let kn = k0 * self.n_ref;
        let dx2 = self.dx * self.dx;

        // Crank-Nicolson coefficient: α = i·dz / (4 k₀ n₀)
        let alpha = Complex64::new(0.0, dz / (4.0 * kn));

        let n = self.nx;

        // Build the tridiagonal H operator diagonal elements
        // H_diag[i] = -2/dx² + 2k0 Δn[i]   (material + kinetic)
        // H_off    = 1/dx²  (constant off-diagonal)
        let h_off = 1.0 / dx2;
        let h_diag: Vec<f64> = (0..n)
            .map(|i| {
                let dn = if let Some(ref np) = self.n_profile {
                    np[i] - self.n_ref
                } else {
                    0.0
                };
                -2.0 / dx2 + 2.0 * k0 * dn
            })
            .collect();

        // Tridiagonal system: (I - α H) E^{n+1} = rhs
        // Diagonal of lhs: 1 - α h_diag[i]
        // Off-diagonal of lhs: -α h_off
        let lhs_diag: Vec<Complex64> = h_diag
            .iter()
            .map(|&h| Complex64::new(1.0, 0.0) - alpha * h)
            .collect();
        let lhs_off: Complex64 = -alpha * h_off;

        // Compute rhs = (I + α H) E^n using the same stencil
        let e = &self.field;
        let mut rhs = vec![Complex64::new(0.0, 0.0); n];
        for i in 0..n {
            let h_e_i = h_diag[i] * e[i]
                + if i > 0 {
                    h_off * e[i - 1]
                } else {
                    Complex64::new(0.0, 0.0)
                }
                + if i + 1 < n {
                    h_off * e[i + 1]
                } else {
                    Complex64::new(0.0, 0.0)
                };
            rhs[i] = e[i] + alpha * h_e_i;
        }

        // Solve tridiagonal system using Thomas algorithm
        self.field = thomas_solve(&lhs_diag, lhs_off, &rhs);
    }

    pub fn propagate(&mut self, dz: f64, n_steps: usize) {
        for _ in 0..n_steps {
            self.step(dz);
        }
    }

    pub fn intensity(&self) -> Vec<f64> {
        self.field.iter().map(|e| e.norm_sqr()).collect()
    }

    pub fn peak_intensity(&self) -> f64 {
        self.intensity().iter().cloned().fold(0.0_f64, f64::max)
    }

    pub fn rms_width(&self) -> f64 {
        let intensity = self.intensity();
        let total: f64 = intensity.iter().sum();
        if total == 0.0 {
            return 0.0;
        }
        let mean = intensity
            .iter()
            .enumerate()
            .map(|(i, &v)| i as f64 * self.dx * v)
            .sum::<f64>()
            / total;
        let var = intensity
            .iter()
            .enumerate()
            .map(|(i, &v)| {
                let x = i as f64 * self.dx - mean;
                x * x * v
            })
            .sum::<f64>()
            / total;
        var.sqrt()
    }
}

/// Result type returned by `FdBpm1d::run_full`.
///
/// Contains the propagated output field, power transmission coefficient,
/// and accumulated phase shift relative to a free-space reference.
#[derive(Debug, Clone)]
pub struct BpmResult {
    /// Complex field profile at the output plane.
    pub output_field: Vec<Complex64>,
    /// Power transmission: P_out / P_in  (0 ≤ T ≤ 1 in lossless media).
    pub power_transmission: f64,
    /// Accumulated phase shift of the beam centroid  \[radians\].
    pub phase_shift: f64,
}

/// Mode field: a guided eigenmode profile of a waveguide cross-section.
#[derive(Debug, Clone)]
pub struct ModeField {
    /// Lateral field profile (normalised so that ∑|E|²·dx = 1).
    pub field: Vec<Complex64>,
    /// Effective index of the mode.
    pub n_eff: f64,
    /// Mode order (0 = fundamental, 1 = first-order, …).
    pub order: usize,
}

/// Extended 1D finite-difference BPM with optional Kerr nonlinearity,
/// mode decomposition, and full-run convenience method.
pub struct FdBpm {
    inner: FdBpm1d,
    /// χ³ Kerr coefficient  \[m²/W\].  Zero → linear propagation.
    chi3: f64,
    /// Number of propagation steps recorded in the last `run_full` call.
    n_steps_last: usize,
    /// Accumulated propagation length  \[m\].
    z_total: f64,
}

impl FdBpm {
    /// Construct an `FdBpm` from an existing `FdBpm1d`.
    pub fn from_bpm1d(bpm: FdBpm1d) -> Self {
        Self {
            inner: bpm,
            chi3: 0.0,
            n_steps_last: 0,
            z_total: 0.0,
        }
    }

    /// Enable Kerr nonlinearity with coefficient `chi3` \[m²/W\].
    ///
    /// During each propagation step the intensity-dependent index correction
    ///   Δn_nl = (χ³ / 2) · |E|²
    /// is added to the refractive-index profile before solving the linear
    /// Crank-Nicolson step.  This implements the scalar nonlinear BPM.
    pub fn with_nonlinear_kerr(mut self, chi3: f64) -> Self {
        self.chi3 = chi3;
        self
    }

    /// Propagate one step `dz` with optional Kerr nonlinearity.
    pub fn step(&mut self, dz: f64) {
        if self.chi3 != 0.0 {
            self.apply_kerr_correction();
        }
        self.inner.step(dz);
        self.z_total += dz;
    }

    /// Apply the intensity-dependent refractive-index perturbation.
    ///
    /// Temporarily adds Δn_nl\[i\] = (χ³ / 2) · |E\[i\]|² to n_profile,
    /// so that the next Crank-Nicolson step sees the nonlinear index.
    fn apply_kerr_correction(&mut self) {
        let chi3 = self.chi3;
        let field = &self.inner.field;
        let n_base = self
            .inner
            .n_profile
            .get_or_insert_with(|| vec![self.inner.n_ref; self.inner.nx]);
        for (ni, ei) in n_base.iter_mut().zip(field.iter()) {
            *ni += (chi3 / 2.0) * ei.norm_sqr();
        }
    }

    /// Run a complete BPM simulation over `total_length` using `n_steps`
    /// uniform propagation steps and return a `BpmResult`.
    ///
    /// The input field **must** be set on `self.inner` before calling this
    /// method (e.g. via `set_gaussian_input` or `set_field`).
    ///
    /// # Arguments
    /// * `total_length` – physical propagation length  \[m\]
    /// * `n_steps`      – number of equal-dz steps
    pub fn run_full(&mut self, total_length: f64, n_steps: usize) -> BpmResult {
        let dz = total_length / n_steps as f64;

        // Capture initial power
        let p_in: f64 = self
            .inner
            .field
            .iter()
            .map(|e| e.norm_sqr() * self.inner.dx)
            .sum();

        // Free-space phase accumulated over the same length
        let k0 = 2.0 * PI / self.inner.wavelength;
        let phi_free = k0 * self.inner.n_ref * total_length;

        for _ in 0..n_steps {
            self.step(dz);
        }
        self.n_steps_last = n_steps;

        let output_field = self.inner.field.clone();

        // Output power
        let p_out: f64 = output_field
            .iter()
            .map(|e| e.norm_sqr() * self.inner.dx)
            .sum();
        let power_transmission = if p_in > 0.0 { p_out / p_in } else { 0.0 };

        // Phase of the field centroid relative to free-space
        let total_intensity: f64 = output_field.iter().map(|e| e.norm_sqr()).sum();
        let centroid_phase = if total_intensity > 0.0 {
            let weighted_phase: f64 = output_field
                .iter()
                .map(|e| e.arg() * e.norm_sqr())
                .sum::<f64>()
                / total_intensity;
            weighted_phase
        } else {
            0.0
        };
        let phase_shift = centroid_phase - phi_free;

        BpmResult {
            output_field,
            power_transmission,
            phase_shift,
        }
    }

    /// Compute approximate guided mode profiles by power-iteration BPM.
    ///
    /// Each mode is found by propagating a trial field with imaginary-step
    /// (lossy) propagation and projecting out lower modes.  This gives the
    /// `n_modes` lowest-order modes of the current index profile.
    ///
    /// # Arguments
    /// * `n_modes` – number of modes to compute (1 = fundamental only)
    pub fn compute_modes(&self, n_modes: usize) -> Vec<ModeField> {
        let nx = self.inner.nx;
        let dx = self.inner.dx;
        let wavelength = self.inner.wavelength;
        let n_ref = self.inner.n_ref;

        // Use imaginary propagation (imag-distance BPM) to find modes.
        // For each mode order, start with a trial Hermite-Gaussian and
        // iterate until convergence.
        let mut modes: Vec<ModeField> = Vec::with_capacity(n_modes);
        let xc = nx as f64 * dx / 2.0;

        for order in 0..n_modes {
            // Trial field: Hermite-Gaussian of the appropriate order
            let mut trial: Vec<Complex64> = (0..nx)
                .map(|i| {
                    let x = i as f64 * dx - xc;
                    let w = dx * nx as f64 / 8.0;
                    let h = hermite_polynomial(order, x / w);
                    let env = (-x * x / (2.0 * w * w)).exp();
                    Complex64::new(h * env, 0.0)
                })
                .collect();

            // Orthogonalise against already-found modes (Gram-Schmidt)
            for prev in &modes {
                let overlap = inner_product(&trial, &prev.field, dx);
                for (ti, pi) in trial.iter_mut().zip(prev.field.iter()) {
                    *ti -= overlap * pi;
                }
            }

            // Normalise
            normalise_field(&mut trial, dx);

            // Imaginary-distance BPM iterations
            let n_iter = 200;
            let idz = -dx * 10.0; // imaginary step
            for _iter in 0..n_iter {
                let mut bpm_iter = FdBpm1d::new(nx, dx, n_ref, wavelength);
                if let Some(ref np) = self.inner.n_profile {
                    bpm_iter.set_index_profile(np.clone());
                }
                bpm_iter.set_field(trial.clone());
                // One imaginary step: multiply diagonal phase by exp(dz·H)
                // For imaginary propagation, use a real dz with sign flipped
                bpm_iter.step(idz.abs());
                trial = bpm_iter.field;

                // Re-orthogonalise
                for prev in &modes {
                    let overlap = inner_product(&trial, &prev.field, dx);
                    for (ti, pi) in trial.iter_mut().zip(prev.field.iter()) {
                        *ti -= overlap * pi;
                    }
                }
                normalise_field(&mut trial, dx);
            }

            // Estimate effective index from phase accumulation
            let k0 = 2.0 * PI / wavelength;
            let phase_avg: f64 = trial.iter().map(|e| e.arg() * e.norm_sqr()).sum::<f64>()
                / trial.iter().map(|e| e.norm_sqr()).sum::<f64>().max(1e-30);
            let n_eff = n_ref + phase_avg / (k0 * idz.abs() * n_iter as f64);
            let n_eff = n_eff.clamp(1.0, 5.0);

            modes.push(ModeField {
                field: trial,
                n_eff,
                order,
            });
        }

        modes
    }

    /// Estimate propagation loss in dB/μm from power decay over `length` \[m\].
    ///
    /// Runs a short simulation and computes the power attenuation rate.
    /// For a lossless waveguide this returns 0; a lossy material or
    /// leaky mode gives a positive value.
    pub fn propagation_loss_db_per_um(&mut self, length: f64) -> f64 {
        let n_steps = 100;
        let p_in: f64 = self
            .inner
            .field
            .iter()
            .map(|e| e.norm_sqr() * self.inner.dx)
            .sum();
        if p_in == 0.0 {
            return 0.0;
        }
        let dz = length / n_steps as f64;
        for _ in 0..n_steps {
            self.inner.step(dz);
        }
        let p_out: f64 = self
            .inner
            .field
            .iter()
            .map(|e| e.norm_sqr() * self.inner.dx)
            .sum();
        if p_out <= 0.0 || p_in <= 0.0 {
            return 0.0;
        }
        let loss_db = -10.0 * (p_out / p_in).log10();
        loss_db / (length * 1e6) // dB per μm
    }

    /// Compute the overlap integral between two complex field vectors.
    ///
    /// Returns the magnitude squared of the normalised overlap:
    ///   η = |⟨a|b⟩|² / (⟨a|a⟩ · ⟨b|b⟩)
    ///
    /// This is the coupling efficiency from mode `b` into mode `a` (0 to 1).
    pub fn overlap_integral(field_a: &[Complex64], field_b: &[Complex64], dx: f64) -> f64 {
        if field_a.len() != field_b.len() {
            return 0.0;
        }
        let ab: Complex64 = field_a
            .iter()
            .zip(field_b.iter())
            .map(|(a, b)| a.conj() * b * dx)
            .sum();
        let aa: f64 = field_a.iter().map(|a| a.norm_sqr() * dx).sum();
        let bb: f64 = field_b.iter().map(|b| b.norm_sqr() * dx).sum();
        if aa == 0.0 || bb == 0.0 {
            return 0.0;
        }
        ab.norm_sqr() / (aa * bb)
    }

    /// Total accumulated propagation length since construction (or reset).
    pub fn z_total(&self) -> f64 {
        self.z_total
    }

    /// Immutable reference to the wrapped `FdBpm1d`.
    pub fn inner(&self) -> &FdBpm1d {
        &self.inner
    }

    /// Mutable reference to the wrapped `FdBpm1d`.
    pub fn inner_mut(&mut self) -> &mut FdBpm1d {
        &mut self.inner
    }
}

// ── Helper functions ─────────────────────────────────────────────────────────

/// Hermite polynomial H_n(x) using recurrence: H_0=1, H_1=2x,
///   H_n = 2x·H_{n-1} − 2(n-1)·H_{n-2}.
fn hermite_polynomial(n: usize, x: f64) -> f64 {
    match n {
        0 => 1.0,
        1 => 2.0 * x,
        _ => {
            let mut h_prev2 = 1.0_f64;
            let mut h_prev1 = 2.0 * x;
            for k in 2..=n {
                let h = 2.0 * x * h_prev1 - 2.0 * (k - 1) as f64 * h_prev2;
                h_prev2 = h_prev1;
                h_prev1 = h;
            }
            h_prev1
        }
    }
}

/// Compute the inner product ⟨a|b⟩ = Σ a*·b·dx.
fn inner_product(a: &[Complex64], b: &[Complex64], dx: f64) -> Complex64 {
    a.iter()
        .zip(b.iter())
        .map(|(ai, bi)| ai.conj() * bi * dx)
        .sum()
}

/// Normalise a field in place so that Σ|E|²·dx = 1.
fn normalise_field(field: &mut [Complex64], dx: f64) {
    let norm: f64 = field.iter().map(|e| e.norm_sqr() * dx).sum::<f64>().sqrt();
    if norm > 0.0 {
        for e in field.iter_mut() {
            *e /= norm;
        }
    }
}

/// Thomas algorithm for solving a tridiagonal complex system.
///
/// Solves: diag\[i\] * x\[i\] + off * x\[i-1\] + off * x\[i+1\] = rhs\[i\]
/// (uniform off-diagonal elements on both sides)
fn thomas_solve(diag: &[Complex64], off: Complex64, rhs: &[Complex64]) -> Vec<Complex64> {
    let n = diag.len();
    let mut c = vec![Complex64::new(0.0, 0.0); n]; // modified off-diagonal
    let mut d = vec![Complex64::new(0.0, 0.0); n]; // modified rhs
    let mut x = vec![Complex64::new(0.0, 0.0); n];

    // Forward sweep
    c[0] = off / diag[0];
    d[0] = rhs[0] / diag[0];
    for i in 1..n {
        let denom = diag[i] - off * c[i - 1];
        c[i] = off / denom;
        d[i] = (rhs[i] - off * d[i - 1]) / denom;
    }

    // Backward substitution
    x[n - 1] = d[n - 1];
    for i in (0..n - 1).rev() {
        x[i] = d[i] - c[i] * x[i + 1];
    }
    x
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fd_bpm_free_space_stable() {
        let mut bpm = FdBpm1d::new(128, 100e-9, 1.0, 1550e-9);
        let xc = 128.0 * 100e-9 / 2.0;
        bpm.set_gaussian_input(1.0, xc, 1e-6);
        bpm.propagate(500e-9, 20);
        // Fields should remain finite
        assert!(bpm.field.iter().all(|e| e.norm().is_finite()));
    }

    #[test]
    fn fd_bpm_beam_spreads() {
        let mut bpm = FdBpm1d::new(256, 50e-9, 1.0, 1550e-9);
        let xc = 256.0 * 50e-9 / 2.0;
        bpm.set_gaussian_input(1.0, xc, 1e-6);
        let w_init = bpm.rms_width();
        bpm.propagate(1e-6, 20);
        let w_final = bpm.rms_width();
        assert!(w_final > w_init, "Beam should spread in free space");
    }

    #[test]
    fn bpm_result_power_transmission_lossless() {
        let nx = 128;
        let dx = 50e-9;
        let mut bpm1d = FdBpm1d::new(nx, dx, 1.5, 1550e-9);
        let xc = nx as f64 * dx / 2.0;
        bpm1d.set_gaussian_input(1.0, xc, 1e-6);

        let mut bpm = FdBpm::from_bpm1d(bpm1d);
        let result = bpm.run_full(5e-6, 50);

        // In a lossless uniform medium, power should be conserved
        assert!(
            result.power_transmission > 0.5,
            "Lossless BPM should conserve power, T = {:.4}",
            result.power_transmission
        );
        assert!(
            result.output_field.len() == nx,
            "Output field length mismatch"
        );
    }

    #[test]
    fn bpm_kerr_enabled_field_finite() {
        let nx = 64;
        let dx = 100e-9;
        let mut bpm1d = FdBpm1d::new(nx, dx, 1.5, 1550e-9);
        let xc = nx as f64 * dx / 2.0;
        bpm1d.set_gaussian_input(1.0, xc, 2e-6);

        let mut bpm = FdBpm::from_bpm1d(bpm1d).with_nonlinear_kerr(1e-20);
        let result = bpm.run_full(2e-6, 20);

        assert!(
            result.output_field.iter().all(|e| e.norm().is_finite()),
            "Kerr BPM should remain finite"
        );
    }

    #[test]
    fn overlap_integral_self_is_unity() {
        let n = 64;
        let dx = 50e-9;
        let field: Vec<Complex64> = (0..n)
            .map(|i| {
                let x = i as f64 * dx - n as f64 * dx / 2.0;
                let env = (-x * x / (1e-6 * 1e-6)).exp();
                Complex64::new(env, 0.0)
            })
            .collect();

        let eta = FdBpm::overlap_integral(&field, &field, dx);
        assert!(
            (eta - 1.0).abs() < 1e-6,
            "Self-overlap should be 1.0, got {eta:.6}"
        );
    }

    #[test]
    fn overlap_integral_orthogonal_fields_small() {
        let n = 128;
        let dx = 50e-9;
        let xc = n as f64 * dx / 2.0;
        let w = 1e-6;

        // Even and odd Hermite-Gaussian modes are orthogonal
        let field_even: Vec<Complex64> = (0..n)
            .map(|i| {
                let x = i as f64 * dx - xc;
                let env = (-x * x / (w * w)).exp();
                Complex64::new(env, 0.0) // H_0 * Gaussian
            })
            .collect();
        let field_odd: Vec<Complex64> = (0..n)
            .map(|i| {
                let x = i as f64 * dx - xc;
                let env = (-x * x / (w * w)).exp();
                Complex64::new(2.0 * x / w * env, 0.0) // H_1 * Gaussian
            })
            .collect();

        let eta = FdBpm::overlap_integral(&field_even, &field_odd, dx);
        assert!(
            eta < 1e-4,
            "Even-odd mode overlap should be near zero, got {eta:.6}"
        );
    }

    #[test]
    fn compute_modes_fundamental_bounded() {
        let nx = 64;
        let dx = 50e-9;
        let n_core = 3.476;
        let n_clad = 1.444;
        let wavelength = 1550e-9;

        // Simple slab: core region in the centre
        let wg_width = 10; // 10 cells of Si core
        let i0 = nx / 2 - wg_width / 2;
        let i1 = nx / 2 + wg_width / 2;
        let n_profile: Vec<f64> = (0..nx)
            .map(|i| if i >= i0 && i < i1 { n_core } else { n_clad })
            .collect();

        let mut bpm1d = FdBpm1d::new(nx, dx, n_clad, wavelength);
        bpm1d.set_index_profile(n_profile);
        let xc = nx as f64 * dx / 2.0;
        bpm1d.set_gaussian_input(1.0, xc, 200e-9);

        let bpm = FdBpm::from_bpm1d(bpm1d);
        let modes = bpm.compute_modes(1);

        assert!(!modes.is_empty(), "Should find at least one mode");
        // n_eff must be physically reasonable (between n_clad and n_core)
        let n_eff = modes[0].n_eff;
        assert!(
            n_eff > 1.0 && n_eff < 5.0,
            "Fundamental mode n_eff should be reasonable: {n_eff:.4}"
        );
    }

    #[test]
    fn hermite_polynomial_values() {
        assert!((hermite_polynomial(0, 1.0) - 1.0).abs() < 1e-12);
        assert!((hermite_polynomial(1, 1.0) - 2.0).abs() < 1e-12);
        // H_2(x) = 4x² - 2: H_2(1) = 2
        assert!((hermite_polynomial(2, 1.0) - 2.0).abs() < 1e-12);
    }

    #[test]
    fn bpm_z_total_accumulates() {
        let mut bpm1d = FdBpm1d::new(32, 50e-9, 1.5, 1550e-9);
        bpm1d.set_gaussian_input(1.0, 32.0 * 50e-9 / 2.0, 500e-9);
        let mut bpm = FdBpm::from_bpm1d(bpm1d);
        bpm.step(1e-6);
        bpm.step(2e-6);
        let z = bpm.z_total();
        assert!(
            (z - 3e-6).abs() < 1e-15,
            "z_total should accumulate: {z:.3e}"
        );
    }
}