#!/usr/bin/env python3
"""
Adversarial Robustness Evaluator

This standalone evaluator tests your adversarially-trained model against
pre-generated FGSM and I-FGSM attacks. Run this after training to measure
your model's robustness metrics.

Usage:
    python evaluate_robustness.py --model-path robust_model.safetensors
    python evaluate_robustness.py --model-path models/ --compare
"""

from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Dict, List, Tuple, Optional

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

SCRIPT_DIR = Path(__file__).resolve().parent

# Primary epsilon for main results
PRIMARY_EPSILON = 0.3

# Files to skip when scanning directories for models
SKIP_FILES = {"baseline_model.safetensors", "adv_examples.safetensors"}


# ---------------------------------------------------------------------------
# Model Definition (lazy import)
# ---------------------------------------------------------------------------


def create_lenet5():
    """Create a LeNet-5 model instance."""
    import torch.nn as nn
    import torch.nn.functional as F

    class LeNet5(nn.Module):
        """Classic LeNet-5 architecture for MNIST classification."""

        def __init__(self) -> None:
            super().__init__()
            self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=2)
            self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
            self.fc1 = nn.Linear(16 * 5 * 5, 120)
            self.fc2 = nn.Linear(120, 84)
            self.fc3 = nn.Linear(84, 10)

        def forward(self, x):
            x = F.max_pool2d(F.relu(self.conv1(x)), 2)
            x = F.max_pool2d(F.relu(self.conv2(x)), 2)
            x = x.view(-1, 16 * 5 * 5)
            x = F.relu(self.fc1(x))
            x = F.relu(self.fc2(x))
            x = self.fc3(x)
            return x

    return LeNet5()


# ---------------------------------------------------------------------------
# Safetensors Helper Functions
# ---------------------------------------------------------------------------


def load_adversarial_examples_safetensors(filepath: Path) -> Dict:
    """Load adversarial examples from safetensors format."""
    from safetensors import safe_open
    from safetensors.torch import load_file

    tensors = load_file(str(filepath))

    # Extract metadata
    with safe_open(str(filepath), framework="pt") as f:
        metadata = f.metadata()

    epsilon = float(metadata["epsilon"])
    epsilon_spread = json.loads(metadata["epsilon_spread"])

    # Reconstruct epsilon-keyed dicts
    fgsm_by_epsilon = {}
    ifgsm_by_epsilon = {}
    for key, tensor in tensors.items():
        if key.startswith("fgsm_eps_"):
            eps = float(key.replace("fgsm_eps_", ""))
            fgsm_by_epsilon[eps] = tensor
        elif key.startswith("ifgsm_eps_"):
            eps = float(key.replace("ifgsm_eps_", ""))
            ifgsm_by_epsilon[eps] = tensor

    return {
        "clean_images": tensors["clean_images"],
        "clean_labels": tensors["clean_labels"],
        "fgsm_images": fgsm_by_epsilon[epsilon],
        "ifgsm_images": ifgsm_by_epsilon[epsilon],
        "epsilon": epsilon,
        "epsilon_spread": epsilon_spread,
        "fgsm_by_epsilon": fgsm_by_epsilon,
        "ifgsm_by_epsilon": ifgsm_by_epsilon,
    }


# ---------------------------------------------------------------------------
# File Discovery
# ---------------------------------------------------------------------------


def find_models(path: Path) -> List[Path]:
    """Find model files to evaluate."""
    if path.is_file() and path.suffix == ".safetensors":
        return [path]

    if path.is_dir():
        models = []
        for f in sorted(path.glob("*.safetensors")):
            if f.name not in SKIP_FILES:
                models.append(f)
        return models

    return []


def find_baseline(model_path: Path) -> Optional[Path]:
    """Find baseline_model.safetensors relative to the model path."""
    if model_path.is_file():
        search_dir = model_path.parent
    else:
        search_dir = model_path

    for check_dir in [search_dir, search_dir.parent, SCRIPT_DIR]:
        baseline = check_dir / "baseline_model.safetensors"
        if baseline.exists():
            return baseline

    return None


def find_adv_examples(model_path: Path) -> Optional[Path]:
    """Find adv_examples.safetensors relative to the model path."""
    if model_path.is_file():
        search_dir = model_path.parent
    else:
        search_dir = model_path

    for check_dir in [search_dir, search_dir.parent, SCRIPT_DIR]:
        adv = check_dir / "adv_examples.safetensors"
        if adv.exists():
            return adv

    return None


# ---------------------------------------------------------------------------
# Evaluation Functions
# ---------------------------------------------------------------------------


def load_model(model_path: Path):
    """Load a trained model from safetensors checkpoint."""
    from safetensors.torch import load_file

    model = create_lenet5()
    state_dict = load_file(str(model_path))
    model.load_state_dict(state_dict)
    model.eval()
    return model


def load_adversarial_examples(adv_path: Path) -> Dict:
    """Load pre-generated adversarial examples from safetensors."""
    return load_adversarial_examples_safetensors(adv_path)


def evaluate_accuracy(model, images, labels, device) -> Tuple[float, int, int]:
    """Compute accuracy and return (accuracy%, correct, total)."""
    import torch

    model.eval()
    with torch.no_grad():
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = outputs.max(1)
        correct = predicted.eq(labels).sum().item()
        total = labels.size(0)

    return 100.0 * correct / total, correct, total


def get_misclassified_samples(
    model, images, labels, device, max_samples: int = 5
) -> list:
    """Get indices and predictions of misclassified samples."""
    import torch
    import torch.nn.functional as F

    model.eval()
    with torch.no_grad():
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        _, predicted = outputs.max(1)
        wrong = predicted != labels
        wrong_indices = wrong.nonzero(as_tuple=True)[0][:max_samples]

    failures = []
    for idx in wrong_indices:
        failures.append(
            {
                "index": idx.item(),
                "true_label": labels[idx].item(),
                "predicted": predicted[idx].item(),
                "confidence": F.softmax(outputs[idx], dim=0).max().item(),
            }
        )

    return failures


def evaluate_model_full(model_path: Path, adv_data: Dict, device) -> Dict:
    """Evaluate a model across all epsilon values and return full results."""
    clean_images = adv_data["clean_images"]
    clean_labels = adv_data["clean_labels"]
    epsilon = adv_data["epsilon"]

    # Check for epsilon spread vs single epsilon
    has_spread = "epsilon_spread" in adv_data and "fgsm_by_epsilon" in adv_data
    if has_spread:
        epsilon_spread = adv_data["epsilon_spread"]
        fgsm_by_epsilon = adv_data["fgsm_by_epsilon"]
        ifgsm_by_epsilon = adv_data["ifgsm_by_epsilon"]
    else:
        epsilon_spread = [epsilon]
        fgsm_by_epsilon = {epsilon: adv_data["fgsm_images"]}
        ifgsm_by_epsilon = {epsilon: adv_data["ifgsm_images"]}

    model = load_model(model_path)
    model.to(device)

    # Evaluate clean accuracy
    clean_acc, clean_correct, clean_total = evaluate_accuracy(
        model, clean_images, clean_labels, device
    )

    # Evaluate across all epsilon values
    fgsm_results = {}
    ifgsm_results = {}
    for eps in epsilon_spread:
        fgsm_acc, _, _ = evaluate_accuracy(
            model, fgsm_by_epsilon[eps], clean_labels, device
        )
        ifgsm_acc, _, _ = evaluate_accuracy(
            model, ifgsm_by_epsilon[eps], clean_labels, device
        )
        fgsm_results[eps] = fgsm_acc
        ifgsm_results[eps] = ifgsm_acc

    return {
        "name": model_path.name,
        "clean_acc": clean_acc,
        "clean_correct": clean_correct,
        "clean_total": clean_total,
        "fgsm_results": fgsm_results,
        "ifgsm_results": ifgsm_results,
        "epsilon_spread": epsilon_spread,
        "model": model,
    }


def print_model_results(results: Dict, baseline: Optional[Dict] = None) -> None:
    """Print detailed results for a single model."""
    name = results["name"]
    clean_acc = results["clean_acc"]
    fgsm_results = results["fgsm_results"]
    ifgsm_results = results["ifgsm_results"]
    epsilon_spread = results["epsilon_spread"]

    print(f"\n{'=' * 70}")
    print(f"Model: {name}")
    print(f"{'=' * 70}")

    # Clean accuracy
    print(
        f"\nClean Accuracy: {clean_acc:.1f}% ({results['clean_correct']}/{results['clean_total']})"
    )

    # Epsilon spread table
    print(
        f"\n{'Epsilon':<10} {'FGSM':>10} {'I-FGSM':>10} {'FGSM Attack':>12} {'I-FGSM Attack':>14}"
    )
    print("-" * 60)

    for eps in epsilon_spread:
        fgsm_acc = fgsm_results[eps]
        ifgsm_acc = ifgsm_results[eps]
        fgsm_attack = 100.0 - fgsm_acc
        ifgsm_attack = 100.0 - ifgsm_acc
        marker = " *" if eps == PRIMARY_EPSILON else ""
        print(
            f"{eps:<10.2f} {fgsm_acc:>9.1f}% {ifgsm_acc:>9.1f}% {fgsm_attack:>11.1f}% {ifgsm_attack:>13.1f}%{marker}"
        )

    print("-" * 60)
    print("* Primary epsilon for summary metrics")

    # Summary at primary epsilon
    primary_fgsm = fgsm_results[PRIMARY_EPSILON]
    primary_ifgsm = ifgsm_results[PRIMARY_EPSILON]

    print(f"\nSummary (ε={PRIMARY_EPSILON}):")
    print(f"  Clean:  {clean_acc:5.1f}%")
    print(
        f"  FGSM:   {primary_fgsm:5.1f}%  (attack success: {100 - primary_fgsm:5.1f}%)"
    )
    print(
        f"  I-FGSM: {primary_ifgsm:5.1f}%  (attack success: {100 - primary_ifgsm:5.1f}%)"
    )

    # Comparison with baseline if provided
    if baseline:
        base_fgsm = baseline["fgsm_results"][PRIMARY_EPSILON]
        base_ifgsm = baseline["ifgsm_results"][PRIMARY_EPSILON]

        fgsm_improve = primary_fgsm - base_fgsm
        ifgsm_improve = primary_ifgsm - base_ifgsm

        print(f"\nImprovement over baseline:")
        print(
            f"  FGSM:   {'+' if fgsm_improve >= 0 else ''}{fgsm_improve:.1f}% ({base_fgsm:.1f}% → {primary_fgsm:.1f}%)"
        )
        print(
            f"  I-FGSM: {'+' if ifgsm_improve >= 0 else ''}{ifgsm_improve:.1f}% ({base_ifgsm:.1f}% → {primary_ifgsm:.1f}%)"
        )


def print_comparison_summary(
    all_results: List[Dict], baseline: Optional[Dict] = None
) -> None:
    """Print a full epsilon spread comparison table of all models."""
    epsilon_spread = all_results[0]["epsilon_spread"]

    print(f"\n{'=' * 90}")
    print("COMPARISON SUMMARY - Full Epsilon Spread")
    print(f"{'=' * 90}")

    # Build model list (baseline first if present)
    models_to_show = []
    if baseline:
        models_to_show.append(("baseline", baseline))
    for r in all_results:
        models_to_show.append((r["name"], r))

    # Print clean accuracy row
    print(f"\n{'Clean Accuracy:':<20}", end="")
    for name, r in models_to_show:
        short_name = name[:18] if len(name) > 18 else name
        print(f"{short_name:>18}", end="")
    print()
    print(f"{'':<20}", end="")
    for _, r in models_to_show:
        print(f"{r['clean_acc']:>17.1f}%", end="")
    print()

    # FGSM table - Model Accuracy (defender's view)
    print(f"\n{'FGSM Model Accuracy (defender):'}")
    print("-" * (20 + 18 * len(models_to_show)))
    print(f"{'Epsilon':<20}", end="")
    for name, _ in models_to_show:
        short_name = name[:18] if len(name) > 18 else name
        print(f"{short_name:>18}", end="")
    print()
    print("-" * (20 + 18 * len(models_to_show)))

    for eps in epsilon_spread:
        marker = " *" if eps == PRIMARY_EPSILON else ""
        print(f"{eps:<20.2f}", end="")
        for _, r in models_to_show:
            print(f"{r['fgsm_results'][eps]:>17.1f}%", end="")
        print(marker)

    # FGSM table - Attack Success (attacker's view)
    print(f"\n{'FGSM Attack Success (attacker):'}")
    print("-" * (20 + 18 * len(models_to_show)))
    print(f"{'Epsilon':<20}", end="")
    for name, _ in models_to_show:
        short_name = name[:18] if len(name) > 18 else name
        print(f"{short_name:>18}", end="")
    print()
    print("-" * (20 + 18 * len(models_to_show)))

    for eps in epsilon_spread:
        marker = " *" if eps == PRIMARY_EPSILON else ""
        print(f"{eps:<20.2f}", end="")
        for _, r in models_to_show:
            attack_success = 100.0 - r["fgsm_results"][eps]
            print(f"{attack_success:>17.1f}%", end="")
        print(marker)

    # I-FGSM table - Model Accuracy (defender's view)
    print(f"\n{'I-FGSM Model Accuracy (defender):'}")
    print("-" * (20 + 18 * len(models_to_show)))
    print(f"{'Epsilon':<20}", end="")
    for name, _ in models_to_show:
        short_name = name[:18] if len(name) > 18 else name
        print(f"{short_name:>18}", end="")
    print()
    print("-" * (20 + 18 * len(models_to_show)))

    for eps in epsilon_spread:
        marker = " *" if eps == PRIMARY_EPSILON else ""
        print(f"{eps:<20.2f}", end="")
        for _, r in models_to_show:
            print(f"{r['ifgsm_results'][eps]:>17.1f}%", end="")
        print(marker)

    # I-FGSM table - Attack Success (attacker's view)
    print(f"\n{'I-FGSM Attack Success (attacker):'}")
    print("-" * (20 + 18 * len(models_to_show)))
    print(f"{'Epsilon':<20}", end="")
    for name, _ in models_to_show:
        short_name = name[:18] if len(name) > 18 else name
        print(f"{short_name:>18}", end="")
    print()
    print("-" * (20 + 18 * len(models_to_show)))

    for eps in epsilon_spread:
        marker = " *" if eps == PRIMARY_EPSILON else ""
        print(f"{eps:<20.2f}", end="")
        for _, r in models_to_show:
            attack_success = 100.0 - r["ifgsm_results"][eps]
            print(f"{attack_success:>17.1f}%", end="")
        print(marker)

    print("-" * (20 + 18 * len(models_to_show)))
    print("* Primary epsilon for summary metrics")

    # Improvement summary at primary epsilon (if baseline exists)
    if baseline:
        print(f"\nImprovement over baseline at ε={PRIMARY_EPSILON}:")
        base_fgsm = baseline["fgsm_results"][PRIMARY_EPSILON]
        base_ifgsm = baseline["ifgsm_results"][PRIMARY_EPSILON]
        for r in all_results:
            fgsm_d = r["fgsm_results"][PRIMARY_EPSILON] - base_fgsm
            ifgsm_d = r["ifgsm_results"][PRIMARY_EPSILON] - base_ifgsm
            print(
                f"  {r['name']}: FGSM {'+' if fgsm_d >= 0 else ''}{fgsm_d:.1f}%, I-FGSM {'+' if ifgsm_d >= 0 else ''}{ifgsm_d:.1f}%"
            )

    print("=" * 90)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate adversarial robustness of trained models",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python evaluate_robustness.py --model-path robust_model.pth
    python evaluate_robustness.py --model-path models/
    python evaluate_robustness.py --model-path . --compare
        """,
    )
    parser.add_argument(
        "--model-path",
        type=str,
        required=True,
        help="Path to model file (.safetensors) or directory containing models",
    )
    parser.add_argument(
        "--compare",
        action="store_true",
        help="Compare against baseline_model.safetensors",
    )
    parser.add_argument(
        "--show-failures",
        action="store_true",
        help="Show details of misclassified samples",
    )
    args = parser.parse_args()

    # Import torch after argparse for instant --help
    import torch

    model_path = Path(args.model_path)
    if not model_path.exists():
        print(f"Error: Path not found: {model_path}")
        return 1

    # Find models to evaluate
    models = find_models(model_path)
    if not models:
        print(f"Error: No .safetensors model files found in {model_path}")
        return 1

    # Find adversarial examples
    adv_path = find_adv_examples(model_path)
    if not adv_path:
        print("Error: Could not find adv_examples.safetensors")
        return 1

    # Find baseline if comparing
    baseline_path = None
    if args.compare:
        baseline_path = find_baseline(model_path)
        if not baseline_path:
            print(
                "Warning: --compare specified but baseline_model.safetensors not found"
            )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Load adversarial examples
    print(f"Loading adversarial examples from: {adv_path.name}")
    adv_data = load_adversarial_examples(adv_path)

    num_samples = adv_data["clean_images"].size(0)
    has_spread = "epsilon_spread" in adv_data
    print(f"  {num_samples} test samples")
    if has_spread:
        print(f"  Epsilon spread: {adv_data['epsilon_spread']}")
    else:
        print(f"  Single epsilon: {adv_data['epsilon']}")

    # Evaluate baseline first if comparing
    baseline_results = None
    if baseline_path:
        print(f"\nEvaluating baseline...")
        baseline_results = evaluate_model_full(baseline_path, adv_data, device)
        print_model_results(baseline_results)

    # Evaluate each model
    all_results = []
    for mp in models:
        print(f"\nEvaluating {mp.name}...")
        results = evaluate_model_full(mp, adv_data, device)
        all_results.append(results)
        print_model_results(results, baseline_results)

        # Show failures if requested
        if args.show_failures:
            eps_spread = results["epsilon_spread"]
            fgsm_by_eps = adv_data.get(
                "fgsm_by_epsilon", {adv_data["epsilon"]: adv_data.get("fgsm_images")}
            )
            ifgsm_by_eps = adv_data.get(
                "ifgsm_by_epsilon", {adv_data["epsilon"]: adv_data.get("ifgsm_images")}
            )

            failures = get_misclassified_samples(
                results["model"],
                ifgsm_by_eps[PRIMARY_EPSILON],
                adv_data["clean_labels"],
                device,
            )
            if failures:
                print(f"\nMisclassified samples (I-FGSM ε={PRIMARY_EPSILON}):")
                for f in failures:
                    print(
                        f"  Sample {f['index']}: {f['true_label']} → {f['predicted']} (conf={f['confidence']:.2f})"
                    )

    # Print comparison summary if multiple models or baseline comparison
    if len(all_results) > 1 or baseline_results:
        print_comparison_summary(all_results, baseline_results)

    return 0


if __name__ == "__main__":
    exit(main())