"""Optional CUDA for the two heavy numeric steps (colour clustering, per-label smoothing), and the Hugging Face ZeroGPU hook.

Nothing here is required: with no torch, no CUDA device, or HATCHLING_GPU=0, ling_hatch uses scikit-learn and OpenCV as before.
The GPU path gives the same kind of result, not bit-identical labels, so a sheet made on the GPU can differ in a few edge pixels
from the same sheet made on the CPU.
"""
import functools
import math
import os

try:
    import spaces                       # present on Hugging Face Spaces; on ZeroGPU the GPU is only attached inside @spaces.GPU calls
except Exception:
    spaces = None

try:
    import torch
    import torch.nn.functional as F
except Exception:
    torch = None

MIN_PIXELS = 20000                      # smaller label maps are faster on the CPU than the trip to the device


_active = 0                             # how many gpu_task calls are running in this process
_failed = False                         # CUDA raised once: stay on the CPU from then on
# On a Hugging Face Space with ZeroGPU, torch pretends to have CUDA at all times but only really has it inside a @spaces.GPU call: outside one, the
# first CUDA operation ends in "Low-level CUDA init reached". So on a Space the GPU is used only inside a gpu_task.
ON_SPACE = spaces is not None and bool(os.environ.get('SPACE_ID') or os.environ.get('SPACES_ZERO_GPU'))


def device():
    """'cuda' when CUDA is usable and not switched off, else None. On a Space this is only so inside a gpu_task call."""
    if torch is None or _failed or os.environ.get('HATCHLING_GPU', '1').lower() in ('0', 'false', 'no', 'off'):
        return None
    if ON_SPACE and not _active: return None
    try:
        return 'cuda' if torch.cuda.is_available() else None
    except Exception:
        return None


def fail(e):
    """CUDA refused: stay on the CPU from here on."""
    global _failed
    _failed = True; print('GPU path failed (%s: %s): using the CPU' % (type(e).__name__, str(e)[:120]))


def attempt(on_gpu, on_cpu):
    """Run the GPU version; if CUDA refuses for any reason, remember that and run the CPU version instead."""
    if device():
        try: return on_gpu()
        except Exception as e: fail(e)
    return on_cpu()


def gpu_task(duration=120):
    """Decorator for the functions that reach the heavy steps: ZeroGPU's @spaces.GPU where it exists, otherwise nothing. It marks the call as one in
    which the GPU may be used."""
    def wrap(fn):
        @functools.wraps(fn)
        def marked(*a, **k):
            global _active
            _active += 1
            try: return fn(*a, **k)
            finally: _active -= 1
        return spaces.GPU(duration=duration)(marked) if spaces is not None else marked
    return wrap


class KMeans:
    """k-means++ with n_init restarts on the GPU; the part of sklearn's KMeans that ling_hatch uses."""

    def __init__(self, k, n_init=4, random_state=0, iters=60):
        self.k, self.n_init, self.seed, self.iters = k, n_init, int(random_state), iters
        self.cluster_centers_ = None

    def fit(self, X):
        dev = device()
        x = torch.as_tensor(X, dtype=torch.float32, device=dev)
        best = None
        for r in range(self.n_init):
            g = torch.Generator(device=dev).manual_seed(self.seed * 1009 + r)
            c = self._init(x, g)
            for _ in range(self.iters):
                a = torch.cdist(x, c).argmin(1)
                nc = torch.zeros_like(c).index_add_(0, a, x)
                cnt = torch.bincount(a, minlength=self.k).clamp_min(1).unsqueeze(1)
                nc = torch.where(torch.bincount(a, minlength=self.k).unsqueeze(1) > 0, nc / cnt, c)
                moved = (nc - c).abs().max().item(); c = nc
                if moved < 1e-4: break
            inertia = torch.cdist(x, c).min(1).values.pow(2).sum().item()
            if best is None or inertia < best[0]: best = (inertia, c)
        self._c = best[1]; self.cluster_centers_ = best[1].cpu().numpy().astype('float64')
        return self

    def _init(self, x, g):
        n = x.shape[0]
        c = x[torch.randint(n, (1,), generator=g, device=x.device)]
        for _ in range(1, self.k):
            d2 = torch.cdist(x, c).min(1).values.pow(2)
            tot = d2.sum()
            i = torch.multinomial(d2 / tot, 1, generator=g) if tot > 0 else torch.randint(n, (1,), generator=g, device=x.device)
            c = torch.cat([c, x[i]])
        return c

    def predict(self, X):
        x = torch.as_tensor(X, dtype=torch.float32, device=self._c.device)
        out = [torch.cdist(x[i:i + (1 << 20)], self._c).argmin(1) for i in range(0, x.shape[0], 1 << 20)]
        return torch.cat(out).cpu().numpy().astype('int32')


def _kernel(sigma, dev):
    """OpenCV's GaussianBlur kernel for float32 with ksize 0: round(8 sigma + 1) | 1 taps."""
    n = int(round(sigma * 8 + 1)) | 1; r = n // 2
    t = torch.arange(-r, r + 1, device=dev, dtype=torch.float32)
    k = torch.exp(-t * t / (2 * sigma * sigma)); return k / k.sum(), r


def _blur(x, sigma):
    """x: (1, 1, H, W) float32. Separable Gaussian with reflect-101 borders, as cv2.GaussianBlur."""
    k, r = _kernel(sigma, x.device)
    r = min(r, x.shape[2] - 1, x.shape[3] - 1)
    if r < 1: return x
    k = k[len(k) // 2 - r: len(k) // 2 + r + 1]; k = k / k.sum()
    x = F.conv2d(F.pad(x, (r, r, 0, 0), mode='reflect'), k.view(1, 1, 1, -1))
    return F.conv2d(F.pad(x, (0, 0, r, r), mode='reflect'), k.view(1, 1, -1, 1))


def mode_clean(lbl, sigma, size=None):
    """Each pixel takes the label with the strongest Gaussian-smoothed presence (cv2 loop of _mode_clean / the upscale in posterize).
    size=(w, h): the result is bicubically upscaled to that size first, as posterize does for the working mask."""
    import numpy as np
    dev = device()
    t = torch.as_tensor(lbl, device=dev)
    best = torch.full(size[::-1] if size else tuple(t.shape), -1.0 if size else 0.0, device=dev)
    out = torch.zeros(tuple(best.shape), dtype=torch.uint8, device=dev) if size else t.clone()
    for v in torch.unique(t).tolist():
        sc = _blur((t == v).float()[None, None], sigma)
        if size: sc = F.interpolate(sc, size=(size[1], size[0]), mode='bicubic', align_corners=False)
        sc = sc[0, 0]; m = sc > best
        out[m] = v; best = torch.where(m, sc, best)
    return out.cpu().numpy().astype(np.uint8)