#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")

import csv
import json
import time
import argparse
import random
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.amp import autocast, GradScaler

from transformers import SegformerForSemanticSegmentation


# ============================================================
# Seed / util
# ============================================================

def set_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)


def load_json(path: str | Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def save_json(path: Path, data: dict):
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


# ============================================================
# Labelmap
# ============================================================

def load_labelmap(labelmap_path: str):
    """
    Tenta usar utils.carregar_labelmap_completo.
    Fallback simples para labelmap com uma classe por linha.
    """
    try:
        from helpers import carregar_labelmap_completo, _infer_ignore_id

        _cor_para_id, _colormap_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)
        ignore_id = _infer_ignore_id(ignore_rgb, 255)

        id2label = {int(k): str(v) for k, v in id_para_nome.items()}
        label2id = {v.lower(): k for k, v in id2label.items()}

        return id2label, label2id, int(ignore_id)

    except Exception as e:
        print(f"[WARN] Não consegui usar utils.carregar_labelmap_completo: {e}")
        print("[WARN] Usando parser simples: uma classe por linha.")

    id2label = {}
    with open(labelmap_path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s or s.startswith("#"):
                continue

            # Aceita:
            #   classe
            #   id classe
            #   id,classe,...
            parts = s.replace(",", " ").split()

            if len(parts) >= 2 and parts[0].isdigit():
                cid = int(parts[0])
                name = parts[1]
            else:
                cid = len(id2label)
                name = parts[0]

            if name.lower() in ("ignore", "void", "background_ignore"):
                continue

            id2label[cid] = name

    label2id = {v.lower(): k for k, v in id2label.items()}
    return id2label, label2id, 255


# ============================================================
# Dataset novo OAK-FCC-3
# ============================================================

class OakFcc3TensorSegDataset(Dataset):
    """
    Lê o contrato pós-normalização/split:

      root/
        group/<grupo>/tensors/<base>.npy  # float32 CHW [5,H,W]
        group/<grupo>/masks/<base>.npy    # int HW
    """

    def __init__(
        self,
        root: str | Path,
        channels: int = 5,
        strict_channels: bool = True,
        resize_hw: Optional[Tuple[int, int]] = None,
        ignore_index: int = 255,
    ):
        self.root = Path(root)
        self.channels = int(channels)
        self.strict_channels = bool(strict_channels)
        self.resize_hw = resize_hw
        self.ignore_index = int(ignore_index)

        self.samples = self._collect_samples()

        if not self.samples:
            raise RuntimeError(f"Nenhuma amostra encontrada em: {self.root}")

    def _collect_samples(self):
        samples = []

        group_root = self.root / "group"
        if not group_root.is_dir():
            raise RuntimeError(f"Pasta group não encontrada em: {self.root}")

        for group_dir in sorted(group_root.iterdir()):
            if not group_dir.is_dir():
                continue

            tensors_dir = group_dir / "tensors"
            masks_dir = group_dir / "masks"
            metas_dir = group_dir / "metas"
            previews_dir = group_dir / "previews"

            if not tensors_dir.is_dir() or not masks_dir.is_dir():
                continue

            for tensor_path in sorted(tensors_dir.glob("*.npy")):
                base = tensor_path.stem
                mask_path = masks_dir / f"{base}.npy"

                if not mask_path.exists():
                    print(f"[WARN] Sem mask para tensor: {tensor_path}")
                    continue

                meta_path = metas_dir / f"{base}.json"
                preview_path = previews_dir / f"{base}.png"

                samples.append({
                    "group": group_dir.name,
                    "base": base,
                    "tensor": tensor_path,
                    "mask": mask_path,
                    "meta": meta_path if meta_path.exists() else None,
                    "preview": preview_path if preview_path.exists() else None,
                })

        return samples

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

    def _resize_tensor_mask(self, x: torch.Tensor, y: torch.Tensor):
        if self.resize_hw is None:
            return x, y

        h, w = self.resize_hw

        if x.shape[-2:] != (h, w):
            x = F.interpolate(
                x.unsqueeze(0),
                size=(h, w),
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)

        if y.shape[-2:] != (h, w):
            y = F.interpolate(
                y.unsqueeze(0).unsqueeze(0).float(),
                size=(h, w),
                mode="nearest",
            ).squeeze(0).squeeze(0).long()

        return x, y

    def __getitem__(self, idx):
        s = self.samples[idx]

        x = np.load(str(s["tensor"])).astype(np.float32)
        y = np.load(str(s["mask"]))

        if x.ndim != 3:
            raise RuntimeError(f"Tensor inválido {s['tensor']}: shape={x.shape}")

        if self.strict_channels and x.shape[0] != self.channels:
            raise RuntimeError(
                f"Channels inválido em {s['tensor']}: veio {x.shape[0]}, esperado {self.channels}"
            )

        if x.shape[0] > self.channels:
            x = x[:self.channels]

        y = y.astype(np.int64)

        xt = torch.from_numpy(np.ascontiguousarray(x)).float()
        yt = torch.from_numpy(np.ascontiguousarray(y)).long()

        xt, yt = self._resize_tensor_mask(xt, yt)

        return {
            "image": xt,
            "mask": yt,
            "group": s["group"],
            "base": s["base"],
        }


def collate_fn(batch):
    imgs = torch.stack([b["image"] for b in batch], dim=0)
    masks = torch.stack([b["mask"] for b in batch], dim=0)
    return imgs, masks


# ============================================================
# Normalização
# ============================================================

class FixedNormalizer(nn.Module):
    def __init__(self, mean: List[float], std: List[float]):
        super().__init__()
        mean_t = torch.tensor(mean, dtype=torch.float32).view(1, -1, 1, 1)
        std_t = torch.tensor(std, dtype=torch.float32).view(1, -1, 1, 1)
        self.register_buffer("mean", mean_t)
        self.register_buffer("std", torch.clamp(std_t, min=1e-6))

    def forward(self, x):
        return (x - self.mean) / self.std


def normalize_per_batch(x: torch.Tensor, eps: float = 1e-6):
    mean = x.mean(dim=(0, 2, 3), keepdim=True)
    std = x.std(dim=(0, 2, 3), keepdim=True).clamp_min(eps)
    return (x - mean) / std


def build_normalizer(config: dict, args, device: torch.device):
    channels = int(config.get("channels", 5))
    model_family = config.get("modelo", "segformer")
    model_name = config.get("model_name", "model")
    stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}")

    candidates = []

    if args.norm_stats:
        candidates.append(Path(args.norm_stats))

    candidates.append(Path("dataset") / f"{config['resolucao'][0]}x{config['resolucao'][1]}" / "group" / "norm_stats.json")
    candidates.append(Path("backup") / model_family / model_name / stats_source_tag / "norm_stats.json")

    for p in candidates:
        if p.exists():
            stats = load_json(p)
            mean = stats.get("mean", [])
            std = stats.get("std", [])
            stat_channels = stats.get("channels", [])

            if len(mean) != channels or len(std) != channels:
                raise RuntimeError(
                    f"norm_stats incompatível: {p} | channels={channels}, mean={len(mean)}, std={len(std)}"
                )

            print(f"[NORM] usando stats fixos: {p}")
            print(f"[NORM] channels={stat_channels}")
            return FixedNormalizer(mean, std).to(device), str(p)

    print("[NORM] norm_stats não encontrado. Usando normalize_per_batch.")
    return None, None


# ============================================================
# Modelo
# ============================================================

def patch_segformer_input_channels(model: nn.Module, in_ch: int):
    """
    Altera o primeiro patch embedding do SegFormer para aceitar C canais.
    Inicializa canais extras pela média dos pesos RGB.
    """
    if in_ch == 3:
        return model

    proj = model.segformer.encoder.patch_embeddings[0].proj

    if proj.in_channels == in_ch:
        return model

    old_weight = proj.weight.data.clone()
    old_bias = proj.bias.data.clone() if proj.bias is not None else None

    new_proj = nn.Conv2d(
        in_channels=in_ch,
        out_channels=proj.out_channels,
        kernel_size=proj.kernel_size,
        stride=proj.stride,
        padding=proj.padding,
        dilation=proj.dilation,
        groups=proj.groups,
        bias=proj.bias is not None,
        padding_mode=proj.padding_mode,
    )

    with torch.no_grad():
        if in_ch <= old_weight.shape[1]:
            new_proj.weight.copy_(old_weight[:, :in_ch, :, :])
        else:
            new_proj.weight[:, :old_weight.shape[1], :, :].copy_(old_weight)
            extra = in_ch - old_weight.shape[1]
            mean_w = old_weight.mean(dim=1, keepdim=True)
            new_proj.weight[:, old_weight.shape[1]:, :, :].copy_(mean_w.repeat(1, extra, 1, 1))

        if old_bias is not None:
            new_proj.bias.copy_(old_bias)

    model.segformer.encoder.patch_embeddings[0].proj = new_proj
    model.config.num_channels = in_ch

    print(f"[MODEL] patch input channels: 3 -> {in_ch}")
    return model


def build_model(
    backbone: str,
    num_classes: int,
    channels: int,
    id2label: Dict[int, str],
    label2id: Dict[str, int],
):
    model = SegformerForSemanticSegmentation.from_pretrained(
        backbone,
        num_labels=num_classes,
        id2label={int(k): str(v) for k, v in id2label.items()},
        label2id={str(k): int(v) for k, v in label2id.items()},
        ignore_mismatched_sizes=True,
    )

    patch_segformer_input_channels(model, channels)

    return model


# ============================================================
# Métricas / loss
# ============================================================

@torch.no_grad()
def update_confusion_matrix(cm, preds, labels, num_classes, ignore_index=255):
    preds = preds.reshape(-1)
    labels = labels.reshape(-1)

    valid = labels != ignore_index
    preds = preds[valid]
    labels = labels[valid]

    valid2 = (labels >= 0) & (labels < num_classes)
    preds = preds[valid2]
    labels = labels[valid2]

    if labels.numel() == 0:
        return

    idx = labels * num_classes + preds
    bins = torch.bincount(idx, minlength=num_classes * num_classes)
    cm += bins.view(num_classes, num_classes)


@torch.no_grad()
def compute_iou_from_cm(cm, eps=1e-6):
    cm = cm.float()
    tp = torch.diag(cm)
    fp = cm.sum(0) - tp
    fn = cm.sum(1) - tp
    denom = tp + fp + fn + eps
    iou = (tp / denom).cpu().tolist()
    miou = float(np.mean(iou))
    return miou, iou


@torch.no_grad()
def compute_pixel_acc_from_cm(cm, eps=1e-6):
    cm = cm.float()
    return float(torch.diag(cm).sum() / (cm.sum() + eps))


def dice_loss(logits, target, num_classes, ignore_index=255, smooth=1.0):
    probs = torch.softmax(logits, dim=1)
    valid = target != ignore_index

    if valid.sum() == 0:
        return logits.new_tensor(0.0)

    target_clamped = target.clone()
    target_clamped[~valid] = 0
    target_clamped = target_clamped.long()

    target_1h = F.one_hot(target_clamped, num_classes=num_classes)
    target_1h = target_1h.permute(0, 3, 1, 2).float()

    valid_f = valid.unsqueeze(1).float()
    probs = probs * valid_f
    target_1h = target_1h * valid_f

    dims = (0, 2, 3)
    inter = (probs * target_1h).sum(dims)
    den = probs.sum(dims) + target_1h.sum(dims)

    dice = (2.0 * inter + smooth) / (den + smooth)
    return 1.0 - dice.mean()


def estimate_class_weights(ds, num_classes, ignore_index=255, max_samples=800, seed=42):
    rng = np.random.default_rng(seed)
    n = min(len(ds), max_samples)
    idxs = rng.choice(len(ds), size=n, replace=False)

    counts = np.zeros(num_classes, dtype=np.float64)

    for i in idxs:
        item = ds[i]
        m = item["mask"].numpy().reshape(-1)
        m = m[m != ignore_index]
        m = m[(m >= 0) & (m < num_classes)]

        if m.size > 0:
            counts += np.bincount(m, minlength=num_classes)[:num_classes]

    freq = counts / (counts.sum() + 1e-12)
    freq = np.clip(freq, 1e-12, 1.0)

    weights = 1.0 / np.log(1.02 + freq)
    weights = weights / weights.mean()

    return torch.tensor(weights, dtype=torch.float32), counts


# ============================================================
# Train / Val
# ============================================================

def run_one_epoch(
    model,
    loader,
    optimizer,
    device,
    num_classes,
    ignore_index,
    criterion,
    amp,
    scaler,
    train,
    grad_accum=1,
    normalizer=None,
):
    model.train(train)

    total_loss = 0.0
    n_batches = 0
    cm = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)

    t0 = time.time()

    with torch.set_grad_enabled(train):
        if train and optimizer is not None:
            optimizer.zero_grad(set_to_none=True)

        for step, (imgs, masks) in enumerate(loader):
            imgs = imgs.to(device, non_blocking=True)
            masks = masks.to(device, non_blocking=True)

            if normalizer is not None:
                imgs = normalizer(imgs)
            else:
                imgs = normalize_per_batch(imgs)

            with autocast(device_type="cuda", enabled=amp and device.type == "cuda"):
                out = model(pixel_values=imgs)
                logits = out.logits

                if logits.shape[-2:] != masks.shape[-2:]:
                    logits = F.interpolate(
                        logits,
                        size=masks.shape[-2:],
                        mode="bilinear",
                        align_corners=False,
                    )

                ce = criterion(logits, masks)
                dice = dice_loss(
                    logits=logits,
                    target=masks,
                    num_classes=num_classes,
                    ignore_index=ignore_index,
                    smooth=1.0,
                )

                loss = 0.7 * ce + 0.3 * dice

                if train and grad_accum > 1:
                    loss = loss / grad_accum

            if train and optimizer is not None:
                if amp and scaler is not None and device.type == "cuda":
                    scaler.scale(loss).backward()

                    if ((step + 1) % grad_accum) == 0 or (step + 1) == len(loader):
                        scaler.step(optimizer)
                        scaler.update()
                        optimizer.zero_grad(set_to_none=True)
                else:
                    loss.backward()

                    if ((step + 1) % grad_accum) == 0 or (step + 1) == len(loader):
                        optimizer.step()
                        optimizer.zero_grad(set_to_none=True)

            loss_value = float(loss.item()) * (grad_accum if train and grad_accum > 1 else 1.0)
            total_loss += loss_value
            n_batches += 1

            with torch.no_grad():
                preds = torch.argmax(logits, dim=1)
                update_confusion_matrix(cm, preds, masks, num_classes, ignore_index)

    avg_loss = total_loss / max(1, n_batches)
    miou, iou_per_class = compute_iou_from_cm(cm)
    acc = compute_pixel_acc_from_cm(cm)

    return {
        "loss": avg_loss,
        "miou": miou,
        "iou_per_class": iou_per_class,
        "acc": acc,
        "time_s": time.time() - t0,
    }


# ============================================================
# Checkpoint / logs
# ============================================================

def save_checkpoint(path, model, optimizer, scaler, epoch, best_miou, best_main_iou, extra=None):
    ckpt = {
        "epoch": epoch,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "best_miou": best_miou,
        "best_main_iou": best_main_iou,
    }

    if scaler is not None:
        ckpt["scaler"] = scaler.state_dict()

    if extra:
        ckpt["extra"] = extra

    torch.save(ckpt, path)


def load_checkpoint(path, model, optimizer=None, scaler=None, map_location="cpu"):
    ckpt = torch.load(path, map_location=map_location, weights_only=False)
    model.load_state_dict(ckpt["model"], strict=True)

    if optimizer is not None and "optimizer" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer"])

    if scaler is not None and "scaler" in ckpt:
        scaler.load_state_dict(ckpt["scaler"])

    return ckpt


def append_train_log(path: Path, row: dict):
    ensure_dir(path.parent)
    exists = path.exists()

    with path.open("a", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(row.keys()))
        if not exists:
            w.writeheader()
        w.writerow(row)


# ============================================================
# Main
# ============================================================

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--config", default="config.json")
    parser.add_argument("--epochs", type=int, default=80)
    parser.add_argument("--batch", type=int, default=1)
    parser.add_argument("--lr", type=float, default=3e-5)
    parser.add_argument("--wd", type=float, default=0.01)
    parser.add_argument("--num_workers", type=int, default=2)

    parser.add_argument("--amp", action="store_true")
    parser.add_argument("--amp_val", action="store_true")
    parser.add_argument("--grad_accum", type=int, default=4)
    parser.add_argument("--grad_ckpt", action="store_true")

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--ignore_index", type=int, default=None)

    parser.add_argument("--class_weights", default="auto")
    parser.add_argument("--main_class", default=None)

    parser.add_argument("--norm_stats", default=None)
    parser.add_argument("--src-root", default="dataset/split")
    parser.add_argument("--save-every", type=int, default=10)

    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--resume-ckpt", default=None)

    parser.add_argument("--early-stop", type=int, default=25)

    args = parser.parse_args()

    set_seed(args.seed)

    config = load_json(args.config)

    if config.get("dual_head", False):
        raise RuntimeError("Este script está preparado para single-head. No config, use dual_head=false.")

    W, H = config["resolucao"]
    channels = int(config.get("channels", 5))
    backbone = config.get("backbone", "nvidia/mit-b1")
    fusion_mode = config.get("fusion_mode", "stacked")

    if fusion_mode != "stacked":
        raise RuntimeError("Este script é para fusion_mode='stacked'.")

    model_family = config.get("modelo", "segformer")
    model_name = config.get("model_name", "test")
    stats_source_tag = config.get("stats_source_tag", f"stacked_raw{channels}")

    main_class_name = str(args.main_class or config.get("main_class_name", "cana")).lower()

    save_dir = Path("backup") / model_family / model_name / f"{fusion_mode}_raw{channels}"
    ensure_dir(save_dir)

    labelmap_path = Path("dataset") / "labelmap.txt"
    if not labelmap_path.exists():
        raise RuntimeError(f"Labelmap não encontrado: {labelmap_path}")

    id2label, label2id, ignore_from_labelmap = load_labelmap(str(labelmap_path))
    ignore_index = int(args.ignore_index if args.ignore_index is not None else ignore_from_labelmap)

    num_classes = len(id2label)

    if num_classes <= 1:
        raise RuntimeError(f"num_classes inválido: {num_classes}")

    main_class_id = label2id.get(main_class_name)

    print("==========================================")
    print("Train SegFormer OAK-FCC-3")
    print(f"Backbone     : {backbone}")
    print(f"Save dir     : {save_dir}")
    print(f"Split root   : {args.src_root}")
    print(f"Resolution   : {W}x{H}")
    print(f"Channels     : {channels}")
    print(f"Classes      : {num_classes} -> {id2label}")
    print(f"Ignore index : {ignore_index}")
    print(f"Main class   : {main_class_name} -> {main_class_id}")
    print("==========================================")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    train_root = Path(args.src_root) / "train"
    val_root = Path(args.src_root) / "val"

    resize_hw = (H, W)

    ds_train = OakFcc3TensorSegDataset(
        train_root,
        channels=channels,
        strict_channels=True,
        resize_hw=resize_hw,
        ignore_index=ignore_index,
    )

    ds_val = OakFcc3TensorSegDataset(
        val_root,
        channels=channels,
        strict_channels=True,
        resize_hw=resize_hw,
        ignore_index=ignore_index,
    )

    print(f"[DATA] train={len(ds_train)} | val={len(ds_val)}")

    dl_train = DataLoader(
        ds_train,
        batch_size=args.batch,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
        drop_last=True if len(ds_train) >= args.batch else False,
    )

    dl_val = DataLoader(
        ds_val,
        batch_size=1,
        shuffle=False,
        num_workers=max(0, args.num_workers // 2),
        pin_memory=True,
        collate_fn=collate_fn,
        drop_last=False,
    )

    normalizer, norm_stats_path = build_normalizer(config, args, device)

    model = build_model(
        backbone=backbone,
        num_classes=num_classes,
        channels=channels,
        id2label=id2label,
        label2id=label2id,
    )

    if args.grad_ckpt:
        try:
            model.gradient_checkpointing_enable()
            print("[MODEL] gradient checkpointing enabled")
        except Exception as e:
            print(f"[WARN] gradient checkpointing não suportado: {e}")

    model.to(device)

    if args.class_weights.lower() == "none":
        weights = None

    elif args.class_weights.lower() == "auto":
        w, counts = estimate_class_weights(
            ds_train,
            num_classes=num_classes,
            ignore_index=ignore_index,
            seed=args.seed,
        )
        weights = w.to(device)
        print("[LOSS] class counts:", counts.astype(int).tolist())
        print("[LOSS] class weights:", w.cpu().numpy().round(3).tolist())

    else:
        parts = [float(x) for x in args.class_weights.split(",")]
        if len(parts) != num_classes:
            raise RuntimeError(f"--class_weights precisa ter {num_classes} valores.")
        weights = torch.tensor(parts, dtype=torch.float32, device=device)

    criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=ignore_index)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.lr,
        weight_decay=args.wd,
    )

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=0.5,
        patience=6,
        threshold=1e-4,
    )

    scaler = GradScaler(enabled=args.amp and device.type == "cuda")

    last_path = save_dir / "last.pt"
    best_miou_path = save_dir / "best_miou.pt"
    best_main_path = save_dir / "best_main.pt"
    train_log_path = save_dir / "train_log.csv"

    save_json(save_dir / "train_config_snapshot.json", {
        "config": config,
        "args": vars(args),
        "id2label": id2label,
        "label2id": label2id,
        "ignore_index": ignore_index,
        "norm_stats_path": norm_stats_path,
    })

    start_epoch = 1
    best_miou = -1.0
    best_main_iou = -1.0
    epochs_without_improve = 0

    resume_path = Path(args.resume_ckpt) if args.resume_ckpt else last_path

    if args.resume and resume_path.exists():
        ckpt = load_checkpoint(
            resume_path,
            model,
            optimizer,
            scaler=scaler,
            map_location="cpu",
        )
        start_epoch = int(ckpt["epoch"]) + 1
        best_miou = float(ckpt.get("best_miou", -1.0))
        best_main_iou = float(ckpt.get("best_main_iou", -1.0))
        print(f"[RESUME] {resume_path} epoch={start_epoch}")

    def pretty_iou(iou_list):
        return " | ".join([
            f"{id2label.get(i, i)}:{v:.3f}"
            for i, v in enumerate(iou_list)
        ])

    for epoch in range(start_epoch, args.epochs + 1):
        lr_now = optimizer.param_groups[0]["lr"]
        print(f"\n==== Epoch {epoch}/{args.epochs} | lr={lr_now:.2e} ====")

        if device.type == "cuda":
            torch.cuda.empty_cache()

        tr = run_one_epoch(
            model=model,
            loader=dl_train,
            optimizer=optimizer,
            device=device,
            num_classes=num_classes,
            ignore_index=ignore_index,
            criterion=criterion,
            amp=args.amp,
            scaler=scaler,
            train=True,
            grad_accum=max(1, args.grad_accum),
            normalizer=normalizer,
        )

        if device.type == "cuda":
            torch.cuda.empty_cache()

        va = run_one_epoch(
            model=model,
            loader=dl_val,
            optimizer=None,
            device=device,
            num_classes=num_classes,
            ignore_index=ignore_index,
            criterion=criterion,
            amp=args.amp_val,
            scaler=None,
            train=False,
            grad_accum=1,
            normalizer=normalizer,
        )

        scheduler.step(va["loss"])

        main_iou = va["miou"] if main_class_id is None else va["iou_per_class"][main_class_id]

        print(
            f"TRAIN: loss={tr['loss']:.4f} acc={tr['acc']:.4f} "
            f"miou={tr['miou']:.4f} t={tr['time_s']:.1f}s"
        )
        print(
            f"VAL  : loss={va['loss']:.4f} acc={va['acc']:.4f} "
            f"miou={va['miou']:.4f} main_iou={float(main_iou):.4f} t={va['time_s']:.1f}s"
        )
        print("IoU:", pretty_iou(va["iou_per_class"]))

        save_checkpoint(
            last_path,
            model,
            optimizer,
            scaler=scaler,
            epoch=epoch,
            best_miou=best_miou,
            best_main_iou=best_main_iou,
            extra={
                "train": tr,
                "val": va,
                "main_iou": float(main_iou),
            },
        )

        if args.save_every > 0 and epoch % args.save_every == 0:
            save_checkpoint(
                save_dir / f"epoch_{epoch:04d}.pt",
                model,
                optimizer,
                scaler=scaler,
                epoch=epoch,
                best_miou=best_miou,
                best_main_iou=best_main_iou,
            )

        improved = False

        if va["miou"] > best_miou:
            best_miou = va["miou"]
            improved = True

            save_checkpoint(
                best_miou_path,
                model,
                optimizer,
                scaler=scaler,
                epoch=epoch,
                best_miou=best_miou,
                best_main_iou=best_main_iou,
                extra={"val": va, "main_iou": float(main_iou)},
            )

            print(f"[BEST mIoU] {best_miou:.4f} -> {best_miou_path}")

        if float(main_iou) > best_main_iou:
            best_main_iou = float(main_iou)
            improved = True

            save_checkpoint(
                best_main_path,
                model,
                optimizer,
                scaler=scaler,
                epoch=epoch,
                best_miou=best_miou,
                best_main_iou=best_main_iou,
                extra={"val": va, "main_iou": float(main_iou)},
            )

            print(f"[BEST MAIN] {best_main_iou:.4f} -> {best_main_path}")

        append_train_log(train_log_path, {
            "epoch": epoch,
            "lr": lr_now,
            "train_loss": tr["loss"],
            "train_acc": tr["acc"],
            "train_miou": tr["miou"],
            "val_loss": va["loss"],
            "val_acc": va["acc"],
            "val_miou": va["miou"],
            "val_main_iou": float(main_iou),
            "best_miou": best_miou,
            "best_main_iou": best_main_iou,
        })

        if improved:
            epochs_without_improve = 0
        else:
            epochs_without_improve += 1

        if args.early_stop > 0 and epochs_without_improve >= args.early_stop:
            print(f"[EARLY STOP] {epochs_without_improve} épocas sem melhora.")
            break

    print("\nTreino finalizado.")
    print(f"Best mIoU     : {best_miou:.4f}")
    print(f"Best main IoU : {best_main_iou:.4f}")
    print(f"Save dir      : {save_dir}")


if __name__ == "__main__":
    main()