#!/usr/bin/env python3
"""Train or fine-tune the signal classifier ONNX model from labeled waterfall captures.

Usage:
    # Train from scratch (default):
    python scripts/train_signal_classifier.py [--epochs N] [--batch-size N]

    # Fine-tune from existing model (preserves existing class knowledge):
    python scripts/train_signal_classifier.py --resume model.onnx [--freeze-features | --freeze-classifier-input]

    # Also accept .pt checkpoint files:
    python scripts/train_signal_classifier.py --resume model.pt [--freeze-features | --freeze-classifier-input]

This script loads all labeled .npz captures from ~/Pycom/signals/captures/,
extracts scope history patches around detected signal regions, trains a
ScopeCNN-Patch model with auxiliary inputs (span, ref level, noise floor,
bandwidth, center freq offset, SNR), and exports
it to ONNX format for use by the SignalClassifier.

With --resume, loads an existing ONNX or .pt model, merges any new labels
found in capture data, expands the classification head for new classes
(Kaiming init for new weights, zero init for new biases), and fine-tunes
while preserving existing class knowledge.

Freezing options:
    --freeze-features         Freeze CNN feature layers (only ~0.3% of params).
                             Allows all classifier layers to adapt.
    --freeze-classifier-input Freeze everything except the final classification
                             head (classifier.3). Preserves ~95% of parameters.
                             Recommended for small datasets with new classes.

Requires PyTorch and ONNX for training (dev-only dependencies, not shipped with the app):
    pip install torch onnx onnxscript torchvision
"""

import argparse
import hashlib
import json
import os
import shutil
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Dict, List, Optional

import numpy as np


SIGNALS_DIR = os.path.join(os.path.expanduser("~"), "Pycom", "signals")
CAPTURES_DIR = os.path.join(SIGNALS_DIR, "captures")
MODEL_DIR = os.path.join(SIGNALS_DIR, "model")
DEFAULT_MODEL_PATH = os.path.join(MODEL_DIR, "signal_classifier.onnx")

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

from model_definition import (
    AUX_DIM,
    SCOPE_HISTORY_COLS,
    SCOPE_PATCH_COLS,
    SCOPE_PATCH_ROWS,
    INPUT_CHANNELS,
    NUM_CLASSES,
    PATCH_STRATEGY,
    SIGNAL_LABELS,
    build_model,
    export_to_onnx,
    prepare_aux_features,
    prepare_derived_features,
    prepare_patch,
    load_onnx_to_state_dict,
    expand_classifier_head,
    merge_labels,
)


HASH_PREFIX_LEN = 16


@dataclass
class TrainResult:
    """Result of training the signal classifier model."""

    model: object
    best_val_accuracy: float
    train_samples: int
    val_samples: int
    active_classes: Dict[str, int] = field(default_factory=dict)
    total_epochs: int = 30
    frozen_layers: Optional[str] = None
    learning_rate: float = 0.001


def compute_model_hash(model_path: str) -> str:
    """Return first HASH_PREFIX_LEN hex chars of SHA256 digest of a model file."""
    h = hashlib.sha256()
    with open(model_path, "rb") as f:
        for chunk in iter(lambda: f.read(65536), b""):
            h.update(chunk)
    return h.hexdigest()[:HASH_PREFIX_LEN]


def load_captures(captures_dir: str, signal_labels: list = None, ai_label_threshold: float = 0.5) -> list:
    """Load all labeled captures from a directory.

    Returns a list of (patch, aux, label_index, is_ai_labeled, ai_confidence, filename) tuples.
    Each detected candidate with a meaningful label becomes its own
    training sample. AI-labeled candidates are flagged with
    is_ai_labeled=True.

    Captures without scope_history are skipped (legacy captures are no
    longer supported).

    Parameters
    ----------
    captures_dir : str
        Directory containing .npz capture files.
    signal_labels : list, optional
        Label list to use for indexing. If None, loads from the model
        directory's signal_labels.json (falling back to SIGNAL_LABELS).
    ai_label_threshold : float
        Minimum AI confidence to use ai_label as ground truth.
    """
    dataset = []
    if signal_labels is None:
        labels_file = os.path.join(SIGNALS_DIR, "model", "signal_labels.json")
        signal_labels = list(SIGNAL_LABELS)
        if os.path.exists(labels_file):
            try:
                with open(labels_file, "r") as f:
                    loaded = json.load(f)
                if isinstance(loaded, dict) and "labels" in loaded:
                    loaded = loaded["labels"]
                if isinstance(loaded, list) and len(loaded) > 0:
                    signal_labels = [lbl for lbl in loaded if lbl != "Unknown"]
                    print(f"Loaded {len(loaded)} labels from {labels_file}")
            except Exception as e:
                print(f"Warning: Could not load labels from {labels_file}: {e}")

    capture_files = [
        f for f in os.listdir(captures_dir) if f.endswith(".npz")
    ]
    print(f"Found {len(capture_files)} capture files in {captures_dir}")

    for filename in capture_files:
        filepath = os.path.join(captures_dir, filename)

        try:
            with np.load(filepath, allow_pickle=True) as data:
                metadata = json.loads(str(data["metadata"]))
                spectrum = data["spectrum"]

                has_scope = "scope_history" in data
                if has_scope:
                    scope_history = data["scope_history"]
                    if scope_history.ndim < 2 or scope_history.size == 0:
                        has_scope = False

                if not has_scope:
                    print(f"Skipping {filename}: no scope_history (legacy capture)")
                    continue

                span_hz = metadata.get("span_hz", metadata.get("effective_span_hz", 100000))
                ref_level_db = metadata.get("reference_level_db", 0.0)
                center_freq_hz = metadata.get("center_freq_hz", 0)
                # Use saved noise_floor from capture metadata (EMA-smoothed
                # detector value) when available; fall back to p25 of spectrum.
                saved_noise_floor = metadata.get("noise_floor")
                if saved_noise_floor is not None:
                    noise_floor = float(saved_noise_floor)
                else:
                    noise_floor = float(np.percentile(spectrum, 25))

                candidates = metadata.get("detected_candidates", [])

                if candidates:
                    for cand in candidates:
                        user_label = cand.get("user_label", "")
                        ai_label = cand.get("ai_label", "")
                        ai_conf = cand.get("ai_confidence", 0.0)

                        if user_label and user_label != "Unknown":
                            patch_label = user_label
                            is_ai_labeled = False
                        elif ai_label and ai_label != "Unknown" and ai_conf > ai_label_threshold:
                            patch_label = ai_label
                            is_ai_labeled = True
                        else:
                            continue

                        if patch_label in signal_labels:
                            patch_label_idx = signal_labels.index(patch_label)
                        else:
                            continue

                        # Map candidate frequency to scope_history column space.
                        # scope_history covers effective_span_hz (2x scope span).
                        cand_freq = cand.get("frequency_hz", 0)
                        cand_bw = cand.get("bandwidth_hz", 0)
                        effective_span = metadata.get("effective_span_hz", 0)
                        if effective_span <= 0 and span_hz > 0:
                            effective_span = span_hz * 2

                        if cand_freq > 0 and effective_span > 0:
                            scope_cols = scope_history.shape[1]
                            bin_start_scope = int(
                                (cand_freq - cand_bw / 2 - center_freq_hz) / effective_span * scope_cols + scope_cols / 2
                            )
                            bin_end_scope = int(
                                (cand_freq + cand_bw / 2 - center_freq_hz) / effective_span * scope_cols + scope_cols / 2
                            )
                        else:
                            bin_start_scope = cand.get("bin_start", cand.get("bin_center", 240)) // 2
                            bin_end_scope = cand.get("bin_end", cand.get("bin_center", 240) + 6) // 2

                        if bin_end_scope <= bin_start_scope:
                            bin_end_scope = bin_start_scope + 3

                        patch = prepare_patch(
                            scope_history, bin_start_scope, bin_end_scope
                        )
                        if patch.shape != (SCOPE_PATCH_ROWS, SCOPE_PATCH_COLS, INPUT_CHANNELS):
                            continue

                        bandwidth_bins = cand.get("bin_end", 0) - cand.get("bin_start", 0)
                        center_freq_offset_hz = cand.get("frequency_hz", 0) - center_freq_hz
                        snr_db = cand.get("snr_estimate", 0.0)

                        derived = prepare_derived_features(
                            scope_history,
                            bin_start_scope,
                            bin_end_scope,
                        )
                        aux = prepare_aux_features(
                            span_hz=span_hz,
                            reference_level_db=ref_level_db,
                            noise_floor=noise_floor,
                            bandwidth_bins=bandwidth_bins,
                            center_freq_offset_hz=center_freq_offset_hz,
                            snr_db=snr_db,
                            derived_features=derived,
                        )
                        dataset.append((patch, aux, patch_label_idx, is_ai_labeled, ai_conf, filename))
                else:
                    whole_label = metadata.get("label", "Unknown")
                    if whole_label not in signal_labels:
                        continue

                    label_idx = signal_labels.index(whole_label)
                    center_scope = SCOPE_HISTORY_COLS // 2
                    patch = prepare_patch(
                        scope_history,
                        center_scope - SCOPE_PATCH_COLS // 2,
                        center_scope + SCOPE_PATCH_COLS // 2,
                    )
                    if patch.shape != (SCOPE_PATCH_ROWS, SCOPE_PATCH_COLS, INPUT_CHANNELS):
                        continue

                    aux = prepare_aux_features(
                        span_hz=span_hz,
                        reference_level_db=ref_level_db,
                        noise_floor=noise_floor,
                        bandwidth_bins=SCOPE_PATCH_COLS,
                        derived_features=prepare_derived_features(
                            scope_history,
                            center_scope - SCOPE_PATCH_COLS // 2,
                            center_scope + SCOPE_PATCH_COLS // 2,
                        ),
                    )
                    dataset.append((patch, aux, label_idx, False, 0.0, filename))

        except Exception as e:
            print(f"Error loading {filename}: {e}")
            continue

    n_ai = sum(1 for _, _, _, is_ai_labeled, _, _ in dataset if is_ai_labeled)
    print(f"Loaded {len(dataset)} training samples ({n_ai} AI-labeled)")
    return dataset


def scan_capture_labels(captures_dir: str, ai_label_threshold: float = 0.5) -> list:
    """Scan capture files for all unique labels (excluding Unknown).

    Returns a sorted list of unique labels found in user_label and
    high-confidence ai_label fields across all captures.
    """
    found_labels = set()
    capture_files = [
        f for f in os.listdir(captures_dir) if f.endswith(".npz")
    ]

    for filename in capture_files:
        filepath = os.path.join(captures_dir, filename)
        try:
            with np.load(filepath, allow_pickle=True) as data:
                metadata = json.loads(str(data["metadata"]))
                candidates = metadata.get("detected_candidates", [])
                for cand in candidates:
                    user_label = cand.get("user_label", "")
                    ai_label = cand.get("ai_label", "")
                    ai_conf = cand.get("ai_confidence", 0.0)
                    if user_label and user_label != "Unknown":
                        found_labels.add(user_label)
                    elif ai_label and ai_label != "Unknown" and ai_conf > ai_label_threshold:
                        found_labels.add(ai_label)

                if not candidates:
                    whole_label = metadata.get("label", "Unknown")
                    if whole_label != "Unknown":
                        found_labels.add(whole_label)
        except Exception as e:
            print(f"Error scanning {filename}: {e}")
            continue

    return sorted(found_labels)


def train_model(
    dataset: list,
    epochs: int = 30,
    batch_size: int = 16,
    learning_rate: float = 0.001,
    num_classes: int = NUM_CLASSES,
    aux_dim: int = AUX_DIM,
    initial_state_dict: dict = None,
    freeze_features: bool = False,
    freeze_classifier_input: bool = False,
    old_num_classes: int = 0,
    merged_labels: List[str] = None,
    difficulty_weights: Optional[List[float]] = None,
):
    """Train or fine-tune the ScopeCNN-Patch model on the dataset.

    Uses class-balanced WeightedRandomSampler to handle imbalanced datasets,
    gradient clipping, and a learning rate scheduler for stable convergence
    on small datasets.

    When initial_state_dict is provided, loads existing weights before
    training (fine-tuning mode). Freezing options control which layers
    are trainable.

    Parameters
    ----------
    dataset : list
        Training data from load_captures().
    epochs : int
        Number of training epochs.
    batch_size : int
        Batch size for training.
    learning_rate : float
        Initial learning rate.
    num_classes : int
        Number of output classes.
    aux_dim : int
        Number of auxiliary input features.
    initial_state_dict : dict, optional
        Pre-trained state dict to load before training.
    freeze_features : bool
        If True, freeze CNN feature layers (features.*).
    freeze_classifier_input : bool
        If True, freeze everything except the final classification head
        (classifier.3). This preserves ~95% of parameters and is
        recommended for small datasets with new classes.
    old_num_classes : int
        Number of classes in the previous model (for gradient hooks).
    merged_labels : list of str, optional
        Label names indexed by label index. Used to build active_classes
        mapping in the returned TrainResult.

    Returns
    -------
    TrainResult
        Trained model and training statistics.
    """
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler

    class SignalDataset(Dataset):
        def __init__(self, samples, weights=None):
            self.samples = samples
            self.weights = weights

        def __len__(self):
            return len(self.samples)

        def __getitem__(self, idx):
            patch, aux, label, is_ai_labeled, ai_conf, filename = self.samples[idx]
            patch_tensor = torch.tensor(patch, dtype=torch.float32).permute(2, 0, 1)
            aux_tensor = torch.tensor(aux, dtype=torch.float32)
            weight = self.weights[idx] if self.weights else 1.0
            return patch_tensor, aux_tensor, label, weight

    if len(dataset) < 2:
        raise ValueError(
            f"Need at least 2 labeled samples for training, got {len(dataset)}. "
            "Capture and label more signals using the app's Capture button."
        )

    model = build_model(num_classes, aux_dim=aux_dim)

    if initial_state_dict is not None:
        model.load_state_dict(initial_state_dict, strict=True)
        print(f"Loaded pre-trained weights ({len(initial_state_dict)} tensors)")

    trainable_params = []
    frozen_names = []
    frozen_layers = None

    if freeze_classifier_input:
        frozen_layers = "classifier.0-2"
        for name, param in model.named_parameters():
            if name.startswith("classifier.3"):
                param.requires_grad = True
                trainable_params.append(param)
            else:
                param.requires_grad = False
                frozen_names.append(name)
        print(f"Frozen {len(frozen_names)} tensors (all except classifier.3)")
        print(f"Trainable parameters: {sum(p.numel() for p in trainable_params):,}")

        if old_num_classes > 0 and old_num_classes < num_classes:
            _frozen_rows = old_num_classes

            def _freeze_old_weight_rows(grad, n=_frozen_rows):
                grad[:n, :] = 0.0
                return grad

            def _freeze_old_bias_rows(grad, n=_frozen_rows):
                grad[:n] = 0.0
                return grad

            model.classifier[3].weight.register_hook(_freeze_old_weight_rows)
            model.classifier[3].bias.register_hook(_freeze_old_bias_rows)
            print(f"Gradient hook: freezing old-class rows 0-{_frozen_rows - 1} in classifier.3")
    elif freeze_features:
        frozen_layers = "features"
        for name, param in model.named_parameters():
            if name.startswith("features."):
                param.requires_grad = False
                frozen_names.append(name)
            else:
                trainable_params.append(param)
                param.requires_grad = True
        print(f"Frozen {len(frozen_names)} feature tensors")
        print(f"Trainable parameters: {sum(p.numel() for p in trainable_params):,}")
    else:
        trainable_params = list(model.parameters())

    from collections import Counter
    label_counts = Counter(s[2] for s in dataset)
    class_weights = torch.zeros(num_classes, dtype=torch.float32)
    total = len(dataset)
    for cls_idx, count in label_counts.items():
        if cls_idx < num_classes:
            class_weights[cls_idx] = total / (len(label_counts) * count) if count > 0 else 0
    print(f"Class weights: {dict(label_counts)}")

    sample_weights = []
    for idx, (_, _, label_idx, _, _, _) in enumerate(dataset):
        w = class_weights[label_idx].item() if label_idx < num_classes else 0.0
        if difficulty_weights is not None:
            w *= difficulty_weights[idx]
        sample_weights.append(w)
    sample_weights_tensor = torch.tensor(sample_weights, dtype=torch.float32)

    # Group-aware split: all candidates from the same capture file go into
    # the same split (train or val) to prevent data leakage. Candidates from
    # the same capture share the same scope_history, so splitting them across
    # train/val inflates validation accuracy.
    capture_to_indices = {}
    capture_to_labels = {}
    for i, (_, _, label_idx, _, _, filename) in enumerate(dataset):
        capture_to_indices.setdefault(filename, []).append(i)
        capture_to_labels.setdefault(filename, set()).add(label_idx)

    # Group capture files by their dominant label for stratification.
    # The dominant label is the one with the most candidates in that file.
    label_to_captures = {}
    for filename in capture_to_indices:
        file_label_counts = {}
        for idx in capture_to_indices[filename]:
            lbl = dataset[idx][2]
            file_label_counts[lbl] = file_label_counts.get(lbl, 0) + 1
        dominant_label = max(file_label_counts, key=file_label_counts.get)
        label_to_captures.setdefault(dominant_label, []).append(filename)

    train_indices = []
    val_indices = []
    _MIN_VAL_CAPTURES_PER_CLASS = 5

    for label_idx, capture_files_for_label in label_to_captures.items():
        np.random.shuffle(capture_files_for_label)
        n_captures = len(capture_files_for_label)
        if n_captures < 2:
            train_indices.extend(
                idx for f in capture_files_for_label for idx in capture_to_indices[f]
            )
        elif n_captures < _MIN_VAL_CAPTURES_PER_CLASS:
            train_indices.extend(
                idx for f in capture_files_for_label for idx in capture_to_indices[f]
            )
        else:
            n_val = max(1, n_captures // 5)
            val_captures = capture_files_for_label[:n_val]
            train_captures = capture_files_for_label[n_val:]
            val_indices.extend(idx for f in val_captures for idx in capture_to_indices[f])
            train_indices.extend(idx for f in train_captures for idx in capture_to_indices[f])

    if not val_indices:
        print("WARNING: No validation samples — insufficient capture files for group-aware split.")
        print("  Falling back to per-candidate split (may have data leakage).")
        label_to_indices_fallback = {}
        for i, (_, _, label_idx, _, _, _) in enumerate(dataset):
            label_to_indices_fallback.setdefault(label_idx, []).append(i)
        for label_idx, indices in label_to_indices_fallback.items():
            np.random.shuffle(indices)
            if len(indices) < 2:
                train_indices.extend(indices)
            else:
                n_val = max(1, len(indices) // 5)
                val_indices.extend(indices[:n_val])
                train_indices.extend(indices[n_val:])

    train_diff_weights = None
    val_diff_weights = None
    if difficulty_weights is not None:
        train_diff_weights = [difficulty_weights[i] for i in train_indices]
        val_diff_weights = [difficulty_weights[i] for i in val_indices]

    train_data = SignalDataset([dataset[i] for i in train_indices], weights=train_diff_weights)
    val_data = SignalDataset([dataset[i] for i in val_indices], weights=val_diff_weights)

    train_sample_weights = sample_weights_tensor[train_indices]
    train_sampler = WeightedRandomSampler(
        train_sample_weights,
        len(train_sample_weights),
        replacement=True,
    )

    train_loader = DataLoader(train_data, batch_size=batch_size, sampler=train_sampler)
    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)

    criterion = nn.CrossEntropyLoss()

    if trainable_params:
        optimizer = optim.Adam(trainable_params, lr=learning_rate)
    else:
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=5, min_lr=1e-6)

    print(f"Training: {len(train_data)} samples, Validation: {len(val_data)} samples")
    print(f"Classes: {num_classes}, Aux dim: {aux_dim}")

    best_val_acc = 0.0
    best_state = None

    grad_params = [p for p in model.parameters() if p.requires_grad]

    for epoch in range(epochs):
        model.train()
        running_loss = 0.0
        correct = 0
        total = 0

        for patches, aux_batch, labels, weights in train_loader:
            optimizer.zero_grad()
            outputs = model(patches, aux=aux_batch)
            if difficulty_weights is not None:
                per_sample = nn.functional.cross_entropy(outputs, labels, reduction="none")
                loss = (per_sample * weights).mean()
            else:
                loss = criterion(outputs, labels)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(grad_params, max_norm=1.0)
            optimizer.step()

            running_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

        train_loss = running_loss / len(train_loader)
        train_acc = correct / total if total > 0 else 0

        model.eval()
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for patches, aux_batch, labels, weights in val_loader:
                outputs = model(patches, aux=aux_batch)
                _, predicted = torch.max(outputs, 1)
                val_total += labels.size(0)
                val_correct += (predicted == labels).sum().item()

        val_acc = val_correct / val_total if val_total > 0 else 0
        scheduler.step(val_acc)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_state = {k: v.clone() for k, v in model.state_dict().items()}

        if (epoch + 1) % 5 == 0 or epoch == 0:
            lr_now = optimizer.param_groups[0]["lr"]
            print(
                f"Epoch {epoch + 1}/{epochs}: "
                f"loss={train_loss:.4f} train_acc={train_acc:.2%} "
                f"val_acc={val_acc:.2%} lr={lr_now:.6f}"
            )

        if train_loss != train_loss:
            print(f"WARNING: NaN loss at epoch {epoch + 1}, stopping early")
            break

    if best_state is not None:
        model.load_state_dict(best_state)
        print(f"Restored best model (val_acc={best_val_acc:.2%})")

    active_classes = {}
    if merged_labels:
        for cls_idx, count in label_counts.items():
            if cls_idx < len(merged_labels):
                active_classes[merged_labels[cls_idx]] = count

    return TrainResult(
        model=model,
        best_val_accuracy=best_val_acc,
        train_samples=len(train_data),
        val_samples=len(val_data),
        active_classes=active_classes,
        total_epochs=epochs,
        frozen_layers=frozen_layers,
        learning_rate=learning_rate,
    )


def main():
    parser = argparse.ArgumentParser(
        description="Train signal classifier from labeled waterfall captures"
    )
    parser.add_argument(
        "--captures-dir",
        default=CAPTURES_DIR,
        help=f"Directory containing .npz capture files (default: {CAPTURES_DIR})",
    )
    parser.add_argument(
        "--output",
        default=DEFAULT_MODEL_PATH,
        help=f"Output ONNX model path (default: {DEFAULT_MODEL_PATH})",
    )
    parser.add_argument("--epochs", type=int, default=30, help="Training epochs (default: 30)")
    parser.add_argument("--batch-size", type=int, default=16, help="Batch size (default: 16)")
    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate (default: 0.001)")
    parser.add_argument(
        "--resume",
        default=None,
        help="Path to existing ONNX or .pt model to fine-tune from. "
             "Loads weights, merges new labels, and expands classification head.",
    )
    parser.add_argument(
        "--labels",
        default=None,
        help="Path to signal_labels.json for the resume model. "
             "By default, looks for signal_labels.json in the same directory "
             "as the resume model. Use this when the labels file is at a "
             "different path (e.g., signal_labels.json.default).",
    )
    parser.add_argument(
        "--freeze-features",
        action="store_true",
        help="Freeze CNN feature layers when fine-tuning (only train classifier).",
    )
    parser.add_argument(
        "--freeze-classifier-input",
        action="store_true",
        help="Freeze everything except the final classification head "
             "(classifier.3). Recommended for small datasets with new classes.",
    )
    parser.add_argument(
        "--checkpoint",
        action="store_true",
        help="Save PyTorch .pt state dict alongside ONNX for future --resume.",
    )
    parser.add_argument(
        "--ai-label-threshold",
        type=float,
        default=None,
        help="Minimum AI confidence to use ai_label as training ground truth. "
             "Default: 0.9 when --resume, 0.5 otherwise. Lower values risk "
             "self-reinforcing misclassification loops.",
    )
    parser.add_argument(
        "--copy-project",
        action="store_true",
        help="Copy trained model and labels to project model/ directory",
    )
    parser.add_argument(
        "--project-dir",
        default=None,
        help="Project repo root for --copy-project (auto-detected if omitted)",
    )
    parser.add_argument(
        "--difficulty-weighting",
        action="store_true",
        default=None,
        help="Enable error-aware difficulty weighting. Evaluates the base model "
             "before training and up-weights misclassified samples. "
             "Default: ON when --resume (fine-tune), OFF otherwise.",
    )
    parser.add_argument(
        "--no-difficulty-weighting",
        action="store_true",
        help="Explicitly disable difficulty weighting (overrides --resume default).",
    )
    parser.add_argument(
        "--misclassification-weight",
        type=float,
        default=2.0,
        help="Multiplier for misclassified samples (default: 2.0). "
             "2.0 is the validated optimum — higher values destabilize training. "
             "Set to 1.0 to disable misclassification factor (span/SNR only).",
    )
    args = parser.parse_args()

    os.makedirs(os.path.dirname(args.output), exist_ok=True)

    merged_labels = list(SIGNAL_LABELS)
    old_num_classes = NUM_CLASSES
    aux_dim = AUX_DIM
    initial_state_dict = None

    if args.resume:
        resume_path = args.resume
        if not os.path.exists(resume_path):
            print(f"Error: Resume model not found: {resume_path}")
            sys.exit(1)

        ai_threshold = args.ai_label_threshold if args.ai_label_threshold is not None else 0.9

        labels_path = args.labels or os.path.join(os.path.dirname(resume_path), "signal_labels.json")
        if os.path.exists(labels_path):
            try:
                with open(labels_path, "r") as f:
                    labels_data = json.load(f)
                if isinstance(labels_data, dict):
                    existing_labels = labels_data.get("labels", [])
                else:
                    existing_labels = labels_data
                print(f"Loaded {len(existing_labels)} existing labels from {labels_path}")
            except Exception as e:
                print(f"Warning: Could not load labels from {labels_path}: {e}")
                existing_labels = list(SIGNAL_LABELS) + ["Unknown"]
        else:
            existing_labels = list(SIGNAL_LABELS) + ["Unknown"]

        capture_labels = scan_capture_labels(args.captures_dir, ai_label_threshold=ai_threshold)
        merged_labels_list, new_labels = merge_labels(existing_labels, capture_labels)
        model_labels = [lbl for lbl in merged_labels_list if lbl != "Unknown"]

        if new_labels:
            print(f"New labels found: {new_labels}")
            print(f"Merged labels: {merged_labels_list}")
        else:
            print(f"No new labels found. Fine-tuning with existing {len(model_labels)} classes.")

        if resume_path.endswith(".pt"):
            import torch
            initial_state_dict = torch.load(resume_path, map_location="cpu", weights_only=True)
            print(f"Loaded checkpoint from {resume_path}")
            old_num_classes = NUM_CLASSES
            for key, value in initial_state_dict.items():
                if key == "classifier.3.weight":
                    old_num_classes = value.shape[0]
                    break
            aux_dim = AUX_DIM
            for key, value in initial_state_dict.items():
                if key == "classifier.0.weight":
                    fc_in = value.shape[1]
                    aux_dim = fc_in - 32 * 12 * 8
                    break

            if old_num_classes < len(model_labels):
                initial_state_dict = expand_classifier_head(
                    initial_state_dict, old_num_classes, len(model_labels), aux_dim=aux_dim
                )
                print(f"Expanded classification head: {old_num_classes} -> {len(model_labels)} classes")
        else:
            state_dict, old_num_classes, aux_dim = load_onnx_to_state_dict(resume_path)
            print(f"Loaded ONNX model: {old_num_classes} classes, aux_dim={aux_dim}")

            if old_num_classes < len(model_labels):
                initial_state_dict = expand_classifier_head(
                    state_dict, old_num_classes, len(model_labels), aux_dim=aux_dim
                )
                print(f"Expanded classification head: {old_num_classes} -> {len(model_labels)} classes")
            else:
                initial_state_dict = state_dict

        merged_labels = model_labels

    else:
        ai_threshold = args.ai_label_threshold if args.ai_label_threshold is not None else 0.5
        capture_labels = scan_capture_labels(args.captures_dir, ai_label_threshold=ai_threshold)
        existing_labels = list(SIGNAL_LABELS)
        merged_labels_list, _ = merge_labels(existing_labels, capture_labels)
        merged_labels = [lbl for lbl in merged_labels_list if lbl != "Unknown"]

    print("Loading labeled captures...")
    dataset = load_captures(args.captures_dir, signal_labels=merged_labels, ai_label_threshold=ai_threshold)
    print(f"AI label confidence threshold: {ai_threshold}")

    if not dataset:
        print("No labeled captures found. Use the Capture button in the app to label signals.")
        sys.exit(1)

    num_classes = len(merged_labels)

    lr = args.lr
    if args.resume and not args.freeze_features and not args.freeze_classifier_input:
        lr = min(lr, 0.0005)
        print(f"Fine-tuning mode: reducing learning rate to {lr}")

    if args.resume and args.freeze_classifier_input:
        lr = min(lr, 0.0003)

    # Determine if difficulty weighting is enabled
    enable_difficulty = False
    if args.no_difficulty_weighting:
        enable_difficulty = False
    elif args.difficulty_weighting:
        enable_difficulty = True
    elif args.resume:
        enable_difficulty = True

    difficulty_weights = None
    if enable_difficulty:
        from difficulty_scorer import compute_difficulty_scores, print_difficulty_report

        print("\n=== Difficulty Weighting Enabled ===")
        misclass_weight = args.misclassification_weight

        eval_model = None
        if initial_state_dict is not None:
            eval_model = build_model(num_classes=num_classes, aux_dim=aux_dim)
            eval_model.load_state_dict(initial_state_dict, strict=False)
            eval_model.eval()
            print(f"Using fine-tune base model for misclassification evaluation")
        else:
            bundled_path = os.path.join(os.path.dirname(SCRIPT_DIR), "model", "signal_classifier.onnx")
            if not os.path.exists(bundled_path):
                bundled_path = os.path.join(os.path.dirname(args.captures_dir), "model", "signal_classifier.onnx")
            if os.path.exists(bundled_path):
                try:
                    sd, nc, ad = load_onnx_to_state_dict(bundled_path)
                    eval_model = build_model(num_classes=nc, aux_dim=ad)
                    eval_model.load_state_dict(sd)
                    eval_model.eval()
                    print(f"Using bundled model for misclassification evaluation: {bundled_path}")
                except Exception as e:
                    print(f"Warning: Could not load bundled model for evaluation: {e}")

        if eval_model is not None:
            print(f"Misclassification weight: {misclass_weight}x")
            difficulty_weights, report = compute_difficulty_scores(
                dataset, eval_model, merged_labels,
                misclassification_weight=misclass_weight,
                enable_misclassification=True,
                enable_span=True,
                enable_snr=True,
            )
            print_difficulty_report(report)
        else:
            print("No model available for evaluation — using physical factors only (span+SNR)")
            difficulty_weights, report = compute_difficulty_scores(
                dataset, None, merged_labels,
                misclassification_weight=1.0,
                enable_misclassification=False,
                enable_span=True,
                enable_snr=True,
            )
            print_difficulty_report(report)

    print(f"\nTraining for {args.epochs} epochs (aux_dim={aux_dim}, classes={num_classes})...")
    result = train_model(
        dataset, epochs=args.epochs, batch_size=args.batch_size,
        learning_rate=lr, num_classes=num_classes, aux_dim=aux_dim,
        initial_state_dict=initial_state_dict,
        freeze_features=args.freeze_features,
        freeze_classifier_input=args.freeze_classifier_input,
        old_num_classes=old_num_classes if args.resume else 0,
        merged_labels=merged_labels,
        difficulty_weights=difficulty_weights,
    )

    print(f"\nExporting ONNX model to {args.output}...")
    export_to_onnx(result.model, args.output, aux_dim=aux_dim)

    file_size = os.path.getsize(args.output)
    print(f"Model exported: {args.output} ({file_size / 1024:.1f} KB)")

    base_model_path = args.resume if args.resume else ""
    base_model_hash = compute_model_hash(args.resume) if args.resume else ""

    training_metadata = {
        "date": datetime.now(timezone.utc).isoformat(),
        "mode": "fine-tune" if args.resume else "scratch",
        "base_model_path": base_model_path,
        "base_model_hash": base_model_hash,
        "total_samples": len(dataset),
        "train_samples": result.train_samples,
        "val_samples": result.val_samples,
        "best_val_accuracy": round(result.best_val_accuracy, 6),
        "epochs": result.total_epochs,
        "active_classes": result.active_classes,
        "frozen_layers": result.frozen_layers,
        "learning_rate": result.learning_rate,
        "difficulty_weighting": enable_difficulty,
        "misclassification_weight": args.misclassification_weight if enable_difficulty else None,
    }
    try:
        import torch
        training_metadata["torch_version"] = torch.__version__
    except ImportError:
        pass

    labels_path = os.path.join(os.path.dirname(args.output), "signal_labels.json")
    labels_data = {
        "labels": merged_labels + ["Unknown"],
        "num_classes": num_classes,
        "aux_dim": aux_dim,
        "patch_strategy": PATCH_STRATEGY,
        "training_metadata": training_metadata,
    }
    with open(labels_path, "w") as f:
        json.dump(labels_data, f, indent=4)
    print(f"Labels exported: {labels_path}")

    if args.checkpoint:
        import torch
        checkpoint_path = os.path.splitext(args.output)[0] + ".pt"
        torch.save(result.model.state_dict(), checkpoint_path)
        print(f"Checkpoint saved: {checkpoint_path}")

    if args.copy_project:
        if args.project_dir:
            project_root = os.path.abspath(args.project_dir)
        else:
            project_root = SCRIPT_DIR
            for _ in range(6):
                project_root = os.path.dirname(project_root)
                if os.path.isdir(os.path.join(project_root, "model")):
                    break
        project_model_dir = os.path.join(project_root, "model")
        if not os.path.isdir(project_model_dir):
            print(f"\nWarning: project model directory not found: {project_model_dir}")
            print("Skipping copy. Run from the project repo for --copy-project.")
        else:
            project_model = os.path.join(project_model_dir, "signal_classifier.onnx")
            project_labels = os.path.join(project_model_dir, "signal_labels.json")
            if os.path.normpath(os.path.abspath(args.output)) == os.path.normpath(os.path.abspath(project_model)):
                print(f"\nModel already at project path: {project_model}")
            else:
                shutil.copy2(args.output, project_model)
                print(f"\nModel copied to project: {project_model}")
            if os.path.normpath(os.path.abspath(labels_path)) == os.path.normpath(os.path.abspath(project_labels)):
                print(f"Labels already at project path: {project_labels}")
            else:
                shutil.copy2(labels_path, project_labels)
                print(f"Labels copied to project: {project_labels}")
    else:
        default_output = os.path.join(MODEL_DIR, "signal_classifier.onnx")
        if os.path.normpath(os.path.abspath(args.output)) == os.path.normpath(os.path.abspath(default_output)):
            print(f"\nModel and labels saved to {MODEL_DIR}/")
            print("Restart the app or click Reload Model in Signal Detection settings to use the new model.")
        else:
            print(f"\nTo use: copy {args.output} and {labels_path} to {MODEL_DIR}/")
            print("Or re-run with --copy-project to copy to the project model/ directory.")


if __name__ == "__main__":
    main()
