#!/usr/bin/env python3
"""Train a small 2D hidden-mask adapter on the frozen accessibility split."""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image, ImageDraw, ImageOps
from torch.utils.data import DataLoader, Dataset


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


CANONICAL_CATEGORIES = ("curb_cut", "ramp", "stairs", "tactile_paving", "walkway")


def row_image_path(row: dict[str, Any]) -> Path:
    if row.get("image_path"):
        path = Path(str(row["image_path"]))
        if path.is_file():
            return path
    sample_dir = Path(row["sample_dir"])
    candidates = [path for path in (sample_dir / "image.jpg", sample_dir / "image.png") if path.is_file()]
    if len(candidates) != 1:
        raise FileNotFoundError(f"Expected exactly one RGB image in {sample_dir}")
    return candidates[0]


def load_rgb(path: Path, size: int) -> np.ndarray:
    # Ignore EXIF display orientation so RGB stays aligned with raw PNG masks.
    image = Image.open(path).convert("RGB")
    return np.asarray(image.resize((size, size), Image.Resampling.BILINEAR), dtype=np.float32) / 255.0


def load_mask(path: Path, size: int | None = None) -> np.ndarray:
    image = Image.open(path).convert("L")
    if size is not None:
        image = image.resize((size, size), Image.Resampling.NEAREST)
    return np.asarray(image) > 127


def category_planes(category: str, size: int) -> np.ndarray:
    planes = np.zeros((len(CANONICAL_CATEGORIES), size, size), dtype=np.float32)
    index = CANONICAL_CATEGORIES.index(category) if category in CANONICAL_CATEGORIES else -1
    if index >= 0:
        planes[index] = 1.0
    return planes


class AccessibilityMaskDataset(Dataset):
    def __init__(self, rows: list[dict[str, Any]], size: int, augment: bool):
        self.rows = rows
        self.size = size
        self.augment = augment

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, index: int) -> dict[str, Any]:
        row = self.rows[index]
        sample_dir = Path(row["sample_dir"])
        rgb = load_rgb(row_image_path(row), self.size)
        visible = load_mask(sample_dir / "target_visible.png", self.size)
        obstacle = load_mask(sample_dir / "obstacle.png", self.size)
        hidden = load_mask(sample_dir / "hidden.png", self.size)
        if self.augment and random.random() < 0.5:
            rgb = rgb[:, ::-1].copy()
            visible = visible[:, ::-1].copy()
            obstacle = obstacle[:, ::-1].copy()
            hidden = hidden[:, ::-1].copy()
        if self.augment:
            gain = random.uniform(0.88, 1.12)
            bias = random.uniform(-0.05, 0.05)
            rgb = np.clip(rgb * gain + bias, 0.0, 1.0)
        inputs = np.concatenate(
            [
                rgb.transpose(2, 0, 1),
                visible[None].astype(np.float32),
                obstacle[None].astype(np.float32),
                category_planes(row["category"], self.size),
            ],
            axis=0,
        )
        return {
            "input": torch.from_numpy(inputs.astype(np.float32)),
            "hidden": torch.from_numpy(hidden[None].astype(np.float32)),
            "obstacle": torch.from_numpy(obstacle[None].astype(np.float32)),
            "sample_id": row["sample_id"],
        }


class ConvBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        groups = min(8, out_channels)
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.GroupNorm(groups, out_channels),
            nn.SiLU(),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.GroupNorm(groups, out_channels),
            nn.SiLU(),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.block(inputs)


class TinyAmodalUNet(nn.Module):
    def __init__(self, base: int = 24, in_channels: int = 5 + len(CANONICAL_CATEGORIES)):
        super().__init__()
        self.enc1 = ConvBlock(in_channels, base)
        self.enc2 = ConvBlock(base, base * 2)
        self.bottleneck = ConvBlock(base * 2, base * 4)
        self.dec2 = ConvBlock(base * 4 + base * 2, base * 2)
        self.dec1 = ConvBlock(base * 2 + base, base)
        self.head = nn.Conv2d(base, 1, 1)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        first = self.enc1(inputs)
        second = self.enc2(F.max_pool2d(first, 2))
        bottleneck = self.bottleneck(F.max_pool2d(second, 2))
        up_second = F.interpolate(bottleneck, size=second.shape[-2:], mode="bilinear", align_corners=False)
        decoded_second = self.dec2(torch.cat([up_second, second], dim=1))
        up_first = F.interpolate(decoded_second, size=first.shape[-2:], mode="bilinear", align_corners=False)
        return self.head(self.dec1(torch.cat([up_first, first], dim=1)))


def training_loss(logits: torch.Tensor, target: torch.Tensor, obstacle: torch.Tensor) -> torch.Tensor:
    positive_weight = torch.tensor(8.0, device=logits.device)
    bce = F.binary_cross_entropy_with_logits(logits, target, pos_weight=positive_weight)
    probabilities = torch.sigmoid(logits)
    intersection = (probabilities * target).sum(dim=(1, 2, 3))
    denominator = probabilities.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3))
    dice_loss = 1.0 - ((2.0 * intersection + 1.0) / (denominator + 1.0)).mean()
    outside_obstacle = (probabilities * (1.0 - obstacle)).mean()
    return bce + dice_loss + 0.20 * outside_obstacle


def iou(left: np.ndarray, right: np.ndarray, empty_value: float = 1.0) -> float:
    union = left | right
    return float((left & right).sum() / union.sum()) if np.any(union) else empty_value


@torch.no_grad()
def predict_probability(
    model: nn.Module, row: dict[str, Any], size: int, device: torch.device
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    sample_dir = Path(row["sample_dir"])
    image_path = row_image_path(row)
    original = Image.open(image_path).convert("RGB")
    width, height = original.size
    rgb = load_rgb(image_path, size)
    visible_small = load_mask(sample_dir / "target_visible.png", size)
    obstacle_small = load_mask(sample_dir / "obstacle.png", size)
    inputs = np.concatenate(
        [
            rgb.transpose(2, 0, 1),
            visible_small[None].astype(np.float32),
            obstacle_small[None].astype(np.float32),
            category_planes(row["category"], size),
        ],
        axis=0,
    )
    logits = model(torch.from_numpy(inputs[None]).to(device)).sigmoid()[0, 0].cpu().numpy()
    probability = np.asarray(
        Image.fromarray(logits.astype(np.float32), mode="F").resize((width, height), Image.Resampling.BILINEAR)
    ).copy()
    visible = load_mask(sample_dir / "target_visible.png")
    obstacle = load_mask(sample_dir / "obstacle.png")
    hidden = load_mask(sample_dir / "hidden.png")
    probability *= (obstacle & ~visible).astype(np.float32)
    return probability, visible, obstacle, hidden


@torch.no_grad()
def evaluate_thresholds(
    model: nn.Module,
    rows: list[dict[str, Any]],
    size: int,
    device: torch.device,
    thresholds: list[float],
) -> tuple[float, dict[str, Any], dict[str, tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]]:
    cached = {row["sample_id"]: predict_probability(model, row, size, device) for row in rows}
    reports = []
    for threshold in thresholds:
        sample_rows = []
        for row in rows:
            probability, visible, _, hidden = cached[row["sample_id"]]
            predicted_hidden = probability >= threshold
            predicted_amodal = visible | predicted_hidden
            target_amodal = visible | hidden
            sample_rows.append(
                {
                    "sample_id": row["sample_id"],
                    "category": row["category"],
                    "hidden_gt_pixels": int(hidden.sum()),
                    "hidden_pred_pixels": int(predicted_hidden.sum()),
                    "hidden_iou": iou(predicted_hidden, hidden),
                    "amodal_iou": iou(predicted_amodal, target_amodal),
                }
            )
        nonempty = [item for item in sample_rows if item["hidden_gt_pixels"] > 0]
        reports.append(
            {
                "threshold": threshold,
                "mean_hidden_iou_nonempty": float(np.mean([item["hidden_iou"] for item in nonempty])) if nonempty else 1.0,
                "mean_amodal_iou": float(np.mean([item["amodal_iou"] for item in sample_rows])),
                "negative_control_false_positive_pixels": int(
                    sum(item["hidden_pred_pixels"] for item in sample_rows if item["hidden_gt_pixels"] == 0)
                ),
                "rows": sample_rows,
            }
        )
    best = max(
        reports,
        key=lambda item: (
            item["mean_hidden_iou_nonempty"],
            item["mean_amodal_iou"],
            -item["negative_control_false_positive_pixels"],
        ),
    )
    return float(best["threshold"]), best, cached


def save_mask(path: Path, mask: np.ndarray) -> None:
    Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(path)


def save_predictions(
    output: Path,
    rows: list[dict[str, Any]],
    cached: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]],
    threshold: float,
) -> None:
    output.mkdir(parents=True, exist_ok=True)
    for row in rows:
        sample_output = output / row["sample_id"]
        sample_output.mkdir(parents=True, exist_ok=True)
        probability, visible, _, _ = cached[row["sample_id"]]
        predicted_hidden = probability >= threshold
        save_mask(sample_output / "target_visible_mask.png", visible)
        save_mask(sample_output / "hidden_completion_mask.png", predicted_hidden)
        save_mask(sample_output / "amodal_accessibility_mask.png", visible | predicted_hidden)
        Image.fromarray(np.clip(probability * 255.0, 0, 255).astype(np.uint8), mode="L").save(
            sample_output / "hidden_probability.png"
        )


def color_overlay(image: Image.Image, masks: list[tuple[np.ndarray, tuple[int, int, int], float]]) -> Image.Image:
    array = np.asarray(image.convert("RGB"), dtype=np.float32).copy()
    for mask, color, alpha in masks:
        array[mask] = array[mask] * (1.0 - alpha) + np.asarray(color) * alpha
    return Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))


def build_contact_sheet(
    path: Path,
    rows: list[dict[str, Any]],
    cached: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]],
    threshold: float,
) -> None:
    panel_size = (320, 180)
    label_height = 30
    sheet = Image.new("RGB", (panel_size[0] * 3, (panel_size[1] + label_height) * len(rows)), "white")
    draw = ImageDraw.Draw(sheet)
    for row_index, row in enumerate(rows):
        image = Image.open(row_image_path(row)).convert("RGB")
        probability, visible, obstacle, hidden = cached[row["sample_id"]]
        predicted = probability >= threshold
        panels = (
            (row["sample_id"], image),
            ("GT: green visible / blue hidden / red obstacle", color_overlay(image, [(visible, (20, 210, 70), 0.38), (hidden, (40, 100, 245), 0.72), (obstacle, (235, 40, 40), 0.25)])),
            ("prediction: green visible / magenta hidden", color_overlay(image, [(visible, (20, 210, 70), 0.38), (predicted, (235, 40, 200), 0.75)])),
        )
        y = row_index * (panel_size[1] + label_height)
        for column, (label, panel) in enumerate(panels):
            x = column * panel_size[0]
            draw.text((x + 5, y + 7), label, fill="black")
            fitted = ImageOps.fit(panel, panel_size, method=Image.Resampling.LANCZOS)
            sheet.paste(fitted, (x, y + label_height))
    sheet.save(path, quality=92, subsampling=0)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--split", default="output/accessibility_training_split_v1")
    parser.add_argument("--output", default="output/accessibility_amodal_adapter_v1")
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--patience", type=int, default=25)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--image-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=3e-4)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--seed", type=int, default=20260623)
    parser.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto")
    args = parser.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)
    split = Path(args.split)
    split_summary = json.loads((split / "summary.json").read_text(encoding="utf-8"))
    train_rows = read_jsonl(split / "train.jsonl") + read_jsonl(split / "auxiliary_train.jsonl")
    validation_rows = read_jsonl(split / "validation.jsonl")
    test_rows = read_jsonl(split / "test.jsonl")
    if any(row["tier"] == "gold" for row in train_rows):
        raise RuntimeError("Gold sample detected in training rows")
    selected_device = (
        "cuda" if args.device == "auto" and torch.cuda.is_available()
        else "cpu" if args.device == "auto"
        else args.device
    )
    device = torch.device(selected_device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA requested but unavailable")

    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    loader = DataLoader(
        AccessibilityMaskDataset(train_rows, args.image_size, augment=True),
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        generator=torch.Generator().manual_seed(args.seed),
    )
    model = TinyAmodalUNet().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=1e-4)
    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")
    thresholds = [round(value, 2) for value in np.arange(0.20, 0.81, 0.05)]
    best_score = (-1.0, -1.0)
    best_epoch = -1
    stale_epochs = 0
    history = []
    for epoch in range(1, args.epochs + 1):
        model.train()
        losses = []
        for batch in loader:
            inputs = batch["input"].to(device, non_blocking=True)
            target = batch["hidden"].to(device, non_blocking=True)
            obstacle = batch["obstacle"].to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=device.type == "cuda"):
                logits = model(inputs)
                loss = training_loss(logits, target, obstacle)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            losses.append(float(loss.detach().cpu()))
        model.eval()
        threshold, validation, _ = evaluate_thresholds(
            model, validation_rows, args.image_size, device, thresholds
        )
        record = {
            "epoch": epoch,
            "train_loss": float(np.mean(losses)),
            "validation_threshold": threshold,
            "validation_hidden_iou_nonempty": validation["mean_hidden_iou_nonempty"],
            "validation_amodal_iou": validation["mean_amodal_iou"],
        }
        history.append(record)
        print(json.dumps(record, ensure_ascii=False), flush=True)
        score = (validation["mean_hidden_iou_nonempty"], validation["mean_amodal_iou"])
        if score > best_score:
            best_score = score
            best_epoch = epoch
            stale_epochs = 0
            torch.save(
                {
                    "model_state": model.state_dict(),
                    "epoch": epoch,
                    "validation_threshold": threshold,
                    "validation_metrics": validation,
                    "split_fingerprint": split_summary["fingerprint"],
                    "architecture": "TinyAmodalUNet",
                    "input_channels": [
                        "rgb",
                        "target_visible",
                        "occluding_obstacle",
                        *CANONICAL_CATEGORIES,
                    ],
                    "model_input_channels": 5 + len(CANONICAL_CATEGORIES),
                    "pseudo_depth_supervision": False,
                },
                output / "best.pt",
            )
        else:
            stale_epochs += 1
        if stale_epochs >= args.patience:
            break

    checkpoint = torch.load(output / "best.pt", map_location=device, weights_only=False)
    model.load_state_dict(checkpoint["model_state"])
    model.eval()
    threshold, validation_report, validation_cache = evaluate_thresholds(
        model, validation_rows, args.image_size, device, thresholds
    )
    _, test_report, test_cache = evaluate_thresholds(
        model, test_rows, args.image_size, device, [threshold]
    )
    save_predictions(output / "predictions" / "validation", validation_rows, validation_cache, threshold)
    save_predictions(output / "predictions" / "test", test_rows, test_cache, threshold)
    build_contact_sheet(output / "validation_contact_sheet.jpg", validation_rows, validation_cache, threshold)
    build_contact_sheet(output / "test_contact_sheet.jpg", test_rows, test_cache, threshold)
    (output / "history.jsonl").write_text(
        "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in history),
        encoding="utf-8",
    )
    baseline = {
        "validation_mean_amodal_iou_visible_only": float(
            np.mean(
                [
                    iou(load_mask(Path(row["sample_dir"]) / "target_visible.png"), load_mask(Path(row["sample_dir"]) / "target_amodal.png"))
                    for row in validation_rows
                ]
            )
        ),
        "validation_mean_hidden_iou_nonempty_visible_only": 0.0,
    }
    summary = {
        "status": "complete",
        "model": "small_2d_hidden_mask_adapter_not_original_amodal3d",
        "best_epoch": best_epoch,
        "epochs_run": len(history),
        "threshold_selected_on_validation": threshold,
        "split_fingerprint": split_summary["fingerprint"],
        "training_real_silver_count": 0,
        "training_strict_gt_count": len(read_jsonl(split / "train.jsonl")),
        "training_annotation_tier": "human_reviewed_strict_gt",
        "strict_gt_used_for_training": True,
        "training_auxiliary_synthetic_silver_count": len(read_jsonl(split / "auxiliary_train.jsonl")),
        "gold_used_for_training": False,
        "pseudo_depth_supervision": False,
        "baseline": baseline,
        "validation": {key: value for key, value in validation_report.items() if key != "rows"},
        "test": {key: value for key, value in test_report.items() if key != "rows"},
        "validation_rows": validation_report["rows"],
        "test_rows": test_report["rows"],
    }
    (output / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())