"""
load.py — mlstac loader for the CHRIS-PROBA1 cloud segmentation model.

This file is executed by mlstac. It must expose two functions that mlstac
calls by contract:

    compiled_model(path, stac_item=None, **kwargs)
        Load the two-checkpoint ensemble (RegNetY + ConvNeXtV2) from `path`
        and return an inference-ready object.

    predict_large(image, model, **kwargs)
        Run tiled inference over a (C, H, W) array using the given model.

The module is self-contained on purpose: it carries the model definition and
the inference logic so it works from Hugging Face without the training repo.

Runtime requirements (install separately):
    torch, segmentation-models-pytorch, pytorch-lightning, timm, numpy
"""
from __future__ import annotations

from pathlib import Path
from typing import Sequence

import numpy as np
import pytorch_lightning as pl
import segmentation_models_pytorch as smp
import torch
import torch.nn.functional as F
from torch import Tensor, nn

# ============================================================================
# Ensemble configuration (CHRIS unified preset)
# ============================================================================
# File name -> timm encoder used to build that checkpoint's network.
ENSEMBLE = [
    ("regnety.ckpt", "tu-regnety_004.pycls_in1k"),
    ("convnextv2.ckpt", "tu-convnextv2_nano.fcmae_ft_in22k_in1k"),
]
IN_CHANNELS = 4          # R, G, B, NIR
NUM_CLASSES = 4          # clear, thick, thin, shadow
PATCH_SIZE = 509
PATCH_OVERLAP = 300


# ============================================================================
# Model definition (must match the trained architecture)
# ============================================================================
class PadToMultiple(nn.Module):
    """Reflect-pad H and W up to a multiple of `multiple`, crop back after."""

    def __init__(self, model: nn.Module, multiple: int = 32):
        super().__init__()
        self.model = model
        self.multiple = multiple

    def forward(self, x: Tensor) -> Tensor:
        h, w = x.shape[-2:]
        ph = (self.multiple - h % self.multiple) % self.multiple
        pw = (self.multiple - w % self.multiple) % self.multiple
        if ph or pw:
            x = F.pad(x, (0, pw, 0, ph), mode="reflect")
        out = self.model(x)
        return out[..., :h, :w]


class OCMSegmenter(pl.LightningModule):
    """U-Net cloud segmenter. Only the inference path is needed here."""

    def __init__(
        self,
        encoder_name: str = "tu-regnety_004.pycls_in1k",
        encoder_weights: str | None = "imagenet",
        in_channels: int = 3,
        num_classes: int = 4,
        ignore_index: int = 99,
        lr: float = 1e-3,
        weight_decay: float = 1e-4,
        pct_start: float = 0.1,
        class_names: Sequence[str] = ("clear", "thick", "thin", "shadow"),
    ):
        super().__init__()
        self.save_hyperparameters()
        base = smp.Unet(
            encoder_name=encoder_name,
            encoder_weights=encoder_weights,
            in_channels=in_channels,
            classes=num_classes,
            activation=None,
        )
        self.model = PadToMultiple(base, multiple=32)
        self.class_names = list(class_names)

    def forward(self, x: Tensor) -> Tensor:
        return self.model(x)


# ============================================================================
# Inference utilities (tiled prediction with gradient blending)
# ============================================================================
def create_gradient_mask(patch_size, patch_overlap, device, dtype):
    if patch_overlap <= 0:
        return torch.ones((patch_size, patch_size), dtype=dtype, device=device)
    if patch_overlap * 2 > patch_size:
        patch_overlap = patch_size // 2
    gradient = torch.ones((patch_size, patch_size), dtype=torch.float32) * patch_overlap
    gradient[:, :patch_overlap] = torch.arange(1, patch_overlap + 1).repeat(patch_size, 1)
    gradient[:, -patch_overlap:] = torch.arange(patch_overlap, 0, -1).repeat(patch_size, 1)
    gradient = gradient / patch_overlap
    combined = torch.rot90(gradient) * gradient
    return combined.to(dtype=dtype, device=device)


def make_patch_indexes(h, w, patch_size, patch_overlap):
    assert patch_size > patch_overlap
    stride = patch_size - patch_overlap
    max_top = h - patch_size
    max_left = w - patch_size
    indexes = []
    for top in range(0, h, stride):
        if top > max_top:
            top = max_top
        bottom = top + patch_size
        for left in range(0, w, stride):
            if left > max_left:
                left = max_left
            right = left + patch_size
            indexes.append((top, bottom, left, right))
    return list(dict.fromkeys(indexes))


def dynamic_zscore(batch, no_data_value=0.0, eps=1e-8):
    valid = batch != no_data_value
    n = valid.sum(dim=(-2, -1), keepdim=True).clamp(min=1)
    mean = (batch * valid).sum(dim=(-2, -1), keepdim=True) / n
    diff_sq = (batch - mean) ** 2 * valid
    std = torch.sqrt(diff_sq.sum(dim=(-2, -1), keepdim=True) / n + eps)
    return torch.where(valid, (batch - mean) / std, torch.zeros_like(batch))


# ============================================================================
# Ensemble wrapper
# ============================================================================
class OCMEnsemble:
    """Holds the loaded models and runs averaged, tiled inference."""

    def __init__(self, models, num_classes, patch_size, patch_overlap,
                 device, dtype, batch_size=4):
        self.models = models
        self.num_classes = num_classes
        self.patch_size = patch_size
        self.patch_overlap = patch_overlap
        self.device = torch.device(device)
        self.dtype = dtype
        self.batch_size = batch_size

    @torch.inference_mode()
    def _predict_batch(self, batch):
        acc = None
        for model in self.models:
            logits = model(batch)
            acc = logits if acc is None else acc + logits
        return acc / len(self.models)

    def predict_array(self, image, apply_nodata_mask=True,
                      return_probs=False, verbose=False):
        assert image.ndim == 3 and image.shape[0] == IN_CHANNELS, (
            f"Expected shape ({IN_CHANNELS}, H, W), got {image.shape}"
        )
        _, H, W = image.shape
        patch_size = min(self.patch_size, H, W)
        overlap = min(self.patch_overlap, patch_size - 1)
        patch_indexes = make_patch_indexes(H, W, patch_size, overlap)
        gradient = create_gradient_mask(patch_size, overlap, self.device, self.dtype)

        pred_tracker = torch.zeros((self.num_classes, H, W), dtype=self.dtype, device=self.device)
        weight_tracker = torch.zeros((H, W), dtype=self.dtype, device=self.device)
        image_tensor = torch.from_numpy(np.ascontiguousarray(image)).to(
            device=self.device, dtype=self.dtype
        )

        iterator = range(0, len(patch_indexes), self.batch_size)
        if verbose:
            from tqdm.auto import tqdm
            iterator = tqdm(iterator, desc="Inference", leave=False)

        for batch_start in iterator:
            batch_idx = patch_indexes[batch_start:batch_start + self.batch_size]
            patches = [image_tensor[:, t:b, l:r] for (t, b, l, r) in batch_idx]
            batch = dynamic_zscore(torch.stack(patches, dim=0))
            probs = torch.softmax(self._predict_batch(batch), dim=1)
            for p, (t, b, l, r) in zip(probs, batch_idx):
                pred_tracker[:, t:b, l:r] += p * gradient[None, :, :]
                weight_tracker[t:b, l:r] += gradient

        weight_tracker = weight_tracker.clamp(min=1e-8)
        pred_tracker = pred_tracker / weight_tracker[None, :, :]

        if return_probs:
            result = pred_tracker.cpu().float().numpy()
        else:
            result = pred_tracker.argmax(dim=0).cpu().numpy().astype(np.uint8)
            if apply_nodata_mask:
                nodata = (image == 0).all(axis=0)
                result[nodata] = 99
        return result


# ============================================================================
# Checkpoint loading
# ============================================================================
def _load_ckpt(ckpt_path, encoder_name, device, dtype):
    # encoder_weights=None: weights come from the checkpoint, so we avoid
    # downloading ImageNet weights at build time.
    # strict=False: the checkpoint also stored loss_fn.weight, which the
    # inference model doesn't have. Ignoring that extra key is safe.
    model = OCMSegmenter.load_from_checkpoint(
        str(ckpt_path),
        encoder_name=encoder_name,
        encoder_weights=None,
        in_channels=IN_CHANNELS,
        num_classes=NUM_CLASSES,
        map_location=device,
        strict=False,
    )
    model.eval().to(device=device, dtype=dtype)
    for p in model.parameters():
        p.requires_grad = False
    return model


# ============================================================================
# mlstac contract
# ============================================================================
def compiled_model(path, stac_item=None, *, device="cuda", dtype=torch.float32,
                   batch_size=4, **kwargs):
    """Load the ensemble from `path` and return an inference-ready object.

    Args:
        path: Local folder holding regnety.ckpt and convnextv2.ckpt.
        stac_item: STAC metadata (unused here, passed by mlstac).
        device: 'cuda', 'cuda:0' or 'cpu'.
        dtype: torch dtype for the weights.
        batch_size: tiles processed per forward pass.

    Returns:
        An OCMEnsemble with a .predict_array(image) method.
    """
    path = Path(path)
    if not torch.cuda.is_available() and str(device).startswith("cuda"):
        device = "cpu"

    models = []
    for filename, encoder_name in ENSEMBLE:
        ckpt = path / filename
        if not ckpt.exists():
            raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
        models.append(_load_ckpt(ckpt, encoder_name, device, dtype))

    return OCMEnsemble(
        models=models,
        num_classes=NUM_CLASSES,
        patch_size=PATCH_SIZE,
        patch_overlap=PATCH_OVERLAP,
        device=device,
        dtype=dtype,
        batch_size=batch_size,
    )


def predict_large(image, model, **kwargs):
    """Run tiled inference over a (C, H, W) RGBN array.

    Args:
        image: np.ndarray of shape (4, H, W), bands ordered R, G, B, NIR.
        model: object returned by compiled_model().
        return_probs: if True, return (num_classes, H, W) probabilities.
        apply_nodata_mask: if True, mark all-zero pixels as 99.
        verbose: show a progress bar.

    Returns:
        (H, W) uint8 label map, or (num_classes, H, W) probabilities.
    """
    return model.predict_array(
        np.asarray(image),
        apply_nodata_mask=kwargs.get("apply_nodata_mask", True),
        return_probs=kwargs.get("return_probs", False),
        verbose=kwargs.get("verbose", False),
    )


# ============================================================================
# CHRIS/PROBA-1 preprocessing (the 6 acquisition modes)
# ============================================================================
# For each CHRIS mode, which raw band indices (1-based) average into each of
# the R, G, B, NIR groups. Taken from the inference pipeline.
BAND_SELECTION = {
    1: {"B4": [23, 24, 25],
        "B3": [13, 14, 15],
        "B2": [4, 5, 6, 7, 8, 9, 10],
        "B8": [43, 44, 45, 46, 47, 48, 49, 50, 51]},
    2: {"B4": [10, 11, 12], "B3": [6, 7], "B2": [3, 4], "B8": [17]},
    3: {"B4": [7, 8],       "B3": [4, 5], "B2": [2],    "B8": [15]},
    4: {"B4": [4],          "B3": [2],    "B2": [1],    "B8": [18]},
    5: {"B4": [7, 8],       "B3": [4, 5], "B2": [2],    "B8": [23, 24, 25, 26]},
    6: {"B2": [1],          "B3": [2],    "B4": [3],    "B8": [4]},
}
BAND_KEYS_RGBN = ["B4", "B3", "B2", "B8"]   # R, G, B, NIR

DN_SCALE = 100_000.0
TOA_SCALE = 10_000.0
CAP = 5.0   # clip ceiling used during training; keep it identical at inference


def _avg_bands(cube, band_list):
    """Average the given 1-based raw bands into a single (H, W) layer."""
    idx = [b - 1 for b in band_list]
    return cube[idx[0]] if len(idx) == 1 else cube[idx].mean(axis=0)


def _infer_source_from_name(tif_path):
    """Guess 'dn' or 'toa' from the file name. Returns None if unclear."""
    name = Path(tif_path).name.lower()
    if "toa" in name:
        return "toa"
    if "dn" in name:
        return "dn"
    return None


def build_rgbn(cube, mode_n, source):
    """Build the (4, H, W) RGBN stack and the nodata mask from a raw cube.

    Args:
        cube: (bands, H, W) float array read from the CHRIS GeoTIFF.
        mode_n: CHRIS acquisition mode, 1 to 6.
        source: 'dn' or 'toa', selects the radiometric scale.

    Returns:
        (stack, nodata) where stack is (4, H, W) float32 and nodata is (H, W) bool.
    """
    sel = BAND_SELECTION.get(int(mode_n))
    if sel is None:
        raise ValueError(f"Unsupported CHRIS mode {mode_n}; expected 1-6.")

    needed_max = max(max(sel[k]) for k in BAND_KEYS_RGBN)
    if cube.shape[0] < needed_max:
        raise ValueError(
            f"Cube has {cube.shape[0]} bands but mode {mode_n} needs {needed_max}."
        )

    # nodata = pixel is 0 in every RGBN group
    H, W = cube.shape[1:]
    nodata = np.ones((H, W), dtype=bool)
    for k in BAND_KEYS_RGBN:
        nodata &= (_avg_bands(cube, sel[k]) == 0)

    scale = DN_SCALE if source == "dn" else TOA_SCALE
    layers = [_avg_bands(cube, sel[k]) for k in BAND_KEYS_RGBN]
    stack = np.stack(layers, axis=0).astype(np.float32) / scale
    stack = np.clip(stack, 0.0, CAP)
    stack[:, nodata] = 0.0
    return stack, nodata


def predict_chris(tif_path, model, mode_n, source=None, **kwargs):
    """Segment a raw CHRIS/PROBA-1 GeoTIFF end to end.

    Reads the cube, builds the RGBN stack for the given mode, runs the
    ensemble, and restores nodata as 99.

    Args:
        tif_path: path to the CHRIS GeoTIFF.
        model: object returned by compiled_model().
        mode_n: CHRIS acquisition mode, 1 to 6.
        source: 'dn' or 'toa'. If None, it is guessed from the file name.
        return_probs: if True, return (num_classes, H, W) probabilities.

    Returns:
        (H, W) uint8 label map (nodata = 99), or probabilities if requested.
    """
    import rasterio as rio

    if source is None:
        source = _infer_source_from_name(tif_path)
        if source is None:
            raise ValueError(
                "Could not tell DN from TOA by the file name; "
                "pass source='dn' or source='toa'."
            )

    with rio.open(tif_path) as src:
        cube = src.read().astype(np.float32)

    stack, nodata = build_rgbn(cube, mode_n, source)

    return_probs = kwargs.get("return_probs", False)
    pred = model.predict_array(
        stack,
        apply_nodata_mask=False,   # we apply the CHRIS nodata mask below
        return_probs=return_probs,
        verbose=kwargs.get("verbose", False),
    )
    if not return_probs:
        pred = np.asarray(pred).astype(np.uint8)
        pred[nodata] = 99
    return pred