#!/usr/bin/env python3
"""Assemble isolated handwriting and Gemini template text in benchmark order."""

from __future__ import annotations

import argparse
import json
import os
import unicodedata
from pathlib import Path
from typing import Any, Callable

from rapidfuzz.distance import Levenshtein

PLACEHOLDER = "{{HANDWRITING}}"
VARIANTS: dict[str, tuple[float | None, bool]] = {
    "simple_top": (None, False),
    "simple_spread_right": (None, True),
    "row055_top": (0.55, False),
    "row070_top": (0.70, False),
    "row055_spread_right": (0.55, True),
    "row070_spread_right": (0.70, True),
}


def normalized(text: str) -> str:
    return " ".join(unicodedata.normalize("NFKC", str(text)).split())


def merge_template_text(template_text: str, handwriting: str) -> tuple[str, int]:
    handwriting = normalized(handwriting)
    count = str(template_text).count(PLACEHOLDER)
    if count:
        merged = str(template_text).replace(PLACEHOLDER, handwriting, 1)
        merged = merged.replace(PLACEHOLDER, "")
    elif normalized(template_text):
        merged = f"{template_text} {handwriting}"
    else:
        merged = handwriting
    return normalized(merged), count


def center(line: dict[str, Any]) -> tuple[float, float]:
    x0, y0, x1, y1 = (float(value) for value in line["bbox"])
    return (x0 + x1) / 2.0, (y0 + y1) / 2.0


def height(line: dict[str, Any]) -> float:
    return max(float(line["bbox"][3]) - float(line["bbox"][1]), 1.0)


def simple(lines: list[dict[str, Any]]) -> list[list[dict[str, Any]]]:
    return [[line] for line in sorted(lines, key=lambda item: (center(item)[1], -center(item)[0]))]


def clustered(lines: list[dict[str, Any]], factor: float) -> list[list[dict[str, Any]]]:
    ordered = sorted(lines, key=lambda item: center(item)[1])
    rows: list[list[dict[str, Any]]] = []
    for line in ordered:
        if not rows:
            rows.append([line])
            continue
        row = rows[-1]
        mean_y = sum(center(item)[1] for item in row) / len(row)
        mean_height = sum(height(item) for item in row) / len(row)
        if abs(center(line)[1] - mean_y) <= factor * max(height(line), mean_height):
            row.append(line)
        else:
            rows.append([line])
    for row in rows:
        row.sort(key=lambda item: center(item)[0], reverse=True)
    return rows


def layout(
    lines: list[dict[str, Any]],
    width: int,
    height_value: int,
    factor: float | None,
    spread: bool,
) -> list[list[dict[str, Any]]]:
    builder: Callable[[list[dict[str, Any]]], list[list[dict[str, Any]]]]
    builder = simple if factor is None else lambda value: clustered(value, factor)
    aspect = width / max(height_value, 1)
    if spread and 1.03 <= aspect <= 1.42:
        midpoint = width / 2.0
        return builder([line for line in lines if center(line)[0] >= midpoint]) + builder(
            [line for line in lines if center(line)[0] < midpoint]
        )
    return builder(lines)


def handwriting_prediction(row: dict[str, Any], variant: str) -> str:
    factor, spread = VARIANTS[variant]
    rows = layout(
        row.get("lines", []),
        int(row.get("width") or 0),
        int(row.get("height") or 0),
        factor,
        spread,
    )
    return "\n".join(
        " ".join(
            str(line.get("text") or "").strip()
            for line in group
            if str(line.get("text") or "").strip()
        )
        for group in rows
        if any(str(line.get("text") or "").strip() for line in group)
    )


def metric(reference: str, hypothesis: str) -> dict[str, int]:
    reference = normalized(reference)
    hypothesis = normalized(hypothesis)
    return {
        "char_errors": int(Levenshtein.distance(reference, hypothesis)),
        "ref_chars": len(reference),
        "word_errors": int(
            Levenshtein.distance(reference.split(), hypothesis.split())
        ),
        "ref_words": len(reference.split()),
    }


def aggregate(rows: list[dict[str, Any]]) -> dict[str, float | int]:
    totals = {key: 0 for key in ("char_errors", "ref_chars", "word_errors", "ref_words")}
    for row in rows:
        for key, value in metric(row["reference"], row["prediction"]).items():
            totals[key] += value
    return {
        **totals,
        "wer": totals["word_errors"] / max(totals["ref_words"], 1),
        "cer": totals["char_errors"] / max(totals["ref_chars"], 1),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--recognized-lines", type=Path, required=True)
    parser.add_argument("--baseline-predictions", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--isolation-variant", required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    recognized = [
        json.loads(line)
        for line in args.recognized_lines.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    selected_ids = {str(row["id"]) for row in recognized}
    baseline = [
        json.loads(line)
        for line in args.baseline_predictions.read_text(encoding="utf-8").splitlines()
        if line.strip() and str(json.loads(line).get("id")) in selected_ids
    ]
    baseline_metrics = aggregate(baseline)
    candidates = {}
    for layout_variant in VARIANTS:
        rows = []
        missing_placeholders = 0
        multiple_placeholders = 0
        for row in recognized:
            handwriting = handwriting_prediction(row, layout_variant)
            prediction, placeholders = merge_template_text(
                str(row.get("template_text") or ""),
                handwriting,
            )
            missing_placeholders += int(placeholders == 0)
            multiple_placeholders += int(placeholders > 1)
            rows.append(
                {
                    "id": row["id"],
                    "split": "handwriting",
                    "reference": row["reference"],
                    "prediction": prediction,
                    "handwriting_prediction": normalized(handwriting),
                    "template_id": row.get("template_id"),
                    "template_placeholder_count": placeholders,
                    "isolation_variant": args.isolation_variant,
                    "layout_variant": layout_variant,
                    "detected_lines": len(row.get("lines", [])),
                    **({"error": row["error"]} if row.get("error") else {}),
                }
            )
        path = args.output_dir / f"predictions-{layout_variant}.jsonl"
        temporary = path.with_suffix(".jsonl.tmp")
        with temporary.open("w", encoding="utf-8") as handle:
            for row in rows:
                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
        os.replace(temporary, path)
        candidates[layout_variant] = {
            **aggregate(rows),
            "missing_placeholders": missing_placeholders,
            "multiple_placeholders": multiple_placeholders,
            "path": str(path),
        }
    best = min(
        candidates,
        key=lambda name: (
            candidates[name]["wer"] + candidates[name]["cer"],
            candidates[name]["cer"],
        ),
    )
    summary = {
        "isolation_variant": args.isolation_variant,
        "rows": len(recognized),
        "baseline": baseline_metrics,
        "candidates": candidates,
        "best_layout_variant": best,
        "best": candidates[best],
    }
    (args.output_dir / "assembly-summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()