"""Canopy-height cue for tree_removal from Meta CHMv2 (DINOv3 ViT-L/16 SAT-493M + DPT head, DINOv3 License).

heights(pre) and heights(post) are canopy-height maps in metres; a pixel is a tree-removal cue when the
pre canopy was tall and the post surface is near the ground.
"""
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForDepthEstimation


class CHM:
    def __init__(self, path, device, flip_tta=True, scale=1.0):
        self.proc = AutoImageProcessor.from_pretrained(path)
        self.m = AutoModelForDepthEstimation.from_pretrained(path).to(device).eval()
        self.device, self.flip_tta, self.scale = device, flip_tta, scale

    @torch.no_grad()
    def _run(self, imgs):
        h, w = imgs[0].shape[:2]
        if self.scale != 1.0:
            sz = (int(round(w * self.scale / 16)) * 16, int(round(h * self.scale / 16)) * 16)
            imgs = [np.asarray(Image.fromarray(im).resize(sz, Image.BICUBIC)) for im in imgs]
        inp = {k: v.to(self.device) for k, v in self.proc(images=list(imgs), return_tensors="pt").items()}
        out = self.m(**inp)
        res = self.proc.post_process_depth_estimation(out, target_sizes=[(h, w)] * len(imgs))
        return torch.stack([r["predicted_depth"].float().reshape(h, w) for r in res])

    @torch.no_grad()
    def heights(self, imgs):
        """list of (H,W,3) uint8 -> (N,H,W) float32 tensor (metres)."""
        h = self._run(imgs)
        if self.flip_tta:
            h = (h + self._run([im[:, ::-1].copy() for im in imgs]).flip(-1)) / 2
        return h


def _smooth(h, k):
    return F.avg_pool2d(h[:, None], k, 1, k // 2, count_include_pad=False)[:, 0]


def cue_abs(hp, hq, k=7, lo=3.0, span=4.0, gmax=4.0, gspan=3.0):
    """tall before (>= lo+span m -> 1) and near ground after (<= gmax-gspan m -> 1)."""
    hp, hq = _smooth(hp, k), _smooth(hq, k)
    return ((hp - lo) / span).clamp(0, 1) * ((gmax - hq) / gspan).clamp(0, 1)


def cue_diff(hp, hq, k=7, lo=2.0, span=4.0):
    """height drop >= lo+span m -> 1."""
    return ((_smooth(hp, k) - _smooth(hq, k) - lo) / span).clamp(0, 1)