"""nika-5: grid ↔ latent-mesh world model of the Sun. Pure torch.

Observations live on a dense grid; state lives on a coarse mesh. One encoder
shared across all AIA channels conditioned on log-T, one for the field vector.
The readout is a function of (cell latent, log-T): the slider is the readout.
"""
from dataclasses import dataclass

import torch
import torch.nn as nn

from lib import AIA_CHANNELS, AIA_SCALE, B_SCALE, LOGT, fourier, mlp, norm_logt, scatter_mean


# ── substrate: regular lattices, analytic connectivity ───────────────

@dataclass
class SolarGraph:
    grid_pos: torch.Tensor   # [G, 2]
    mu: torch.Tensor         # [G]     cos heliocentric angle (0 off-disk)
    g2m: torch.Tensor        # [Eg, 2] (grid, mesh): 4 surrounding mesh nodes
    g2m_feat: torch.Tensor   # [Eg, 3] relative geometry
    m2m: torch.Tensor        # [Em, 2] directed: 8-neighbor + strides 2, 4
    m2m_feat: torch.Tensor   # [Em, 3]
    n_mesh: int

    def to(self, device):
        moved = {f: getattr(self, f).to(device)
                 for f in ("grid_pos", "mu", "g2m", "g2m_feat", "m2m", "m2m_feat")}
        return SolarGraph(**moved, n_mesh=self.n_mesh)


def _lattice(side):
    ax = (torch.arange(side) + 0.5) / side
    y, x = torch.meshgrid(ax, ax, indexing="ij")
    return torch.stack([x.flatten(), y.flatten()], -1)


def _mesh_edges(side, offsets):
    idx = torch.arange(side * side).view(side, side)
    src, dst = [], []
    for dy, dx in offsets:
        s = idx[max(0, -dy): side - max(0, dy), max(0, -dx): side - max(0, dx)]
        d = idx[max(0, dy): side + min(0, dy), max(0, dx): side + min(0, dx)]
        src.append(s.flatten())
        dst.append(d.flatten())
    e = torch.stack([torch.cat(src), torch.cat(dst)], -1)
    return torch.cat([e, e.flip(-1)])


def build_graph(res=256, mesh_side=48):
    grid = _lattice(res)
    r = ((grid - 0.5).square().sum(-1)).sqrt() / 0.5
    mu = (1 - r.square()).clamp(min=0).sqrt()

    mesh = _lattice(mesh_side)
    near = [(0, 1), (1, 0), (1, 1), (1, -1)]
    far = [(0, 2), (2, 0), (0, 4), (4, 0)]
    m2m = _mesh_edges(mesh_side, near + far)
    scale = 1.0 / mesh_side
    rel = (mesh[m2m[:, 0]] - mesh[m2m[:, 1]]) / scale
    m2m_feat = torch.cat([rel, rel.norm(dim=-1, keepdim=True)], -1)

    cell = grid * mesh_side - 0.5
    base = cell.floor().long()
    corners = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]])
    m_yx = (base[:, None] + corners[None]).clamp(0, mesh_side - 1)      # [G,4,2]
    m_idx = (m_yx[..., 1] * mesh_side + m_yx[..., 0]).flatten()
    g_idx = torch.arange(len(grid)).repeat_interleave(4)
    g2m = torch.stack([g_idx, m_idx], -1)
    rel = (grid[g_idx] - mesh[m_idx]) / scale
    g2m_feat = torch.cat([rel, rel.norm(dim=-1, keepdim=True)], -1)

    return SolarGraph(grid, mu, g2m, g2m_feat, m2m, m2m_feat, n_mesh=mesh_side**2)


# ── observations ─────────────────────────────────────────────────────

@dataclass
class Observation:
    aia: torch.Tensor       # [B, 9, G] asinh-normalized
    field: torch.Tensor     # [B, 3, G] tanh-normalized
    present: torch.Tensor   # [B, 9] bool

    @staticmethod
    def pack(frames, res=256):
        pool = lambda x: torch.as_tensor(x).view(res, 512 // res, res, 512 // res).mean((1, 3))
        aia = torch.stack([
            torch.stack([pool(f[c]) if c in f else torch.zeros(res, res) for c in AIA_CHANNELS])
            for f in frames]).flatten(2)
        field = torch.stack([
            torch.stack([pool(f[c]) for c in ("Bx", "By", "Bz")]) for f in frames]).flatten(2)
        present = torch.tensor([[c in f for c in AIA_CHANNELS] for f in frames])
        return Observation(torch.asinh(aia / AIA_SCALE[:, None]), torch.tanh(field / B_SCALE), present)

    def masked(self, hide):
        vis = self.present.clone()
        vis[:, hide] = False
        return Observation(self.aia * vis[:, :, None], self.field, vis)

    def to(self, device):
        return Observation(self.aia.to(device), self.field.to(device), self.present.to(device))


# ── layers ───────────────────────────────────────────────────────────

class ChannelEncoder(nn.Module):
    """One encoder for every AIA channel: (intensity, log-T) -> d. Sum over present."""

    def __init__(self, d):
        super().__init__()
        self.net = mlp(2, d, d)

    def forward(self, aia, present):
        logt = norm_logt(LOGT.to(aia.device))[:, None].expand_as(aia)
        return (self.net(torch.stack([aia, logt], -1)) * present[:, :, None, None]).sum(1)


class FieldEncoder(nn.Module):
    """Magnetic vector -> d. The causal modality gets its own pathway."""

    def __init__(self, d):
        super().__init__()
        self.net = mlp(3, d, d)

    def forward(self, field):
        return self.net(field.transpose(1, 2))


class GeometryEmbed(nn.Module):
    """μ + Fourier position -> d. Static, computed once per graph."""

    def __init__(self, d):
        super().__init__()
        self.net = nn.Linear(17, d)

    def forward(self, graph):
        return self.net(torch.cat([graph.mu[:, None], fourier(graph.grid_pos)], -1))


class GridToMesh(nn.Module):
    """Lift cell tokens onto the latent mesh along g2m edges."""

    def __init__(self, d):
        super().__init__()
        self.edge = mlp(d + 3, d, d)

    def forward(self, cells, graph):
        g, m = graph.g2m[:, 0], graph.g2m[:, 1]
        feat = graph.g2m_feat.expand(cells.shape[0], -1, -1)
        return scatter_mean(self.edge(torch.cat([cells[:, g], feat], -1)), m, graph.n_mesh)


class DynamicsRound(nn.Module):
    """One tick: field-gated messages update a velocity; state integrates it."""

    def __init__(self, d, damp=0.9, dt=0.2):
        super().__init__()
        self.msg = mlp(2 * d + 9, d, d)
        self.gate = mlp(2 * d + 9, d // 2, 1)
        self.upd = mlp(2 * d, d, d)
        self.damp, self.dt = damp, dt

    def forward(self, h, v, edge_field, graph):
        src, dst = graph.m2m[:, 0], graph.m2m[:, 1]
        feat = graph.m2m_feat.expand(h.shape[0], -1, -1)
        pair = torch.cat([h[:, dst], h[:, src], feat, edge_field], -1)
        gate = torch.sigmoid(self.gate(pair))
        agg = scatter_mean(gate * self.msg(pair), dst, graph.n_mesh)
        v = self.damp * v + self.upd(torch.cat([h, agg], -1))
        return h + self.dt * v, v, gate.squeeze(-1)


class MeshDynamics(nn.Module):
    """Second-order rollout over the mesh; exposes mean gates for physics probes."""

    def __init__(self, d, rounds):
        super().__init__()
        self.rounds = nn.ModuleList(DynamicsRound(d) for _ in range(rounds))

    def forward(self, h, field_mesh, graph):
        src, dst = graph.m2m[:, 0], graph.m2m[:, 1]
        edge_field = torch.cat([field_mesh[:, src], field_mesh[:, dst]], -1)
        v = torch.zeros_like(h)
        gates = []
        for tick in self.rounds:
            h, v, g = tick(h, v, edge_field, graph)
            gates.append(g)
        return h, torch.stack(gates).mean(0)


class MeshToGrid(nn.Module):
    """Project mesh state back to cell latents along the same edges."""

    def __init__(self, d):
        super().__init__()
        self.edge = mlp(d + 3, d, d)

    def forward(self, h, n_cells, graph):
        g, m = graph.g2m[:, 0], graph.g2m[:, 1]
        feat = graph.g2m_feat.expand(h.shape[0], -1, -1)
        return scatter_mean(self.edge(torch.cat([h[:, m], feat], -1)), g, n_cells)


class TemperatureReadout(nn.Module):
    """Intensity as a function of (cell latent, log-T). Query any temperature."""

    def __init__(self, d_in, d):
        super().__init__()
        self.net = mlp(d_in + 1, d, 1, depth=3)

    def forward(self, cells, logt_q):
        q = norm_logt(logt_q.to(cells.device))[None, :, None, None]
        z = cells[:, None].expand(-1, len(logt_q), -1, -1)
        return self.net(torch.cat([z, q.expand(*z.shape[:3], 1)], -1)).squeeze(-1)


# ── model ────────────────────────────────────────────────────────────

@dataclass
class Config:
    dim: int = 128
    rounds: int = 6
    res: int = 256
    mesh_side: int = 48


class Nika5(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        d = cfg.dim
        self.channels = ChannelEncoder(d)
        self.field = FieldEncoder(d)
        self.geometry = GeometryEmbed(d)
        self.norm = nn.LayerNorm(d)
        self.lift = GridToMesh(d)
        self.dynamics = MeshDynamics(d, cfg.rounds)
        self.project = MeshToGrid(d)
        # mesh context ‖ grid skip ‖ raw visible intensities — concurrent
        # channels alone carry real signal (the linear baseline proves it),
        # so the readout keeps a direct pathway to them.
        self.readout = TemperatureReadout(2 * d + 9, d)

    def forward(self, obs: Observation, graph: SolarGraph, logt_q):
        cells = self.norm(self.channels(obs.aia, obs.present)
                          + self.field(obs.field)
                          + self.geometry(graph))
        state = self.lift(cells, graph)
        field_mesh = scatter_mean(obs.field.transpose(1, 2)[:, graph.g2m[:, 0]],
                                  graph.g2m[:, 1], graph.n_mesh)
        state, self.gates = self.dynamics(state, field_mesh, graph)
        context = self.project(state, cells.shape[1], graph)
        raw = (obs.aia * obs.present[:, :, None]).transpose(1, 2)
        return self.readout(torch.cat([context, cells, raw], -1), logt_q)