#!/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_patch,
    waterfall_to_scope_history,
    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_legacy, is_ai_labeled) tuples.
    Each detected candidate with a meaningful label becomes its own
    training sample. Legacy captures (without scope_history) are flagged
    with is_legacy=True. AI-labeled candidates are flagged with
    is_ai_labeled=True.

    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:
            data = np.load(filepath, allow_pickle=True)
            metadata = json.loads(str(data["metadata"]))
            waterfall = data["waterfall"]
            spectrum = data["spectrum"]

            has_scope = "scope_history" in data
            if has_scope:
                scope_history = data["scope_history"]
                if scope_history.shape == (0, 0) or scope_history.size == 0:
                    has_scope = False

            if not has_scope:
                color_range_pct = metadata.get("color_range_percentage", 75)
                scope_history = waterfall_to_scope_history(
                    waterfall, color_range_percentage=color_range_pct
                )

            is_legacy = not has_scope

            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)
            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

                    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)

                    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,
                    )
                    dataset.append((patch, aux, patch_label_idx, is_legacy, is_ai_labeled))
            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,
                )
                dataset.append((patch, aux, label_idx, is_legacy, False))

            data.close()

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

    n_legacy = sum(1 for _, _, _, is_legacy, _ in dataset if is_legacy)
    n_ai = sum(1 for _, _, _, _, is_ai_labeled in dataset if is_ai_labeled)
    print(f"Loaded {len(dataset)} training samples ({n_legacy} legacy, {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:
            data = np.load(filepath, allow_pickle=True)
            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)

            data.close()
        except Exception:
            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,
):
    """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.

    Legacy samples (is_legacy=True) are down-weighted by 0.5 in the loss
    function to reduce the impact of approximate amplitude data
    reverse-engineered from waterfall RGB.

    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):
            self.samples = samples

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

        def __getitem__(self, idx):
            patch, aux, label, is_legacy, is_ai_labeled = self.samples[idx]
            patch_tensor = torch.tensor(patch, dtype=torch.float32).permute(2, 0, 1)
            aux_tensor = torch.tensor(aux, dtype=torch.float32)
            weight = 0.5 if is_legacy 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 _, _, label_idx, is_legacy, is_ai_labeled in dataset:
        w = class_weights[label_idx].item() if label_idx < num_classes else 0.0
        if is_legacy:
            w *= 0.5
        sample_weights.append(w)
    sample_weights_tensor = torch.tensor(sample_weights, dtype=torch.float32)

    label_to_indices = {}
    for i, (_, _, label_idx, _, _) in enumerate(dataset):
        label_to_indices.setdefault(label_idx, []).append(i)

    train_indices = []
    val_indices = []
    for label_idx, indices in label_to_indices.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_data = SignalDataset([dataset[i] for i in train_indices])
    val_data = SignalDataset([dataset[i] for i in val_indices])

    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)
            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)",
    )
    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)

    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,
    )

    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,
    }
    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()
