"""
PlantDoc benchmark using SimpleShot for plant disease classification

Paper "PlantDoc: A Dataset for Visual Plant Disease Detection",
https://arxiv.org/abs/1911.10317
"""

import argparse
import logging
import os
import time
from pathlib import Path
from typing import Any
from typing import Optional

import numpy as np
import numpy.typing as npt
import polars as pl
from rich.console import Console
from rich.table import Table

from birder.common import cli
from birder.common import lib
from birder.conf import settings
from birder.data.datasets.directory import make_image_dataset
from birder.datahub.evaluation import PlantDoc
from birder.eval._embeddings import load_embeddings
from birder.eval.methods.simpleshot import evaluate_simpleshot
from birder.eval.methods.simpleshot import sample_k_shot

logger = logging.getLogger(__name__)


def _results_filename(args: argparse.Namespace) -> str:
    k_values = "-".join(str(k) for k in args.k_shot)
    return f"plantdoc_simpleshot_k{k_values}_runs{args.runs}.csv"


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

    table = Table(show_header=True, header_style="bold dark_magenta")
    table.add_column("PlantDoc (SimpleShot)", style="dim")
    for k in k_shots:
        table.add_column(f"{k}-shot", justify="right")

    table.add_column("Runs", justify="right")

    for result in sorted(results, key=lambda result: Path(result["embeddings_file"]).name):
        row = [Path(result["embeddings_file"]).name]
        for k in k_shots:
            acc = result["accuracies"].get(k)
            row.append(f"{acc:.4f}" if acc is not None else "-")

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

    console.print(table)


def _result_to_row(result: dict[str, Any], k_shots: list[int]) -> dict[str, Any]:
    row: dict[str, Any] = {
        "embeddings_file": result["embeddings_file"],
        "method": result["method"],
        "num_runs": result["num_runs"],
    }
    for k in k_shots:
        row[f"{k}_shot_acc"] = result["accuracies"].get(k)
        row[f"{k}_shot_std"] = result["accuracies_std"].get(k)

    return row


def _append_result_csv(result: dict[str, Any], k_shots: list[int], output_path: Path) -> None:
    row_df = pl.DataFrame([_result_to_row(result, k_shots)])
    file_exists = output_path.exists()
    mode = "a" if file_exists is True else "w"
    with open(output_path, mode, encoding="utf-8") as handle:
        row_df.write_csv(handle, include_header=file_exists is False)

    logger.info(f"Results updated at {output_path}")


def _row_to_result(row: dict[str, Any], k_shots: list[int]) -> dict[str, Any]:
    accuracies = {k: float(value) for k in k_shots if (value := row.get(f"{k}_shot_acc")) is not None}
    accuracies_std = {k: float(value) for k in k_shots if (value := row.get(f"{k}_shot_std")) is not None}

    return {
        "embeddings_file": str(row["embeddings_file"]),
        "method": str(row["method"]),
        "num_runs": int(row["num_runs"]),
        "accuracies": accuracies,
        "accuracies_std": accuracies_std,
    }


def _load_cached_result(
    existing_df: Optional[pl.DataFrame], embeddings_path: str, k_shots: list[int]
) -> Optional[dict[str, Any]]:
    if existing_df is None or existing_df.is_empty() is True:
        return None

    match_df = existing_df.filter(pl.col("embeddings_file") == embeddings_path)
    if match_df.is_empty() is True:
        return None

    return _row_to_result(match_df.row(0, named=True), k_shots)


def _summary_results(
    existing_df: Optional[pl.DataFrame], current_results: list[dict[str, Any]], k_shots: list[int]
) -> list[dict[str, Any]]:
    if existing_df is None or existing_df.is_empty() is True:
        return current_results

    return [_row_to_result(row, k_shots) for row in existing_df.iter_rows(named=True)]


def _load_plantdoc_metadata(dataset: PlantDoc) -> pl.DataFrame:
    """
    Load metadata using make_image_dataset with a fixed class_to_idx

    Returns DataFrame with columns: id (filename stem), label, split
    """

    # Build unified class_to_idx from the union of both splits
    all_classes: set[str] = set()
    for split_dir in [dataset.train_dir, dataset.test_dir]:
        all_classes.update(entry.name for entry in os.scandir(str(split_dir)) if entry.is_dir())

    class_to_idx = {cls_name: idx for idx, cls_name in enumerate(sorted(all_classes))}

    rows: list[dict[str, Any]] = []
    for split, split_dir in [("train", dataset.train_dir), ("test", dataset.test_dir)]:
        image_dataset = make_image_dataset([str(split_dir)], class_to_idx)
        for i in range(len(image_dataset)):
            path = image_dataset.paths[i].decode("utf-8")
            label = image_dataset.labels[i].item()
            rows.append({"id": Path(path).stem, "label": label, "split": split})

    return pl.DataFrame(rows)


def _load_embeddings_with_split(
    embeddings_path: str, metadata_df: pl.DataFrame
) -> tuple[npt.NDArray[np.float32], npt.NDArray[np.int_], npt.NDArray[np.float32], npt.NDArray[np.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"])
    test_meta = metadata_df.filter(pl.col("split") == "test").select(["id", "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
    dropped_total = dropped_train + dropped_test
    if dropped_total > 0:
        logger.warning(
            f"Join dropped {dropped_total} samples (missing embeddings): train={dropped_train}, test={dropped_test}"
        )

    x_train = train_join.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_train = train_join.get_column("label").to_numpy().astype(np.int_)
    x_test = test_join.get_column("embedding").to_numpy().astype(np.float32, copy=False)
    y_test = test_join.get_column("label").to_numpy().astype(np.int_)

    num_classes = metadata_df.get_column("label").max() + 1  # type: ignore[operator]
    total_samples = x_train.shape[0] + x_test.shape[0]
    embedding_dim = x_train.shape[1]
    logger.info(f"Loaded {total_samples} samples with {embedding_dim} dimensions, {num_classes} classes")

    logger.info(f"Train: {len(y_train)} samples, Test: {len(y_test)} samples")

    return (x_train, y_train, x_test, y_test)


def evaluate_plantdoc_single(
    x_train: npt.NDArray[np.float32],
    y_train: npt.NDArray[np.int_],
    x_test: npt.NDArray[np.float32],
    y_test: npt.NDArray[np.int_],
    k_shot: int,
    num_runs: int,
    seed: int,
    embeddings_path: str,
) -> dict[str, Any]:
    logger.info(f"Evaluating {k_shot}-shot")

    scores: list[float] = []
    for run in range(num_runs):
        run_seed = seed + run
        rng = np.random.default_rng(run_seed)

        # Sample k examples per class
        x_train_k_shot, y_train_k_shot = sample_k_shot(x_train, y_train, k_shot, rng)

        # Evaluate using SimpleShot
        y_pred, y_true = evaluate_simpleshot(x_train_k_shot, y_train_k_shot, x_test, y_test)

        acc = float(np.mean(y_pred == y_true))
        scores.append(acc)
        logger.info(f"Run {run + 1}/{num_runs} - Accuracy: {acc:.4f}")

    scores_arr = np.array(scores)
    mean_acc = float(scores_arr.mean())
    std_acc = float(scores_arr.std(ddof=1)) if len(scores) > 1 else 0.0

    logger.info(f"Mean accuracy over {num_runs} runs: {mean_acc:.4f} +/- {std_acc:.4f} (std)")

    return {
        "method": "simpleshot",
        "k_shot": k_shot,
        "accuracy": mean_acc,
        "accuracy_std": std_acc,
        "num_runs": num_runs,
        "embeddings_file": str(embeddings_path),
    }


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

    logger.info(f"Loading PlantDoc dataset from {args.dataset_path}")
    dataset = PlantDoc(args.dataset_path)
    existing_df: Optional[pl.DataFrame] = None
    output_path: Optional[Path] = None

    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(_results_filename(args))
        if output_path.exists() is True:
            logger.info(f"Loading existing results from {output_path}")
            existing_df = pl.read_csv(output_path)

    if args.embeddings is None:
        summary_results = _summary_results(existing_df, [], args.k_shot)
        if len(summary_results) == 0:
            raise cli.ValidationError("--embeddings is required when no cached results are available")

        logger.info("No embeddings provided, showing cached results")
        _print_summary_table(summary_results, args.k_shot)
        toc = time.time()
        logger.info(f"PlantDoc benchmark completed in {lib.format_duration(toc - tic)}")
        return

    metadata_df = _load_plantdoc_metadata(dataset)
    logger.info(f"Loaded metadata for {metadata_df.height} images")

    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}")

        cached_result = _load_cached_result(existing_df, embeddings_path, args.k_shot)
        if cached_result is not None:
            logger.info(f"Using cached result for {embeddings_path}")
            results.append(cached_result)
            continue

        if Path(embeddings_path).exists() is False:
            logger.warning(f"Embeddings file not found, skipping: {embeddings_path}")
            continue

        x_train, y_train, x_test, y_test = _load_embeddings_with_split(embeddings_path, metadata_df)

        accuracies: dict[int, float] = {}
        accuracies_std: dict[int, float] = {}
        for k_shot in args.k_shot:
            single_result = evaluate_plantdoc_single(
                x_train, y_train, x_test, y_test, k_shot, args.runs, args.seed, embeddings_path
            )
            accuracies[k_shot] = single_result["accuracy"]
            accuracies_std[k_shot] = single_result["accuracy_std"]

        results.append(
            {
                "embeddings_file": str(embeddings_path),
                "method": "simpleshot",
                "num_runs": args.runs,
                "accuracies": accuracies,
                "accuracies_std": accuracies_std,
            }
        )
        if output_path is not None:
            result = results[-1]
            _append_result_csv(result, args.k_shot, output_path)
            row_df = pl.DataFrame([_result_to_row(result, args.k_shot)])
            if existing_df is None:
                existing_df = row_df
            else:
                existing_df = pl.concat([existing_df, row_df], how="diagonal_relaxed")

    _print_summary_table(_summary_results(existing_df, results, args.k_shot), args.k_shot)

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


def set_parser(subparsers: Any) -> None:
    subparser = subparsers.add_parser(
        "plantdoc",
        allow_abbrev=False,
        help="run PlantDoc benchmark - 27 class plant disease classification using SimpleShot",
        description="run PlantDoc benchmark - 27 class plant disease classification using SimpleShot",
        epilog=(
            "Usage examples:\n"
            "python -m birder.eval plantdoc --embeddings "
            "results/plantdoc_embeddings.parquet --dataset-path ~/Datasets/PlantDoc --dry-run\n"
            "python -m birder.eval plantdoc --embeddings results/plantdoc/*.parquet "
            "--dataset-path ~/Datasets/PlantDoc\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 PlantDoc dataset root")
    subparser.add_argument(
        "--k-shot", type=int, nargs="+", default=[2, 5], help="number of examples per class for few-shot learning"
    )
    subparser.add_argument("--runs", type=int, default=5, help="number of evaluation runs")
    subparser.add_argument("--seed", type=int, default=0, help="base random seed")
    subparser.add_argument(
        "--dir", type=str, default="plantdoc", 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.dataset_path is None:
        raise cli.ValidationError("--dataset-path is required")


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