"""Template-conditioned fingerprint RENDER generator (G_render).

Goal (see design discussion): generate a NEW clean fingerprint image that is forensic and
hallucination-free WITHOUT relying on raw pixels. The trick: decouple *identity information*
(a TEMPLATE extracted once from the observation) from *image rendering* (a generator that only
sees the template, never the raw image). Forensic == the template is faithful; no-hallucination ==
the generator can only render what the template specifies (extracted minutiae + deterministic
ridge flow), and blank/low-coverage regions stay blank.

TEMPLATE conditioning channels (NO raw pixels here — this is the whole point):
  0   minutiae heatmap   (Gaussian blobs at extracted minutiae; identity anchors)
  1-2 orientation field  [cos 2θ, sin 2θ] * coherence  (deterministic ridge flow)
  3   coherence / coverage  (where the template HAS evidence; 0 => stay blank)
  4   ridge-frequency map    (ridge spacing prior for realistic texture)

Self-supervised training: target = a clean rolled print R; cond = template(R). The generator
learns to re-render R from R's own template. At inference, feed a template extracted from a LATENT
=> a clean synthetic print of that identity, with gaps left blank (never invented).
"""
from __future__ import annotations

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

from g_render.models.structure.heatmap import compute_orientation_field, render_gaussian_points
from g_render.models.frequency.ridge_freq import RidgeFrequencyNormalizer

TEMPLATE_CHANNELS = 5


@torch.no_grad()
def build_template_conditioning(
    image: torch.Tensor,
    minutiae: list | None,
    size: int,
    rn: RidgeFrequencyNormalizer | None = None,
    minu_sigma: float = 4.0,
) -> torch.Tensor:
    """image: [1,1,H,W] in [0,1] (used ONLY to extract the template, discarded after).
    minutiae: list of (x,y[,theta]) in image pixels, or None. Returns cond [1, TEMPLATE_CHANNELS, size, size].
    NOTE: the returned tensor contains NO raw pixels — only extracted structure."""
    rn = rn or RidgeFrequencyNormalizer()
    dev = image.device
    H, W = image.shape[-2:]
    of = compute_orientation_field(image)                       # [1,2,H,W] = [cos2,sin2]*coh
    coh = (of.pow(2).sum(1, keepdim=True) + 1e-12).sqrt().clamp(0, 1)
    fmap, _ = rn.estimate_frequency_map(image)                  # [1,1,gh,gw]
    fmap = F.interpolate(fmap, size=(H, W), mode="bilinear", align_corners=False)
    if minutiae:
        pts = [(float(m[0]), float(m[1])) for m in minutiae]
        minu = render_gaussian_points(pts, out_size=size, src_size=(float(W), float(H)),
                                      sigma=minu_sigma, device=dev)[None]   # [1,1,size,size]
    else:
        minu = torch.zeros(1, 1, size, size, device=dev)
    def rs(t):
        return F.interpolate(t, size=(size, size), mode="bilinear", align_corners=False)
    cond = torch.cat([minu, rs(of), rs(coh), rs(fmap.clamp(0, 1))], dim=1)  # [1,5,size,size]
    return cond.clamp(-1, 1)


class RenderGenerator(nn.Module):
    """template conditioning -> clean fingerprint image, ATTENTION backbone.

    Wraps the in-repo UformerCondGenerator (window self-attention U-Net): attention enforces the
    long-range ridge-flow / singular-point coherence that a plain conv U-Net misses, and propagates
    the sparse minutiae/OF template globally (esp. across gap regions). No raw pixels in the input.
    """
    def __init__(self, in_ch: int = TEMPLATE_CHANNELS, dims=(32, 64, 128, 256),
                 window: int = 8, embed_dim: int = 128, cover_tau: float = 0.15):
        super().__init__()
        from g_render.models.generators.uformer_gen import UformerCondGenerator
        # skip_mode="lowpass": template has no raw speckle, but keep it for ridge-scale skips.
        self.body = UformerCondGenerator(in_channels=in_ch, dims=dims, embed_dim=embed_dim,
                                         window=window, skip_mode="lowpass")
        self.head = nn.Conv2d(2, 1, 1)                      # Uformer emits 2ch -> 1ch image
        nn.init.zeros_(self.head.bias)
        self.cover_tau = float(cover_tau)

    def forward(self, cond: torch.Tensor) -> torch.Tensor:
        y = self.body(cond)                                 # [B,2,H,W] window-attention output
        img = torch.sigmoid(self.head(y))                   # -> [B,1,H,W] in [0,1]
        # coverage gate: outside template evidence (coh ch3 ~0) -> white (blank). The generator
        # CANNOT paint ridges where the template has no evidence (structural anti-hallucination).
        cover = (cond[:, 3:4] / self.cover_tau).clamp(0, 1)
        return (img * cover + 1.0 * (1 - cover)).clamp(0, 1)