"""Train local damage-classification models on hand-crafted image features.

Five sklearn classifiers are compared (Dummy, Logistic Regression, SVM RBF,
Random Forest, Gradient Boosting) using 5-fold stratified cross-validation.
All features are extracted locally — no images are sent to any external API.

This script represents the hand-crafted feature baseline. For the CNN
transfer learning approach (EfficientNet B0, 1280-dim features) see
train_cnn_cv_model.py.

IMPORTANT: This script trains on a local representative subset only.
Do NOT add more than ~320 images. Use prepare_local_cv_dataset.py
--max-per-class 80 to prepare the balanced 240-image subset.

Expected folder structure:

data/cv_damage/
  no visible damage/
    image_001.jpg
  minor damage/
    image_002.jpg
  moderate damage/
    image_003.jpg
  severe damage/
    image_004.jpg
"""

from __future__ import annotations

import argparse
import json
import sys
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.ensemble import GradientBoostingClassifier, RandomForestClassifier
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 app.damage_model import LOCAL_DAMAGE_MODEL_PATH, extract_local_cv_features


DEFAULT_INPUT = PROJECT_ROOT / "data/cv_damage"
REPORT_PATH = PROJECT_ROOT / "reports/local_cv_model_report.md"
METRICS_PATH = PROJECT_ROOT / "reports/local_cv_model_metrics.json"
CONFUSION_PATH = PROJECT_ROOT / "reports/local_cv_confusion_matrix.csv"
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp"}


def collect_labeled_images(input_dir: Path) -> list[tuple[Path, str]]:
    """Collect images from class-named subfolders."""
    if not input_dir.exists():
        return []
    rows: list[tuple[Path, str]] = []
    for label_dir in sorted(path for path in input_dir.iterdir() if path.is_dir()):
        label = label_dir.name.strip()
        for image_path in sorted(label_dir.rglob("*")):
            if image_path.suffix.lower() in IMAGE_EXTENSIONS:
                rows.append((image_path, label))
    return rows


def build_feature_matrix(rows: list[tuple[Path, str]]) -> tuple[np.ndarray, np.ndarray, list[str]]:
    """Extract deterministic image features for all labeled images."""
    features = []
    labels = []
    paths = []
    for image_path, label in rows:
        features.append(extract_local_cv_features(image_path))
        labels.append(label)
        paths.append(str(image_path.relative_to(PROJECT_ROOT)))
    return np.vstack(features), np.asarray(labels), paths


def split_data(
    x: np.ndarray,
    y: np.ndarray,
    test_size: float,
    random_state: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Split data, stratifying only when class counts allow it."""
    counts = Counter(y)
    stratify = y if len(counts) > 1 and min(counts.values()) >= 2 else None
    return train_test_split(
        x,
        y,
        test_size=test_size,
        random_state=random_state,
        stratify=stratify,
    )


def train_models(
    x_train: np.ndarray,
    y_train: np.ndarray,
) -> dict[str, Pipeline]:
    """Train and return all candidate classifiers (fitted on x_train/y_train)."""
    candidates: dict[str, Pipeline] = {
        "dummy_most_frequent": Pipeline([
            ("scaler", StandardScaler()),
            ("model", DummyClassifier(strategy="most_frequent")),
        ]),
        "logistic_regression": Pipeline([
            ("scaler", StandardScaler()),
            ("model", LogisticRegression(
                max_iter=1000, class_weight="balanced", random_state=42,
            )),
        ]),
        "svm_rbf": Pipeline([
            ("scaler", StandardScaler()),
            ("model", SVC(
                kernel="rbf", C=5.0, gamma="scale",
                class_weight="balanced", random_state=42,
                probability=True,
            )),
        ]),
        "random_forest": Pipeline([
            ("scaler", StandardScaler()),
            ("model", RandomForestClassifier(
                n_estimators=100, class_weight="balanced",
                random_state=42, n_jobs=-1,
            )),
        ]),
        "gradient_boosting": Pipeline([
            ("scaler", StandardScaler()),
            ("model", GradientBoostingClassifier(
                n_estimators=80, learning_rate=0.1,
                max_depth=3, random_state=42,
            )),
        ]),
    }
    for clf in candidates.values():
        clf.fit(x_train, y_train)
    return candidates


def run_cross_validation(
    candidates: dict[str, Pipeline],
    x_train: np.ndarray,
    y_train: np.ndarray,
    n_splits: int = 5,
    random_state: int = 42,
) -> dict[str, dict[str, float]]:
    """Evaluate all classifiers with stratified k-fold CV on the training set."""
    import warnings
    cv = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state)
    cv_results: dict[str, dict[str, float]] = {}
    for name, clf in candidates.items():
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            scores = cross_validate(clf, x_train, y_train, cv=cv,
                                    scoring=["f1_macro", "accuracy"])
        cv_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()),
        }
    return cv_results


def evaluate_model(model: Pipeline, x_test: np.ndarray, y_test: np.ndarray) -> dict[str, float]:
    """Calculate classification metrics."""
    predictions = model.predict(x_test)
    return {
        "accuracy": accuracy_score(y_test, predictions),
        "precision_macro": precision_score(y_test, predictions, average="macro", zero_division=0),
        "recall_macro": recall_score(y_test, predictions, average="macro", zero_division=0),
        "f1_macro": f1_score(y_test, predictions, average="macro", zero_division=0),
    }


def write_empty_report(input_dir: Path) -> None:
    """Write an actionable report when no images are available yet."""
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
    lines = [
        "# Local CV Model Report",
        "",
        "No local CV model was trained because no labeled images were found.",
        "",
        f"Expected folder: `{input_dir.relative_to(PROJECT_ROOT)}`",
        "",
        "Create class folders such as:",
        "",
        "```text",
        "data/cv_damage/",
        "  no visible damage/",
        "  minor damage/",
        "  moderate damage/",
        "  severe damage/",
        "```",
        "",
        "Then run:",
        "",
        "```bash",
        "python scripts/train_local_cv_model.py",
        "```",
    ]
    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")
    METRICS_PATH.write_text(json.dumps({"trained": False, "reason": "no labeled images"}, indent=2), encoding="utf-8")


def write_report(
    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,
    cv_results: dict[str, dict[str, float]] | None = None,
    n_splits: int = 5,
) -> None:
    """Write markdown, JSON metrics, and confusion matrix outputs."""
    REPORT_PATH.parent.mkdir(parents=True, exist_ok=True)
    METRICS_PATH.write_text(
        json.dumps(
            {
                "trained": True,
                "dataset_size": dataset_size,
                "class_counts": dict(class_counts),
                "selected_model": selected_model_name,
                "metrics": metrics,
                "cv_results": cv_results,
            },
            indent=2,
        ),
        encoding="utf-8",
    )

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

    lines = [
        "# Local CV Model Report (Hand-Crafted Features)",
        "",
        "This report documents the reproducible hand-crafted feature baseline for local",
        "vehicle damage classification. Five sklearn classifiers are compared with",
        "5-fold stratified cross-validation. No images are sent to any external API.",
        "",
        f"- Dataset size: {dataset_size} images (representative subset)",
        "- Feature type: hand-crafted (colour histograms, edge statistics) — 53 dimensions",
        f"- Evaluation: {n_splits}-fold stratified CV + held-out test set",
        f"- Selected model: `{selected_model_name}`",
        f"- Saved artifact: `{LOCAL_DAMAGE_MODEL_PATH.relative_to(PROJECT_ROOT)}`",
        "",
        "## Class Distribution",
        "",
        "| Class | Images |",
        "|---|---:|",
    ]
    for label, count in sorted(class_counts.items()):
        lines.append(f"| {label} | {count} |")

    if cv_results:
        lines.extend([
            "",
            f"## {n_splits}-Fold Stratified Cross-Validation (Training Set)",
            "",
            "| Model | 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",
            "",
            "| Model | Accuracy | Precision macro | Recall macro | F1 macro |",
            "|---|---:|---:|---:|---:|",
        ]
    )
    for model_name, values in metrics.items():
        lines.append(
            f"| {model_name} | {values['accuracy']:.3f} | {values['precision_macro']:.3f} | "
            f"{values['recall_macro']:.3f} | {values['f1_macro']:.3f} |"
        )

    lines.extend(
        [
            "",
            "## Classification Report",
            "",
            "```text",
            report_text,
            "```",
            "",
            "## Interpretation",
            "",
            "The local model uses hand-crafted image statistics as features. It provides a",
            "reproducible training/evaluation baseline and is intentionally less expressive",
            "than the OpenAI Vision path used in the deployed app. It estimates image-level",
            "damage classes only; it does not estimate bounding boxes or repair costs.",
            "Outputs are converted to the same transparent `damage_score` logic used by the",
            "deployed app.",
        ]
    )
    REPORT_PATH.write_text("\n".join(lines), encoding="utf-8")


def parse_args() -> argparse.Namespace:
    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("--random-state", type=int, default=42)
    parser.add_argument("--cv-splits", type=int, default=5)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    rows = collect_labeled_images(args.input_dir)
    if len(rows) < 8 or len({label for _, label in rows}) < 2:
        write_empty_report(args.input_dir)
        print(f"Not enough labeled images. Wrote {REPORT_PATH}")
        return

    x, y, _ = build_feature_matrix(rows)
    x_train, x_test, y_train, y_test = split_data(
        x, y, test_size=args.test_size, random_state=args.random_state,
    )

    models = train_models(x_train, y_train)

    print(f"Running {args.cv_splits}-fold stratified CV on training set...")
    cv_results = run_cross_validation(
        models, x_train, y_train, n_splits=args.cv_splits, random_state=args.random_state,
    )
    for name, res in cv_results.items():
        print(f"  {name}: F1={res['f1_macro_mean']:.3f} ± {res['f1_macro_std']:.3f}")

    metrics = {name: evaluate_model(model, x_test, y_test) for name, model in models.items()}
    selected_model_name = max(cv_results, key=lambda name: cv_results[name]["f1_macro_mean"])
    selected_model = models[selected_model_name]
    predictions = selected_model.predict(x_test)
    labels = sorted(set(y))

    LOCAL_DAMAGE_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    joblib.dump(
        {
            "pipeline": selected_model,
            "labels": labels,
            "feature_extractor": "extract_local_cv_features",
            "selected_model": selected_model_name,
        },
        LOCAL_DAMAGE_MODEL_PATH,
    )
    write_report(
        metrics=metrics,
        selected_model_name=selected_model_name,
        y_test=y_test,
        predictions=predictions,
        labels=labels,
        dataset_size=len(rows),
        class_counts=Counter(y),
        cv_results=cv_results,
        n_splits=args.cv_splits,
    )
    print(f"Wrote {LOCAL_DAMAGE_MODEL_PATH}")
    print(f"Wrote {REPORT_PATH}")


if __name__ == "__main__":
    main()