import torch
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
 
BASE_CHECKPOINT = "nvidia/mit-b2"
NUM_CLASSES     = 7
WATER_CLASS_ID  = 1
 
CLASS_NAMES = {
    0: "background", 1: "water", 2: "sky",
    3: "vegetation", 4: "building", 5: "vehicle", 6: "person",
}
 
 
def build_processor():
    
    return SegformerImageProcessor.from_pretrained(
        BASE_CHECKPOINT,
        do_resize=True,
        size={"height": 512, "width": 512},
        do_normalize=True,
    )
 
 
def build_model_architecture():
    """Instantiates the 7-class architecture. decode_head starts randomly
    initialized here; the trained weights get loaded on top right after."""
    return SegformerForSemanticSegmentation.from_pretrained(
        BASE_CHECKPOINT,
        num_labels=NUM_CLASSES,
        ignore_mismatched_sizes=True,
        id2label={str(k): v for k, v in CLASS_NAMES.items()},
        label2id={v: str(k) for k, v in CLASS_NAMES.items()},
    )
 
 
def load_trained_model(ckpt_path: str):
    model = build_model_architecture()
    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
 
    missing, unexpected = model.load_state_dict(state, strict=False)
    real_missing = [k for k in missing if not k.endswith("num_batches_tracked")]
    total_keys   = len(model.state_dict())
    ratio        = (total_keys - len(real_missing)) / total_keys
 
    print(f"[LOAD] Loaded {ratio*100:.1f}% of params "
          f"({len(real_missing)} missing, {len(unexpected)} unexpected).")
 
    if ratio < 0.999:
        raise RuntimeError(
            "Checkpoint keys do not match the model architecture. This is "
            "likely the transformers version mismatch known to rename "
            "decode_head submodules. Verify transformers==4.45.2 is installed. "
            f"First missing: {real_missing[:8]} | First unexpected: {unexpected[:8]}"
        )
 
    model.eval()
    
    return model