#!/usr/bin/env python3
"""Train a wall, paint, and skirting segmentation model.

Layout: DATA/{train,val}/{images,masks}/<matching-stem>.(jpg|png)
Indexed masks: 0=other/non-wall, 1=wall-unpainted, 2=wall-painted,
3=wall-uncertain, 4=skirting/trim, 5=switch/outlet, 6=AC, 7=door,
8=window, 9=wall-attached obstacle, 255=ignore/void. Independent material masks use 0=other wall material, 1=drywall,
2=wallpaper, 255=unknown/non-wall.

Validation scores every class and material, with an explicit status for any the
split has no ground truth for (see validation.py), and breaks scores down by
data source so synthetic scenes (make_synthetic_rooms.py) never inflate a
real-data number.
"""
import argparse
import copy
import csv
import json
import math
import random
from pathlib import Path

import numpy as np
import torch
from PIL import Image, ImageEnhance, ImageFilter
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
from data_balance import image_sampling_weights
from label_schema import (CLASSES, CONDITION_CLASSES, CONDITION_VALUES, IGNORE, MATERIAL_CLASSES,
                          MATERIAL_DRYWALL_OR_OTHER, MATERIAL_VALUES, mask_to_labels)
from models import DEFAULT_ARCH, available_architectures, build_from_checkpoint, build_model, warm_start
from validation import ScoreAccumulator, split_coverage
from wallpaper import paste_wallpaper

MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
def letterbox(image, size, fill):
    width, height = image.size
    scale = size / max(width, height)
    out_w, out_h = max(1, round(width * scale)), max(1, round(height * scale))
    resized = image.resize((out_w, out_h), Image.Resampling.BILINEAR if image.mode == "RGB" else Image.Resampling.NEAREST)
    canvas = Image.new(image.mode, (size, size), fill)
    left, top = (size - out_w) // 2, (size - out_h) // 2
    canvas.paste(resized, (left, top))
    return canvas


def ignore_fill(mask):
    # Letterbox/rotation padding is not scene content: label it ignore, never
    # "other", or the model learns "grey border = other" and val IoU is inflated.
    return (255, 255, 255) if mask.mode == "RGB" else IGNORE


def scale_crop(image, mask, material, size, crop, scale_range, condition=None):
    """Multi-scale training view: resize the long side to ``size * s`` and take a random ``crop``.

    Sources are often much larger than the training canvas (the facade set is
    1024 px), so plain letterboxing to ``size`` throws away the fine detail that
    wall edges and mullions depend on. Random scale + crop shows the model both
    full context and near-native-resolution detail. ``condition`` (optional) is
    cropped with the masks; it is then returned as a fourth item.
    """
    s = random.uniform(*scale_range)
    long_side = max(crop, round(size * s))
    width, height = image.size
    k = long_side / max(width, height)
    out = max(1, round(width * k)), max(1, round(height * k))
    image = image.resize(out, Image.Resampling.BILINEAR)
    mask = mask.resize(out, Image.Resampling.NEAREST)
    material = material.resize(out, Image.Resampling.NEAREST)
    layers = [(image, (114, 114, 114)), (mask, ignore_fill(mask)), (material, 255)]
    if condition is not None:
        layers.append((condition.resize(out, Image.Resampling.NEAREST), 255))
    canvas_w, canvas_h = max(out[0], crop), max(out[1], crop)
    if (canvas_w, canvas_h) != out:
        padded = []
        for im, fill in layers:
            canvas = Image.new(im.mode, (canvas_w, canvas_h), fill)
            canvas.paste(im, (0, 0))
            padded.append((canvas, fill))
        layers = padded
    left, top = random.randint(0, canvas_w - crop), random.randint(0, canvas_h - crop)
    box = (left, top, left + crop, top + crop)
    return tuple(im.crop(box) for im, _ in layers)


def paired_mask(mask_dir, image_path):
    return next((mask_dir / (image_path.stem + s) for s in (".png", ".jpg", ".jpeg")
                 if (mask_dir / (image_path.stem + s)).is_file()), None)


class WallDataset(Dataset):
    def __init__(self, root, split, size, augment=False, crop=None, scale_range=(1.0, 1.0), wallpaper_paste=0.0):
        self.image_dir = Path(root) / split / "images"
        self.mask_dir = Path(root) / split / "masks"
        self.material_dir = Path(root) / "drywall_masks" / split
        # Optional: images without a condition mask leave the condition head unsupervised.
        self.condition_dir = Path(root) / "condition_masks" / split
        if not self.image_dir.is_dir() or not self.mask_dir.is_dir():
            raise ValueError(f"expected directories {self.image_dir} and {self.mask_dir}")
        if not self.material_dir.is_dir():
            raise ValueError(f"expected drywall material masks in {self.material_dir}")
        self.files = sorted(p for p in self.image_dir.glob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png"})
        missing = [p.name for p in self.files if paired_mask(self.mask_dir, p) is None or not (self.material_dir / (p.stem + ".png")).is_file()]
        if not self.files or missing:
            raise ValueError(f"{self.image_dir}: no images or missing masks (examples: {missing[:5]})")
        self.size, self.augment = size, augment
        self.crop, self.scale_range = crop or size, scale_range
        self.wallpaper_paste = wallpaper_paste

    def __len__(self):
        return len(self.files)

    def condition_path(self, image_path):
        path = self.condition_dir / (image_path.stem + ".png")
        return path if path.is_file() else None

    def __getitem__(self, i):
        image_path = self.files[i]
        image = Image.open(image_path).convert("RGB")
        mask = Image.open(paired_mask(self.mask_dir, image_path))
        material = Image.open(self.material_dir / (image_path.stem + ".png"))
        condition_file = self.condition_path(image_path)
        condition = (Image.open(condition_file) if condition_file else
                     Image.new("L", image.size, IGNORE))
        if condition.size != image.size or condition.mode not in ("L", "P") \
                or not np.isin(np.unique(np.asarray(condition)), CONDITION_VALUES).all():
            raise ValueError(f"condition mask for {image_path.name} must be a same-size indexed PNG "
                             f"using only {list(CONDITION_VALUES)}")
        if image.size != mask.size:
            raise ValueError(f"image and mask dimensions differ for {image_path.name}: {image.size} vs {mask.size}")
        if image.size != material.size or material.mode not in ("L", "P"):
            raise ValueError(f"drywall mask for {image_path.name} must be a same-size indexed/grayscale PNG")
        material_values = np.asarray(material)
        if not np.isin(np.unique(material_values), MATERIAL_VALUES).all():
            raise ValueError(f"material mask {image_path.stem}.png must use only {list(MATERIAL_VALUES)}")
        if self.augment and random.random() < self.wallpaper_paste:
            # Re-surface part of the wall with procedural wallpaper under the photo's own lighting.
            rng = np.random.default_rng(random.getrandbits(32))
            before = material_values
            pixels, labels, material_values = paste_wallpaper(np.asarray(image), mask_to_labels(mask),
                                                              material_values, rng)
            image, mask, material = Image.fromarray(pixels), Image.fromarray(labels), Image.fromarray(material_values)
            # Fresh paper covers whatever condition the wall was in: unknown there.
            condition_values = np.asarray(condition).copy()
            condition_values[material_values != before] = IGNORE
            condition = Image.fromarray(condition_values)
        if self.augment:
            image, mask, material, condition = scale_crop(image, mask, material, self.size, self.crop,
                                                          self.scale_range, condition)
            if random.random() < 0.5:
                image, mask = image.transpose(Image.Transpose.FLIP_LEFT_RIGHT), mask.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
                material = material.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
                condition = condition.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
            # Keep vertical orientation: skirting belongs at the bottom of the wall.
            if random.random() < 0.5:
                angle = random.uniform(-3.0, 3.0)
                image = image.rotate(angle, resample=Image.Resampling.BILINEAR, fillcolor=(114, 114, 114))
                mask = mask.rotate(angle, resample=Image.Resampling.NEAREST, fillcolor=ignore_fill(mask))
                material = material.rotate(angle, resample=Image.Resampling.NEAREST, fillcolor=255)
                condition = condition.rotate(angle, resample=Image.Resampling.NEAREST, fillcolor=255)
            # Lighting varies far more on site than between dataset photos; hue is
            # left alone because paint colour is itself the signal.
            for enhancer in (ImageEnhance.Brightness, ImageEnhance.Contrast, ImageEnhance.Color):
                image = enhancer(image).enhance(random.uniform(0.7, 1.3))
            if random.random() < 0.2:
                image = image.filter(ImageFilter.GaussianBlur(random.uniform(0.3, 1.2)))
        else:
            image = letterbox(image, self.size, (114, 114, 114))
            mask = letterbox(mask, self.size, ignore_fill(mask))
            material = letterbox(material, self.size, 255)
            condition = letterbox(condition, self.size, 255)
        x = np.asarray(image, dtype=np.float32).transpose(2, 0, 1) / 255.0
        x = torch.from_numpy(((x - MEAN) / STD).copy())
        y = torch.from_numpy(mask_to_labels(mask).astype(np.int64).copy())
        m = torch.from_numpy(np.asarray(material).astype(np.int64).copy())
        material_known = m != IGNORE
        semantic_wall = (y == 1) | (y == 2) | (y == 3)
        if (material_known & ~semantic_wall).any():
            raise ValueError(f"drywall material mask for {image_path.name} labels pixels outside the wall region")
        c = torch.from_numpy(np.asarray(condition).astype(np.int64).copy())
        c[~semantic_wall] = IGNORE          # condition is only defined on wall
        return x, y, m, c

    def class_counts(self):
        counts = np.zeros(len(CLASSES), dtype=np.int64)
        for image_path in self.files:
            mask = Image.open(paired_mask(self.mask_dir, image_path))
            labels = mask_to_labels(mask)
            counts += np.bincount(labels[labels != IGNORE], minlength=len(CLASSES))[:len(CLASSES)]
        return counts

    def material_counts(self):
        counts = np.zeros(len(MATERIAL_CLASSES), dtype=np.int64)
        for image_path in self.files:
            with Image.open(self.material_dir / (image_path.stem + ".png")) as material:
                values = np.asarray(material)
            if values.ndim != 2 or not np.isin(np.unique(values), MATERIAL_VALUES).all():
                raise ValueError(f"invalid material mask: {image_path.stem}.png")
            counts += np.bincount(values[values != 255], minlength=len(MATERIAL_CLASSES))[:len(MATERIAL_CLASSES)]
        return counts

    def condition_counts(self):
        counts = np.zeros(len(CONDITION_CLASSES), dtype=np.int64)
        for image_path in self.files:
            path = self.condition_path(image_path)
            if path:
                with Image.open(path) as condition:
                    values = np.asarray(condition)
                counts += np.bincount(values[values != IGNORE], minlength=len(CONDITION_CLASSES))[:len(CONDITION_CLASSES)]
        return counts

    def material_presence(self):
        """Boolean (image, material) matrix: which materials each image labels."""
        presence = np.zeros((len(self.files), len(MATERIAL_CLASSES)), dtype=bool)
        for i, image_path in enumerate(self.files):
            with Image.open(self.material_dir / (image_path.stem + ".png")) as material:
                values = np.asarray(material)
            present = np.unique(values[values != 255])
            presence[i, present[present < len(MATERIAL_CLASSES)].astype(np.int64)] = True
        return presence

    def coverage_report(self):
        """What this split can score at all (classes/materials with ground truth)."""
        return split_coverage(self.class_presence(), self.material_presence())

    def groups(self, metadata_path):
        """Score groups per image: ``source:real|synthetic`` and ``environment:<value>``."""
        metadata_path = Path(metadata_path)
        rows = {}
        if metadata_path.is_file():
            with metadata_path.open(newline="", encoding="utf-8-sig") as stream:
                for row in csv.DictReader(stream):
                    rows[Path(row.get("image") or "").name] = row
        out = []
        for path in self.files:
            row = rows.get(path.name, {})
            source = (row.get("source") or "").strip().lower() or ("synthetic" if path.name.startswith("synth_") else "real")
            keys = [f"source:{source}"]
            environment = (row.get("environment") or "").strip().lower()
            if environment:
                keys.append(f"environment:{environment}")
            out.append(keys)
        return out

    def class_presence(self):
        """Boolean (image, class) matrix: which classes each image contains."""
        presence = np.zeros((len(self.files), len(CLASSES)), dtype=bool)
        for i, image_path in enumerate(self.files):
            labels = mask_to_labels(Image.open(paired_mask(self.mask_dir, image_path)))
            present = np.unique(labels[labels != IGNORE])
            presence[i, present.astype(np.int64)] = True
        return presence

    def wall_ids(self, metadata_path):
        """Wall/group id per image from metadata, falling back to the stem.

        Grouping by physical wall keeps a heavily cropped wall from dominating
        sampling just because it produced more frames.
        """
        metadata_path = Path(metadata_path)
        mapping = {}
        if metadata_path.is_file():
            with metadata_path.open(newline="", encoding="utf-8-sig") as stream:
                for row in csv.DictReader(stream):
                    name = Path(row.get("image") or "").name
                    if name:
                        mapping[name] = (row.get("wall_id") or "").strip()
        return [mapping.get(path.name) or f"__wall__{path.stem}" for path in self.files]


def class_balanced_weights(counts):
    counts = np.asarray(counts, dtype=np.float64)
    weights = np.zeros_like(counts)
    present = counts > 0
    if present.any():
        weights[present] = np.sqrt(counts[present].sum() / (present.sum() * counts[present]))
        weights[present] /= weights[present].mean()
        weights = np.minimum(weights, 5.0)
    return torch.tensor(weights, dtype=torch.float32)


def dice_loss(logits, target, class_weights):
    """Soft Dice over the classes present in training (weight > 0).

    Cross-entropy is dominated by easy interior pixels; Dice scores the overlap
    of each region as a whole, which tracks IoU and penalises ragged wall edges
    and fragmented windows that CE barely notices.
    """
    active = class_weights > 0
    valid = (target != IGNORE).unsqueeze(1)
    prob = torch.softmax(logits.float(), dim=1) * valid
    onehot = torch.nn.functional.one_hot(target.clamp(0, logits.shape[1] - 1), logits.shape[1]).permute(0, 3, 1, 2)
    onehot = onehot.float() * valid
    intersection = (prob * onehot).sum(dim=(0, 2, 3))
    denominator = prob.sum(dim=(0, 2, 3)) + onehot.sum(dim=(0, 2, 3))
    dice = (2 * intersection + 1.0) / (denominator + 1.0)
    return (1.0 - dice[active]).mean() if active.any() else logits.sum() * 0.0


class ModelEma:
    """Exponential moving average of weights; evaluated and saved instead of the raw model.

    With a few hundred training images the raw weights swing between epochs
    (val IoU jumped +/-0.02 run to run); the average is smoother and usually better.
    """

    def __init__(self, model, decay):
        self.decay = decay
        self.module = copy.deepcopy(model).eval()
        for p in self.module.parameters():
            p.requires_grad_(False)

    @torch.no_grad()
    def update(self, model, step):
        # Warm the decay up so early, poor weights do not linger in the average.
        d = min(self.decay, (1 + step) / (10 + step))
        for ema_v, v in zip(self.module.state_dict().values(), model.state_dict().values()):
            if ema_v.dtype.is_floating_point:
                ema_v.mul_(d).add_(v.detach(), alpha=1 - d)
            else:
                ema_v.copy_(v)


def parameter_groups(model, lr, head_lr_mult, weight_decay):
    """Pretrained encoder at ``lr``; freshly initialised heads at ``lr * head_lr_mult``."""
    head_keys = ("drywall_head", "condition_head", "segmenter.decode_head", "segmenter.classifier", "segmenter.aux_classifier")
    encoder, head = [], []
    for name, p in model.named_parameters():
        if p.requires_grad:
            (head if name.startswith(head_keys) else encoder).append(p)
    return [{"params": encoder, "lr": lr, "weight_decay": weight_decay},
            {"params": head, "lr": lr * head_lr_mult, "weight_decay": weight_decay}]


def evaluate(model, loader, device, flip_tta=False, groups=None, temperature=(1.0, 1.0)):
    """Score ``model`` on ``loader`` (unshuffled); ``groups`` lists group keys per dataset item.

    ``temperature`` (semantic, material) divides the logits before softmax; it
    changes confidence (ECE, uncertainty thresholds) but never the argmax.
    """
    model.eval()
    t_sem, t_mat = temperature
    scores = ScoreAccumulator()
    index = 0
    with torch.inference_mode():
        for x, y, material, condition in loader:
            x = x.to(device, non_blocking=True)
            output = model(x)
            semantic = (output["semantic"].float() / t_sem).softmax(1)
            material_prob = (output["drywall"].float() / t_mat).softmax(1)
            if flip_tta:
                flipped = model(torch.flip(x, dims=[3]))
                semantic = 0.5 * (semantic + (flipped["semantic"].float() / t_sem).softmax(1).flip(dims=[3]))
                material_prob = 0.5 * (material_prob + (flipped["drywall"].float() / t_mat).softmax(1).flip(dims=[3]))
            confidence, pred = semantic.max(1)
            material_pred = material_prob.argmax(1)
            painted = semantic[:, CLASSES.index("wall_painted")]
            condition_pred = [None] * x.shape[0]
            if "condition" in output:
                condition_logits = output["condition"].float()
                if flip_tta:
                    condition_logits = 0.5 * (condition_logits + flipped["condition"].float().flip(dims=[3]))
                condition_pred = condition_logits.argmax(1).cpu().numpy()
            for p, t, mp, mt, c, pp, cp, ct in zip(pred.cpu().numpy(), y.numpy(), material_pred.cpu().numpy(),
                                                   material.numpy(), confidence.cpu().numpy(), painted.cpu().numpy(),
                                                   condition_pred, condition.numpy()):
                scores.add(p, t, mp, mt, confidence=c, groups=groups[index] if groups else (),
                           painted_probability=pp, condition_pred=cp, condition_target=ct)
                index += 1
    return scores.result()


def fit_temperature(model, loader, device, pixels_per_image=4096, seed=0):
    """Fit one temperature per head on held-out pixels by minimising NLL (Guo et al. 2017).

    Segmentation nets trained with CE are overconfident (the first checkpoint
    scored ECE 0.51 on unseen indoor scenes); a single scalar fixes most of it
    without touching the argmax, so it is safe to fit on the selection split.
    """
    model.eval()
    generator = torch.Generator().manual_seed(seed)
    collected = {"semantic": ([], []), "drywall": ([], [])}
    with torch.inference_mode():
        for x, y, material, _ in loader:
            output = model(x.to(device))
            for head, target in (("semantic", y), ("drywall", material)):
                logits = output[head].float().cpu().permute(0, 2, 3, 1).reshape(-1, output[head].shape[1])
                labels = target.reshape(-1)
                valid = torch.nonzero(labels < output[head].shape[1]).squeeze(1)   # skips 254/255
                take = valid[torch.randperm(valid.numel(), generator=generator)[:pixels_per_image * x.shape[0]]]
                collected[head][0].append(logits[take])
                collected[head][1].append(labels[take])
    temperatures = []
    for head in ("semantic", "drywall"):
        logits, labels = torch.cat(collected[head][0]), torch.cat(collected[head][1])
        if labels.numel() == 0:
            temperatures.append(1.0)
            continue
        log_t = torch.zeros(1, requires_grad=True)
        optimizer = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)

        def closure():
            optimizer.zero_grad()
            loss = torch.nn.functional.cross_entropy(logits / log_t.exp(), labels)
            loss.backward()
            return loss

        optimizer.step(closure)
        temperatures.append(round(float(log_t.exp().clamp(0.05, 20.0)), 4))
    return tuple(temperatures)


def require_coverage(report, split, strict):
    """Print what ``split`` cannot score; exit when ``strict`` and anything is missing."""
    missing = report["classes_missing"] + [f"material:{m}" for m in report.get("materials_missing", [])]
    print(json.dumps({"split_coverage": {"split": split, **report}}), flush=True)
    if missing:
        message = (f"{split} split has no ground truth for {missing}: those scores will be reported as "
                   f"no_ground_truth. Add labelled data or synthetic scenes (make_synthetic_rooms.py --into).")
        if strict:
            raise SystemExit(message)
        print(f"warning: {message}", flush=True)


def load_checkpoint_model(path, device):
    checkpoint = torch.load(path, map_location=device, weights_only=True)
    model = build_from_checkpoint(checkpoint, len(CLASSES))
    return model.to(device).eval(), checkpoint


def checkpoint_temperature(checkpoint):
    return float(checkpoint.get("temperature", 1.0)), float(checkpoint.get("material_temperature", 1.0))


def calibrate(model, checkpoint, val_loader, device, val_groups=None):
    """Fit temperatures, then the completion threshold, on val; both are written into ``checkpoint``.

    The threshold is chosen after temperature scaling because it is a threshold on
    the calibrated P(painted) that predict.py computes. ``completion_threshold`` is
    None when no threshold keeps false completes under the target, and predict.py
    then never reports a wall as complete.
    """
    temperature = fit_temperature(model, val_loader, device)
    checkpoint["temperature"], checkpoint["material_temperature"] = temperature
    scores = evaluate(model, val_loader, device, groups=val_groups, temperature=temperature)
    choice = scores.get("completion_threshold") or {"threshold": None, "reason": "val has no known wall"}
    checkpoint["completion_threshold"] = choice["threshold"]
    checkpoint["metrics"] = {**checkpoint.get("metrics", {}), "val_ece_semantic": scores["ece_semantic"],
                             "val_completion": choice,
                             "val_completion_operating_points": scores.get("completion_operating_points")}
    return temperature, {"temperature": temperature[0], "material_temperature": temperature[1],
                         "val_ece_semantic": scores["ece_semantic"], "completion": choice}


def at_checkpoint_threshold(scores, checkpoint):
    """The operating point predict.py will actually use, read off a held-out score dict."""
    threshold = checkpoint.get("completion_threshold")
    for point in scores.get("completion_operating_points") or []:
        if threshold is not None and abs(point["threshold"] - threshold) < 1e-9:
            return {"completion_threshold": threshold, "false_complete_rate_at_threshold": point["false_complete_rate"],
                    "complete_recall_at_threshold": point["complete_recall"]}
    return {"completion_threshold": threshold}


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--data", type=Path, required=True)
    ap.add_argument("--output", type=Path, default=Path("artifacts"))
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--size", type=int, default=640,
                    help="evaluation/inference canvas (long side); saved in the checkpoint for predict.py")
    ap.add_argument("--crop", type=int, help="training crop size (default: --size)")
    ap.add_argument("--scale-min", type=float, default=1.0,
                    help="random training scale range, relative to --size (e.g. 0.75 1.5)")
    ap.add_argument("--scale-max", type=float, default=1.0)
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--head-lr-mult", type=float, default=1.0,
                    help="learning-rate multiplier for the freshly initialised heads (10 is typical)")
    ap.add_argument("--warmup-epochs", type=float, default=0.0, help="linear LR warmup before the cosine decay")
    ap.add_argument("--weight-decay", type=float, default=1e-4)
    ap.add_argument("--dice-weight", type=float, default=0.0, help="weight of the soft-Dice term added to cross-entropy")
    ap.add_argument("--ema", type=float, default=0.0, help="EMA decay for evaluated/saved weights (0 disables; 0.999 typical)")
    ap.add_argument("--edge-boost", type=float, default=2.0, help="extra cross-entropy weight near wall contours")
    ap.add_argument("--drywall-loss-weight", type=float, default=0.5)
    ap.add_argument("--condition-loss-weight", type=float, default=0.5,
                    help="weight of the surface-condition loss (paint failure / rust); needs condition_masks/")
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--arch", choices=available_architectures(), default=DEFAULT_ARCH,
                    help="backbone/architecture; segformer_b2+ and deeplabv3_resnet50+ are stronger than the default")
    ap.add_argument("--sampling", choices=("uniform", "class", "wall", "hybrid"), default="uniform",
                    help="class rare classes more often, wall diverse physical walls, or hybrid both")
    ap.add_argument("--window-loss-boost", type=float, default=1.0,
                    help="extra loss weight for the window class (windows are small, rare, and safety-critical)")
    ap.add_argument("--metadata", type=Path, help="CSV with wall_id per image; defaults to DATA/metadata.csv")
    ap.add_argument("--manifest", type=Path,
                    help="provenance manifest; training refuses to start if the dataset does not match")
    ap.add_argument("--registry", type=Path, help="append a run record to this hash-linked ledger on completion")
    ap.add_argument("--run-id", help="optional run identifier recorded in the registry")
    ap.add_argument("--no-pretrained", action="store_true")
    ap.add_argument("--init-checkpoint", type=Path,
                    help="warm-start from this checkpoint (same --arch); a pre-wallpaper material head is grown to 3 classes")
    ap.add_argument("--no-calibrate", action="store_true",
                    help="skip fitting softmax temperatures on val after training")
    ap.add_argument("--eval-checkpoint", type=Path, action="append", default=[],
                    help="skip training; score this checkpoint (repeatable) on --eval-split and exit")
    ap.add_argument("--calibrate-checkpoint", type=Path,
                    help="skip training; fit temperatures and the completion threshold on val into this checkpoint "
                         "(rewritten in place) and exit")
    ap.add_argument("--eval-split", default="test", help="split for --eval-checkpoint and the final best.pt score")
    ap.add_argument("--eval-size", type=int, help="canvas for --eval-checkpoint (default: the checkpoint's own size)")
    ap.add_argument("--wallpaper-paste", type=float, default=0.0,
                    help="probability of re-surfacing a training image's wall with procedural wallpaper (0.25 typical)")
    ap.add_argument("--require-full-validation", action="store_true",
                    help="refuse to run when the val/eval split lacks any class or material (no unscored heads)")
    ap.add_argument("--report", type=Path, help="write the final held-out score dict(s) here as JSON")
    args = ap.parse_args()

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    metadata_path = args.metadata or args.data / "metadata.csv"
    reports = []
    if args.eval_checkpoint:
        for path in args.eval_checkpoint:
            model, checkpoint = load_checkpoint_model(path, device)
            size = args.eval_size or int(checkpoint["size"])
            eval_ds = WallDataset(args.data, args.eval_split, size)
            require_coverage(eval_ds.coverage_report(), args.eval_split, args.require_full_validation)
            loader = DataLoader(eval_ds, batch_size=args.batch_size, num_workers=args.workers)
            for tta in (False, True):
                scores = evaluate(model, loader, device, flip_tta=tta, groups=eval_ds.groups(metadata_path),
                                  temperature=checkpoint_temperature(checkpoint))
                record = {"checkpoint": str(path), "split": args.eval_split, "size": size, "flip_tta": tta,
                          "material_classes": list(checkpoint.get("material_classes", [])),
                          "temperature": checkpoint_temperature(checkpoint),
                          **at_checkpoint_threshold(scores, checkpoint), **scores}
                reports.append(record)
                print(json.dumps(record), flush=True)
        if args.report:
            args.report.write_text(json.dumps(reports, indent=2) + "\n")
        return
    if args.calibrate_checkpoint:
        model, checkpoint = load_checkpoint_model(args.calibrate_checkpoint, device)
        val_ds = WallDataset(args.data, "val", int(checkpoint["size"]))
        require_coverage(val_ds.coverage_report(), "val", args.require_full_validation)
        loader = DataLoader(val_ds, batch_size=args.batch_size, num_workers=args.workers)
        _, summary = calibrate(model, checkpoint, loader, device, val_ds.groups(metadata_path))
        torch.save(checkpoint, args.calibrate_checkpoint)
        print(json.dumps({"calibration": summary}), flush=True)
        return

    dataset_id = None
    if args.manifest:
        from provenance import verify_manifest
        manifest = json.loads(args.manifest.read_text(encoding="utf-8"))
        ok, report = verify_manifest(args.data, manifest)
        if not ok:
            raise SystemExit(f"dataset does not match manifest {args.manifest}: {report}")
        dataset_id = manifest["dataset_id"]
        print(json.dumps({"provenance": {"dataset_id": dataset_id, "verified": True}}), flush=True)
    if not 0 < args.scale_min <= args.scale_max:
        ap.error("need 0 < --scale-min <= --scale-max")
    if not 0.0 <= args.wallpaper_paste <= 1.0:
        ap.error("--wallpaper-paste must be between 0 and 1")
    train_ds = WallDataset(args.data, "train", args.size, augment=True, crop=args.crop,
                           scale_range=(args.scale_min, args.scale_max), wallpaper_paste=args.wallpaper_paste)
    val_ds = WallDataset(args.data, "val", args.size)
    require_coverage(val_ds.coverage_report(), "val", args.require_full_validation)
    val_groups = val_ds.groups(metadata_path)
    counts = train_ds.class_counts()
    class_weights = class_balanced_weights(counts).to(device)
    if args.window_loss_boost != 1.0:
        class_weights[CLASSES.index("window")] *= args.window_loss_boost
    material_counts = train_ds.material_counts()
    material_weights = class_balanced_weights(material_counts).to(device)
    condition_counts = train_ds.condition_counts()
    condition_weights = class_balanced_weights(condition_counts).to(device) if condition_counts.sum() else None
    sampler, sampling_summary = None, {"mode": "uniform"}
    if args.sampling != "uniform":
        wall_ids = train_ds.wall_ids(metadata_path)
        presence = train_ds.class_presence() if args.sampling in ("class", "hybrid") else None
        weights = image_sampling_weights(presence, wall_ids, args.sampling)
        # Keep every image sampleable: a rare-class weighting must not starve
        # the plenty-class frames that teach the common background.
        weights = np.maximum(weights, 0.1 * weights.mean()) if weights.mean() else np.ones_like(weights)
        sampler = WeightedRandomSampler(torch.as_tensor(weights, dtype=torch.double),
                                        num_samples=len(train_ds), replacement=True)
        sampling_summary = {"mode": args.sampling, "metadata": str(metadata_path),
                            "walls_in_train": len(set(wall_ids)),
                            "image_weight_min": round(float(weights.min()), 4),
                            "image_weight_max": round(float(weights.max()), 4),
                            "image_weight_mean": round(float(weights.mean()), 4)}
    print(json.dumps({"device": str(device), "arch": args.arch, "train_images": len(train_ds), "val_images": len(val_ds),
                      "sampling": sampling_summary,
                      "crop": train_ds.crop, "scale_range": [args.scale_min, args.scale_max],
                      "dice_weight": args.dice_weight, "ema": args.ema, "head_lr_mult": args.head_lr_mult,
                      "window_loss_boost": args.window_loss_boost, "wallpaper_paste": args.wallpaper_paste,
                      "train_pixels_by_class": dict(zip(CLASSES, counts.tolist())),
                      "semantic_loss_weights": dict(zip(CLASSES, class_weights.cpu().tolist())),
                      "material_pixels": dict(zip(MATERIAL_CLASSES, material_counts.tolist())),
                      "drywall_loss_weights": material_weights.cpu().tolist(),
                      "train_pixels_by_condition": dict(zip(CONDITION_CLASSES, condition_counts.tolist())),
                      "condition_loss_weights": (condition_weights.cpu().tolist() if condition_weights is not None
                                                 else None)}), flush=True)
    train = DataLoader(train_ds, batch_size=args.batch_size, shuffle=sampler is None, sampler=sampler,
                       num_workers=args.workers, pin_memory=device.type == "cuda",
                       persistent_workers=args.workers > 0, drop_last=len(train_ds) > args.batch_size)
    val = DataLoader(val_ds, batch_size=args.batch_size, num_workers=args.workers,
                     pin_memory=device.type == "cuda", persistent_workers=args.workers > 0)
    model = build_model(args.arch, len(CLASSES), pretrained_backbone=not args.no_pretrained,
                        num_material_classes=len(MATERIAL_CLASSES))
    if args.init_checkpoint:
        init = torch.load(args.init_checkpoint, map_location="cpu", weights_only=True)
        if str(init.get("arch", DEFAULT_ARCH)) != args.arch:
            ap.error(f"--init-checkpoint is {init.get('arch')!r}, not --arch {args.arch!r}")
        skipped = warm_start(model, init)
        print(json.dumps({"init_checkpoint": str(args.init_checkpoint), "tensors_not_copied": skipped}), flush=True)
    model = model.to(device)
    optimizer = torch.optim.AdamW(parameter_groups(model, args.lr, args.head_lr_mult, args.weight_decay))
    total_steps = max(1, args.epochs * len(train))
    warmup_steps = int(args.warmup_epochs * len(train))

    def lr_factor(step):
        if step < warmup_steps:
            return (step + 1) / warmup_steps
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1 + math.cos(math.pi * min(progress, 1.0)))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_factor)
    scaler = torch.cuda.amp.GradScaler(enabled=device.type == "cuda")
    ema = ModelEma(model, args.ema) if args.ema > 0 else None
    args.output.mkdir(parents=True, exist_ok=True)
    best = -1.0
    best_metrics = None
    step = 0

    for epoch in range(args.epochs):
        model.train()
        loss_sum = 0.0
        for x, y, material, condition in train:
            x, y, material = x.to(device, non_blocking=True), y.to(device, non_blocking=True), material.to(device, non_blocking=True)
            condition = condition.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"):
                output = model(x)
                pixel_loss = torch.nn.functional.cross_entropy(output["semantic"], y, weight=class_weights,
                                                               ignore_index=IGNORE, reduction="none")
                edge = torch.zeros_like(y, dtype=torch.bool)
                valid_h = (y[:, 1:, :] != IGNORE) & (y[:, :-1, :] != IGNORE) & (y[:, 1:, :] != y[:, :-1, :])
                valid_v = (y[:, :, 1:] != IGNORE) & (y[:, :, :-1] != IGNORE) & (y[:, :, 1:] != y[:, :, :-1])
                edge[:, 1:, :] |= valid_h
                edge[:, :-1, :] |= valid_h
                edge[:, :, 1:] |= valid_v
                edge[:, :, :-1] |= valid_v
                weights = 1.0 + args.edge_boost * edge.float()
                valid = y != IGNORE
                semantic_loss = (pixel_loss * weights * valid).sum() / (weights * valid).sum().clamp_min(1)
                if args.dice_weight:
                    semantic_loss = semantic_loss + args.dice_weight * dice_loss(output["semantic"], y, class_weights)
                hard = material < len(MATERIAL_CLASSES)
                material_loss_map = torch.nn.functional.cross_entropy(output["drywall"], material.masked_fill(~hard, IGNORE),
                                                                      weight=material_weights, ignore_index=IGNORE,
                                                                      reduction="none")
                # Smooth painted board/plaster: soft target 50/50 drywall vs other material.
                either = material == MATERIAL_DRYWALL_OR_OTHER
                log_p = torch.log_softmax(output["drywall"].float(), dim=1)
                either_loss_map = -0.5 * (log_p[:, MATERIAL_CLASSES.index("other_wall_material")]
                                          + log_p[:, MATERIAL_CLASSES.index("drywall")])
                material_valid = hard | either
                material_loss = ((material_loss_map * hard).sum() + (either_loss_map * either).sum()) \
                    / material_valid.sum().clamp_min(1)
                loss = semantic_loss + args.drywall_loss_weight * material_loss
                condition_valid = condition != IGNORE
                if condition_weights is not None and condition_valid.any():
                    condition_loss_map = torch.nn.functional.cross_entropy(output["condition"], condition,
                                                                           weight=condition_weights,
                                                                           ignore_index=IGNORE, reduction="none")
                    condition_loss = (condition_loss_map * condition_valid).sum() / condition_valid.sum().clamp_min(1)
                    loss = loss + args.condition_loss_weight * condition_loss
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            step += 1
            if ema:
                ema.update(model, step)
            loss_sum += loss.item()

        eval_model = ema.module if ema else model
        scores = evaluate(eval_model, val, device, groups=val_groups)
        metrics = {"epoch": epoch + 1, "train_loss": round(loss_sum / max(len(train), 1), 5),
                   **{f"val_{k}": v for k, v in scores.items()},
                   "learning_rate": optimizer.param_groups[0]["lr"]}
        print(json.dumps(metrics), flush=True)
        if scores["mean_iou"] > best:
            best = scores["mean_iou"]
            best_metrics = metrics
            torch.save({"model": eval_model.state_dict(), "classes": CLASSES, "size": args.size,
                        "arch": args.arch,
                        "material_classes": MATERIAL_CLASSES,
                        "condition_classes": CONDITION_CLASSES,
                        "mean": MEAN[:, 0, 0].tolist(), "std": STD[:, 0, 0].tolist(),
                        "metrics": metrics}, args.output / "best.pt")

    # Score the selected checkpoint once on the held-out split (never used for selection).
    best_path = args.output / "best.pt"
    temperature = (1.0, 1.0)
    if best_path.is_file() and not args.no_calibrate:
        best_model, best_checkpoint = load_checkpoint_model(best_path, device)
        temperature, summary = calibrate(best_model, best_checkpoint, val, device, val_groups)
        torch.save(best_checkpoint, best_path)
        print(json.dumps({"calibration": summary}), flush=True)
    if best_path.is_file() and (args.data / args.eval_split / "images").is_dir():
        best_model, best_checkpoint = load_checkpoint_model(best_path, device)
        eval_ds = WallDataset(args.data, args.eval_split, args.size)
        require_coverage(eval_ds.coverage_report(), args.eval_split, False)
        loader = DataLoader(eval_ds, batch_size=args.batch_size, num_workers=args.workers)
        for tta in (False, True):
            scores = evaluate(best_model, loader, device, flip_tta=tta, groups=eval_ds.groups(metadata_path),
                              temperature=temperature)
            record = {"checkpoint": str(best_path), "split": args.eval_split, "size": args.size, "flip_tta": tta,
                      "temperature": temperature, **at_checkpoint_threshold(scores, best_checkpoint), **scores}
            reports.append(record)
            print(json.dumps(record), flush=True)
        if args.report:
            args.report.write_text(json.dumps(reports, indent=2) + "\n")

    if args.registry:
        from provenance import sha256_file
        from run_registry import append as registry_append
        record = registry_append(args.registry, "train", dataset_id=dataset_id,
                                 checkpoint_sha256=sha256_file(best_path) if best_path.is_file() else None,
                                 arch=args.arch,
                                 config={"epochs": args.epochs, "size": args.size, "crop": train_ds.crop,
                                         "scale_range": [args.scale_min, args.scale_max],
                                         "batch_size": args.batch_size, "lr": args.lr,
                                         "head_lr_mult": args.head_lr_mult, "dice_weight": args.dice_weight,
                                         "ema": args.ema, "sampling": args.sampling,
                                         "window_loss_boost": args.window_loss_boost,
                                         "wallpaper_paste": args.wallpaper_paste,
                                         "init_checkpoint": str(args.init_checkpoint) if args.init_checkpoint else None,
                                         "temperature": list(temperature)},
                                 metrics=best_metrics or {}, run_id=args.run_id)
        print(json.dumps({"registry": record["hash"], "seq": record["seq"]}), flush=True)


if __name__ == "__main__":
    main()