"""
Caltech Camera Traps benchmark for cross-location category classification and empty-vs-non-empty scoring.

Dataset: https://beerys.github.io/CaltechCameraTraps/

Generated by gpt-5.3-codex xhigh.
"""

import argparse
import json
import logging
import time
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt
import polars as pl
import torch
from rich.console import Console
from rich.table import Table
from sklearn.metrics import f1_score

from birder.common import cli
from birder.common import lib
from birder.conf import settings
from birder.datahub.evaluation import CaltechCameraTraps
from birder.eval._embeddings import load_embeddings
from birder.eval.methods.knn import evaluate_knn
from birder.eval.methods.simpleshot import evaluate_simpleshot

logger = logging.getLogger(__name__)

SplitData = tuple[npt.NDArray[np.float32], npt.NDArray[np.int_], npt.NDArray[np.float32], npt.NDArray[np.int_]]


def _print_summary_table(results: list[dict[str, Any]], k_values: list[int]) -> None:
    console = Console()

    table = Table(show_header=True, header_style="bold dark_magenta")
    table.add_column("CCT (KNN and SimpleShot)", style="dim")
    for k in k_values:
        table.add_column(f"Cross-location k={k}", justify="right")

    table.add_column("Non-empty F1", justify="right")

    for result in results:
        row = [Path(result["embeddings_file"]).name]
        for k in k_values:
            acc = result["cross_location_accuracies"].get(k)
            row.append(f"{acc:.4f}" if acc is not None else "-")

        row.append(f"{result['non_empty_f1']:.4f}")
        table.add_row(*row)

    console.print(table)


def _write_results_csv(results: list[dict[str, Any]], k_values: list[int], output_path: Path) -> None:
    rows: list[dict[str, Any]] = []
    for result in results:
        row: dict[str, Any] = {
            "embeddings_file": result["embeddings_file"],
            "method": result["method"],
            "split_set": result["split_set"],
            "non_empty_f1": result["non_empty_f1"],
            "train_locations": result["train_locations"],
            "test_locations": result["test_locations"],
            "category_train_samples": result["category_train_samples"],
            "category_test_samples": result["category_test_samples"],
            "non_empty_train_samples": result["non_empty_train_samples"],
            "non_empty_test_samples": result["non_empty_test_samples"],
            "category_train_classes": result["category_train_classes"],
            "category_test_classes": result["category_test_classes"],
            "category_unseen_test_classes": result["category_unseen_test_classes"],
        }
        for k in k_values:
            row[f"k_{k}_cross_location_acc"] = result["cross_location_accuracies"].get(k)

        rows.append(row)

    pl.DataFrame(rows).write_csv(output_path)
    logger.info(f"Results saved to {output_path}")


def _load_split_locations(dataset: CaltechCameraTraps, split_set: str) -> tuple[list[str], list[str]]:
    with open(dataset.splits_path, encoding="utf-8") as f:
        split_payload = json.load(f)

    split_map = {
        "main": ("train", "val"),
        "eccv": ("eccv_train", "eccv_val"),
    }
    train_split_name, test_split_name = split_map[split_set]

    split_dict = split_payload["splits"]
    train_raw = [str(loc) for loc in split_dict[train_split_name]]
    test_raw = [str(loc) for loc in split_dict[test_split_name]]

    train_locations = sorted(set(train_raw))
    test_locations = sorted(set(test_raw))

    dropped_train_duplicates = len(train_raw) - len(train_locations)
    dropped_test_duplicates = len(test_raw) - len(test_locations)
    if dropped_train_duplicates > 0 or dropped_test_duplicates > 0:
        logger.info(
            f"Deduplicated split locations: train removed={dropped_train_duplicates}, "
            f"test removed={dropped_test_duplicates}"
        )

    overlap = set(train_locations).intersection(test_locations)
    if len(overlap) > 0:
        raise RuntimeError(f"Split leakage detected ({len(overlap)} overlapping locations) for split set '{split_set}'")

    return (train_locations, test_locations)


def _load_cct_metadata(dataset: CaltechCameraTraps, split_set: str) -> tuple[pl.DataFrame, list[str], list[str]]:
    logger.info(f"Loading CCT annotations from {dataset.annotations_path}")
    with open(dataset.annotations_path, encoding="utf-8") as f:
        payload = json.load(f)

    categories = payload["categories"]
    empty_category = [row for row in categories if row["name"] == "empty"]
    if len(empty_category) != 1:
        raise RuntimeError(f"Expected exactly one 'empty' category, found {len(empty_category)}")

    empty_category_id = int(empty_category[0]["id"])

    images_df = pl.DataFrame(payload["images"]).select(
        pl.col("id").cast(pl.Utf8),
        pl.col("location").cast(pl.Utf8),
    )
    annotations_df = pl.DataFrame(payload["annotations"]).select(
        pl.col("image_id").cast(pl.Utf8),
        pl.col("category_id").cast(pl.Int64),
    )
    labels_df = annotations_df.group_by("image_id").agg(
        pl.col("category_id").first().alias("label"),
        pl.col("category_id").n_unique().alias("num_labels"),
        (pl.col("category_id") != empty_category_id).any().alias("is_non_empty"),
    )

    metadata_df = images_df.join(labels_df, left_on="id", right_on="image_id", how="inner")

    train_locations, test_locations = _load_split_locations(dataset, split_set)
    metadata_df = metadata_df.with_columns(
        pl.when(pl.col("location").is_in(train_locations))
        .then(pl.lit("train"))
        .when(pl.col("location").is_in(test_locations))
        .then(pl.lit("test"))
        .otherwise(pl.lit("ignore"))
        .alias("split"),
        (pl.col("num_labels") == 1).alias("is_single_label"),
    )

    ignored_rows = metadata_df.filter(pl.col("split") == "ignore").height
    if ignored_rows > 0:
        logger.warning(f"Ignoring {ignored_rows} images from locations not present in selected split set")

    metadata_df = metadata_df.filter(pl.col("split") != "ignore")
    logger.info(
        f"Metadata loaded: images={metadata_df.height}, train_locations={len(train_locations)}, "
        f"test_locations={len(test_locations)}"
    )

    return (metadata_df, train_locations, test_locations)


def _load_embeddings_with_split(
    embeddings_path: str, metadata_df: pl.DataFrame
) -> tuple[SplitData, SplitData, dict[str, int]]:
    logger.info(f"Loading embeddings from {embeddings_path}")
    emb_df = load_embeddings(embeddings_path)

    train_meta = metadata_df.filter(pl.col("split") == "train").select(
        ["id", "label", "is_non_empty", "is_single_label"]
    )
    test_meta = metadata_df.filter(pl.col("split") == "test").select(["id", "label", "is_non_empty", "is_single_label"])

    train_join = train_meta.join(emb_df, on="id", how="inner")
    test_join = test_meta.join(emb_df, on="id", how="inner")

    dropped_train = train_meta.height - train_join.height
    dropped_test = test_meta.height - test_join.height
    if dropped_train > 0 or dropped_test > 0:
        logger.warning(
            f"Join dropped {dropped_train + dropped_test} samples (missing embeddings): "
            f"train={dropped_train}, test={dropped_test}"
        )

    train_category = train_join.filter(pl.col("is_single_label"))
    test_category = test_join.filter(pl.col("is_single_label"))

    x_train_category = train_category.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_train_category = train_category.get_column("label").to_numpy().astype(np.int_, copy=False)
    x_test_category = test_category.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_test_category = test_category.get_column("label").to_numpy().astype(np.int_, copy=False)

    x_train_non_empty = train_join.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_train_non_empty = train_join.get_column("is_non_empty").cast(pl.Int64).to_numpy().astype(np.int_, copy=False)
    x_test_non_empty = test_join.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_test_non_empty = test_join.get_column("is_non_empty").cast(pl.Int64).to_numpy().astype(np.int_, copy=False)

    embedding_dim = x_train_category.shape[1] if x_train_category.size > 0 else x_train_non_empty.shape[1]
    total_samples = train_join.height + test_join.height
    logger.info(f"Loaded {total_samples} samples with {embedding_dim} dimensions")
    logger.info(
        f"Category classification (single-label only) - Train: {train_category.height} samples, "
        f"Test: {test_category.height} samples | "
        f"Non-empty classification - Train: {train_join.height} samples, Test: {test_join.height} samples"
    )

    stats = {
        "category_train_samples": int(train_category.height),
        "category_test_samples": int(test_category.height),
        "non_empty_train_samples": int(train_join.height),
        "non_empty_test_samples": int(test_join.height),
        "category_train_classes": int(np.unique(y_train_category).shape[0]) if y_train_category.size > 0 else 0,
        "category_test_classes": int(np.unique(y_test_category).shape[0]) if y_test_category.size > 0 else 0,
        "category_unseen_test_classes": (
            int(np.setdiff1d(np.unique(y_test_category), np.unique(y_train_category)).shape[0])
            if y_test_category.size > 0 and y_train_category.size > 0
            else 0
        ),
    }

    category_data: SplitData = (x_train_category, y_train_category, x_test_category, y_test_category)
    non_empty_data: SplitData = (x_train_non_empty, y_train_non_empty, x_test_non_empty, y_test_non_empty)

    return (category_data, non_empty_data, stats)


def evaluate_cct_single(
    category_data: SplitData,
    non_empty_data: SplitData,
    embeddings_path: str,
    split_set: str,
    train_locations: int,
    test_locations: int,
    stats: dict[str, int],
    k_values: list[int],
    device: torch.device,
) -> dict[str, Any]:
    x_train_category, y_train_category, x_test_category, y_test_category = category_data
    x_train_non_empty, y_train_non_empty, x_test_non_empty, y_test_non_empty = non_empty_data

    if stats["category_unseen_test_classes"] > 0:
        logger.warning(
            f"Found {stats['category_unseen_test_classes']} test classes not present in train locations; "
            "these classes will be counted as errors in cross-location category accuracy"
        )

    cross_location_accuracies: dict[int, float] = {}
    for k in k_values:
        logger.info(f"Evaluating KNN with k={k}")
        y_pred_category, y_true_category = evaluate_knn(
            x_train_category, y_train_category, x_test_category, y_test_category, k=k, device=device
        )
        acc = float(np.mean(y_pred_category == y_true_category))
        cross_location_accuracies[k] = acc
        logger.info(f"k={k} - Cross-location category accuracy: {acc:.4f}")

    y_pred_non_empty, y_true_non_empty = evaluate_simpleshot(
        x_train_non_empty, y_train_non_empty, x_test_non_empty, y_test_non_empty
    )
    non_empty_f1 = float(f1_score(y_true_non_empty, y_pred_non_empty, average="binary", zero_division=0.0, pos_label=1))
    logger.info(f"Non-empty F1 (SimpleShot): {non_empty_f1:.4f}")

    return {
        "method": "knn+simpleshot",
        "split_set": split_set,
        "cross_location_accuracies": cross_location_accuracies,
        "non_empty_f1": non_empty_f1,
        "embeddings_file": str(embeddings_path),
        "train_locations": train_locations,
        "test_locations": test_locations,
        **stats,
    }


def evaluate_cct(args: argparse.Namespace) -> None:
    tic = time.time()

    if args.gpu is True:
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    if args.gpu_id is not None:
        torch.cuda.set_device(args.gpu_id)

    logger.info(f"Using device {device}")
    logger.info(f"Loading Caltech Camera Traps dataset from {args.dataset_path}")
    dataset = CaltechCameraTraps(args.dataset_path, require_images=False)
    metadata_df, train_locations, test_locations = _load_cct_metadata(dataset, args.split_set)

    results: list[dict[str, Any]] = []
    total = len(args.embeddings)
    for idx, embeddings_path in enumerate(args.embeddings, start=1):
        logger.info(f"Processing embeddings {idx}/{total}: {embeddings_path}")
        category_data, non_empty_data, stats = _load_embeddings_with_split(embeddings_path, metadata_df)
        result = evaluate_cct_single(
            category_data,
            non_empty_data,
            embeddings_path=embeddings_path,
            split_set=args.split_set,
            train_locations=len(train_locations),
            test_locations=len(test_locations),
            stats=stats,
            k_values=args.k,
            device=device,
        )
        results.append(result)

    _print_summary_table(results, args.k)

    if args.dry_run is False:
        output_dir = settings.RESULTS_DIR.joinpath(args.dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir.joinpath("cct.csv")
        _write_results_csv(results, args.k, output_path)

    toc = time.time()
    logger.info(f"CCT benchmark completed in {lib.format_duration(toc - tic)}")


def set_parser(subparsers: Any) -> None:
    subparser = subparsers.add_parser(
        "cct",
        allow_abbrev=False,
        help=(
            "run CCT benchmark - cross-location category classification (KNN) + "
            "empty-vs-non-empty detection (SimpleShot)"
        ),
        description=(
            "run CCT benchmark - cross-location category classification (KNN) + "
            "empty-vs-non-empty detection (SimpleShot)"
        ),
        epilog=(
            "Usage examples:\n"
            "python -m birder.eval cct --embeddings "
            "results/cct_embeddings.parquet --dataset-path ~/Datasets/CaltechCameraTraps --dry-run\n"
            "python -m birder.eval cct --embeddings results/cct/*.parquet "
            "--dataset-path ~/Datasets/CaltechCameraTraps --gpu\n"
        ),
        formatter_class=cli.ArgumentHelpFormatter,
    )
    subparser.add_argument(
        "--embeddings", type=str, nargs="+", metavar="FILE", help="paths to embeddings parquet files"
    )
    subparser.add_argument("--dataset-path", type=str, metavar="PATH", help="path to Caltech Camera Traps dataset root")
    subparser.add_argument(
        "--split-set",
        type=str,
        choices=["main", "eccv"],
        default="main",
        help="which official location split set to evaluate",
    )
    subparser.add_argument(
        "--k", type=int, nargs="+", default=[10, 20, 100], help="number of nearest neighbors for KNN"
    )
    subparser.add_argument("--gpu", default=False, action="store_true", help="use gpu")
    subparser.add_argument("--gpu-id", type=int, metavar="ID", help="gpu id to use")
    subparser.add_argument(
        "--dir", type=str, default="cct", help="place all outputs in a sub-directory (relative to results)"
    )
    subparser.add_argument("--dry-run", default=False, action="store_true", help="skip saving results to file")
    subparser.set_defaults(func=main)


def validate_args(args: argparse.Namespace) -> None:
    if args.embeddings is None:
        raise cli.ValidationError("--embeddings is required")
    if args.dataset_path is None:
        raise cli.ValidationError("--dataset-path is required")


def main(args: argparse.Namespace) -> None:
    validate_args(args)
    evaluate_cct(args)
