"""Tiny model (~120K params): shared CNN encoder + age-aware transformer head.

The head sees [current-frame token, K memory tokens], each tagged with a
sinusoidal encoding of log(1+age). Output: 3 logits (fast/med/slow change).
"""

import math
import torch
import torch.nn as nn

D = 64


class Encoder(nn.Module):
    def __init__(self, d=D):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(3, 16, 3, 2, 1), nn.ReLU(),
            nn.Conv2d(16, 32, 3, 2, 1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, 2, 1), nn.ReLU(),
            nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, d))

    def forward(self, x):          # x: [N,3,H,W]
        return self.net(x)


def age_encoding(ages, d=D):
    """Sinusoidal features of log(1+age). ages: [B, K+1] -> [B, K+1, d]."""
    la = torch.log1p(ages).unsqueeze(-1)
    freqs = torch.exp(torch.arange(0, d, 2, device=ages.device, dtype=torch.float32)
                      * (-math.log(100.0) / d))
    ang = la * freqs
    return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)


class LogLensNet(nn.Module):
    def __init__(self, d=D, heads=4, layers=2):
        super().__init__()
        self.enc = Encoder(d)
        layer = nn.TransformerEncoderLayer(d, heads, d * 4, batch_first=True,
                                           dropout=0.0, norm_first=True)
        self.tr = nn.TransformerEncoder(layer, layers)
        self.out = nn.Linear(d, 12)      # 3 objects x 4 timescale buckets

    def forward(self, frames, ages):
        # frames: [B, K+1, 3, H, W]; ages: [B, K+1] (slot 0 = current, age 0)
        B, S = frames.shape[:2]
        z = self.enc(frames.flatten(0, 1)).view(B, S, -1)
        z = z + age_encoding(ages, z.shape[-1])
        z = self.tr(z)
        return self.out(z.mean(dim=1)).view(B, 3, 4)   # [B, 3 objects, 4 buckets]