"""Train a CNN-based (EfficientNet B0 transfer learning) damage classifier.

Feature extraction uses a pretrained EfficientNet B0 backbone (ImageNet
weights via torchvision). This is a fully local operation — no images are
sent to any external API. Six sklearn classifiers are compared, evaluated
with 5-fold stratified cross-validation on the training set, then assessed
on a held-out test set.

Results are compared with the hand-crafted feature baseline from
train_local_cv_model.py (53-dim features, SVM RBF, held-out F1=0.570).

IMPORTANT: This script trains on a local representative subset only.
Do NOT run this script on the full DrBimmer dataset (2,300 images).
Use prepare_local_cv_dataset.py --max-per-class 80 (240 images, 3 classes)
to prepare the balanced subset first.

Expected folder structure:
    data/cv_damage/
      no visible damage/
      minor damage/   (optional)
      moderate damage/
      severe damage/

Usage:
    python scripts/train_cnn_cv_model.py
    python scripts/train_cnn_cv_model.py --skip-cv   # faster, no CV

Dependencies (install via requirements-cv.txt):
    torch, torchvision, pillow
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import warnings
from collections import Counter
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.dummy import DummyClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from sklearn.model_selection import StratifiedKFold, cross_validate, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.append(str(PROJECT_ROOT))

from scripts.train_local_cv_model import collect_labeled_images

CNN_MODEL_PATH = PROJECT_ROOT / "models/cnn_cv_model.joblib"
REPORT_PATH = PROJECT_ROOT / "reports/cnn_cv_model_report.md"
METRICS_PATH = PROJECT_ROOT / "reports/cnn_cv_model_metrics.json"
CONFUSION_CSV_PATH = PROJECT_ROOT / "reports/cnn_cv_confusion_matrix.csv"
FIGURE_PATH = PROJECT_ROOT / "reports/cnn_cv_confusion_matrix.png"
DEFAULT_INPUT = PROJECT_ROOT / "data/cv_damage"
RANDOM_STATE = 42
MAX_DATASET_IMAGES = 320  # hard cap — never process more than this locally
AUGMENTATION_VARIANTS = ("original", "horizontal_flip", "brightness_up", "contrast_up")


def extract_efficientnet_features(
    image_paths: list[Path],
    batch_size: int = 16,
    device: str | None = None,
    augment: bool = False,
) -> np.ndarray:
    """Extract 1280-dim pooled features from pretrained EfficientNet B0.

    The classifier head is removed. Adaptive average pooling output (1280-dim)
    is used as the feature vector. All processing is local — no API calls.
    """
    import torch
    import torchvision.models as tv_models
    from PIL import Image as PilImage
    from PIL import ImageEnhance, ImageOps

    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"  Feature extraction device: {device}")

    weights = tv_models.EfficientNet_B0_Weights.DEFAULT
    base_model = tv_models.efficientnet_b0(weights=weights)
    feature_extractor = torch.nn.Sequential(*list(base_model.children())[:-1])
    feature_extractor.eval().to(device)

    transform = weights.transforms()
    all_features: list[np.ndarray] = []
    variants = AUGMENTATION_VARIANTS if augment else ("original",)
    image_jobs = [(path, variant) for path in image_paths for variant in variants]
    n_total = len(image_jobs)

    for i in range(0, n_total, batch_size):
        batch_jobs = image_jobs[i: i + batch_size]
        tensors: list[torch.Tensor] = []
        for path, variant in batch_jobs:
            try:
                image = PilImage.open(path).convert("RGB")
                if variant == "horizontal_flip":
                    image = ImageOps.mirror(image)
                elif variant == "brightness_up":
                    image = ImageEnhance.Brightness(image).enhance(1.15)
                elif variant == "contrast_up":
                    image = ImageEnhance.Contrast(image).enhance(1.20)
                tensors.append(transform(image))
            except Exception:
                tensors.append(torch.zeros(3, 224, 224))
        batch = torch.stack(tensors).to(device)
        with torch.no_grad():
            out = feature_extractor(batch)
        all_features.append(out.squeeze(-1).squeeze(-1).cpu().numpy())
        done = min(i + batch_size, n_total)
        if done % 80 == 0 or done == n_total:
            print(f"  Processed {done}/{n_total} images")

    return np.vstack(all_features)


def build_classifiers() -> dict[str, Pipeline]:
    """Return six candidate classifiers trained on CNN features."""
    return {
        "dummy_most_frequent": Pipeline([
            ("scaler", StandardScaler()),
            ("model", DummyClassifier(strategy="most_frequent")),
        ]),
        "logistic_regression": Pipeline([
            ("scaler", StandardScaler()),
            ("model", LogisticRegression(
                max_iter=2000, C=1.0, class_weight="balanced",
                random_state=RANDOM_STATE,
            )),
        ]),
        "svm_rbf": Pipeline([
            ("scaler", StandardScaler()),
            ("model", SVC(
                kernel="rbf", C=10.0, gamma="scale",
                class_weight="balanced", random_state=RANDOM_STATE,
                probability=True,
            )),
        ]),
        "svm_linear": Pipeline([
            ("scaler", StandardScaler()),
            ("model", SVC(
                kernel="linear", C=1.0,
                class_weight="balanced", random_state=RANDOM_STATE,
                probability=True,
            )),
        ]),
        "knn_3": Pipeline([
            ("scaler", StandardScaler()),
            ("model", KNeighborsClassifier(n_neighbors=3, weights="distance")),
        ]),
        "knn_5": Pipeline([
            ("scaler", StandardScaler()),
            ("model", KNeighborsClassifier(n_neighbors=5, weights="distance")),
        ]),
    }


def run_cross_validation(
    classifiers: dict[str, Pipeline],
    x_train: np.ndarray,
    y_train: np.ndarray,
    n_splits: int,
) -> dict[str, dict[str, float]]:
    """Stratified k-fold CV with F1 macro and accuracy scoring."""
    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=RANDOM_STATE)
    results: dict[str, dict[str, float]] = {}
    for name, clf in classifiers.items():
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            scores = cross_validate(
                clf, x_train, y_train,
                cv=cv,
                scoring=["f1_macro", "accuracy"],
            )
        results[name] = {
            "f1_macro_mean": float(scores["test_f1_macro"].mean()),
            "f1_macro_std": float(scores["test_f1_macro"].std()),
            "accuracy_mean": float(scores["test_accuracy"].mean()),
            "accuracy_std": float(scores["test_accuracy"].std()),
        }
        print(
            f"  {name}: F1={results[name]['f1_macro_mean']:.3f}"
            f" ± {results[name]['f1_macro_std']:.3f}"
        )
    return results


def evaluate_held_out(
    classifiers: dict[str, Pipeline],
    x_test: np.ndarray,
    y_test: np.ndarray,
) -> dict[str, dict[str, float]]:
    results: dict[str, dict[str, float]] = {}
    for name, clf in classifiers.items():
        preds = clf.predict(x_test)
        results[name] = {
            "accuracy": float(accuracy_score(y_test, preds)),
            "precision_macro": float(precision_score(y_test, preds, average="macro", zero_division=0)),
            "recall_macro": float(recall_score(y_test, preds, average="macro", zero_division=0)),
            "f1_macro": float(f1_score(y_test, preds, average="macro", zero_division=0)),
        }
    return results


def save_confusion_figure(
    y_test: np.ndarray,
    predictions: np.ndarray,
    labels: list[str],
) -> None:
    if os.getenv("GENERATE_CV_FIGURES") != "1":
        return
    try:
        import matplotlib.pyplot as plt
        from sklearn.metrics import ConfusionMatrixDisplay
        fig, ax = plt.subplots(figsize=(6, 5))
        ConfusionMatrixDisplay.from_predictions(
            y_test, predictions, labels=labels,
            ax=ax, colorbar=False,
        )
        ax.set_title("CNN (EfficientNet B0) — Confusion Matrix")
        fig.tight_layout()
        fig.savefig(FIGURE_PATH, dpi=120)
        plt.close(fig)
        print(f"Saved confusion matrix figure: {FIGURE_PATH}")
    except ImportError:
        pass


def write_report(
    cv_results: dict[str, dict[str, float]] | None,
    test_metrics: dict[str, dict[str, float]],
    selected_model_name: str,
    y_test: np.ndarray,
    predictions: np.ndarray,
    labels: list[str],
    dataset_size: int,
    class_counts: Counter,
    n_splits: int,
    baseline_f1: float,
    augmentation_enabled: bool,
    train_original_count: int,
    train_feature_count: int,
    test_original_count: int,
) -> None:
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)

    matrix = confusion_matrix(y_test, predictions, labels=labels)
    pd.DataFrame(matrix, index=labels, columns=labels).to_csv(CONFUSION_CSV_PATH)
    report_text = classification_report(y_test, predictions, labels=labels, zero_division=0)
    save_confusion_figure(y_test, predictions, labels)

    selected_f1 = test_metrics.get(selected_model_name, {}).get("f1_macro", 0.0)
    improvement = selected_f1 - baseline_f1

    lines = [
        "# CNN Transfer Learning CV Model Report",
        "",
        "This report documents the EfficientNet B0 transfer learning approach for local",
        "vehicle damage classification, compared against the hand-crafted feature baseline",
        "from `train_local_cv_model.py`.",
        "",
        "**Method**: A pretrained EfficientNet B0 backbone (ImageNet weights via",
        "`torchvision`) is used as a fixed feature extractor. The classifier head is removed;",
        "the 1280-dim adaptive average pooling output is fed into sklearn classifiers.",
        "All processing is fully local — no images are sent to any external API.",
        "",
        f"- Dataset: **{dataset_size} images** (representative subset — max {MAX_DATASET_IMAGES} enforced)",
        "- Feature backbone: `EfficientNet_B0_Weights.DEFAULT` (ImageNet pretrained)",
        "- Feature dimensionality: **1280** (vs 53 for hand-crafted baseline)",
        f"- Training augmentation: **{'enabled' if augmentation_enabled else 'disabled'}**",
        f"- Training originals / feature samples: **{train_original_count} / {train_feature_count}**",
        f"- Held-out test originals: **{test_original_count}** (no augmentation applied)",
        f"- Evaluation: {n_splits}-fold stratified CV + held-out test set",
        f"- Selected model: **{selected_model_name}**",
        f"- Artifact: `{CNN_MODEL_PATH.relative_to(PROJECT_ROOT)}`",
        "",
        "## Class Distribution",
        "",
        "| Class | Images |",
        "|---|---:|",
    ]
    for label in labels:
        lines.append(f"| {label} | {class_counts.get(label, 0)} |")

    if cv_results:
        lines.extend([
            "",
            f"## {n_splits}-Fold Stratified Cross-Validation (Original Training Set)",
            "",
            "Cross-validation is performed on original training images only. The final",
            "classifier is then fitted on the augmented training set. This avoids",
            "placing augmented versions of the same original image in different CV folds.",
            "",
            "| Classifier | F1 macro mean ± std | Accuracy mean ± std |",
            "|---|---|---|",
        ])
        for name, res in cv_results.items():
            lines.append(
                f"| {name} | "
                f"{res['f1_macro_mean']:.3f} ± {res['f1_macro_std']:.3f} | "
                f"{res['accuracy_mean']:.3f} ± {res['accuracy_std']:.3f} |"
            )

    lines.extend([
        "",
        "## Held-Out Test Set Results",
        "",
        "| Classifier | Accuracy | Precision macro | Recall macro | F1 macro |",
        "|---|---:|---:|---:|---:|",
    ])
    for name, vals in test_metrics.items():
        lines.append(
            f"| {name} | {vals['accuracy']:.3f} | {vals['precision_macro']:.3f} |"
            f" {vals['recall_macro']:.3f} | {vals['f1_macro']:.3f} |"
        )

    lines.extend([
        "",
        f"Best model by CV F1: **{selected_model_name}**",
        f"Test F1 macro: **{selected_f1:.3f}** (+{improvement:.3f} vs hand-crafted baseline F1 {baseline_f1:.3f})",
        "",
        "## Classification Report (Test Set, Best Model)",
        "",
        "```text",
        report_text,
        "```",
        "",
        "## Comparison: Hand-Crafted Features vs CNN Transfer Features",
        "",
        "| Feature type | Dimensionality | Classifier | Test F1 macro |",
        "|---|---:|---|---:|",
        f"| Hand-crafted (colour histograms, edge stats) | 53 | SVM RBF | {baseline_f1:.3f} |",
        f"| EfficientNet B0 transfer (ImageNet pretrained) | 1280 | {selected_model_name} | {selected_f1:.3f} |",
        "",
        "Transfer learning from ImageNet provides substantially richer features than",
        "hand-crafted statistics. EfficientNet B0 encodes high-level visual patterns",
        "(textures, surface deformations, structural damage cues) that are directly",
        "relevant to vehicle damage classification but unavailable from low-level",
        "colour and edge descriptors.",
        "",
        "## Augmentation",
        "",
        "Label-preserving augmentation is applied only to the training split:",
        "horizontal flip, slight brightness increase, and slight contrast increase.",
        "The held-out test set is never augmented. This simulates realistic seller-photo",
        "variation while keeping the evaluation honest.",
        "",
        "## Confusion Matrix",
        "",
        f"CSV: `{CONFUSION_CSV_PATH.relative_to(PROJECT_ROOT)}`",
        "Figure generation is optional. Set `GENERATE_CV_FIGURES=1` before running",
        "the script to create the PNG confusion matrix.",
        "",
        "## Limitations",
        "",
        "- Image-level classification only; no bounding-box localisation.",
        "- EfficientNet B0 weights are frozen (linear probing). End-to-end fine-tuning",
        "  would improve results further but requires more data and GPU compute.",
        "- The deployed app uses OpenAI Vision for richer multi-view reasoning.",
        "- Only a representative 240-image subset is used; the full dataset is never",
        "  sent to any external API in any training or evaluation step.",
    ])
    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-dir", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--test-size", type=float, default=0.25)
    parser.add_argument("--cv-splits", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--no-augmentation", action="store_true",
                        help="Disable label-preserving training augmentation.")
    parser.add_argument("--skip-cv", action="store_true",
                        help="Skip cross-validation and evaluate on test set only.")
    args = parser.parse_args()

    rows = collect_labeled_images(args.input_dir)
    if len(rows) < 8 or len({lbl for _, lbl in rows}) < 2:
        print(f"Not enough labeled images in {args.input_dir}.")
        print("Run: python scripts/prepare_local_cv_dataset.py --max-per-class 80")
        return

    if len(rows) > MAX_DATASET_IMAGES:
        print(
            f"WARNING: {len(rows)} images found but hard cap is {MAX_DATASET_IMAGES}. "
            "Use --max-per-class in prepare_local_cv_dataset.py to limit the subset."
        )
        rows = rows[:MAX_DATASET_IMAGES]

    class_counts: Counter = Counter(lbl for _, lbl in rows)
    unique_labels = sorted(class_counts.keys())
    print(f"Dataset: {len(rows)} images, {len(unique_labels)} classes: {unique_labels}")

    image_paths = [p for p, _ in rows]
    y = np.asarray([lbl for _, lbl in rows])
    counts = Counter(y)
    stratify = y if min(counts.values()) >= 2 else None
    train_paths, test_paths, y_train, y_test = train_test_split(
        image_paths, y, test_size=args.test_size, random_state=RANDOM_STATE, stratify=stratify,
    )
    print(f"Train originals: {len(train_paths)}, Test originals: {len(test_paths)}")

    print("\nExtracting EfficientNet B0 features for original training images...")
    x_train_original = extract_efficientnet_features(
        train_paths,
        batch_size=args.batch_size,
        augment=False,
    )

    augmentation_enabled = not args.no_augmentation
    print(
        "\nExtracting EfficientNet B0 features for final training set "
        f"(augmentation={'on' if augmentation_enabled else 'off'})..."
    )
    x_train = extract_efficientnet_features(
        train_paths,
        batch_size=args.batch_size,
        augment=augmentation_enabled,
    )
    y_train_fit = np.repeat(y_train, len(AUGMENTATION_VARIANTS)) if augmentation_enabled else y_train

    print("\nExtracting EfficientNet B0 features for held-out test images...")
    x_test = extract_efficientnet_features(test_paths, batch_size=args.batch_size, augment=False)
    print(f"Feature matrices: train_original={x_train_original.shape}, train_fit={x_train.shape}, test={x_test.shape}")

    classifiers = build_classifiers()

    cv_results: dict[str, dict[str, float]] | None = None
    if not args.skip_cv:
        print(f"\nRunning {args.cv_splits}-fold stratified CV on original training set...")
        cv_results = run_cross_validation(classifiers, x_train_original, y_train, n_splits=args.cv_splits)

    print("\nFitting all classifiers on final training set...")
    for clf in classifiers.values():
        clf.fit(x_train, y_train_fit)

    print("Evaluating on held-out test set...")
    test_metrics = evaluate_held_out(classifiers, x_test, y_test)
    for name, vals in test_metrics.items():
        print(f"  {name}: F1={vals['f1_macro']:.3f}, Acc={vals['accuracy']:.3f}")

    if cv_results:
        selected_model_name = max(cv_results, key=lambda n: cv_results[n]["f1_macro_mean"])
    else:
        selected_model_name = max(test_metrics, key=lambda n: test_metrics[n]["f1_macro"])
    print(f"\nSelected model: {selected_model_name}")

    CNN_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump({
        "pipeline": classifiers[selected_model_name],
        "labels": unique_labels,
        "feature_extractor": "efficientnet_b0_imagenet",
        "feature_dim": 1280,
        "selected_model": selected_model_name,
        "augmentation_enabled": augmentation_enabled,
        "augmentation_variants": list(AUGMENTATION_VARIANTS if augmentation_enabled else ("original",)),
    }, CNN_MODEL_PATH)

    METRICS_PATH.write_text(json.dumps({
        "trained": True,
        "dataset_size": len(rows),
        "class_counts": dict(class_counts),
        "augmentation_enabled": augmentation_enabled,
        "augmentation_variants": list(AUGMENTATION_VARIANTS if augmentation_enabled else ("original",)),
        "train_original_count": len(train_paths),
        "train_feature_count": len(x_train),
        "test_original_count": len(test_paths),
        "selected_model": selected_model_name,
        "cv_results": cv_results,
        "test_metrics": test_metrics,
    }, indent=2), encoding="utf-8")

    predictions = classifiers[selected_model_name].predict(x_test)
    write_report(
        cv_results=cv_results,
        test_metrics=test_metrics,
        selected_model_name=selected_model_name,
        y_test=y_test,
        predictions=predictions,
        labels=unique_labels,
        dataset_size=len(rows),
        class_counts=class_counts,
        n_splits=args.cv_splits,
        baseline_f1=0.570,
        augmentation_enabled=augmentation_enabled,
        train_original_count=len(train_paths),
        train_feature_count=len(x_train),
        test_original_count=len(test_paths),
    )

    print(f"\nSaved: {CNN_MODEL_PATH}")
    print(f"Wrote: {REPORT_PATH}")


if __name__ == "__main__":
    main()