"""Codelor Color v2 - Chroma key. Autonome (torch uniquement)."""
import torch, torch.nn as nn, torch.nn.functional as F
from pathlib import Path
import json

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__(); self.eps = eps; self.w = nn.Parameter(torch.ones(dim))
    def forward(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w

class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__(); self.w12 = nn.Linear(dim, hidden * 2, bias=False); self.w3 = nn.Linear(hidden, dim, bias=False)
    def forward(self, x):
        a, b = self.w12(x).chunk(2, dim=-1); return self.w3(F.silu(a) * b)

class LayerScale(nn.Module):
    def __init__(self, dim, init=1e-6):
        super().__init__(); self.gamma = nn.Parameter(torch.ones(dim) * init)
    def forward(self, x): return x * self.gamma

class DropPath(nn.Module):
    def __init__(self, p=0.0):
        super().__init__(); self.p = p
    def forward(self, x):
        if self.p == 0.0 or not self.training: return x
        keep = 1 - self.p; mask = torch.rand(x.shape[0], 1, 1, device=x.device) < keep
        return x * mask / keep

class Attention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__(); assert dim % heads == 0
        self.heads = heads; self.dh = dim // heads
        self.qkv = nn.Linear(dim, dim * 3, bias=False); self.q_norm = RMSNorm(self.dh); self.k_norm = RMSNorm(self.dh)
        self.proj = nn.Linear(dim, dim, bias=False)
    def forward(self, x, h, w):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.heads, self.dh).permute(2, 0, 3, 1, 4)
        q, k, v = qkv.unbind(0); q, k = self.q_norm(q), self.k_norm(k); q, k = rope_2d(q, k, h, w)
        out = F.scaled_dot_product_attention(q, k, v)
        return self.proj(out.transpose(1, 2).reshape(B, N, C))

class ColorCrossAttention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__(); assert dim % heads == 0
        self.heads = heads; self.dh = dim // heads
        self.q  = nn.Linear(dim, dim, bias=False); self.kv = nn.Linear(dim, dim * 2, bias=False)
        self.q_norm = RMSNorm(self.dh); self.k_norm = RMSNorm(self.dh); self.proj = nn.Linear(dim, dim, bias=False)
    def forward(self, spatial, cond):
        B, N, C = spatial.shape; M = cond.shape[1]
        q = self.q(spatial).reshape(B, N, self.heads, self.dh).transpose(1, 2)
        kv = self.kv(cond).reshape(B, M, 2, self.heads, self.dh).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0); q, k = self.q_norm(q), self.k_norm(k)
        out = F.scaled_dot_product_attention(q, k, v)
        return self.proj(out.transpose(1, 2).reshape(B, N, C))

class Block(nn.Module):
    def __init__(self, dim, heads, mlp_ratio=3.0, drop_path=0.0):
        super().__init__(); self.norm1 = RMSNorm(dim); self.attn  = Attention(dim, heads)
        self.ls1 = LayerScale(dim); self.dp1 = DropPath(drop_path); self.norm2 = RMSNorm(dim)
        self.ffn = SwiGLU(dim, int(dim * mlp_ratio)); self.ls2 = LayerScale(dim); self.dp2 = DropPath(drop_path)
    def forward(self, x, h, w):
        x = x + self.dp1(self.ls1(self.attn(self.norm1(x), h, w)))
        return x + self.dp2(self.ls2(self.ffn(self.norm2(x))))

class AntiAliasDownsample(nn.Module):
    def __init__(self, c_in, c_out):
        super().__init__()
        self.blur = nn.Conv2d(c_in, c_in, 3, padding=1, groups=c_in, bias=False)
        kernel = torch.tensor([1., 2., 1.]); kernel = kernel / kernel.sum(); bk = kernel[:, None] * kernel[None, :]
        with torch.no_grad(): self.blur.weight.copy_(bk.expand(c_in, 1, 3, 3))
        self.blur.weight.requires_grad_(False)
        self.conv = nn.Conv2d(c_in, c_out, 1, stride=2, bias=False); self.bn = nn.BatchNorm2d(c_out); self.act = nn.SiLU(inplace=True)
    def forward(self, x): return self.act(self.bn(self.conv(self.blur(x))))

class ConvNeXtStem(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.patch = nn.Sequential(nn.Conv2d(3, 32, kernel_size=4, stride=4, bias=False), nn.BatchNorm2d(32), nn.SiLU(inplace=True))
        self.block1 = nn.Sequential(nn.Conv2d(32, 32, 7, padding=3, groups=32, bias=False), nn.BatchNorm2d(32), nn.SiLU(inplace=True), nn.Conv2d(32, 32, 1, bias=False), nn.BatchNorm2d(32))
        self.down8  = AntiAliasDownsample(32, 64); self.down16 = AntiAliasDownsample(64, dim)
    def forward(self, x):
        x = self.patch(x); x = x + self.block1(x); f4 = x; f8 = self.down8(f4); f16 = self.down16(f8)
        return f4, f8, f16

class SobelEdge(nn.Module):
    def __init__(self):
        super().__init__()
        sx = torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]]).float().view(1,1,3,3); sy = torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]]).float().view(1,1,3,3)
        self.register_buffer('sx', sx); self.register_buffer('sy', sy)
    def forward(self, gray):
        ex = F.conv2d(gray, self.sx, padding=1); ey = F.conv2d(gray, self.sy, padding=1)
        return (ex**2 + ey**2 + 1e-6).sqrt()

class GatedDecoder(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.proj16 = nn.Conv2d(dim, dim, 1, bias=False); self.aux_head = nn.Conv2d(dim, 1, 1)
        self.up8 = nn.Sequential(nn.Conv2d(dim+64, dim, 3, padding=1, bias=False), nn.BatchNorm2d(dim), nn.SiLU(inplace=True))
        self.up4 = nn.Sequential(nn.Conv2d(dim+32, 64, 3, padding=1, bias=False), nn.BatchNorm2d(64), nn.SiLU(inplace=True))
        self.up1 = nn.Sequential(nn.Conv2d(64+3+1+1, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.SiLU(inplace=True), nn.Conv2d(32, 32, 3, padding=1, bias=False), nn.BatchNorm2d(32), nn.SiLU(inplace=True))
        self.head = nn.Conv2d(32, 1, 1); self.gate8 = nn.Parameter(torch.zeros(1)); self.gate4 = nn.Parameter(torch.zeros(1)); self.edge = SobelEdge()
    def forward(self, latent, f4, f8, img, ref_rgb):
        B, _, H, W = img.shape; x = self.proj16(latent)
        aux = self.aux_head(x); aux = F.interpolate(aux, size=f4.shape[-2:], mode='bilinear', align_corners=False)
        x = F.interpolate(x, size=f8.shape[-2:], mode='bilinear', align_corners=False)
        x = self.up8(torch.cat([x, torch.sigmoid(self.gate8) * f8], 1))
        x = F.interpolate(x, size=f4.shape[-2:], mode='bilinear', align_corners=False)
        x = self.up4(torch.cat([x, torch.sigmoid(self.gate4) * f4], 1))
        x = F.interpolate(x, size=(H, W), mode='bilinear', align_corners=False)
        ref_map = ref_rgb.view(B, 3, 1, 1).expand(B, 3, H, W); cos = F.cosine_similarity(img, ref_map, dim=1, eps=1e-6).unsqueeze(1)
        gray = img.mean(1, keepdim=True); edge = self.edge(gray); x = self.up1(torch.cat([x, img, cos, edge], 1))
        return self.head(x), aux

class CodelorColor(nn.Module):
    def __init__(self, dim=128, depth=8, heads=8, tied_after=4, mlp_ratio=3.0, max_tokens=1024, drop_path_rate=0.1):
        super().__init__(); self.dim, self.depth, self.heads = dim, depth, heads; self.tied_after, self.mlp_ratio = tied_after, mlp_ratio; self.max_tokens = max_tokens
        self.stem = ConvNeXtStem(dim); self.color_mlp = nn.Sequential(nn.Linear(3, dim), nn.SiLU(inplace=True), nn.Linear(dim, dim))
        self.pos = nn.Parameter(torch.zeros(1, max_tokens, dim)); nn.init.trunc_normal_(self.pos, std=0.02)
        n_unique = depth if tied_after <= 0 else min(depth, tied_after)
        dp = [drop_path_rate * i / max(1, depth - 1) for i in range(depth)]
        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio, drop_path=dp[i] if i < len(dp) else 0.0) for i in range(n_unique)])
        self.color_xattn = nn.ModuleList([ColorCrossAttention(dim, heads) for _ in range(n_unique)])
        self.norm_out = RMSNorm(dim); self.decoder = GatedDecoder(dim)
    def _get(self, i):
        if self.tied_after > 0 and i >= self.tied_after: return self.blocks[i - self.tied_after], self.color_xattn[i - self.tied_after]
        return self.blocks[i], self.color_xattn[i]
    def forward(self, img, ref_rgb, return_aux=False):
        B = img.shape[0]; f4, f8, f16 = self.stem(img); _, _, h, w = f16.shape; N = h * w
        tokens = f16.flatten(2).transpose(1, 2)
        if N <= self.pos.shape[1]: tokens = tokens + self.pos[:, :N]
        c_tok = self.color_mlp(ref_rgb).unsqueeze(1); x = tokens
        for i in range(self.depth):
            blk, xa = self._get(i); x = blk(x, h, w); x = x + xa(x, c_tok)
        x = self.norm_out(x); latent = x.transpose(1, 2).reshape(B, -1, h, w)
        main, aux = self.decoder(latent, f4, f8, img, ref_rgb)
        return (main, aux) if return_aux else main
    def config(self): return dict(dim=self.dim, depth=self.depth, heads=self.heads, tied_after=self.tied_after, mlp_ratio=self.mlp_ratio, max_tokens=self.max_tokens)

def rope_2d(q, k, h, w, base=100.0):
    B, Hh, N, D = q.shape; d4 = D // 4
    freqs = 1.0 / (base ** (torch.arange(0, d4, device=q.device).float() / d4))
    xs = torch.arange(w, device=q.device).float(); ys = torch.arange(h, device=q.device).float()
    ax = torch.outer(xs, freqs)[None, :, :].expand(h, w, d4).reshape(h * w, d4)
    ay = torch.outer(ys, freqs)[:, None, :].expand(h, w, d4).reshape(h * w, d4)
    ang = torch.cat([ax, ay], dim=-1); cos, sin = ang.cos()[None, None], ang.sin()[None, None]
    def rot(x):
        x1, x2 = x[..., :D//2], x[..., D//2:]
        return torch.cat([x1*cos - x2*sin, x1*sin + x2*cos], dim=-1)
    return rot(q), rot(k)



def load_from_pretrained(path_or_repo, device="cpu"):
    p = Path(path_or_repo)
    if not p.exists():
        from huggingface_hub import snapshot_download
        p = Path(snapshot_download(path_or_repo))
    cfg = json.loads((p / "config.json").read_text())
    model = CodelorColor(**cfg["hyperparameters"])
    from safetensors.torch import load_file
    model.load_state_dict(load_file(p / "model.safetensors"))
    return model.eval().to(device)

def predict_mask(model, image_tensor, color_rgb=None, threshold=0.5, device="cpu"):
    if color_rgb is None: color_rgb = (0.20, 0.85, 0.30)
    if image_tensor.ndim == 3: image_tensor = image_tensor.unsqueeze(0)
    image_tensor = image_tensor.to(device)
    ref = torch.tensor(color_rgb, dtype=torch.float32, device=device).view(1, 3)
    with torch.no_grad():
        prob = torch.sigmoid(model(image_tensor, ref))[0, 0].cpu().numpy()
    return prob > threshold, prob

def remove_green(image_tensor, model, color_rgb=None, threshold=0.5, device="cpu"):
    import numpy as np
    _, prob = predict_mask(model, image_tensor, color_rgb, threshold, device)
    alpha = torch.from_numpy((1.0 - prob.astype("float32")).clip(0, 1)).unsqueeze(0)
    return torch.cat([image_tensor.cpu(), alpha], dim=0)