"""ScopeCNN-Patch model definition for signal classification.

This module defines the PyTorch model architecture used for training
the ONNX signal classifier. It is NOT imported at runtime by the main
application — it is only used by the training pipeline script
(train_signal_classifier.py).

Architecture (with auxiliary inputs):
    Scope history input: 100x64x1 (raw amplitude patch, normalized 0-1)
    Aux input: 6 floats (span, ref_level, noise_floor, bandwidth_bins,
                          center_freq_offset_hz, snr_db)
    Conv2d(1->8,   5x5, pad=2) -> ReLU -> MaxPool2d(2)   -> 50x32x8
    Conv2d(8->16,  3x3, pad=1) -> ReLU -> MaxPool2d(2)   -> 25x16x16
    Conv2d(16->32, 3x3, pad=1) -> ReLU -> MaxPool2d(2)  -> 12x8x32
    Flatten(3072) + Concat(aux 6) -> 3078
    Linear(3078->128) -> ReLU -> Dropout(0.3)
    Linear(128->num_classes)

~400K parameters, ~1.6 MB ONNX export, ~3-5 ms CPU inference.

When aux_dim=0, falls back to image-only architecture (Flatten -> 3072 -> 128).

Constants and prepare_aux_features are imported from patch_config so they
stay synchronised with the runtime classifier. In the deployed environment
(~/Pycom/signals/scripts/), patch_config.py is copied alongside this file
and imported directly.
"""

import os

import numpy as np

try:
    from lib.signal.patch_config import (
        AUX_DIM,
        AUX_FEATURE_NAMES,
        INPUT_CHANNELS,
        PATCH_MARGIN_FRACTION,
        PATCH_MIN_MARGIN_BINS,
        PATCH_STRATEGY,
        SCOPE_AMP_MAX,
        SCOPE_HISTORY_COLS,
        SCOPE_HISTORY_ROWS,
        SCOPE_PATCH_COLS,
        SCOPE_PATCH_ROWS,
        SIGNAL_LABELS,
        prepare_aux_features,
        prepare_derived_features,
    )
except ImportError:
    from patch_config import (
        AUX_DIM,
        AUX_FEATURE_NAMES,
        INPUT_CHANNELS,
        PATCH_MARGIN_FRACTION,
        PATCH_MIN_MARGIN_BINS,
        PATCH_STRATEGY,
        SCOPE_AMP_MAX,
        SCOPE_HISTORY_COLS,
        SCOPE_HISTORY_ROWS,
        SCOPE_PATCH_COLS,
        SCOPE_PATCH_ROWS,
        SIGNAL_LABELS,
        prepare_aux_features,
        prepare_derived_features,
    )

NUM_CLASSES = len(SIGNAL_LABELS)


def build_model(num_classes: int = NUM_CLASSES, aux_dim: int = AUX_DIM):
    """Build and return the ScopeCNN-Patch PyTorch model.

    Parameters
    ----------
    num_classes : int
        Number of output classes.
    aux_dim : int
        Number of auxiliary input features. Set to 0 for image-only model
        (backward compatible with models trained without aux inputs).

    Requires torch to be installed (dev-only dependency).
    """
    try:
        import torch
        import torch.nn as nn
    except ImportError:
        raise ImportError(
            "PyTorch is required for model definition. "
            "Install with: pip install torch"
        )

    class ScopeCNN(nn.Module):
        def __init__(self, num_classes: int, aux_dim: int):
            super().__init__()
            self.aux_dim = aux_dim
            self.features = nn.Sequential(
                nn.Conv2d(INPUT_CHANNELS, 8, kernel_size=5, padding=2),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=2, stride=2),
                nn.Conv2d(8, 16, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=2, stride=2),
                nn.Conv2d(16, 32, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=2, stride=2),
            )
            fc_input_dim = 32 * 12 * 8 + aux_dim
            self.classifier = nn.Sequential(
                nn.Linear(fc_input_dim, 128),
                nn.ReLU(inplace=True),
                nn.Dropout(0.3),
                nn.Linear(128, num_classes),
            )

        def forward(self, x, aux=None):
            x = self.features(x)
            x = x.view(x.size(0), -1)
            if self.aux_dim > 0 and aux is not None:
                x = torch.cat([x, aux], dim=1)
            x = self.classifier(x)
            return x

    return ScopeCNN(num_classes, aux_dim)


def export_to_onnx(
    model,
    output_path: str,
    input_height: int = SCOPE_PATCH_ROWS,
    input_width: int = SCOPE_PATCH_COLS,
    aux_dim: int = AUX_DIM,
):
    """Export a trained PyTorch model to ONNX format.

    Parameters
    ----------
    model : torch.nn.Module
        Trained PyTorch model.
    output_path : str
        Path to write the ONNX model file.
    input_height : int
        Input height (number of scope history rows).
    input_width : int
        Input width (number of scope bins in patch).
    aux_dim : int
        Number of auxiliary input features. 0 for image-only model.
    """
    import torch
    import warnings

    model.eval()
    dummy_input = torch.randn(1, INPUT_CHANNELS, input_height, input_width)

    batch_dim = torch.export.Dim("batch_size", min=1)

    if aux_dim > 0:
        dummy_aux = torch.randn(1, aux_dim)
        input_names = ["input", "aux"]

        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message=".*axis name.*will not be used.*",
                category=UserWarning,
            )
            warnings.filterwarnings(
                "ignore",
                message=".*isinstance.*treespec.*is deprecated.*",
                category=FutureWarning,
            )
            torch.onnx.export(
                model,
                (dummy_input, dummy_aux),
                output_path,
                dynamo=True,
                dynamic_shapes=({0: batch_dim}, {0: batch_dim}),
                export_params=True,
                opset_version=18,
                do_constant_folding=True,
                input_names=input_names,
                output_names=["output"],
            )
    else:
        input_names = ["input"]

        with warnings.catch_warnings():
            warnings.filterwarnings(
                "ignore",
                message=".*axis name.*will not be used.*",
                category=UserWarning,
            )
            warnings.filterwarnings(
                "ignore",
                message=".*isinstance.*treespec.*is deprecated.*",
                category=FutureWarning,
            )
            torch.onnx.export(
                model,
                dummy_input,
                output_path,
                dynamo=True,
                dynamic_shapes=({0: batch_dim},),
                export_params=True,
                opset_version=18,
                do_constant_folding=True,
                input_names=input_names,
                output_names=["output"],
            )

    import onnx
    onnx_model = onnx.load(output_path, load_external_data=True)
    onnx.save_model(
        onnx_model,
        output_path,
        save_as_external_data=False,
    )
    data_file = output_path + ".data"
    if os.path.exists(data_file):
        os.remove(data_file)


_VALID_STATE_DICT_PREFIXES = ("features.", "classifier.")


def load_onnx_to_state_dict(onnx_path: str):
    """Load weights from an ONNX model file into a PyTorch state dict.

    Extracts float32 initializer tensors from the ONNX model and maps them
    to PyTorch state dict keys. Non-parameter constants (e.g., reshape
    constants like val_7) and non-float tensors are excluded.

    Parameters
    ----------
    onnx_path : str
        Path to the ONNX model file.

    Returns
    -------
    dict
        State dict compatible with build_model(num_classes, aux_dim).

    Raises
    ------
    ValueError
        If the ONNX model has no output shape info or the architecture
        does not match the expected ScopeCNN structure.
    """
    import onnx
    import torch

    onnx_model = onnx.load(onnx_path, load_external_data=True)

    output_nodes = [n for n in onnx_model.graph.node if n.op_type == "Reshape" or n.output]
    output_shape = None
    for output_info in onnx_model.graph.output:
        shape = [d.dim_value for d in output_info.type.tensor_type.shape.dim]
        if shape and shape[-1] > 1:
            output_shape = shape
            break

    if output_shape is None:
        for initializer in onnx_model.graph.initializer:
            if initializer.name.startswith("classifier.3."):
                arr = np.frombuffer(initializer.raw_data, dtype=np.float32)
                output_dim = arr.shape[0] if len(arr.shape) == 0 else arr.shape[0]
                output_shape = [0, output_dim]
                break

    if output_shape is None or output_shape[-1] < 2:
        raise ValueError(
            f"Cannot determine num_classes from ONNX model at {onnx_path}. "
            "The model output shape could not be inferred."
        )

    num_classes = output_shape[-1]

    aux_dim = 0
    for inp in onnx_model.graph.input:
        if inp.name == "aux":
            dims = [d.dim_value for d in inp.type.tensor_type.shape.dim]
            aux_dim = dims[-1] if dims else 0
            break

    expected_model = build_model(num_classes=num_classes, aux_dim=aux_dim if aux_dim > 0 else AUX_DIM)
    expected_state = expected_model.state_dict()

    state_dict = {}
    for initializer in onnx_model.graph.initializer:
        name = initializer.name
        if not name.startswith(_VALID_STATE_DICT_PREFIXES):
            continue

        if initializer.data_type != 1:
            continue

        if name not in expected_state:
            continue

        arr = np.frombuffer(initializer.raw_data, dtype=np.float32).copy()
        expected_shape = expected_state[name].shape

        if arr.shape != expected_shape and arr.size == np.prod(expected_shape):
            arr = arr.reshape(expected_shape)

        if arr.shape != expected_shape:
            raise ValueError(
                f"Shape mismatch for '{name}': ONNX has {arr.shape}, "
                f"expected {expected_shape}. The model architecture does not match."
            )

        state_dict[name] = torch.from_numpy(arr)

    missing = set(expected_state.keys()) - set(state_dict.keys())
    if missing:
        raise ValueError(
            f"Missing keys in ONNX model: {sorted(missing)}. "
            "The ONNX model architecture does not match ScopeCNN-Patch."
        )

    return state_dict, num_classes, aux_dim if aux_dim > 0 else AUX_DIM


def expand_classifier_head(
    state_dict: dict,
    old_num_classes: int,
    new_num_classes: int,
    aux_dim: int = AUX_DIM,
):
    """Expand the classification head of a state dict to support more classes.

    Copies all shared weights exactly, preserves existing class rows, and
    initializes new class rows using small values (0.01x Kaiming uniform)
    with zero-initialized bias. Small initialization ensures new classes
    start with near-zero logits, preventing them from competing with
    existing class predictions before training.

    Parameters
    ----------
    state_dict : dict
        Existing model state dict (from load_onnx_to_state_dict or checkpoint).
    old_num_classes : int
        Number of output classes in the existing model.
    new_num_classes : int
        Number of output classes in the expanded model (must be >= old_num_classes).
    aux_dim : int
        Number of auxiliary input features.

    Returns
    -------
    dict
        Expanded state dict compatible with build_model(new_num_classes, aux_dim).
    """
    import torch
    import torch.nn as nn

    if new_num_classes < old_num_classes:
        raise ValueError(
            f"new_num_classes ({new_num_classes}) must be >= old_num_classes "
            f"({old_num_classes}). Reducing classes is not supported."
        )

    if new_num_classes == old_num_classes:
        return dict(state_dict)

    new_model = build_model(num_classes=new_num_classes, aux_dim=aux_dim)
    new_state = new_model.state_dict()

    for key, value in state_dict.items():
        if key not in new_state:
            continue

        if key == "classifier.3.weight":
            old_weight = value
            new_state[key][:old_num_classes] = old_weight
            nn.init.kaiming_uniform_(
                new_state[key][old_num_classes:],
                a=np.sqrt(5),
            )
            new_state[key][old_num_classes:] *= 0.01
        elif key == "classifier.3.bias":
            old_bias = value
            new_state[key][:old_num_classes] = old_bias
            new_state[key][old_num_classes:] = 0.0
        else:
            new_state[key] = value

    return new_state


def merge_labels(existing_labels: list, capture_labels: list):
    """Merge capture labels into existing labels, preserving order.

    Appends any new labels from capture_labels that are not already in
    existing_labels. "Unknown" is always kept at the end. Returns the
    merged labels list and the list of new labels that were added.

    Duplicate labels in capture_labels are ignored — only unique new
    labels that are not already in existing_labels are appended.

    Parameters
    ----------
    existing_labels : list
        Current label list (e.g., from signal_labels.json "labels" field).
        May include "Unknown" as the last entry.
    capture_labels : list
        Labels found in the user's capture data (may contain duplicates).

    Returns
    -------
    tuple of (list, list)
        (merged_labels, new_labels) where merged_labels has new labels
        appended before "Unknown", and new_labels is the list of labels
        that were added (unique, deduplicated).
    """
    model_labels = [l for l in existing_labels if l != "Unknown"]
    has_unknown = "Unknown" in existing_labels

    unique_capture = list(dict.fromkeys(capture_labels))
    new_labels = [l for l in unique_capture if l not in model_labels and l != "Unknown"]

    merged = list(model_labels) + new_labels
    if has_unknown or "Unknown" in existing_labels:
        merged.append("Unknown")

    return merged, new_labels


def prepare_patch(
    scope_history: np.ndarray,
    bin_start_scope: int,
    bin_end_scope: int,
    patch_rows: int = SCOPE_PATCH_ROWS,
    patch_cols: int = SCOPE_PATCH_COLS,
) -> np.ndarray:
    """Extract a bandwidth-aware patch from scope history for model input.

    For narrow signals (bandwidth < patch_cols), the patch is centered:
    the signal's center is placed in the middle of the patch, providing
    equal context on both sides. For wide signals (bandwidth >= patch_cols),
    the patch is centered and resampled.

    Parameters
    ----------
    scope_history : np.ndarray
        Scope history (ROWS, COLS) float32, amplitude values 0-160.
    bin_start_scope, bin_end_scope : int
        Start/end bins in 475-bin scope space.
    patch_rows, patch_cols : int
        Output patch dimensions.

    Returns
    -------
    np.ndarray
        (patch_rows, patch_cols, 1) float32 normalized to [0, 1].
    """
    rows, cols = scope_history.shape[:2]

    if bin_end_scope <= bin_start_scope:
        bin_end_scope = bin_start_scope + 1

    signal_bw = bin_end_scope - bin_start_scope

    if signal_bw < patch_cols:
        center = (bin_start_scope + bin_end_scope) // 2
        left = max(0, center - patch_cols // 2)
        right = left + patch_cols
        if right > cols:
            right = cols
            left = max(0, right - patch_cols)
    else:
        center = (bin_start_scope + bin_end_scope) // 2
        half_w = (signal_bw + 1) // 2
        left = max(0, center - half_w)
        right = min(cols, center + half_w)

    region = scope_history[:, left:right].astype(np.float32) / float(SCOPE_AMP_MAX)

    if rows >= patch_rows:
        start_row = rows - patch_rows
        region = region[start_row:start_row + patch_rows, :]
    else:
        padded = np.zeros((patch_rows, region.shape[1]), dtype=np.float32)
        padded[patch_rows - rows:, :] = region
        region = padded

    if region.shape[1] != patch_cols:
        src_cols = np.arange(region.shape[1])
        dst_cols = np.linspace(0, region.shape[1] - 1, patch_cols)
        resampled = np.zeros((patch_rows, patch_cols), dtype=np.float32)
        for r in range(patch_rows):
            resampled[r, :] = np.interp(dst_cols, src_cols, region[r, :])
        region = resampled

    return region[:, :, np.newaxis]
