"""CorridorKey Real-Time Edge Refiner: standalone reference implementation.

Licensed under the Corridor Key Licence (CC BY-NC-SA 4.0 + additional terms), see LICENSE.
Built for use with CorridorKey 2.0 by Corridor Digital LLC.

The refiner restores native-resolution edge detail (hair strands, motion blur, soft edges) on top of a
CorridorKey 2.0 matte that was computed at a reduced resolution. Inputs per native pixel:

  I      native RGB frame (sRGB, [0, 1])
  B      smooth background estimate (sRGB), from pixels CorridorKey calls background
  F      CorridorKey foreground transferred onto the native frame (I + upsampled (F_low - I_low))
  a_up   CorridorKey alpha, bilinear upsampled
  a_gf   alpha of a colour guided filter (He et al.) driven by CorridorKey's low-res alpha

It predicts residuals for alpha (on top of a_gf) and foreground (on top of F). The result is only used
inside an uncertainty band around CorridorKey's edges; elsewhere the guided-filter result is kept.

`refine()` below processes a whole frame at once (simple, a few ms on a modern GPU at 1080p). The real-time
engine evaluates exactly the same maths on 32x32 band tiles only.
"""
from __future__ import annotations

import json

import torch
import torch.nn as nn
import torch.nn.functional as F

IN_CHANNELS = 11


class EdgeRefiner(nn.Module):
    def __init__(self, width: int = 24, dilations=(1, 2, 4, 1)):
        super().__init__()
        self.dilations = tuple(int(d) for d in dilations)
        self.inp = nn.Conv2d(IN_CHANNELS, width, 3, padding=1)
        self.blocks = nn.ModuleList(nn.Conv2d(width, width, 3, padding=d, dilation=d) for d in self.dilations)
        self.out = nn.Conv2d(width, 4, 1)

    def receptive_radius(self) -> int:
        return 1 + sum(self.dilations)

    def forward(self, x: torch.Tensor):
        """x: (B, 11, H, W) = [I(3), B(3), F(3), a_up(1), a_gf(1)]. Returns (alpha, fg_srgb) in [0, 1]."""
        h = F.gelu(self.inp(x))
        for blk in self.blocks:
            h = h + F.gelu(blk(h))
        d = self.out(h)
        alpha = (x[:, 10:11] + d[:, 0:1]).clamp(0.0, 1.0)
        fg = (x[:, 6:9] + d[:, 1:4]).clamp(0.0, 1.0)
        return alpha, fg


def load_refiner(path: str, device="cuda", dtype=torch.float16) -> EdgeRefiner:
    """Load a .safetensors checkpoint of this repository."""
    from safetensors import safe_open
    from safetensors.torch import load_file
    with safe_open(path, framework="pt") as f:
        cfg = json.loads(f.metadata()["config"])
    net = EdgeRefiner(width=cfg["width"], dilations=tuple(cfg["dilations"]))
    net.load_state_dict(load_file(path))
    return net.to(device=device, dtype=dtype).eval()


# ---------------------------------------------------------------------------------------------------------
# Input construction (identical to the real-time engine)
# ---------------------------------------------------------------------------------------------------------

def srgb_to_linear(x):
    return torch.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4)


def linear_to_srgb(x):
    x = x.clamp(min=0.0)
    return torch.where(x <= 0.0031308, x * 12.92, 1.055 * x.clamp(min=0.0031308) ** (1 / 2.4) - 0.055)


def guided_coef(rgb_low, alpha_low, radius: int = 2, eps: float = 5e-4):
    """Colour guided filter (He et al.) at low resolution: smoothed (a_r, a_g, a_b, b) coefficients."""
    k = 2 * radius + 1
    box = lambda z: F.avg_pool2d(z, k, stride=1, padding=radius, count_include_pad=False)
    I, p = rgb_low, alpha_low
    _, _, h, w = p.shape
    mean_I, mean_p = box(I), box(p)
    cov_Ip = box(I * p) - mean_I * mean_p
    mean_II = box((I.unsqueeze(1) * I.unsqueeze(2)).reshape(1, 9, h, w)).reshape(1, 3, 3, h, w)
    s = mean_II - mean_I.unsqueeze(1) * mean_I.unsqueeze(2)
    s = s + eps * torch.eye(3, device=s.device, dtype=s.dtype).view(1, 3, 3, 1, 1)
    a = torch.linalg.solve(s.permute(0, 3, 4, 1, 2), cov_Ip.permute(0, 2, 3, 1).unsqueeze(-1))
    a = a.squeeze(-1).permute(0, 3, 1, 2)
    b = mean_p - (a * mean_I).sum(1, keepdim=True)
    return torch.cat([box(a), box(b)], 1)


def push_pull(num, den):
    """Hole-filling interpolation: weighted average num/den, holes (den ~ 0) filled from coarser levels."""
    levels = [(num, den)]
    while min(levels[-1][0].shape[-2:]) > 4:
        n, d = levels[-1]
        levels.append((F.avg_pool2d(n, 2, ceil_mode=True), F.avg_pool2d(d, 2, ceil_mode=True)))
    n, d = levels[-1]
    est = (n.sum((-2, -1), keepdim=True) / d.sum((-2, -1), keepdim=True).clamp(min=1e-6)).expand_as(n)
    for n, d in reversed(levels):
        est = F.interpolate(est, size=n.shape[-2:], mode="bilinear", align_corners=False)
        d = d.clamp(0.0, 1.0)
        est = (n + (1.0 - d) * est) / (d + (1.0 - d)).clamp(min=1e-6)
    return est


def background_low(rgb_low, alpha_low):
    w = (1.0 - alpha_low).clamp(0.0, 1.0) ** 4
    return linear_to_srgb(push_pull(srgb_to_linear(rgb_low.clamp(0.0, 1.0)) * w, w).clamp(0.0, 1.0))


def refiner_band(alpha_low, gain: float = 10.0, lo: float = 0.004):
    """Where the refiner may change the guided result: uncertain alpha, dilated, weighted by nearby
    foreground (gain > 0 keeps it out of noisy/dark screen areas) and feathered."""
    unc = ((alpha_low > lo) & (alpha_low < 0.996)).float()
    unc = F.max_pool2d(unc, 5, stride=1, padding=2)
    if gain > 0:
        near_fg = (F.max_pool2d(alpha_low, 5, stride=1, padding=2) * gain).clamp(0.0, 1.0)
        unc = F.avg_pool2d(unc * near_fg, 3, stride=1, padding=1, count_include_pad=False)
    return unc


@torch.inference_mode()
def refine(net: EdgeRefiner, rgb, rgb_low, alpha_low, fg_low, band_gain: float = 10.0):
    """
    rgb:       (1, 3, H, W) native frame, sRGB in [0, 1]
    rgb_low:   (1, 3, h, w) the frame as CorridorKey saw it (downscaled), sRGB
    alpha_low: (1, 1, h', w') CorridorKey alpha (any low resolution, e.g. 256x256 from a 512x512 pass)
    fg_low:    (1, 3, h', w') CorridorKey straight foreground, sRGB
    Returns (alpha, fg) at native resolution, float32. Apply despill / alpha contrast afterwards as needed.
    """
    rgb, alpha_low, fg_low = rgb.float(), alpha_low.float(), fg_low.float()
    rgb_low = F.interpolate(rgb_low.float(), size=alpha_low.shape[-2:], mode="bilinear", align_corners=False)
    H, W = rgb.shape[-2:]
    up = lambda z: F.interpolate(z, size=(H, W), mode="bilinear", align_corners=False)
    coef = up(guided_coef(rgb_low, alpha_low))
    a_gf = ((coef[:, :3] * rgb).sum(1, keepdim=True) + coef[:, 3:4]).clamp(0.0, 1.0)
    fg_in = (rgb + up(fg_low - rgb_low)).clamp(0.0, 1.0)
    x = torch.cat([rgb, up(background_low(rgb_low, alpha_low)), fg_in, up(alpha_low), a_gf], 1)
    p = next(net.parameters())
    a_r, f_r = net(x.to(p.dtype))
    band = up(refiner_band(alpha_low, band_gain))
    alpha = band * a_r.float() + (1.0 - band) * a_gf
    fg = band * f_r.float() + (1.0 - band) * fg_in
    return alpha, fg