core.py

from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

APP_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = APP_DIR.parent
LOCAL_MODAL_PACKAGE_ROOT = PROJECT_ROOT / "WaveguideModalBPM1D"

if LOCAL_MODAL_PACKAGE_ROOT.exists() and str(LOCAL_MODAL_PACKAGE_ROOT) not in sys.path:
    # Prioritize the local package in this workspace over any older installed version.
    sys.path.insert(0, str(LOCAL_MODAL_PACKAGE_ROOT))
if str(APP_DIR) not in sys.path:
    sys.path.insert(0, str(APP_DIR))

from WaveguideModalBPM1D import (  # noqa: E402
    CouplerTestParams,
    estimate_first_relevant_peak,
    nm,
    sweep_gap_to_lc_coupler,
    simulate_coupler,
    summarize_coupler_output,
    sweep_coupler_wavelength_response,
    um,
)
from WaveguideModalBPM1D.test_cases import CouplerSimulationResult  # noqa: E402


@dataclass(frozen=True)
class SpectrumResult:
    lambda_nm: np.ndarray
    upper_pct: np.ndarray
    lower_pct: np.ndarray
    loss_pct: np.ndarray
    dz_um: np.ndarray
    rows: list[dict[str, float]]
    backend: str


@dataclass(frozen=True)
class CouplerViewModel:
    params: CouplerTestParams
    result: CouplerSimulationResult
    lc_idx: int
    lc_z_value: float
    lc_local_value: float
    l3db_value: float
    beat_length_value: float | None
    lc_theoretical_value: float | None
    power_ref: float
    power_total_norm: np.ndarray
    power_upper_norm: np.ndarray
    power_lower_norm: np.ndarray
    input_profile_norm: np.ndarray
    output_profile_norm: np.ndarray
    output_upper_pct: float
    output_lower_pct: float
    output_loss_pct: float
    sweep_backend: str
    spectrum: SpectrumResult | None
    gap_sweep_results: list[dict[str, Any]] | None


def estimate_coupler_lc(
    z_vec: np.ndarray,
    power_lower_norm: np.ndarray,
    params: CouplerTestParams,
) -> tuple[int, float, float]:
    start_idx = int(np.searchsorted(z_vec, params.bend_length, side="left"))
    end_idx = int(
        np.searchsorted(
            z_vec,
            params.bend_length + params.length_coupling,
            side="right",
        )
    ) - 1
    end_idx = max(start_idx, min(end_idx, len(z_vec) - 1))
    peak_info: dict[str, Any] = estimate_first_relevant_peak(
        z_vec,
        power_lower_norm,
        z_start=params.bend_length,
        z_end=params.bend_length + params.length_coupling,
        min_fraction=0.01,
        min_power=0.01,
    )

    peak_idx_global = peak_info.get("peak_index_global")
    peak_z_value = peak_info.get("z_peak_absolute")
    if peak_idx_global is None or not np.isfinite(float(peak_z_value)):
        segment = power_lower_norm[start_idx : end_idx + 1]
        lc_idx = start_idx + int(np.argmax(segment))
        lc_z_value = float(z_vec[lc_idx])
    else:
        lc_idx = int(peak_idx_global)
        lc_z_value = float(peak_z_value)

    lc_local_value = max(0.0, lc_z_value - params.bend_length)
    return lc_idx, lc_z_value, lc_local_value


def estimate_coupler_beat_length(
    z_vec: np.ndarray,
    power_lower_norm: np.ndarray,
    params: CouplerTestParams,
    *,
    min_peak_height: float = 0.05,
) -> tuple[float | None, float | None]:
    start_idx = int(np.searchsorted(z_vec, params.bend_length, side="left"))
    end_idx = int(np.searchsorted(z_vec, params.bend_length + params.length_coupling, side="right")) - 1
    end_idx = max(start_idx, min(end_idx, len(z_vec) - 1))

    segment = power_lower_norm[start_idx : end_idx + 1]
    if len(segment) < 3:
        return None, None

    delta = np.diff(segment)
    local_peaks = np.where((delta[:-1] > 0) & (delta[1:] <= 0))[0] + 1
    valid_peaks = [idx for idx in local_peaks if segment[idx] >= min_peak_height]
    if len(valid_peaks) < 2:
        return None, None

    min_separation = max(10.0 * (z_vec[1] - z_vec[0]), 0.1 * params.length_coupling)
    chosen_pair = None
    for first_idx, second_idx in zip(valid_peaks[:-1], valid_peaks[1:]):
        z_first = float(z_vec[start_idx + first_idx])
        z_second = float(z_vec[start_idx + second_idx])
        if (z_second - z_first) >= min_separation:
            chosen_pair = (z_first, z_second)
            break

    if chosen_pair is None:
        return None, None

    z_first, z_second = chosen_pair
    beat_length = z_second - z_first
    lc_theoretical = 0.5 * beat_length
    return beat_length, lc_theoretical


def default_coupler_params() -> CouplerTestParams:
    return CouplerTestParams(
        n_core=1.50,
        n_clad=1.42,
        window=40.0 * um,
        lambda_launch=1.55 * um,
        length_total=1764.0 * um,
        length_coupling=882.0 * um,
        port_offset=10.0 * um,
        core_width=1.0 * um,
        center_offset=1.0 * um,
        sigma=0.45 * um,
        n_points=401,
        mode_search_count=8,
    )


def build_coupler_params(values: dict[str, float | int | bool]) -> CouplerTestParams:
    return CouplerTestParams(
        n_core=float(values["n_core"]),
        n_clad=float(values["n_clad"]),
        window=float(values["window_um"]) * um,
        lambda_launch=float(values["lambda_um"]) * um,
        length_total=float(values["length_total_um"]) * um,
        length_coupling=float(values["length_coupling_um"]) * um,
        port_offset=float(values["port_offset_um"]) * um,
        core_width=float(values["core_width_um"]) * um,
        center_offset=float(values["center_offset_um"]) * um,
        sigma=float(values["sigma_um"]) * um,
        n_points=int(values["n_points"]),
        mode_search_count=int(values["mode_search_count"]),
    )


def run_coupler_simulation(
    params: CouplerTestParams,
    *,
    run_sweep: bool,
    sweep_min_um: float,
    sweep_max_um: float,
    sweep_points: int,
    run_gap_sweep: bool,
    gap_sweep_min_um: float,
    gap_sweep_max_um: float,
    gap_sweep_points: int,
) -> CouplerViewModel:
    result = simulate_coupler(params=params)
    summary = summarize_coupler_output(result)

    power_ref = float(result.power_total[0])
    power_total_norm = result.power_total / power_ref
    power_upper_norm = result.power_upper / power_ref
    power_lower_norm = result.power_lower / power_ref

    lc_idx, lc_z_value, lc_local_value = estimate_coupler_lc(result.z_vec, power_lower_norm, params)
    beat_length_value, lc_theoretical_value = estimate_coupler_beat_length(
        result.z_vec, power_lower_norm, params
    )
    l3db_value = 0.5 * lc_local_value

    input_profile = np.abs(result.e_evol[0, :])
    output_profile = np.abs(result.e_evol[-1, :])
    input_profile_norm = input_profile / np.max(input_profile)
    output_profile_norm = output_profile / np.max(output_profile)

    output_upper_pct = summary["upper_out_percent"]
    output_lower_pct = summary["lower_out_percent"]
    output_loss_pct = summary["outside_cores_percent"]

    spectrum = None
    sweep_backend = "disabled"
    if run_sweep:
        lambda_vec = np.linspace(sweep_min_um, sweep_max_um, int(sweep_points)) * um
        sweep_rows = sweep_coupler_wavelength_response(params, lambda_vec)

        upper_pct = np.array([row["upper"] for row in sweep_rows], dtype=float)
        lower_pct = np.array([row["lower"] for row in sweep_rows], dtype=float)
        loss_pct = np.array([row["loss"] for row in sweep_rows], dtype=float)
        dz_um = np.array([row["dz_um"] for row in sweep_rows], dtype=float)
        sweep_backend = "package-workflow"

        spectrum = SpectrumResult(
            lambda_nm=lambda_vec / nm,
            upper_pct=upper_pct,
            lower_pct=lower_pct,
            loss_pct=loss_pct,
            dz_um=dz_um,
            rows=sweep_rows,
            backend=sweep_backend,
        )

    gap_sweep_results = None
    if run_gap_sweep:
        gap_values = np.linspace(gap_sweep_min_um, gap_sweep_max_um, int(gap_sweep_points)) * um
        gap_sweep_results = sweep_gap_to_lc_coupler(params=params, gap_values=gap_values)

    return CouplerViewModel(
        params=params,
        result=result,
        lc_idx=int(lc_idx),
        lc_z_value=float(lc_z_value),
        lc_local_value=float(lc_local_value),
        l3db_value=float(l3db_value),
        beat_length_value=beat_length_value,
        lc_theoretical_value=lc_theoretical_value,
        power_ref=power_ref,
        power_total_norm=power_total_norm,
        power_upper_norm=power_upper_norm,
        power_lower_norm=power_lower_norm,
        input_profile_norm=input_profile_norm,
        output_profile_norm=output_profile_norm,
        output_upper_pct=float(output_upper_pct),
        output_lower_pct=float(output_lower_pct),
        output_loss_pct=float(output_loss_pct),
        sweep_backend=sweep_backend,
        spectrum=spectrum,
        gap_sweep_results=gap_sweep_results,
    )