"""
EchoCoder — FSI coding specialist.

A small decoder-only Transformer (Llama-flavored: RMSNorm + RoPE + SwiGLU + causal
MHA) trained FROM SCRATCH on a code corpus. Char-level vocab => zero tokenizer
dependencies, fully self-contained, runs locally on CPU.

Why this is "our own model":
  * architecture is written here from scratch (no pretrained weights, no HuggingFace
    base model).
  * trained on a corpus we generate (TinyCode) so the whole pipeline is reproducible.
  * exports to TorchScript (portable CPU runtime) AND to GGUF (llama.cpp compatible)
    so it fits the sovereign / off-grid / private-first stack.

Tensor names intentionally match llama.cpp so the GGUF export is loadable.
"""

import math
import random
import struct
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ----------------------------- config --------------------------------------

class Cfg:
    d_model = 128
    n_layers = 4
    n_heads = 4
    ctx = 192
    ffn_mult = 2
    vocab = 256          # raw bytes / chars
    rope_base = 10000.0


# ----------------------------- model ----------------------------------------

class RMSNorm(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.w = nn.Parameter(torch.ones(d))

    def forward(self, x):
        return F.rms_norm(x, (x.shape[-1],), self.w, 1e-6)


def rope(x, base):
    # x: B, nh, T, hd
    T = x.shape[2]
    hd = x.shape[3]
    inv = 1.0 / (base ** (torch.arange(0, hd, 2, dtype=torch.float32) / hd))
    t = torch.arange(T, dtype=torch.float32)
    freqs = torch.outer(t, inv)              # T, hd/2
    cos = torch.cos(freqs).unsqueeze(0).unsqueeze(0)   # 1,1,T,hd/2
    sin = torch.sin(freqs).unsqueeze(0).unsqueeze(0)
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    rot1 = x1 * cos - x2 * sin
    rot2 = x1 * sin + x2 * cos
    out = torch.empty_like(x)
    out[..., 0::2] = rot1
    out[..., 1::2] = rot2
    return out


class Attn(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.q = nn.Linear(c.d_model, c.d_model, bias=False)
        self.k = nn.Linear(c.d_model, c.d_model, bias=False)
        self.v = nn.Linear(c.d_model, c.d_model, bias=False)
        self.o = nn.Linear(c.d_model, c.d_model, bias=False)
        self.nh = c.n_heads
        self.hd = c.d_model // c.n_heads
        self.d_model = c.d_model
        self.base = c.rope_base

    def forward(self, x):
        B, T, _ = x.shape
        q = self.q(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        k = self.k(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        v = self.v(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        q, k = rope(q, self.base), rope(k, self.base)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).reshape(B, T, self.d_model)
        return self.o(out)


class SwiGLU(nn.Module):
    def __init__(self, c):
        super().__init__()
        h = c.d_model * c.ffn_mult
        self.gate = nn.Linear(c.d_model, h, bias=False)
        self.up = nn.Linear(c.d_model, h, bias=False)
        self.down = nn.Linear(h, c.d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class Block(nn.Module):
    def __init__(self, c):
        super().__init__()
        self.attn_norm = RMSNorm(c.d_model)
        self.attn = Attn(c)
        self.ffn_norm = RMSNorm(c.d_model)
        self.ffn = SwiGLU(c)

    def forward(self, x):
        x = x + self.attn(self.attn_norm(x))
        x = x + self.ffn(self.ffn_norm(x))
        return x


class EchoCoder(nn.Module):
    def __init__(self, c=Cfg()):
        super().__init__()
        self.c = c
        self.tok = nn.Parameter(torch.zeros(c.vocab, c.d_model))   # token_embd
        self.blocks = nn.ModuleList([Block(c) for _ in range(c.n_layers)])
        self.norm = RMSNorm(c.d_model)                            # output_norm
        self.head = nn.Linear(c.d_model, c.vocab, bias=False)     # output

    def forward(self, idx):
        # idx: B,T of ints in [0,vocab)
        B, T = idx.shape
        x = self.tok[idx]                       # B,T,d
        for blk in self.blocks:
            x = blk(x)
        x = self.norm(x)
        return self.head(x)                     # B,T,vocab

    # ---- naming for GGUF/llama.cpp ----
    def state_dict_llama(self):
        sd = {}
        sd["token_embd.weight"] = self.tok
        for i, b in enumerate(self.blocks):
            p = b.attn
            sd[f"blk.{i}.attn_norm.weight"] = b.attn_norm.w
            sd[f"blk.{i}.attn_q.weight"] = p.q.weight
            sd[f"blk.{i}.attn_k.weight"] = p.k.weight
            sd[f"blk.{i}.attn_v.weight"] = p.v.weight
            sd[f"blk.{i}.attn_output.weight"] = p.o.weight
            sd[f"blk.{i}.ffn_norm.weight"] = b.ffn_norm.w
            sd[f"blk.{i}.ffn_gate.weight"] = b.ffn.gate.weight
            sd[f"blk.{i}.ffn_up.weight"] = b.ffn.up.weight
            sd[f"blk.{i}.ffn_down.weight"] = b.ffn.down.weight
        sd["output_norm.weight"] = self.norm.w
        sd["output.weight"] = self.head.weight
        return sd


def count_params(m):
    return sum(p.numel() for p in m.parameters())


# ----------------------------- corpus ---------------------------------------

def gen_code_snippet(rng):
    styles = [
        lambda: f"def {rng.choice(['add','sub','mul','div','max','min'])}"
                f"(a, b):\n    return a {rng.choice(['+','-','*','/'])} b\n",
        lambda: f"def factorial(n):\n    if n <= 1:\n        return 1\n"
                f"    return n * factorial(n - 1)\n",
        lambda: f"def sum_list(xs):\n    total = 0\n    for x in xs:\n"
                f"        total += x\n    return total\n",
        lambda: f"def is_prime(n):\n    if n < 2:\n        return False\n"
                f"    for i in range(2, int(n ** 0.5) + 1):\n"
                f"        if n % i == 0:\n            return False\n    return True\n",
        lambda: f"def greet(name):\n    return f'hello {{name}}'\n",
        lambda: f"class {rng.choice(['Stack','Queue','Node','Calc'])}:\n"
                f"    def __init__(self):\n        self.items = []\n"
                f"    def push(self, v):\n        self.items.append(v)\n"
                f"    def pop(self):\n        return self.items.pop()\n",
        lambda: f"def fib(n):\n    a, b = 0, 1\n    for _ in range(n):\n"
                f"        a, b = b, a + b\n    return a\n",
        lambda: f"def map_double(xs):\n    return [x * 2 for x in xs]\n",
        lambda: f"def clamp(v, lo, hi):\n    return max(lo, min(hi, v))\n",
        lambda: f"def load_config(path):\n    with open(path) as f:\n"
                f"        return f.read().splitlines()\n",
    ]
    s = rng.choice(styles)()
    # pad with a blank line so the model learns structure
    return s + "\n"


def build_corpus(path, n=4000, seed=0):
    rng = random.Random(seed)
    with open(path, "w") as f:
        for _ in range(n):
            f.write(gen_code_snippet(rng))
    size = sum(len(gen_code_snippet(rng)) for _ in range(0))  # placeholder
    total = 0
    with open(path) as f:
        total = len(f.read())
    return total


# ----------------------------- data -----------------------------------------

class CharData:
    def __init__(self, path, cfg, seq):
        data = open(path, "rb").read()
        self.ids = [b if b < cfg.vocab else ord("?") for b in data]
        self.seq = seq
        self.n = len(self.ids)

    def batch(self, batch_size):
        idxs = [random.randint(0, self.n - self.seq - 1) for _ in range(batch_size)]
        x = torch.tensor([self.ids[i:i + self.seq] for i in idxs], dtype=torch.long)
        y = torch.tensor([self.ids[i + 1:i + self.seq + 1] for i in idxs], dtype=torch.long)
        return x, y


# ----------------------------- train ----------------------------------------

def train(corpus_path, out_dir, steps=3000, batch=32, lr=3e-3, seed=0):
    torch.manual_seed(seed)
    random.seed(seed)
    cfg = Cfg()
    build_if_needed(corpus_path)
    data = CharData(corpus_path, cfg, cfg.ctx)
    model = EchoCoder(cfg)
    print(f"[EchoCoder] params={count_params(model):,}", flush=True)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    best = 1e9
    t0 = time.time()
    for step in range(1, steps + 1):
        x, y = data.batch(batch)
        loss = F.cross_entropy(model(x).view(-1, cfg.vocab), y.view(-1))
        opt.zero_grad()
        loss.backward()
        opt.step()
        if step % 100 == 0:
            print(f"step {step}/{steps} loss={loss.item():.3f} "
                  f"({time.time()-t0:.0f}s)", flush=True)
        if loss.item() < best:
            best = loss.item()
            torch.save(model.state_dict_llama(), f"{out_dir}/echocoder_best.pt")
    torch.save(model.state_dict_llama(), f"{out_dir}/echocoder.pt")
    # torchscript
    model.eval()
    traced = torch.jit.trace(model.eval(), torch.zeros(1, 1, dtype=torch.long))
    traced.save(f"{out_dir}/echocoder_ts.pt")
    print(f"[EchoCoder] DONE best_loss={best:.3f} -> {out_dir}", flush=True)
    return model, cfg


def build_if_needed(path):
    import os
    if not os.path.exists(path):
        build_corpus(path)


# ----------------------------- generate -------------------------------------

@torch.no_grad()
def generate(model, cfg, prompt="def ", length=120, temp=0.8):
    model.eval()
    ids = [ord(c) if ord(c) < cfg.vocab else ord("?") for c in prompt]
    for _ in range(length):
        ctx = torch.tensor([ids[-cfg.ctx:]], dtype=torch.long)
        logits = model(ctx)[0, -1]
        if temp > 0:
            logits = logits / temp
            p = torch.softmax(logits, -1)
            nxt = torch.multinomial(p, 1).item()
        else:
            nxt = int(logits.argmax())
        ids.append(nxt)
        if nxt == ord("\n") and ids[-2] == ord("\n"):
            break
    return bytes(ids).decode("utf-8", "replace")


# ----------------------------- GGUF export ----------------------------------

def export_gguf(state_dict, cfg, path):
    # minimal f32 GGUF, llama.cpp-compatible tensor names
    def gguf_str(s):
        b = s.encode()
        return struct.pack("i", len(b)) + b

    n_tensors = len(state_dict)
    # hyperparams
    hp = {
        "general.architecture": ("str", "llama"),
        "llama.context_length": ("uint32", cfg.ctx),
        "llama.embedding_length": ("uint32", cfg.d_model),
        "llama.block_count": ("uint32", cfg.n_layers),
        "llama.attention.head_count": ("uint32", cfg.n_heads),
        "llama.feed_forward_length": ("uint32", cfg.d_model * cfg.ffn_mult),
        "llama.attention.layer_norm_rms_epsilon": ("float32", 1e-6),
        "general.file_type": ("uint32", 0),
    }
    # build metadata + tensor info
    meta_buf = b""
    for k, (t, v) in hp.items():
        if t == "str":
            meta_buf += gguf_str(k) + struct.pack("i", 0) + gguf_str(v)
        elif t == "uint32":
            meta_buf += gguf_str(k) + struct.pack("i", 4) + struct.pack("I", v)
        elif t == "float32":
            meta_buf += gguf_str(k) + struct.pack("i", 5) + struct.pack("f", v)
    # tensor info: name, n_dims, dims..., ggml type (0=f32)
    ti = b""
    data = b""
    type_f32 = 0
    for name, t in state_dict.items():
        arr = t.detach().cpu().numpy().astype("float32")
        ti += gguf_str(name)
        ti += struct.pack("i", len(arr.shape))
        for d in arr.shape:
            ti += struct.pack("I", d)
        ti += struct.pack("i", type_f32)
        data += arr.tobytes()
    header = b"gguf" + struct.pack("i", 3) + struct.pack("Q", len(hp)) \
        + struct.pack("Q", n_tensors) + struct.pack("Q", len(meta_buf) + len(ti)) \
        + meta_buf + ti
    with open(path, "wb") as f:
        f.write(header)
        f.write(data)
    print(f"[EchoCoder] wrote GGUF -> {path} ({len(header)+len(data)} bytes)", flush=True)
    return path


if __name__ == "__main__":
    import os
    d = "/tmp/echocoder"
    os.makedirs(d, exist_ok=True)
    cp = f"{d}/tinycode.txt"
    if not os.path.exists(cp):
        n = build_corpus(cp)
        print(f"[corpus] {n} chars", flush=True)
    model, cfg = train(cp, d, steps=1200, batch=64)
    # GGUF
    sd = torch.load(f"{d}/echocoder.pt")
    export_gguf(sd, cfg, f"{d}/echocoder.f32.gguf")
    # demo generation
    print("=== sample generation ===")
    print(generate(model, cfg, "def ", length=140))