"""
K-R Pilot-Identifiable WiFi Receiver
======================================
Complete reference implementation for:
"Pilot-Identifiable Single-Parameter Compensation for
 Power-Amplifier Nonlinearity in IEEE 802.11ax Receivers"

Author: Ramakrishna Pasupuleti
        Independent Researcher, Suryapet, Telangana, India
ORCID:  0009-0008-8418-1430

Usage:
    from kr_wifi_receiver import KRReceiver, build_lut
    lut = build_lut()                      # calibrate LUT once
    rx  = KRReceiver(lut=lut)              # instantiate receiver
    r   = rx.receive(y, h_est, sigma_n,    # per-packet call
                     x_pilot, snr_db)
"""

import numpy as np
from typing import Optional, Tuple, Dict
import json

# ── Constants ────────────────────────────────────────────────────────
PILOT_IDX       = np.array([7, 21, 43, 57])   # IEEE 802.11ax HE-LTF 52-tone RU
TAU1            = 0.15   # gate lower threshold
TAU2            = 0.30   # gate upper threshold
A_LO            = 0.3    # saturation search lower bound
A_HI            = 1.2    # saturation search upper bound
RAPP_P          = 2.0    # Rapp smoothness exponent (fixed, known)
GSS_N_ITER      = 12     # golden-section iterations (validated)
SNR_THRESHOLD   = 20.0   # dB: below this K-R is disabled (SNR-aware gate)
GOLDEN_RATIO    = (np.sqrt(5) - 1) / 2


# ── Rapp model ───────────────────────────────────────────────────────
def pa_rapp(x: np.ndarray, A_sat: float, p: float = RAPP_P) -> np.ndarray:
    """Rapp PA model: f(x; A_sat) = x / (1 + (|x|/A_sat)^(2p))^(1/(2p))"""
    return x / (1.0 + (np.abs(x) / A_sat) ** (2*p)) ** (1.0 / (2*p))


def pa_rapp_inverse(y: np.ndarray, A_sat: float, p: float = RAPP_P) -> np.ndarray:
    """Inverse Rapp model: f^{-1}(y; A_sat)"""
    mag = np.minimum(np.abs(y), A_sat * 0.999)
    denom = np.maximum(1.0 - (mag / A_sat) ** (2*p), 1e-6)
    return (mag / denom ** (1.0 / (2*p))) * np.exp(1j * np.angle(y))


# ── Bounded golden-section search ────────────────────────────────────
def fit_A_sat_gss(r_pilot: np.ndarray, x_pilot: np.ndarray,
                  n_iter: int = GSS_N_ITER) -> float:
    """
    Estimate A_sat from 4 pilot observations via bounded golden-section search.

    Args:
        r_pilot: MMSE-equalised pilot symbols (4-element complex array)
        x_pilot: known transmitted pilot symbols (4-element complex array)
        n_iter:  number of golden-section iterations (default 12)

    Returns:
        A_hat: estimated saturation amplitude in [A_LO, A_HI]
    """
    a, b = A_LO, A_HI

    def loss(A: float) -> float:
        diff = pa_rapp_inverse(r_pilot, A) - x_pilot
        return float(np.real(np.vdot(diff, diff)))

    c = b - GOLDEN_RATIO * (b - a)
    d = a + GOLDEN_RATIO * (b - a)
    fc, fd = loss(c), loss(d)

    for _ in range(n_iter):
        if fc < fd:
            b, d, fd = d, c, fc
            c = b - GOLDEN_RATIO * (b - a)
            fc = loss(c)
        else:
            a, c, fc = c, d, fd
            d = a + GOLDEN_RATIO * (b - a)
            fd = loss(d)

    return (a + b) / 2.0


# ── LUT estimator (deployment-grade, pure-Python speed) ─────────────
def build_lut(n_cal: int = 1200, n_lut: int = 128,
              seed: int = 500000) -> Dict:
    """
    Build a lookup table mapping distortion metric D -> A_hat.

    The LUT is calibrated from a swept set of (D, A_hat) pairs obtained
    by running the golden-section estimator on packets with random A_sat
    and SNR values drawn from the practical operating range.  Call this
    ONCE offline; store the result and pass it to KRReceiver.

    Args:
        n_cal: number of calibration packets (default 1200)
        n_lut: number of LUT entries (default 128)
        seed:  random seed for calibration packets (must not overlap
               with validation seeds 0..n_cal to avoid data leakage)

    Returns:
        dict with keys 'D_sorted', 'A_sorted', 'D_range'
    """
    from simulation_utils import make_packet  # defined in simulation_utils.py

    D_cal, A_cal = [], []
    for s in range(n_cal):
        rng = np.random.RandomState(seed + s)
        asat = rng.uniform(A_LO, A_HI)
        snr  = rng.uniform(15.0, 40.0)
        y, h, sigma_n, x, _ = make_packet(snr, asat, seed + s)
        r = mmse_equalise(y, h, sigma_n)
        rp, yp, xp, hp = r[PILOT_IDX], y[PILOT_IDX], x[PILOT_IDX], h[PILOT_IDX]
        D = float(distortion_metric(yp, hp, xp))
        alpha = soft_gate(D)
        if alpha > 0:
            A_fit = fit_A_sat_gss(rp, xp)
            D_cal.append(D)
            A_cal.append(A_fit)

    pairs  = sorted(zip(D_cal, A_cal))
    D_sort = np.array([d for d, _ in pairs])
    A_sort = np.array([a for _, a in pairs])

    return {'D_sorted': D_sort.tolist(), 'A_sorted': A_sort.tolist(),
            'n_lut': n_lut, 'D_range': [float(D_sort[0]), float(D_sort[-1])]}


def lut_lookup(D: float, lut: Dict) -> float:
    """Linear interpolation in a pre-sorted LUT."""
    # Support both key conventions
    D_arr = np.asarray(lut.get('D_sorted', lut.get('D_LUT')))
    A_arr = np.asarray(lut.get('A_sorted', lut.get('A_LUT')))
    idx = int(np.searchsorted(D_arr, D))
    idx = max(1, min(idx, len(D_arr) - 1))
    t = (D - D_arr[idx-1]) / (D_arr[idx] - D_arr[idx-1] + 1e-12)
    return float(A_arr[idx-1] + t * (A_arr[idx] - A_arr[idx-1]))


# ── Core per-packet operations ────────────────────────────────────────
def mmse_equalise(y: np.ndarray, h: np.ndarray,
                  sigma_n: float) -> np.ndarray:
    """Standard frequency-domain MMSE equalisation."""
    return (np.conj(h) / (np.abs(h)**2 + sigma_n**2 + 1e-12)) * y


def distortion_metric(y_pilot: np.ndarray, h_pilot: np.ndarray,
                      x_pilot: np.ndarray) -> float:
    """Pilot-domain distortion metric D (Eq. 2 in paper)."""
    e = y_pilot - h_pilot * x_pilot
    return float(np.sum(np.abs(e)**2) / (np.sum(np.abs(y_pilot)**2) + 1e-12))


def soft_gate(D: float, tau1: float = TAU1,
              tau2: float = TAU2) -> float:
    """Soft gate alpha(D) (Eq. 3 in paper)."""
    if D <= tau1: return 0.0
    if D >= tau2: return 1.0
    return (D - tau1) / (tau2 - tau1)


# ── Main receiver class ───────────────────────────────────────────────
class KRReceiver:
    """
    K-R pilot-identifiable WiFi receiver (IEEE 802.11ax).

    Supports three A_sat estimation back-ends:
        'gss'  : fixed-iteration golden-section search (reference)
        'lut'  : lookup-table interpolation (deployment-grade speed)
        'auto' : use LUT if available, fall back to GSS

    Also implements the SNR-aware gate: K-R is disabled when the
    estimated SNR falls below snr_threshold (default 20 dB), consistent
    with the estimator's high variance at low SNR (Fig. 6a in paper).
    """

    def __init__(self,
                 lut: Optional[Dict] = None,
                 estimator: str = 'auto',
                 snr_threshold: float = SNR_THRESHOLD,
                 tau1: float = TAU1,
                 tau2: float = TAU2,
                 a_lo: float = A_LO,
                 a_hi: float = A_HI,
                 n_iter: int = GSS_N_ITER):

        self.lut           = lut
        self.estimator     = estimator if (lut or estimator == 'gss') else 'gss'
        self.snr_threshold = snr_threshold
        self.tau1, self.tau2 = tau1, tau2
        self.a_lo, self.a_hi = a_lo, a_hi
        self.n_iter        = n_iter

    def receive(self,
                y_he: np.ndarray,
                h_est: np.ndarray,
                sigma_n: float,
                x_pilot: np.ndarray,
                snr_db: float) -> Tuple[np.ndarray, Dict]:
        """
        Process one 64-subcarrier HE-LTF OFDM symbol.

        Args:
            y_he:    received symbol after ADC (64 complex samples)
            h_est:   channel estimate for all 64 subcarriers (from L-LTF)
            sigma_n: estimated noise standard deviation
            x_pilot: known transmitted pilot symbols (4 complex values,
                     ordered by PILOT_IDX = [7, 21, 43, 57])
            snr_db:  estimated SNR in dB (from preamble)

        Returns:
            (r_out, meta)
            r_out: corrected output, 64 complex samples
            meta:  dict with diagnostic fields:
                   alpha, A_hat, method, D, snr_db
        """
        # Step 1: MMSE equalisation
        r = mmse_equalise(y_he, h_est, sigma_n)

        # Step 2: SNR-aware gate
        if snr_db < self.snr_threshold:
            return r, {'alpha': 0.0, 'A_hat': None, 'D': 0.0,
                       'method': 'MMSE_LOW_SNR', 'snr_db': snr_db}

        # Step 3: Distortion metric on 4 pilots
        rp = r[PILOT_IDX]
        yp = y_he[PILOT_IDX]
        hp = h_est[PILOT_IDX]
        D  = distortion_metric(yp, hp, x_pilot)

        # Step 4: Soft gate
        alpha = soft_gate(D, self.tau1, self.tau2)
        if alpha == 0.0:
            return r, {'alpha': 0.0, 'A_hat': None, 'D': D,
                       'method': 'MMSE_GATE_CLOSED', 'snr_db': snr_db}

        # Step 5: A_sat estimation
        use_lut = (self.estimator == 'lut' or
                   (self.estimator == 'auto' and self.lut is not None))

        if use_lut and self.lut is not None:
            A_hat  = lut_lookup(D, self.lut)
            method = 'KR_LUT'
        else:
            A_hat  = fit_A_sat_gss(rp, x_pilot, self.n_iter)
            method = 'KR_GSS'

        # Step 6: Gated nonlinear inversion
        r_inv = pa_rapp_inverse(r, A_hat)
        r_out = r + alpha * (r_inv - r)

        return r_out, {'alpha': alpha, 'A_hat': A_hat, 'D': D,
                       'method': method, 'snr_db': snr_db}

    def save_lut(self, path: str) -> None:
        """Save LUT to JSON file for reuse."""
        if self.lut is None:
            raise ValueError("No LUT to save; build one with build_lut() first.")
        with open(path, 'w') as f:
            json.dump(self.lut, f, indent=2)

    @classmethod
    def load_lut(cls, path: str, **kwargs) -> 'KRReceiver':
        """Load LUT from JSON file and return a configured KRReceiver."""
        with open(path) as f:
            lut = json.load(f)
        return cls(lut=lut, estimator='lut', **kwargs)