from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np


GRID_SIZE = 17


@dataclass(frozen=True)
class MatrixMetrics:
    stem: str
    true_positive: int
    false_positive: int
    false_negative: int
    manual_cells: int
    generated_cells: int
    manual_2x2: int
    generated_2x2: int

    @property
    def precision(self) -> float:
        denominator = self.true_positive + self.false_positive
        return self.true_positive / denominator if denominator else 1.0

    @property
    def recall(self) -> float:
        denominator = self.true_positive + self.false_negative
        return self.true_positive / denominator if denominator else 1.0

    @property
    def f1(self) -> float:
        denominator = self.precision + self.recall
        return 2.0 * self.precision * self.recall / denominator if denominator else 0.0


def read_matrix(path: Path) -> np.ndarray:
    rows: list[str] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#"):
            continue
        rows.append(stripped)
        if len(rows) == GRID_SIZE:
            break

    if len(rows) != GRID_SIZE:
        raise ValueError(f"{path} has {len(rows)} matrix rows, expected {GRID_SIZE}")

    matrix = np.zeros((GRID_SIZE, GRID_SIZE), dtype=np.uint8)
    for row_index, row_text in enumerate(rows):
        if len(row_text) < GRID_SIZE:
            raise ValueError(
                f"{path} row {row_index + 1} has {len(row_text)} columns, "
                f"expected {GRID_SIZE}"
            )
        for col_index, char in enumerate(row_text[:GRID_SIZE]):
            matrix[row_index, col_index] = 1 if char == "1" else 0

    return matrix


def count_2x2_blocks(matrix: np.ndarray) -> int:
    total = 0
    for row in range(GRID_SIZE - 1):
        for col in range(GRID_SIZE - 1):
            total += int(matrix[row : row + 2, col : col + 2].sum() == 4)
    return total


def diff_image(manual: np.ndarray, generated: np.ndarray, scale: int) -> np.ndarray:
    both = (manual == 1) & (generated == 1)
    false_positive = (manual == 0) & (generated == 1)
    false_negative = (manual == 1) & (generated == 0)

    image = np.full((GRID_SIZE, GRID_SIZE, 3), 255, dtype=np.uint8)
    image[both] = (0, 0, 0)
    image[false_positive] = (40, 40, 220)
    image[false_negative] = (220, 80, 40)
    return cv2.resize(image, (GRID_SIZE * scale, GRID_SIZE * scale), interpolation=cv2.INTER_NEAREST)


def write_image(path: Path, image: np.ndarray) -> None:
    ok, encoded = cv2.imencode(path.suffix, image)
    if not ok:
        raise ValueError(f"Cannot encode image: {path}")
    encoded.tofile(str(path))


def metrics_for_pair(manual_path: Path, generated_path: Path) -> MatrixMetrics:
    manual = read_matrix(manual_path)
    generated = read_matrix(generated_path)
    true_positive = int(((manual == 1) & (generated == 1)).sum())
    false_positive = int(((manual == 0) & (generated == 1)).sum())
    false_negative = int(((manual == 1) & (generated == 0)).sum())
    stem = manual_path.name.removesuffix("_matrix17.txt")
    return MatrixMetrics(
        stem=stem,
        true_positive=true_positive,
        false_positive=false_positive,
        false_negative=false_negative,
        manual_cells=int(manual.sum()),
        generated_cells=int(generated.sum()),
        manual_2x2=count_2x2_blocks(manual),
        generated_2x2=count_2x2_blocks(generated),
    )


def print_table(metrics: list[MatrixMetrics]) -> None:
    header = (
        "stem",
        "P",
        "R",
        "F1",
        "TP",
        "FP",
        "FN",
        "manual",
        "gen",
        "manual_2x2",
        "gen_2x2",
    )
    print(
        f"{header[0]:42} {header[1]:>6} {header[2]:>6} {header[3]:>6} "
        f"{header[4]:>4} {header[5]:>4} {header[6]:>4} {header[7]:>6} "
        f"{header[8]:>6} {header[9]:>11} {header[10]:>8}"
    )
    for item in metrics:
        print(
            f"{item.stem:42} {item.precision:6.3f} {item.recall:6.3f} {item.f1:6.3f} "
            f"{item.true_positive:4d} {item.false_positive:4d} {item.false_negative:4d} "
            f"{item.manual_cells:6d} {item.generated_cells:6d} "
            f"{item.manual_2x2:11d} {item.generated_2x2:8d}"
        )

    total_tp = sum(item.true_positive for item in metrics)
    total_fp = sum(item.false_positive for item in metrics)
    total_fn = sum(item.false_negative for item in metrics)
    total_manual = sum(item.manual_cells for item in metrics)
    total_generated = sum(item.generated_cells for item in metrics)
    total_manual_2x2 = sum(item.manual_2x2 for item in metrics)
    total_generated_2x2 = sum(item.generated_2x2 for item in metrics)
    total = MatrixMetrics(
        stem="TOTAL",
        true_positive=total_tp,
        false_positive=total_fp,
        false_negative=total_fn,
        manual_cells=total_manual,
        generated_cells=total_generated,
        manual_2x2=total_manual_2x2,
        generated_2x2=total_generated_2x2,
    )
    print("-" * 113)
    print(
        f"{total.stem:42} {total.precision:6.3f} {total.recall:6.3f} {total.f1:6.3f} "
        f"{total.true_positive:4d} {total.false_positive:4d} {total.false_negative:4d} "
        f"{total.manual_cells:6d} {total.generated_cells:6d} "
        f"{total.manual_2x2:11d} {total.generated_2x2:8d}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compare manual 17x17 labels against generated matrix output."
    )
    parser.add_argument(
        "--manual-dir",
        type=Path,
        default=Path("PreProcessing/manual_matrices"),
    )
    parser.add_argument(
        "--generated-dir",
        type=Path,
        default=Path("PreProcessing/processed_characters"),
    )
    parser.add_argument(
        "--diff-dir",
        type=Path,
        default=Path("PreProcessing/manual_diffs"),
    )
    parser.add_argument("--scale", type=int, default=24)
    parser.add_argument(
        "--include-empty-manual",
        action="store_true",
        help="Include manual matrices with no black cells instead of treating them as unlabelled.",
    )
    args = parser.parse_args()

    manual_paths = sorted(args.manual_dir.glob("*_matrix17.txt"))
    if not manual_paths:
        raise SystemExit(f"No manual matrices found in {args.manual_dir}")

    args.diff_dir.mkdir(parents=True, exist_ok=True)
    metrics: list[MatrixMetrics] = []
    missing: list[Path] = []
    skipped_empty: list[Path] = []

    for manual_path in manual_paths:
        manual = read_matrix(manual_path)
        if not args.include_empty_manual and int(manual.sum()) == 0:
            skipped_empty.append(manual_path)
            continue

        generated_path = args.generated_dir / manual_path.name
        if not generated_path.exists():
            missing.append(generated_path)
            continue

        generated = read_matrix(generated_path)
        metrics.append(metrics_for_pair(manual_path, generated_path))
        stem = manual_path.name.removesuffix("_matrix17.txt")
        write_image(args.diff_dir / f"{stem}_diff.png", diff_image(manual, generated, args.scale))

    if missing:
        print("Missing generated matrices:")
        for path in missing:
            print(f"  {path}")

    if skipped_empty:
        print("Skipped empty manual matrices:")
        for path in skipped_empty:
            print(f"  {path}")

    if not metrics:
        raise SystemExit("No comparable matrix pairs found.")

    print_table(metrics)
    print(f"\nDiff images written to: {args.diff_dir}")
    print("Diff colors: black=match, red=generated extra, blue=manual missing.")


if __name__ == "__main__":
    main()