"""Verify a trained signal classifier ONNX model and labels file.

Checks model dimensions, aux input shape, ONNX structure, inference,
label validity, training metadata, and optionally file sync between
project and user dirs.

Usage:
    python verify_model.py [--model-dir DIR] [--scripts-dir DIR]
                           [--project-scripts DIR] [--project-model DIR]

    # Compare two models' training metadata:
    python verify_model.py --compare DIR1 DIR2
"""

import argparse
import json
import os
import sys


def _load_labels_json(model_dir: str) -> dict:
    """Load signal_labels.json from a model directory, returning empty dict on failure."""
    labels_path = os.path.join(model_dir, "signal_labels.json")
    if not os.path.isfile(labels_path):
        return {}
    try:
        with open(labels_path, "r") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def display_metadata(labels_data: dict) -> None:
    """Print training_metadata from a labels JSON dict if present."""
    metadata = labels_data.get("training_metadata")
    if metadata is None:
        print("  Training metadata: (none)")
        return

    mode = metadata.get("mode", "unknown")
    date = metadata.get("date", "unknown")
    total = metadata.get("total_samples", "?")
    train = metadata.get("train_samples", "?")
    val = metadata.get("val_samples", "?")
    acc = metadata.get("best_val_accuracy", 0)
    epochs = metadata.get("epochs", "?")
    active = metadata.get("active_classes", {})
    frozen = metadata.get("frozen_layers")
    lr = metadata.get("learning_rate", "?")
    torch_ver = metadata.get("torch_version", "?")
    base_hash = metadata.get("base_model_hash", "")
    base_path = metadata.get("base_model_path", "")

    acc_str = f"{acc:.2%}" if isinstance(acc, (int, float)) and acc > 0 else str(acc)
    print(f"  Training mode: {mode}")
    print(f"  Date: {date}")
    print(f"  Samples: {total} (train={train}, val={val})")
    print(f"  Best val accuracy: {acc_str}")
    print(f"  Epochs: {epochs}, LR: {lr}")
    if frozen:
        print(f"  Frozen layers: {frozen}")
    if active:
        active_str = ", ".join(f"{k}={v}" for k, v in sorted(active.items()))
        print(f"  Active classes ({len(active)}): {active_str}")
    print(f"  Torch version: {torch_ver}")
    if base_hash:
        print(f"  Base model hash: {base_hash}")
    if base_path:
        print(f"  Base model path: {base_path}")


def compare_models(dir_a: str, dir_b: str) -> int:
    """Compare training metadata between two model directories.

    Returns 0 if no regression detected, 1 if regression found.
    Only requires signal_labels.json — no torch/onnx needed.
    """
    data_a = _load_labels_json(dir_a)
    data_b = _load_labels_json(dir_b)

    if not data_a:
        print(f"Error: No labels JSON found in {dir_a}")
        return 1
    if not data_b:
        print(f"Error: No labels JSON found in {dir_b}")
        return 1

    meta_a = data_a.get("training_metadata", {})
    meta_b = data_b.get("training_metadata", {})
    labels_a = data_a.get("labels", [])
    labels_b = data_b.get("labels", [])

    mode_a = meta_a.get("mode", "unknown") if meta_a else "unknown"
    mode_b = meta_b.get("mode", "unknown") if meta_b else "unknown"
    date_a = meta_a.get("date", "?") if meta_a else "?"
    date_b = meta_b.get("date", "?") if meta_b else "?"
    samples_a = meta_a.get("total_samples", 0) if meta_a else 0
    samples_b = meta_b.get("total_samples", 0) if meta_b else 0
    acc_a = meta_a.get("best_val_accuracy", 0) if meta_a else 0
    acc_b = meta_b.get("best_val_accuracy", 0) if meta_b else 0

    print("Comparing:")
    print(f"  A: {dir_a} ({mode_a}, {date_a}, {samples_a} samples, val_acc={acc_a:.2%})")
    print(f"  B: {dir_b} ({mode_b}, {date_b}, {samples_b} samples, val_acc={acc_b:.2%})")
    print()

    regression = False

    aux_a = data_a.get("aux_dim")
    aux_b = data_b.get("aux_dim")
    if aux_a is not None and aux_b is not None:
        if aux_a != aux_b:
            print(f"  aux_dim: A={aux_a}, B={aux_b} ⚠ CRITICAL MISMATCH")
            regression = True
        else:
            print(f"  aux_dim: A={aux_a}, B={aux_b} ✓")
    elif aux_a is not None or aux_b is not None:
        print(f"  aux_dim: A={aux_a}, B={aux_b} (cannot compare)")

    nc_a = len([lbl for lbl in labels_a if lbl != "Unknown"])
    nc_b = len([lbl for lbl in labels_b if lbl != "Unknown"])
    if nc_a != nc_b:
        print(f"  num_classes: A={nc_a}, B={nc_b} (different, not an error)")

    if samples_a > 0 and samples_b > 0:
        if samples_a < samples_b:
            pct = (1 - samples_a / samples_b) * 100
            print(f"  Sample regression: {samples_b} → {samples_a} ({pct:.1f}% decrease) ⚠")
            regression = True
        else:
            print(f"  Samples: A={samples_a}, B={samples_b} ✓")

    if acc_a > 0 and acc_b > 0:
        if acc_a < acc_b:
            print(f"  Accuracy regression: {acc_b:.2%} → {acc_a:.2%} ⚠")
            regression = True
        else:
            print(f"  Accuracy: A={acc_a:.2%}, B={acc_b:.2%} ✓")

    active_a = len(meta_a.get("active_classes", {})) if meta_a else 0
    active_b = len(meta_b.get("active_classes", {})) if meta_b else 0
    if active_a > 0 or active_b > 0:
        print(f"  Active classes: A={active_a}, B={active_b}")

    print()
    if regression:
        print("WARNING: Model A shows regression compared to model B.")
        print("Consider using model B or fine-tuning from model B instead.")
        return 1

    print("No regression detected.")
    return 0


def verify(model_dir, scripts_dir, project_scripts=None, project_model=None):
    sys.path.insert(0, scripts_dir)

    from model_definition import (
        AUX_DIM,
        SCOPE_PATCH_ROWS,
        SCOPE_PATCH_COLS,
        INPUT_CHANNELS,
        SCOPE_HISTORY_COLS,
        PATCH_STRATEGY,
        NUM_CLASSES,
        SIGNAL_LABELS,
        prepare_patch,
        prepare_aux_features,
    )
    import numpy as np

    onnx_path = os.path.join(model_dir, "signal_classifier.onnx")
    labels_path = os.path.join(model_dir, "signal_labels.json")

    print(f"  AUX_DIM: {AUX_DIM}")
    print(f"  PATCH_STRATEGY: {PATCH_STRATEGY}")
    print(f"  NUM_CLASSES: {NUM_CLASSES}")
    print(f"  INPUT_CHANNELS: {INPUT_CHANNELS}")

    aux = prepare_aux_features(
        span_hz=10000, reference_level_db=0.0, noise_floor=15.0,
        bandwidth_bins=10, center_freq_offset_hz=500, snr_db=12.0,
    )
    print(f"  Aux shape: {aux.shape}, Aux dim: {AUX_DIM}")
    assert aux.shape == (AUX_DIM,), f"Expected ({AUX_DIM},), got {aux.shape}"

    scope_history = np.random.randint(0, 160, (100, SCOPE_HISTORY_COLS)).astype(np.float32)
    patch = prepare_patch(scope_history, 220, 260)
    print(f"  Patch shape: {patch.shape}")
    assert patch.shape == (SCOPE_PATCH_ROWS, SCOPE_PATCH_COLS, INPUT_CHANNELS), \
        f"Expected ({SCOPE_PATCH_ROWS}, {SCOPE_PATCH_COLS}, {INPUT_CHANNELS}), got {patch.shape}"

    try:
        import onnx
    except ImportError:
        print("  onnx not installed — skipping ONNX verification")
        onnx = None

    if onnx is not None and os.path.exists(onnx_path):
        onnx_model = onnx.load(onnx_path)
        inputs = {
            inp.name: [d.dim_value for d in inp.type.tensor_type.shape.dim]
            for inp in onnx_model.graph.input
        }
        print(f"  ONNX inputs: {inputs}")
        graph_output = onnx_model.graph.output[0].name
        print(f"  ONNX output: {graph_output}")

        aux_input = [inp for inp in onnx_model.graph.input if inp.name == "aux"]
        if aux_input:
            aux_dims = [d.dim_value for d in aux_input[0].type.tensor_type.shape.dim]
            print(f"  ONNX aux input dims: {aux_dims}")
            aux_size = aux_dims[-1] if aux_dims[-1] else "?"
            if aux_size != "?":
                assert int(aux_size) == AUX_DIM, \
                    f"ONNX aux dim mismatch: expected {AUX_DIM}, got {aux_size}"
                print(f"  ONNX aux dim verified: {aux_size} == {AUX_DIM}")
        else:
            print("  WARNING: No aux input found in ONNX model")

        nan_count = sum(
            1 for init in onnx_model.graph.initializer
            if init.data_type == 1 and np.isnan(np.frombuffer(init.raw_data, dtype=np.float32)).any()
        )
        print(f"  ONNX NaN params (model weights): {nan_count}")

        try:
            import onnxruntime as ort
            sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
            dummy_patch = np.random.randint(0, 160, (1, INPUT_CHANNELS, SCOPE_PATCH_ROWS, SCOPE_PATCH_COLS)).astype(np.float32)
            dummy_aux = np.random.rand(1, AUX_DIM).astype(np.float32)
            result = sess.run(None, {"input": dummy_patch, "aux": dummy_aux})
            pred_class = int(np.argmax(result[0][0]))
            label = SIGNAL_LABELS[pred_class] if pred_class < len(SIGNAL_LABELS) else "Unknown"
            print(f"  ONNX Runtime inference OK: class {pred_class} ({label})")
        except ImportError:
            print("  onnxruntime not installed — skipping inference verification")
        except Exception as e:
            print(f"  WARNING: ONNX Runtime inference failed: {e}")
    elif onnx is not None:
        print(f"  WARNING: ONNX model not found at {onnx_path}")

    if os.path.exists(labels_path):
        with open(labels_path) as f:
            labels_data = json.load(f)
        if isinstance(labels_data, dict):
            labels = labels_data.get("labels", [])
            strategy = labels_data.get("patch_strategy", "unknown")
            labels_aux_dim = labels_data.get("aux_dim")
        else:
            labels = labels_data
            strategy = "unknown (old format)"
            labels_aux_dim = None
        print(f"  Labels ({len(labels)}): {labels[:5]}... + {len(labels) - 5} more")
        print(f"  Patch strategy: {strategy}")
        if labels_aux_dim is not None:
            print(f"  Labels aux_dim: {labels_aux_dim}")
        assert strategy == PATCH_STRATEGY, \
            f"Patch strategy mismatch: model={strategy}, runtime={PATCH_STRATEGY}"
        assert "Unknown" in labels, "Unknown label missing from labels file"

        if isinstance(labels_data, dict) and labels_data.get("training_metadata"):
            print()
            display_metadata(labels_data)
    else:
        print(f"  WARNING: Labels file not found at {labels_path}")

    if project_scripts and project_model:
        sync_ok = True
        for fname in ["model_definition.py", "patch_config.py", "train_signal_classifier.py", "retrain_model.py", "verify_model.py"]:
            p = os.path.join(project_scripts, fname)
            u = os.path.join(scripts_dir, fname)
            if os.path.isfile(p) and os.path.isfile(u):
                if open(p, "rb").read() != open(u, "rb").read():
                    print(f"  MISMATCH: {fname}")
                    sync_ok = False

        for fname in ["signal_classifier.onnx", "signal_labels.json"]:
            p = os.path.join(project_model, fname)
            u = os.path.join(model_dir, fname)
            if os.path.isfile(p) and os.path.isfile(u):
                if open(p, "rb").read() != open(u, "rb").read():
                    print(f"  MISMATCH: {fname}")
                    sync_ok = False

        if sync_ok:
            print("  All files in sync ✓")
        else:
            print("  Some files out of sync — re-run to fix")

    print("Verification passed.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Verify signal classifier ONNX model and labels")
    parser.add_argument(
        "--model-dir",
        default=os.path.join(os.path.expanduser("~"), "Pycom", "signals", "model"),
        help="Path to model directory (default: ~/Pycom/signals/model)",
    )
    parser.add_argument(
        "--scripts-dir",
        default=os.path.join(os.path.expanduser("~"), "Pycom", "signals", "scripts"),
        help="Path to scripts directory for imports (default: ~/Pycom/signals/scripts)",
    )
    parser.add_argument(
        "--project-scripts",
        default=None,
        help="Project scripts dir for sync check (optional)",
    )
    parser.add_argument(
        "--project-model",
        default=None,
        help="Project model dir for sync check (optional)",
    )
    parser.add_argument(
        "--compare",
        nargs=2,
        metavar=("DIR_A", "DIR_B"),
        help="Compare training metadata between two model directories",
    )
    args = parser.parse_args()

    if args.compare:
        sys.exit(compare_models(args.compare[0], args.compare[1]))

    print("=== Verifying model ===\n")
    verify(args.model_dir, args.scripts_dir, args.project_scripts, args.project_model)
