from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np

from .io import load_image, read_json, save_image, write_jsonl


def _draw_segment(mask: np.ndarray, t0: float, x0: float, t1: float, x1: float, thickness: int) -> None:
    dt = float(t1 - t0)
    dx = float(x1 - x0)
    steps = int(max(abs(dt), abs(dx))) + 1
    if steps < 2:
        steps = 2
    rr = max(0, int(thickness // 2))

    for i in range(steps):
        a = i / (steps - 1)
        t = t0 + a * dt
        x = x0 + a * dx
        ti = int(round(t))
        xj = int(round(x))
        tmin = max(0, ti - rr)
        tmax = min(mask.shape[0], ti + rr + 1)
        xmin = max(0, xj - rr)
        xmax = min(mask.shape[1], xj + rr + 1)
        if tmin >= tmax or xmin >= xmax:
            continue
        yy, xx = np.ogrid[tmin:tmax, xmin:xmax]
        disk = (yy - t) ** 2 + (xx - x) ** 2 <= rr**2 if rr > 0 else np.ones((tmax - tmin, xmax - xmin), dtype=bool)
        mask[tmin:tmax, xmin:xmax][disk] = 1


def rasterize_one(sample_dir: Path, thickness: int) -> dict | None:
    kymo_path = sample_dir / "kymo.png"
    label_json_path = sample_dir / "label.json"
    if not kymo_path.exists() or not label_json_path.exists():
        return None

    kymo = load_image(kymo_path)
    if kymo.ndim != 2:
        raise ValueError(f"Expected grayscale kymo at {kymo_path}, got shape {kymo.shape}")
    h, w = kymo.shape

    label_json = read_json(label_json_path)
    tracks = label_json.get("tracks", [])

    mask = np.zeros((h, w), dtype=np.uint8)
    for track in tracks:
        anchors = track.get("anchors", [])
        if len(anchors) < 2:
            continue
        ordered = sorted(anchors, key=lambda a: float(a.get("t", 0)))
        for a0, a1 in zip(ordered[:-1], ordered[1:]):
            t0 = float(a0.get("t", 0))
            x0 = float(a0.get("x", 0))
            t1 = float(a1.get("t", 0))
            x1 = float(a1.get("x", 0))
            _draw_segment(mask, t0=t0, x0=x0, t1=t1, x1=x1, thickness=thickness)

    label_png = sample_dir / "label.png"
    save_image(label_png, mask)

    return {
        "image": str(kymo_path),
        "label": str(label_png),
        "meta": str(sample_dir / "meta.json") if (sample_dir / "meta.json").exists() else None,
        "sample_id": sample_dir.name,
    }


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Rasterize real anchor labels into dense masks")
    ap.add_argument("--real_root", required=True, help="Root containing data/real/<sample_id>/")
    ap.add_argument("--out_manifest", required=True, help="Output JSONL manifest path")
    ap.add_argument("--thickness", type=int, default=3, help="Line thickness in pixels")
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    real_root = Path(args.real_root)
    rows = []
    for sample_dir in sorted(real_root.iterdir()):
        if not sample_dir.is_dir():
            continue
        row = rasterize_one(sample_dir, thickness=args.thickness)
        if row is not None:
            rows.append(row)

    write_jsonl(args.out_manifest, rows)
    print(f"Wrote {len(rows)} samples to {args.out_manifest}")


if __name__ == "__main__":
    main()
