"""PixelModel v3 - architecture + PNG weight codec.

A text-to-image CPPN whose entire weight set is stored as pixels in a PNG.

v3 changes relative to v1/v2 (see README):
  1. Learned word-embedding table (mean-pooled) replaces the 0-param hash embed.
  2. SIREN (sine) decoder with the real SIREN init instead of tanh.
  3. Wider/deeper decoder (128-256 wide, 4-5 layers).
  4. FiLM conditioning (per-layer gamma/beta from z) replaces concat conditioning.
  5. PNG weight-encoding scheme extended to the new param count.
  6/7. Fully batched coordinate decoding (no per-pixel python loop).

The PNG codec is deliberately identical in spirit to v1: one weight per pixel,
16-bit precision spread across the R (high byte) and G (low byte) channels,
B reserved, values linearly mapped from a fixed [-2, 2] range.
"""

from __future__ import annotations

import json
import math
import os
import time
from dataclasses import dataclass, asdict, field

import numpy as np

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



WEIGHT_RANGE = (-2.0, 2.0)

PAD_ID = 0
UNK_ID = 1


@dataclass
class ModelConfig:
    vocab_size: int = 8192
    embed_dim: int = 64
    max_tokens: int = 20
    text_hidden: int = 256
    z_dim: int = 192

    decoder_width: int = 256
    num_sine_layers: int = 4
    num_freq: int = 8
    w0_first: float = 30.0
    w0_hidden: float = 30.0

    weight_range: tuple = WEIGHT_RANGE

    @property
    def coord_dim(self) -> int:
        return 2 + 4 * self.num_freq

    def to_json(self) -> dict:
        d = asdict(self)
        d["weight_range"] = list(self.weight_range)
        d["coord_dim"] = self.coord_dim
        return d

    @staticmethod
    def from_json(d: dict) -> "ModelConfig":
        keys = {f for f in ModelConfig.__dataclass_fields__}
        clean = {k: v for k, v in d.items() if k in keys}
        if "weight_range" in clean:
            clean["weight_range"] = tuple(clean["weight_range"])
        return ModelConfig(**clean)



import re as _re

_TOKEN_RE = _re.compile(r"[a-z0-9]+")


def simple_tokenize(text: str) -> list[str]:
    """Lowercase, split on non-alphanumeric runs. Deterministic and reversible
    enough for a bag-of-words embedding."""
    return _TOKEN_RE.findall(text.lower())


def build_vocab(captions, vocab_size: int, min_freq: int = 1) -> dict:
    """captions: iterable of strings -> {token: id}. Ids 0/1 are PAD/UNK."""
    from collections import Counter
    counter = Counter()
    for cap in captions:
        counter.update(simple_tokenize(cap))
    vocab = {"<pad>": PAD_ID, "<unk>": UNK_ID}
    for tok, freq in counter.most_common():
        if len(vocab) >= vocab_size:
            break
        if freq < min_freq:
            break
        if tok not in vocab:
            vocab[tok] = len(vocab)
    return vocab


def load_vocab(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def encode_caption(text: str, vocab: dict, max_tokens: int) -> np.ndarray:
    """text -> int32 array of length max_tokens (PAD-filled, UNK-mapped)."""
    ids = [vocab.get(t, UNK_ID) for t in simple_tokenize(text)][:max_tokens]
    out = np.full(max_tokens, PAD_ID, dtype=np.int32)
    out[:len(ids)] = ids
    return out



def fourier_features(coords, num_freq: int):
    """coords: (..., 2) in [-1, 1] -> (..., 2 + 4*num_freq).

    One explicit batched tensor op over the whole pixel set; no python loop
    over pixels. Frequencies are the octaves pi * 2^k, k = 0..num_freq-1.
    """
    freqs = (2.0 ** torch.arange(num_freq, device=coords.device, dtype=coords.dtype)) * math.pi
    scaled = coords.unsqueeze(-1) * freqs
    sin = torch.sin(scaled)
    cos = torch.cos(scaled)
    enc = torch.cat([coords, sin.flatten(-2), cos.flatten(-2)], dim=-1)
    return enc



class FiLMSineLayer(nn.Module):
    """Linear -> FiLM(gamma, beta) -> sin(w0 * .).

    FiLM is applied *after* the linear and *before* the sine, so the
    conditioning warps the pre-activation the same way a learned bias would,
    while the SIREN frequency scaling still governs the activation spectrum.
    """

    def __init__(self, in_features: int, out_features: int, w0: float, is_first: bool):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.w0 = w0
        self.is_first = is_first
        self.linear = nn.Linear(in_features, out_features)
        self.siren_init()

    def siren_init(self):
        with torch.no_grad():
            if self.is_first:
                bound = 1.0 / self.in_features
            else:
                bound = math.sqrt(6.0 / self.in_features) / self.w0
            self.linear.weight.uniform_(-bound, bound)
            if self.linear.bias is not None:
                self.linear.bias.uniform_(-bound, bound)

    def forward(self, x, gamma, beta):
        h = self.linear(x)
        h = gamma * h + beta
        return torch.sin(self.w0 * h)



class PixelModelV3(nn.Module):
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg

        self.embed = nn.Embedding(cfg.vocab_size, cfg.embed_dim, padding_idx=PAD_ID)
        nn.init.normal_(self.embed.weight, std=0.05)
        with torch.no_grad():
            self.embed.weight[PAD_ID].zero_()

        self.text_fc1 = nn.Linear(cfg.embed_dim, cfg.text_hidden)
        self.text_fc2 = nn.Linear(cfg.text_hidden, cfg.z_dim)

        self.film = nn.Linear(cfg.z_dim, cfg.num_sine_layers * 2 * cfg.decoder_width)
        nn.init.zeros_(self.film.weight)
        nn.init.zeros_(self.film.bias)

        layers = []
        in_dim = cfg.coord_dim
        for i in range(cfg.num_sine_layers):
            is_first = (i == 0)
            w0 = cfg.w0_first if is_first else cfg.w0_hidden
            layers.append(FiLMSineLayer(in_dim, cfg.decoder_width, w0=w0, is_first=is_first))
            in_dim = cfg.decoder_width
        self.sine_layers = nn.ModuleList(layers)

        self.head = nn.Linear(cfg.decoder_width, 3)
        with torch.no_grad():
            bound = math.sqrt(6.0 / cfg.decoder_width) / cfg.w0_hidden
            self.head.weight.uniform_(-bound, bound)
            nn.init.zeros_(self.head.bias)

        self.num_freq = cfg.num_freq

        self.grad_checkpoint = False

    def encode_text(self, token_ids):
        """token_ids: (B, L) long -> z: (B, z_dim)."""
        emb = self.embed(token_ids)
        mask = (token_ids != PAD_ID).float().unsqueeze(-1)
        summed = (emb * mask).sum(dim=1)
        count = mask.sum(dim=1).clamp(min=1.0)
        pooled = summed / count
        h = torch.sin(self.text_fc1(pooled))
        z = self.text_fc2(h)
        return z

    def film_params(self, z):
        """z: (B, z_dim) -> list of (gamma, beta), each (B, 1, width)."""
        raw = self.film(z)
        raw = raw.view(-1, self.cfg.num_sine_layers, 2, self.cfg.decoder_width)
        out = []
        for i in range(self.cfg.num_sine_layers):
            gamma = 1.0 + raw[:, i, 0, :].unsqueeze(1)
            beta = raw[:, i, 1, :].unsqueeze(1)
            out.append((gamma, beta))
        return out

    def decode(self, coords, z, return_stats: bool = False):
        """coords: (B, N, 2) in [-1, 1]; z: (B, z_dim) -> rgb (B, N, 3).

        The whole pixel set N is decoded as one batched tensor op per layer.
        """
        x = fourier_features(coords, self.num_freq)
        films = self.film_params(z)
        use_ckpt = self.grad_checkpoint and self.training and not return_stats and x.requires_grad
        stats = []
        for layer, (gamma, beta) in zip(self.sine_layers, films):
            if use_ckpt:
                x = torch.utils.checkpoint.checkpoint(layer, x, gamma, beta, use_reentrant=False)
            else:
                x = layer(x, gamma, beta)
            if return_stats:
                stats.append((float(x.mean()), float(x.std())))
        rgb = torch.sigmoid(self.head(x))
        if return_stats:
            return rgb, stats
        return rgb

    def forward(self, token_ids, coords, return_stats: bool = False):
        z = self.encode_text(token_ids)
        return self.decode(coords, z, return_stats=return_stats)



def make_coord_grid(height: int, width: int, device=None, dtype=None):
    """-> (H*W, 2) coordinate grid in [-1, 1]. One tensor op, no python loop."""
    ys = torch.linspace(-1.0, 1.0, height, device=device, dtype=dtype)
    xs = torch.linspace(-1.0, 1.0, width, device=device, dtype=dtype)
    gy, gx = torch.meshgrid(ys, xs, indexing="ij")
    return torch.stack([gx, gy], dim=-1).reshape(-1, 2)



def param_breakdown(model: "PixelModelV3") -> dict:
    """Per-parameter-tensor breakdown, in the exact order used by the PNG codec."""
    breakdown = []
    total = 0
    for name, p in model.named_parameters():
        n = int(p.numel())
        breakdown.append({"name": name, "shape": list(p.shape), "params": n})
        total += n
    return {"total_parameters": total, "param_breakdown": breakdown}



def png_dims_for(n_params: int) -> tuple[int, int]:
    """Roughly-square PNG that holds n_params pixels (one weight per pixel)."""
    width = int(math.ceil(math.sqrt(n_params)))
    height = int(math.ceil(n_params / width))
    return width, height


def flatten_params(model: "PixelModelV3") -> np.ndarray:
    chunks = [p.detach().cpu().float().reshape(-1).numpy() for p in model.parameters()]
    return np.concatenate(chunks)


def load_flat_params(model: "PixelModelV3", flat: np.ndarray) -> None:
    i = 0
    for p in model.parameters():
        n = p.numel()
        chunk = torch.from_numpy(np.ascontiguousarray(flat[i:i + n])).to(p.dtype).view_as(p)
        p.data.copy_(chunk)
        i += n
    if i != flat.shape[0]:
        raise ValueError(f"param/flat size mismatch: consumed {i}, have {flat.shape[0]}")


def weights_to_png(flat: np.ndarray, path: str, weight_range=WEIGHT_RANGE):
    """Write a float weight vector to a PNG. Returns (width, height, n_clamped)."""
    from PIL import Image

    n = int(flat.size)
    width, height = png_dims_for(n)
    total = width * height
    padded = np.zeros(total, dtype=np.float64)
    padded[:n] = flat.astype(np.float64)

    lo, hi = weight_range
    n_clamped = int(np.count_nonzero((padded[:n] < lo) | (padded[:n] > hi)))
    clamped = np.clip(padded, lo, hi)
    u16 = np.round((clamped - lo) / (hi - lo) * 65535.0).astype(np.uint16)

    img = np.zeros((height, width, 3), dtype=np.uint8)
    img[..., 0] = (u16 >> 8).astype(np.uint8).reshape(height, width)
    img[..., 1] = (u16 & 0xFF).astype(np.uint8).reshape(height, width)
    Image.fromarray(img, "RGB").save(path, optimize=True)
    return width, height, n_clamped


def png_to_weights(path: str, n_params: int, weight_range=WEIGHT_RANGE) -> np.ndarray:
    """Decode a PNG back into the first n_params weight values (float32)."""
    from PIL import Image

    arr = np.asarray(Image.open(path).convert("RGB"))
    hi_byte = arr[..., 0].astype(np.uint16)
    lo_byte = arr[..., 1].astype(np.uint16)
    u16 = ((hi_byte << 8) | lo_byte).reshape(-1)[:n_params]
    lo, hi = weight_range
    vals = u16.astype(np.float64) / 65535.0 * (hi - lo) + lo
    return vals.astype(np.float32)



def save_model_png(model: "PixelModelV3", png_path: str, config_path: str | None = None,
                   verbose: bool = True) -> dict:
    """Write model.png (and optionally config.json). Returns a stats dict."""
    flat = flatten_params(model)
    n = int(flat.size)
    t0 = time.time()
    width, height, n_clamped = weights_to_png(flat, png_path, model.cfg.weight_range)
    write_s = time.time() - t0
    size_bytes = os.path.getsize(png_path)

    info = param_breakdown(model)
    info.update({
        "png_width": width,
        "png_height": height,
        "png_pixels": width * height,
        "png_bytes": size_bytes,
        "png_write_seconds": round(write_s, 4),
        "weights_clamped": n_clamped,
        "config": model.cfg.to_json(),
    })

    if config_path is not None:
        with open(config_path, "w") as f:
            json.dump(info, f, indent=2)

    if verbose:
        print(f"[png] total params      : {n:,}")
        print(f"[png] resolution        : {width} x {height}  ({width*height:,} pixels)")
        print(f"[png] file size on disk : {size_bytes/1024:.1f} KiB ({size_bytes/1e6:.3f} MB)")
        print(f"[png] write time        : {write_s*1000:.1f} ms")
        if n_clamped:
            frac = 100.0 * n_clamped / n
            print(f"[png] WARNING clamped   : {n_clamped:,} weights ({frac:.3f}%) outside {model.cfg.weight_range}")
        if size_bytes > 4_000_000:
            print(f"[png] WARNING size      : {size_bytes/1e6:.2f} MB is larger than expected")
    return info


def load_model_png(png_path: str, cfg: ModelConfig, map_location="cpu",
                   verbose: bool = True) -> "PixelModelV3":
    """Rebuild the model architecture from cfg, then fill weights from the PNG."""
    model = PixelModelV3(cfg)
    n_params = sum(p.numel() for p in model.parameters())
    t0 = time.time()
    flat = png_to_weights(png_path, n_params, cfg.weight_range)
    load_flat_params(model, flat)
    model.to(map_location)
    model.eval()
    if verbose:
        print(f"[png] loaded {n_params:,} weights from {png_path} in {(time.time()-t0)*1000:.1f} ms")
    return model


def load_config(config_path: str) -> ModelConfig:
    with open(config_path) as f:
        info = json.load(f)
    cfg_dict = info.get("config", info)
    return ModelConfig.from_json(cfg_dict)



def roundtrip_error(model: "PixelModelV3", tmp_png: str) -> float:
    """Max abs weight error after a PNG encode/decode round trip."""
    before = flatten_params(model)
    weights_to_png(before, tmp_png, model.cfg.weight_range)
    after = png_to_weights(tmp_png, before.size, model.cfg.weight_range)
    return float(np.max(np.abs(before - after)))