import numpy as np

from vna_system.core.logging.logger import get_component_logger
from vna_system.core.acquisition.sweep_buffer import SweepData
from vna_system.core.settings.preset_manager import VNAMode
from vna_system.core.settings.calibration_manager import CalibrationSet, CalibrationStandard

logger = get_component_logger(__file__)


class CalibrationProcessor:
    """
    Apply VNA calibration to raw sweeps.

    Supports:
        - S11 (reflection) using OSL (Open–Short–Load) error model
        - S21 (transmission) using THRU reference

    All operations are vectorized with NumPy and return data as a list of (real, imag) tuples.
    """

    def apply_calibration(self, sweep_data: SweepData, calibration_set: CalibrationSet) -> list[tuple[float, float]]:
        """
        Calibrate a sweep and return corrected complex points.

        Parameters
        ----------
        sweep_data
            Raw sweep as (real, imag) tuples with metadata.
        calibration_set
            A complete set of standards for the current VNA mode.

        Returns
        -------
        list[tuple[float, float]]
            Calibrated complex points as (real, imag) pairs.

        Raises
        ------
        ValueError
            If the calibration set is incomplete or the VNA mode is unsupported.
        """
        if not calibration_set.is_complete():
            raise ValueError("Calibration set is incomplete")

        raw = self._to_complex_array(sweep_data)

        if calibration_set.preset.mode == VNAMode.S21:
            logger.debug("Applying S21 calibration", sweep_number=sweep_data.sweep_number, points=sweep_data.total_points)
            calibrated = self._apply_s21(raw, calibration_set)
        elif calibration_set.preset.mode == VNAMode.S11:
            logger.debug("Applying S11 calibration (OSL)", sweep_number=sweep_data.sweep_number, points=sweep_data.total_points)
            calibrated = self._apply_s11_osl(raw, calibration_set)
        else:
            raise ValueError(f"Unsupported measurement mode: {calibration_set.preset.mode}")

        return [(float(z.real), float(z.imag)) for z in calibrated]

    # --------------------------------------------------------------------- #
    # Helpers
    # --------------------------------------------------------------------- #
    @staticmethod
    def _to_complex_array(sweep: SweepData) -> np.ndarray:
        """Convert `SweepData.points` to a 1-D complex NumPy array."""
        # Using vectorized construction for speed and clarity
        if not sweep.points:
            return np.empty(0, dtype=np.complex64)
        arr = np.asarray(sweep.points, dtype=np.float32)
        return arr[:, 0].astype(np.float32) + 1j * arr[:, 1].astype(np.float32)

    @staticmethod
    def _safe_divide(num: np.ndarray, den: np.ndarray, eps: float = 1e-12) -> np.ndarray:
        """
        Elementwise complex-safe division with small epsilon guard.

        Avoids division by zero by clamping |den|<eps to eps (preserves phase).
        """
        mask = np.abs(den) < eps
        if np.any(mask):
            den = den.copy()
            # Scale only magnitude; keep angle of denominator
            den[mask] = eps * np.exp(1j * np.angle(den[mask]))
        return num / den

    # --------------------------------------------------------------------- #
    # S21 calibration (THRU)
    # --------------------------------------------------------------------- #
    def _apply_s21(self, raw: np.ndarray, calib: CalibrationSet) -> np.ndarray:
        """
        S21: normalize the DUT response by the THRU reference.

        Calibrated = Raw / Through
        """
        through = self._to_complex_array(calib.standards[CalibrationStandard.THROUGH])
        if raw.size != through.size:
            raise ValueError("Signal and THRU reference have different lengths")
        return self._safe_divide(raw, through)

    # --------------------------------------------------------------------- #
    # S11 calibration (OSL)
    # --------------------------------------------------------------------- #
    def _apply_s11_osl(self, raw: np.ndarray, calib: CalibrationSet) -> np.ndarray:
        """
        S11 OSL correction using a 3-term error model.

        Error terms
        -----------
        Ed (directivity)          := Load
        Es (source match)         := (Open + Short - 2*Load) / (Open - Short)
        Er (reflection tracking)  := -2*(Open - Load)*(Short - Load) / (Open - Short)

        Final correction
        ----------------
        S11 = (Raw - Ed) / (Er + Es * (Raw - Ed))
        """
        open_ref = self._to_complex_array(calib.standards[CalibrationStandard.OPEN])
        short_ref = self._to_complex_array(calib.standards[CalibrationStandard.SHORT])
        load_ref = self._to_complex_array(calib.standards[CalibrationStandard.LOAD])

        n = raw.size
        if not (open_ref.size == short_ref.size == load_ref.size == n):
            raise ValueError("Signal and OSL standards have different lengths")

        # Ed
        ed = load_ref

        # Es, Er with guarded denominators
        denom = open_ref - short_ref
        denom = np.where(np.abs(denom) < 1e-12, 1e-12 * np.exp(1j * np.angle(denom)), denom)

        es = (open_ref + short_ref - 2.0 * load_ref) / denom
        er = -2.0 * (open_ref - load_ref) * (short_ref - load_ref) / denom

        num = raw - ed
        den = er + es * num
        return self._safe_divide(num, den)
