"""Shared constants and functions for scope patch extraction and signal labels.

This module is the single source of truth for parameters that must stay
synchronised between the training pipeline (scripts/model_definition.py)
and the runtime classifier (lib/signal/signal_classifier.py).

scripts/model_definition.py imports from this module via a try/except
fallback so that it works both in the project repo (where lib/ is on
sys.path) and in the deployed environment (~/Pycom/signals/scripts/,
where patch_config.py is copied alongside model_definition.py).
"""

from typing import Optional

import numpy as np


# ---------------------------------------------------------------------------
# Scope / patch geometry
# ---------------------------------------------------------------------------

SCOPE_HISTORY_ROWS = 100
SCOPE_HISTORY_COLS = 475
SCOPE_PATCH_ROWS = 100
SCOPE_PATCH_COLS = 64
INPUT_CHANNELS = 1
SCOPE_AMP_MAX = 160
PATCH_MIN_MARGIN_BINS = 20
PATCH_MARGIN_FRACTION = 0.15
PATCH_STRATEGY = "centered"

# Waterfall scale factor: the waterfall/effective span is this many times
# the radio's raw scope_span. Used by WaterfallManager, capture window
# display, and radio context construction. Single source of truth.
WATERFALL_SCALE_FACTOR = 2

# Auxiliary input dimension: 6 base features + 7 derived spectral features
# (span, ref_level, noise_floor, bandwidth_bins, center_freq_offset_hz, snr_db,
#  spectral_flatness, spectral_kurtosis, peak_count, spectral_periodicity,
#  spectral_bw_ratio, centroid_offset, temporal_env_cv)
AUX_DIM = 13

AUX_FEATURE_NAMES = [
    "span_hz",
    "reference_level_db",
    "noise_floor",
    "bandwidth_bins",
    "center_freq_offset_hz",
    "snr_db",
    "spectral_flatness",
    "spectral_kurtosis",
    "peak_count",
    "spectral_periodicity",
    "spectral_bw_ratio",
    "centroid_offset",
    "temporal_env_cv",
]

assert len(AUX_FEATURE_NAMES) == AUX_DIM, (
    f"AUX_FEATURE_NAMES length ({len(AUX_FEATURE_NAMES)}) must match AUX_DIM ({AUX_DIM})"
)

# ---------------------------------------------------------------------------
# Signal labels
# ---------------------------------------------------------------------------

# Model labels: the 15 classes the classifier is trained on.
# This list MUST match model/signal_labels.json at training time.
SIGNAL_LABELS = [
    "CW",
    "FM",
    "AM",
    "SSB",
    "RTTY",
    "FT8",
    "FT4",
    "FT2",
    "Q65",
    "JT65",
    "AFSK",
    "PSK",
    "FSK",
    "MFSK",
    "MSK",
]

# Capture labels: model labels plus "Unknown" for human annotation
# in the capture manager UI.
CAPTURE_LABELS = SIGNAL_LABELS + ["Unknown"]

# Fallback used at runtime when the labels JSON is missing or corrupt.
SIGNAL_LABELS_FALLBACK = ["Unknown"]


# ---------------------------------------------------------------------------
# Auxiliary feature normalisation
# ---------------------------------------------------------------------------

def prepare_aux_features(
    span_hz: float = 100000,
    reference_level_db: float = 0,
    noise_floor: float = 10,
    bandwidth_bins: int = 5,
    center_freq_offset_hz: float = 0,
    snr_db: float = 0,
    derived_features: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Normalize auxiliary features to [0, 1] or [-1, 1] for the model.

    Parameters
    ----------
    span_hz : float
        Scope span in Hz. Log-scaled to [0, 1] (1 Hz -> 0, 1 MHz -> 1).
    reference_level_db : float
        Scope reference level in dB. Maps [-20, 60] to [0, 1].
    noise_floor : float
        Noise floor in raw scope units (0-160). Maps to [0, 1].
    bandwidth_bins : int
        Signal width in detector bins (0-950). Maps to [0, 1].
    center_freq_offset_hz : float
        Offset from center frequency in Hz. Candidate frequencies span
        +/- radio_span (the detector uses effective_span = 2 * radio_span).
        Normalized by radio_span to [-1, 1].
    snr_db : float
        Signal-to-noise ratio in dB. Maps [0, 40] to [0, 1].
    derived_features : np.ndarray, optional
        Pre-computed derived spectral features (shape (7,) float32 from
        prepare_derived_features). If None, padded with zeros.

    Returns
    -------
    np.ndarray
        Shape (AUX_DIM,) float32 with normalized features.
    """
    norm_span = min(max(np.log10(max(span_hz, 1)) / 6.0, 0.0), 1.0)
    norm_ref = min(max((reference_level_db + 20) / 80.0, 0.0), 1.0)
    norm_noise = min(max(noise_floor / float(SCOPE_AMP_MAX), 0.0), 1.0)
    norm_bw = min(max(bandwidth_bins / 950.0, 0.0), 1.0)
    norm_offset = min(max(center_freq_offset_hz / max(span_hz, 1), -1.0), 1.0)
    norm_snr = min(max(snr_db / 40.0, 0.0), 1.0)
    base = np.array(
        [norm_span, norm_ref, norm_noise, norm_bw, norm_offset, norm_snr],
        dtype=np.float32,
    )
    if derived_features is not None:
        derived = np.asarray(derived_features, dtype=np.float32)
        if derived.shape != (DERIVED_FEATURE_DIM,):
            if derived.size >= DERIVED_FEATURE_DIM:
                derived = derived[:DERIVED_FEATURE_DIM]
            else:
                padded = np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)
                padded[:derived.size] = derived
                derived = padded
    else:
        derived = np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)
    return np.concatenate([base, derived], axis=0)


# ---------------------------------------------------------------------------
# Derived spectral features (7 descriptors from scope_history signal region)
# ---------------------------------------------------------------------------

DERIVED_FEATURE_NAMES = [
    "spectral_flatness",
    "spectral_kurtosis",
    "peak_count",
    "spectral_periodicity",
    "spectral_bw_ratio",
    "centroid_offset",
    "temporal_env_cv",
]

DERIVED_FEATURE_DIM = 7


def prepare_derived_features(
    scope_history: np.ndarray,
    bin_start: int,
    bin_end: int,
    noise_level: Optional[float] = None,
) -> np.ndarray:
    """Compute 7 spectral descriptors from the signal region of scope_history.

    These features capture modulation characteristics that the CNN cannot
    extract from waterfall image patches alone, particularly at typical
    scope spans where individual tones of multi-tone modes (FT8, FT4, JT65)
    are sub-pixel width. See docs/DERIVED_SPECTRAL_FEATURES.md for the
    full analysis and separability measurements.

    Parameters
    ----------
    scope_history : np.ndarray
        Scope history buffer (100, 475) float32, amplitude values 0-160.
        Bin coordinates are 475-bin scope space (NOT 950-bin detector space).
    bin_start : int
        Signal start bin in 475-bin scope coordinates.
    bin_end : int
        Signal end bin in 475-bin scope coordinates.
    noise_level : float, optional
        Noise floor in raw scope units (0-160). If None, computed internally
        as p25 of the last frame's amplitudes — ensuring identical computation
        at training time and runtime (train/serve consistency).

    Returns
    -------
    np.ndarray
        Shape (7,) float32 with normalized features in [0, 1]:
        spectral_flatness, spectral_kurtosis, peak_count,
        spectral_periodicity, spectral_bw_ratio, centroid_offset,
        temporal_env_cv.
    """
    if scope_history is None or scope_history.size == 0:
        return np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)

    if bin_end <= bin_start:
        return np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)

    rows, cols = scope_history.shape[:2]

    # Extract signal region with margin for context
    margin = max(3, (bin_end - bin_start))
    left = max(0, bin_start - margin)
    right = min(cols, bin_end + margin)
    if right <= left:
        return np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)

    region = scope_history[:, left:right].astype(np.float32)
    n_frames, n_bins = region.shape

    if n_bins < 2:
        return np.zeros(DERIVED_FEATURE_DIM, dtype=np.float32)

    # Average spectrum across frames (stable estimate)
    avg_spectrum = region.mean(axis=0)

    # Compute noise level if not provided (self-contained for train/serve parity)
    if noise_level is None:
        noise_level = float(np.percentile(scope_history[-1], 25))

    # 1. Spectral flatness (Wiener entropy)
    #    Tone-like (CW) -> ~0, noise-like (FM) -> higher
    power = np.maximum(avg_spectrum - noise_level, 0.01) ** 2
    log_power = np.log(power + 1e-10)
    geo_mean = float(np.exp(np.mean(log_power)))
    arith_mean = float(np.mean(power))
    flatness = geo_mean / (arith_mean + 1e-10)
    flatness = min(max(flatness, 0.0), 1.0)

    # 2. Spectral kurtosis (peakedness)
    #    Single spike (CW) -> high, flat spread (FM) -> low
    spec_mean = float(np.mean(avg_spectrum))
    spec_std = float(np.std(avg_spectrum)) + 1e-10
    kurtosis = float(np.mean(((avg_spectrum - spec_mean) / spec_std) ** 4) - 3)
    kurtosis_norm = min(max(kurtosis / (abs(kurtosis) + 5.0), 0.0), 1.0)

    # 3. Peak count (contiguous regions above noise + 3 dB)
    #    CW=1, RTTY=2, multi-tone=varies
    threshold = noise_level + 3
    above = (avg_spectrum > threshold).astype(np.int8)
    if above[0] == 1:
        peaks = 1
    else:
        peaks = 0
    peaks += int(np.sum(np.diff(above) == 1))
    peak_count_norm = min(peaks / 10.0, 1.0)

    # 4. Spectral periodicity (FFT of magnitude spectrum -> comb detector)
    #    Multi-tone MFSK produces comb pattern even when tones aren't resolved
    spectrum_centered = avg_spectrum - spec_mean
    if spec_std > 0.1 and n_bins > 4:
        spectrum_fft = np.abs(np.fft.rfft(spectrum_centered))[1:]
        if len(spectrum_fft) > 2:
            peak_val = float(np.max(spectrum_fft[1:]))
            total_val = float(np.sum(spectrum_fft)) + 1e-10
            periodicity = peak_val / total_val
        else:
            periodicity = 0.0
    else:
        periodicity = 0.0
    periodicity_norm = min(max(periodicity, 0.0), 1.0)

    # 5. Spectral bandwidth ratio (energy spread around centroid)
    total_energy = float(np.sum(avg_spectrum)) + 1e-10
    freq_axis = np.arange(n_bins, dtype=np.float32)
    centroid = float(np.sum(freq_axis * avg_spectrum) / total_energy)
    spectral_bw = float(np.sqrt(np.sum(((freq_axis - centroid) ** 2) * avg_spectrum) / total_energy))
    spectral_bw_norm = min(spectral_bw / max(n_bins, 1), 1.0)

    # 6. Centroid offset from geometric center (spectral asymmetry)
    #    AM centered -> ~0, SSB asymmetric -> higher
    center = n_bins / 2.0
    centroid_offset = float(abs(centroid - center) / max(n_bins / 2.0, 1.0))
    centroid_offset_norm = min(centroid_offset * 2.0, 1.0)

    # 7. Temporal envelope coefficient of variation (duty cycle proxy)
    #    CW/AFSK bursty -> high, FM/PSK continuous -> low
    energy_per_frame = region.sum(axis=1)
    env_mean = float(np.mean(energy_per_frame))
    env_std = float(np.std(energy_per_frame))
    env_cv = env_std / (env_mean + 1e-10)
    env_cv_norm = min(env_cv / (env_cv + 1.0), 1.0)

    return np.array([
        flatness,
        kurtosis_norm,
        peak_count_norm,
        periodicity_norm,
        spectral_bw_norm,
        centroid_offset_norm,
        env_cv_norm,
    ], dtype=np.float32)
