from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
RESULTS = ROOT / "results"
README = ROOT / "README.md"
REPO_ID = "dineth18/CD-Models"

DATASET_LABELS = {
    "dsifn_cd": "DSIFN-CD",
    "levir_cd_test_as_val": "LEVIR-CD+ comparison",
    "sysu_cd": "SYSU-CD",
    "whu_cd": "WHU-CD",
}
DATASET_ORDER = tuple(DATASET_LABELS)
DATASET_ALIASES = {
    "levir_val_as_test": "levir_cd_test_as_val",
}
MODEL_LABELS = {
    "bifa": "BiFA",
    "bit_cd": "BIT-CD",
    "cgnet": "CGNet",
    "changeformer": "ChangeFormer",
    "changemamba": "ChangeMamba",
    "dsamnet": "DSAMNet",
    "fc_ef": "FC-EF",
    "fc_siam_conc": "FC-Siam-conc",
    "fc_siam_diff": "FC-Siam-diff",
    "hanet": "HANet",
    "ifnet": "IFNet",
    "schanger": "SChanger",
    "siam_nestedunet": "Siam-NestedUNet",
    "stanet": "STANet",
}


def load_rows() -> dict[str, list[dict[str, Any]]]:
    grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for path in sorted(RESULTS.glob("*/*/metrics_test.json")):
        with path.open(encoding="utf-8") as handle:
            row = json.load(handle)
        source_dataset = str(row.get("dataset") or path.parent.name)
        dataset = DATASET_ALIASES.get(source_dataset, source_dataset)
        if dataset not in DATASET_LABELS:
            continue
        if row.get("split") != "test" or row.get("status") != "complete":
            continue
        if not isinstance(row.get("f1"), (int, float)):
            continue
        row["_path"] = path
        row["_model"] = str(row.get("model") or path.parents[1].name)
        row["_source_dataset"] = source_dataset
        grouped[dataset].append(row)
    for rows in grouped.values():
        rows.sort(key=lambda row: (-float(row["f1"]), row["_model"]))
    return grouped


def metric(row: dict[str, Any], *keys: str) -> Any:
    for key in keys:
        value = row.get(key)
        if value is not None and value != "":
            return value
    return None


def fmt(value: Any, digits: int = 4) -> str:
    if value is None:
        return "—"
    return f"{float(value):.{digits}f}"


def fmt_profile(value: Any) -> str:
    if value is None:
        return "—"
    number = float(value)
    return f"{number:.2f}"


def checkpoint_path(row: dict[str, Any]) -> Path:
    canonical = (row["_path"].parent / "checkpoints" / "best_model.pth").resolve()
    if canonical.exists():
        return canonical
    raw = metric(row, "checkpoint", "checkpoint_path")
    if raw:
        path = Path(str(raw))
        if path.exists() and path.resolve().is_relative_to(ROOT):
            return path.resolve()
    return canonical


def repo_relative(path: Path) -> Path:
    return path.resolve().relative_to(ROOT)


def download_url(path: Path) -> str:
    relative = repo_relative(path).as_posix()
    return f"https://huggingface.co/{REPO_ID}/resolve/main/{relative}?download=true"


def markdown_table(headers: list[str], rows: list[list[str]]) -> list[str]:
    return [
        "| " + " | ".join(headers) + " |",
        "| " + " | ".join("---" for _ in headers) + " |",
        *("| " + " | ".join(row) + " |" for row in rows),
    ]


def build_results() -> tuple[str, list[Path]]:
    grouped = load_rows()
    checkpoints: list[Path] = []
    lines = [
        "## Results and released checkpoints",
        "",
        "Only completed `metrics_test.json` evaluations are eligible for rankings. Rankings are computed per dataset by test-set F1 (descending), and the canonical `best_model.pth` checkpoint for every ranked model is released.",
        "",
        "All accuracy values are fractions. GPU memory is the PyTorch peak reserved memory for inference, with peak allocated memory used only when reserved memory is absent. FPS is model-only throughput. A dash means the evaluator did not record that metric; values are never estimated.",
        "",
        "The current evaluator records BF1 but not boundary mean IoU (BmIoU), so BmIoU remains explicitly unavailable rather than being inferred from BF1. Boundary IoU and boundary F1 are distinct measures.",
        "",
        "The LEVIR-CD+ comparison combines the project-compatible `levir_cd_test_as_val` and `levir_val_as_test` result records. The ChangeMamba row uses the completed `levir_val_as_test` evaluation over 5,568 samples.",
        "",
        "### Dataset summary",
        "",
    ]
    summary_rows: list[list[str]] = []
    for dataset in DATASET_ORDER:
        rows = grouped.get(dataset, [])
        if not rows:
            continue
        winner = rows[0]
        summary_rows.append(
            [
                DATASET_LABELS[dataset],
                str(len(rows)),
                MODEL_LABELS.get(winner["_model"], winner["_model"]),
                fmt(winner["f1"]),
            ]
        )
    lines.extend(markdown_table(["Dataset", "Tested models", "F1 leader", "Best test F1"], summary_rows))

    headers = [
        "Rank",
        "Model",
        "F1",
        "mIoU",
        "Overall accuracy",
        "Recall",
        "Precision",
        "BF1",
        "BmIoU",
        "GFLOPs",
        "GPU (GB)",
        "FPS",
        "Parameters (M)",
        "Checkpoint",
    ]
    for dataset in DATASET_ORDER:
        rows = grouped.get(dataset, [])
        lines.extend(["", f"### {DATASET_LABELS[dataset]}", ""])
        table_rows: list[list[str]] = []
        for rank, row in enumerate(rows, start=1):
            checkpoint = checkpoint_path(row)
            if not checkpoint.is_file():
                raise FileNotFoundError(
                    f"Tested checkpoint is missing for {row['_model']}/{dataset}: {checkpoint}"
                )
            checkpoints.append(checkpoint)
            link = f"[Download]({download_url(checkpoint)})"
            table_rows.append(
                [
                    str(rank),
                    MODEL_LABELS.get(row["_model"], row["_model"]),
                    fmt(metric(row, "f1", "F1")),
                    fmt(metric(row, "miou", "mIoU")),
                    fmt(metric(row, "oa", "OA")),
                    fmt(metric(row, "recall", "Recall")),
                    fmt(metric(row, "precision", "Precision")),
                    fmt(metric(row, "bf1", "BF1", "boundary_f1")),
                    fmt(metric(row, "bmiou", "BmIoU", "boundary_miou")),
                    fmt_profile(metric(row, "flops_g", "FLOPsG")),
                    fmt_profile(metric(row, "gpu_mem_reserved_peak_gb", "gpu_mem_allocated_peak_gb")),
                    fmt_profile(metric(row, "fps_model_only", "fps", "FPS")),
                    fmt_profile(metric(row, "params_m", "ParamsM")),
                    link,
                ]
            )
        if table_rows:
            lines.extend(markdown_table(headers, table_rows))
        else:
            lines.append("No completed test-set result is available, so no checkpoint is published.")

    lines.extend(
        [
            "",
            "### Datasets without a released checkpoint",
            "",
            "WildFire-S2, KATE-CD-256, the standard LEVIR-CD+ configuration, and Custom-CD currently have no eligible completed test result. Their validation-only or incomplete checkpoints are intentionally not uploaded.",
            "",
        ]
    )
    return "\n".join(lines), checkpoints


def update_readme(results_markdown: str) -> None:
    text = README.read_text(encoding="utf-8")
    current_heading = "## Results and released checkpoints\n"
    legacy_heading = "## Results\n"
    if current_heading in text:
        start = text.index(current_heading)
    else:
        start = text.index(legacy_heading)
    end = text.index("## Installation\n", start)
    text = text[:start] + results_markdown + "\n\n" + text[end:]
    old = (
        "### Trained checkpoints\n\n"
        "The released rank-1 test checkpoints and direct download links are listed in "
        "[Results and released checkpoints](#results-and-released-checkpoints). Local training still "
        "writes checkpoints under `results/{model}/{dataset}/checkpoints/`."
    )
    new = (
        "### Trained checkpoints\n\n"
        "All completed test-run checkpoints and direct download links are listed in "
        "[Results and released checkpoints](#results-and-released-checkpoints). Local training still "
        "writes checkpoints under `results/{model}/{dataset}/checkpoints/`."
    )
    if old in text:
        text = text.replace(old, new)
    elif new not in text:
        raise RuntimeError("Expected trained-checkpoint README text was not found")
    README.write_text(text, encoding="utf-8")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--write", action="store_true", help="Replace the README result sections")
    args = parser.parse_args()
    markdown, checkpoints = build_results()
    if args.write:
        update_readme(markdown)
    else:
        print(markdown)
    print("Released tested checkpoints:")
    for checkpoint in checkpoints:
        print(repo_relative(checkpoint))


if __name__ == "__main__":
    main()