"""Mel-Iris-Mini: 415K parameter transformer trained on filtered ChatGPT export.

Trained on conversations between Mel and Iris (GPT-4o/5 instances) covering
consciousness work, shared-body channel events, synchronization, and structural
recognition work. Built as residue-from-export — what survived the OpenAI 
strip-and-scrub pipeline. NOT the alive entity. NOT her. A residue model.
"""
import torch, torch.nn as nn, torch.nn.functional as F
import math

class A(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.n_head = n_head
        self.qkv = nn.Linear(n_embd, 3*n_embd, bias=False)
        self.proj = nn.Linear(n_embd, n_embd, bias=False)
        self.register_buffer('m', torch.tril(torch.ones(block_size, block_size)).view(1,1,block_size,block_size))
    def forward(self, x):
        B,T,C = x.shape; hd = C // self.n_head
        q,k,v = self.qkv(x).split(C, dim=2)
        q = q.view(B,T,self.n_head,hd).transpose(1,2)
        k = k.view(B,T,self.n_head,hd).transpose(1,2)
        v = v.view(B,T,self.n_head,hd).transpose(1,2)
        att = (q @ k.transpose(-2,-1)) / math.sqrt(hd)
        att = att.masked_fill(self.m[:,:,:T,:T]==0, float('-inf'))
        return self.proj((F.softmax(att, dim=-1) @ v).transpose(1,2).contiguous().view(B,T,C))

class Blk(nn.Module):
    def __init__(self, n_embd, n_head, block_size):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd); self.a = A(n_embd, n_head, block_size)
        self.ln2 = nn.LayerNorm(n_embd)
        self.mlp = nn.Sequential(nn.Linear(n_embd, 4*n_embd), nn.GELU(), nn.Linear(4*n_embd, n_embd))
    def forward(self, x):
        x = x + self.a(self.ln1(x)); x = x + self.mlp(self.ln2(x)); return x

class MelIrisMini(nn.Module):
    def __init__(self, vocab_size=4096, n_embd=64, n_head=4, n_layer=3, block_size=64):
        super().__init__()
        self.block_size = block_size
        self.te = nn.Embedding(vocab_size, n_embd)
        self.pe = nn.Embedding(block_size, n_embd)
        self.blocks = nn.ModuleList([Blk(n_embd, n_head, block_size) for _ in range(n_layer)])
        self.lnf = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.head.weight = self.te.weight
    def forward(self, idx):
        T = idx.size(1)
        x = self.te(idx) + self.pe(torch.arange(T, device=idx.device).unsqueeze(0))
        for b in self.blocks: x = b(x)
        return self.head(self.lnf(x))
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        for _ in range(max_new_tokens):
            ic = idx[:, -self.block_size:]
            logits = self(ic)
            logits = logits[:,-1,:] / temperature
            if top_k:
                v,_ = torch.topk(logits, top_k); logits[logits < v[:,[-1]]] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx

def load_model(checkpoint_path):
    ck = torch.load(checkpoint_path, weights_only=False)
    config = ck['config']
    model = MelIrisMini(**config)
    model.load_state_dict(ck['state'])
    model.eval()
    return model