"""Retrain or train the signal classifier ONNX model from the command line.

Cross-platform replacement for retrain_model.sh. Works on macOS, Linux,
and Windows with Python 3.10+.

Usage:
    # Train from scratch (default):
    python retrain_model.py

    # Fine-tune from the distributed base model:
    python retrain_model.py --fine-tune --freeze-classifier-input

    # With custom epochs:
    python retrain_model.py --fine-tune --freeze-classifier-input --epochs 50

Requires PyTorch and ONNX for training (dev-only dependencies):
    pip install torch onnx onnxscript torchvision

Training modes:
    From scratch (no flags): Creates a new model with random initial weights
        and trains on all labeled captures. Use when building a completely
        custom model.

    Fine-tune (--fine-tune): Loads the distributed base model weights and
        trains on all labeled captures. Use when adding your signals to the
        proven model. Always fine-tunes from the protected .default model,
        never from the active model.
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
from typing import Dict, List, Optional


SIGNALS_DIR = os.path.join(os.path.expanduser("~"), "Pycom", "signals")
MODEL_DIR = os.path.join(SIGNALS_DIR, "model")
MODEL_PATH = os.path.join(MODEL_DIR, "signal_classifier.onnx")
DEFAULT_MODEL_PATH = MODEL_PATH + ".default"
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

_SAMPLE_REGRESSION_FACTOR = 0.5
_CLASS_COVERAGE_STRONG = 0.33
_CLASS_COVERAGE_MODERATE = 0.5
_CLASS_IMBALANCE_MAX = 0.5
_SMALL_MODEL_THRESHOLD = 50


_LABELS_DEFAULT_PATH = os.path.join(MODEL_DIR, "signal_labels.json.default")

_UNRELIABLE_VAL_THRESHOLD = 10


def load_training_metadata(model_dir: str) -> Optional[dict]:
    """Load training_metadata from signal_labels.json in model_dir.

    Returns None if the file doesn't exist, has no training_metadata,
    or is malformed.
    """
    labels_path = os.path.join(model_dir, "signal_labels.json")
    if not os.path.isfile(labels_path):
        return None
    try:
        with open(labels_path, "r") as f:
            data = json.load(f)
        if isinstance(data, dict):
            return data.get("training_metadata")
    except Exception:
        pass
    return None


def check_training_quality(
    existing_metadata: Optional[dict],
    new_sample_count: int,
    active_class_count: int,
    num_classes: int,
    class_distribution: Dict[str, int],
    fine_tune: bool = False,
) -> List[str]:
    """Check for training quality issues and return warning messages.

    Pure function for testability. Returns an empty list if no warnings.

    Parameters
    ----------
    existing_metadata : dict or None
        training_metadata from the existing model's labels JSON, or None
        if no previous model exists.
    new_sample_count : int
        Total number of training samples available for this run.
    active_class_count : int
        Number of classes that have at least one training sample.
    num_classes : int
        Total number of output classes in the model.
    class_distribution : dict
        Mapping of class name to sample count for this run.
    fine_tune : bool
        Whether this run is a fine-tune from the distributed base model.
        Suppresses the "consider --fine-tune" recommendation when already
        fine-tuning.
    """
    warnings = []

    if existing_metadata and existing_metadata.get("total_samples", 0) > _SMALL_MODEL_THRESHOLD:
        prev = existing_metadata["total_samples"]
        if new_sample_count < prev * _SAMPLE_REGRESSION_FACTOR:
            warnings.append(
                f"Fine-tuning with significantly fewer samples ({new_sample_count}) "
                f"than the base model was trained on ({prev}). "
                "This may degrade classification quality for underrepresented classes."
            )

    if num_classes > 0:
        if active_class_count < num_classes * _CLASS_COVERAGE_STRONG:
            warnings.append(
                f"Only {active_class_count} of {num_classes} classes have training data. "
                f"The other {num_classes - active_class_count} classes will lose accuracy."
            )
        elif active_class_count < num_classes * _CLASS_COVERAGE_MODERATE:
            warnings.append(
                f"{active_class_count} of {num_classes} classes have training data. "
                "Classes without training samples may lose accuracy."
            )

    total = sum(class_distribution.values())
    if total > 0:
        for label, count in class_distribution.items():
            if count > total * _CLASS_IMBALANCE_MAX:
                pct = count * 100 // total
                warnings.append(
                    f"One class ({label}) has {count}/{total} ({pct}%) of training samples. "
                    "The model will be biased toward this class. "
                    "Consider capturing more varied signal types."
                )
                break

    if not fine_tune and existing_metadata and existing_metadata.get("mode") == "scratch":
        prev_samples = existing_metadata.get("total_samples", 0)
        if prev_samples < _SMALL_MODEL_THRESHOLD:
            warnings.append(
                f"The current model was trained from scratch with only {prev_samples} samples. "
                "Consider using --fine-tune to add your data to the distributed "
                "base model instead."
            )

    return warnings


def main():
    parser = argparse.ArgumentParser(
        description="Train or fine-tune the signal classifier ONNX model",
    )
    parser.add_argument(
        "--fine-tune",
        action="store_true",
        help="Fine-tune from the distributed base model. Always loads from "
             "the protected .default model, not the active model.",
    )
    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. "
             "Recommended for small datasets with new classes.",
    )
    parser.add_argument(
        "--checkpoint",
        action="store_true",
        help="Save PyTorch .pt state dict alongside ONNX for future "
             "training via train_signal_classifier.py --resume",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=30,
        help="Training epochs (default: 30)",
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=None,
        help="Learning rate override",
    )
    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 --fine-tune, OFF otherwise.",
    )
    parser.add_argument(
        "--no-difficulty-weighting",
        action="store_true",
        help="Explicitly disable difficulty weighting (overrides --fine-tune 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.",
    )
    args = parser.parse_args()

    if args.fine_tune and not os.path.isfile(DEFAULT_MODEL_PATH):
        print(
            f"Error: No base model found at {DEFAULT_MODEL_PATH}. "
            "Start the app at least once to deploy the default model, "
            "or train from scratch without --fine-tune."
        )
        sys.exit(1)

    if args.fine_tune and not os.path.isfile(_LABELS_DEFAULT_PATH):
        print(
            f"Error: No base labels found at {_LABELS_DEFAULT_PATH}. "
            "Start the app at least once to deploy the default labels, "
            "or train from scratch without --fine-tune."
        )
        sys.exit(1)

    os.makedirs(MODEL_DIR, exist_ok=True)

    existing_metadata = load_training_metadata(MODEL_DIR)

    if args.fine_tune:
        if os.path.isfile(MODEL_PATH):
            backup = MODEL_PATH + ".bak"
            shutil.copy2(MODEL_PATH, backup)
            print(f"=== Backed up existing model to {backup} ===")
            labels_path = os.path.join(MODEL_DIR, "signal_labels.json")
            if os.path.isfile(labels_path):
                shutil.copy2(labels_path, labels_path + ".bak")
                print(f"=== Backed up labels to {labels_path}.bak ===")
        print(f"=== Fine-tuning signal classifier from {DEFAULT_MODEL_PATH} ===")
    else:
        print("=== Training signal classifier from scratch (ScopeCNN-Patch) ===")
        if existing_metadata and existing_metadata.get("total_samples", 0) > _SMALL_MODEL_THRESHOLD:
            prev = existing_metadata["total_samples"]
            print(
                f"WARNING: Training from scratch but an existing model was trained "
                f"on {prev} samples. Consider --fine-tune to add your data to the "
                f"distributed base model instead."
            )

    sys.path.insert(0, SCRIPT_DIR)
    try:
        from train_signal_classifier import scan_capture_labels
        from model_definition import SIGNAL_LABELS, merge_labels
    except ImportError as e:
        print(f"Warning: cannot import training modules for quality checks: {e}")
        print("Quality checks will be skipped.")
        scan_capture_labels = None

    if scan_capture_labels is not None:
        captures_dir = os.path.join(SIGNALS_DIR, "captures")
        if os.path.isdir(captures_dir):
            capture_labels = scan_capture_labels(captures_dir, ai_label_threshold=0.9 if args.fine_tune else 0.5)
            existing_labels = list(SIGNAL_LABELS)
            merged, _ = merge_labels(existing_labels, capture_labels)
            model_labels = [label for label in merged if label != "Unknown"]
            num_classes = len(model_labels)

            from collections import Counter
            label_counts = Counter(capture_labels)
            active_class_count = len(label_counts)
            total_samples = sum(label_counts.values())

            warnings = check_training_quality(
                existing_metadata=existing_metadata,
                new_sample_count=total_samples,
                active_class_count=active_class_count,
                num_classes=num_classes,
                class_distribution=dict(label_counts),
                fine_tune=args.fine_tune,
            )
            for w in warnings:
                print(f"WARNING: {w}")

    cmd = [sys.executable, os.path.join(SCRIPT_DIR, "train_signal_classifier.py")]
    cmd.append(f"--epochs={args.epochs}")

    if args.fine_tune:
        cmd.extend(["--resume", DEFAULT_MODEL_PATH])
        cmd.extend(["--labels", _LABELS_DEFAULT_PATH])

    if args.freeze_features:
        cmd.append("--freeze-features")

    if args.freeze_classifier_input:
        cmd.append("--freeze-classifier-input")

    if args.checkpoint:
        cmd.append("--checkpoint")

    if args.lr is not None:
        cmd.append(f"--lr={args.lr}")

    if args.copy_project:
        cmd.append("--copy-project")
        if args.project_dir:
            cmd.extend(["--project-dir", args.project_dir])

    if args.no_difficulty_weighting:
        cmd.append("--no-difficulty-weighting")
    elif args.difficulty_weighting:
        cmd.append("--difficulty-weighting")

    if args.misclassification_weight != 2.0:
        cmd.append(f"--misclassification-weight={args.misclassification_weight}")

    print(f"  Running: {' '.join(cmd)}")
    print()

    result = subprocess.run(cmd, cwd=SCRIPT_DIR)
    if result.returncode != 0:
        print(f"\nERROR: Training failed with exit code {result.returncode}")
        sys.exit(result.returncode)

    new_metadata = load_training_metadata(MODEL_DIR)
    if existing_metadata and new_metadata:
        prev_acc = existing_metadata.get("best_val_accuracy", 0.0)
        new_acc = new_metadata.get("best_val_accuracy", 0.0)
        if prev_acc > 0 and new_acc < prev_acc * 0.8:
            print(
                f"\nWARNING: Model accuracy regressed: "
                f"{prev_acc:.2%} → {new_acc:.2%}. "
                "The new model may perform worse than the previous one."
            )

    if new_metadata:
        val_samples = new_metadata.get("val_samples", 0)
        if val_samples > 0 and val_samples < _UNRELIABLE_VAL_THRESHOLD:
            val_acc = new_metadata.get("best_val_accuracy", 0.0)
            print(
                f"\nNOTE: Validation accuracy ({val_acc:.2%}) is based on only "
                f"{val_samples} sample{'s' if val_samples != 1 else ''}. "
                "Results may not generalize reliably to new data. "
                "Consider capturing more signals before training."
            )

    print("\n=== Verifying model ===\n")
    verify_cmd = [
        sys.executable,
        os.path.join(SCRIPT_DIR, "verify_model.py"),
        "--model-dir", MODEL_DIR,
        "--scripts-dir", SCRIPT_DIR,
    ]
    if args.copy_project and args.project_dir:
        verify_cmd.extend([
            "--project-scripts", os.path.join(args.project_dir, "scripts"),
            "--project-model", os.path.join(args.project_dir, "model"),
        ])

    subprocess.run(verify_cmd, cwd=SCRIPT_DIR)

    print("\n=== Done ===")
    for fname in ["signal_classifier.onnx", "signal_labels.json"]:
        fpath = os.path.join(MODEL_DIR, fname)
        if os.path.isfile(fpath):
            size_kb = os.path.getsize(fpath) / 1024
            print(f"  {fname}: {size_kb:.1f} KB")


if __name__ == "__main__":
    main()
