#!/usr/bin/env python3
"""
FSI_ECHO - Morphing Code Swarm
Novel architecture: token morph embedding + nanobot swarm + assembly blocks + self-verification.
2.6M params — fits in 1.3MB at q4, runs on any phone.
"""
import os, sys, json, time, math, random, re, struct
from typing import List, Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F

# =============================================================================
# 1. TOKENIZER
# =============================================================================
class CodeTokenizer:
    SPECIAL = {
        '<PAD>': 0, '<EOS>': 1, '<BOS>': 2, '<UNK>': 3,
        '<BUG>': 4, '<FIX>': 5, '<CODE>': 6, '<EXPLAIN>': 7,
        '<MORPH>': 8, '<ASSEMBLE>': 9, '<SCOUT>': 10, '<COMBAT>': 11,
    }
    def __init__(self, vocab_size: int = 4096):
        self.vocab_size = vocab_size
        self.vocab = dict(self.SPECIAL)
        self.inverse = {v: k for k, v in self.SPECIAL.items()}
        self.next_id = len(self.SPECIAL)
        self._build()
    def _build(self):
        for i in range(32, 127):
            self._add(chr(i))
        for t in ['def','class','return','if','else','elif','for','while','in','not',
            'and','or','import','from','as','try','except','finally','raise','with',
            'pass','break','continue','yield','lambda','self','None','True','False',
            'async','await','global','nonlocal','assert','del','print','len','range',
            'int','str','float','list','dict','set','tuple','type','is','isinstance',
            'hasattr','getattr','setattr','super','open','Exception','ValueError',
            'TypeError','KeyError','IndexError','AttributeError','ImportError',
            'Error','Warning','property','staticmethod','classmethod']:
            self._add(t)
        for s in ['==','!=','<=','>=','->','+=','-=','*=','/=','//=','**=','%=',
                  '<<','>>','**','//','::','=>','++','--','...']:
            self._add(s)
        for t in ['fn','func','function','const','let','var','this','typeof','void',
            'null','undefined','prototype','module','exports','require','new','delete',
            'throw','catch','switch','case','default','do','while','interface','enum',
            'implements','private','public','protected','abstract','final','static',
            'package','boolean','byte','char','double','float','int','long','short',
            'printf','scanf','malloc','free','sizeof','typedef','struct','union',
            'include','define','template','typename','namespace','using','virtual',
            'override','friend','operator','inline','explicit','string','vector',
            'map','set','auto','decltype','noexcept','constexpr','std','cout','cin',
            'endl','printf','scanf','NULL','nullptr','true','false','bool']:
            self._add(t)
        while self.next_id < self.vocab_size:
            self._add(f'v{self.next_id}')
    def _add(self, t):
        if t not in self.vocab and self.next_id < self.vocab_size:
            self.vocab[t] = self.next_id
            self.inverse[self.next_id] = t
            self.next_id += 1
    def encode(self, text: str, bos: bool = True, eos: bool = False) -> List[int]:
        ids = []
        if bos:
            ids.append(2)
        for token in re.findall(r'<[^>]+>|[A-Za-z_][A-Za-z0-9_]*|\.\.\.|==|!=|<=|>=|->|\*\*|//|::|=>|\d+\.\d*|\d+|\S', text):
            if token in self.vocab:
                ids.append(self.vocab[token])
            elif token.lower() in self.vocab:
                ids.append(self.vocab[token.lower()])
            else:
                for ch in token:
                    if ch in self.vocab:
                        ids.append(self.vocab[ch])
                    else:
                        ids.append(3)
        if eos:
            ids.append(1)
        return ids[:2048]
    def decode(self, ids: List[int], skip_special: bool = True) -> str:
        tokens = []
        for i in ids:
            if i in self.inverse:
                t = self.inverse[i]
                if skip_special and t.startswith('<') and t.endswith('>'):
                    continue
                tokens.append(t)
            else:
                tokens.append(' ')
        return ''.join(tokens)
    @property
    def pad_id(self): return 0
    @property
    def eos_id(self): return 1
    @property
    def bos_id(self): return 2
    @property
    def vocab_size_(self): return len(self.vocab)

# =============================================================================
# 2. MORPH EMBEDDING
# =============================================================================
class MorphEmbedding(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, morph_width: int = 3):
        super().__init__()
        self.d_model = d_model
        self.morph_width = morph_width
        self.base_embed = nn.Embedding(vocab_size, d_model)
        # Causal morph: only current + past (morph_width tokens, no future leakage)
        self.morph_net = nn.Sequential(
            nn.Linear(d_model * morph_width, d_model), nn.Tanh(),
            nn.Linear(d_model, d_model),
        )
        self.gate = nn.Linear(d_model * 2, d_model)

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        B, T = tokens.shape
        base = self.base_embed(tokens)
        # Causal sliding window: each position only sees itself and preceding tokens
        padded = F.pad(base, (0, 0, self.morph_width - 1, 0), mode='replicate')
        contexts = []
        for i in range(self.morph_width):
            contexts.append(padded[:, i:i+T, :])
        context = torch.cat(contexts, dim=-1)
        morph = self.morph_net(context)
        gate = torch.sigmoid(self.gate(torch.cat([base, morph], dim=-1)))
        return base + gate * morph

# =============================================================================
# 3. NANOBOT SWARM
# =============================================================================
class NanobotSwarm(nn.Module):
    def __init__(self, d_model: int, n_nanobots: int = 512, scout_dim: int = 64, combat_dim: int = 128):
        super().__init__()
        self.d_model = d_model
        self.n_nanobots = n_nanobots
        self.nano_keys = nn.Parameter(torch.randn(n_nanobots, d_model // 4))
        self.nano_vals = nn.Parameter(torch.randn(n_nanobots, d_model))
        self.scout_ffn = nn.Sequential(
            nn.Linear(d_model, scout_dim), nn.GELU(), nn.Linear(scout_dim, d_model),
        )
        self.combat_ffn = nn.Sequential(
            nn.Linear(d_model, combat_dim), nn.GELU(), nn.Linear(combat_dim, d_model),
        )
        self.mode_router = nn.Linear(d_model, 2)
        self.assembly = nn.Linear(d_model, d_model)
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        residual = x
        x = self.norm(x)
        router_query = x[..., :self.d_model // 4]
        nano_scores = torch.matmul(router_query, self.nano_keys.T)
        nano_weights = F.softmax(nano_scores / math.sqrt(self.d_model // 4), dim=-1)
        mode_logits = self.mode_router(x)
        mode_w = F.softmax(mode_logits, dim=-1)
        scout_out = self.scout_ffn(x)
        combat_out = self.combat_ffn(x)
        mode_out = mode_w[:, :, 0:1] * scout_out + mode_w[:, :, 1:2] * combat_out
        nano_out = torch.matmul(nano_weights, self.nano_vals)
        out = residual + self.dropout(self.assembly(mode_out + nano_out))
        return out

# =============================================================================
# 4. ASSEMBLY BLOCK
# =============================================================================
class AssemblyBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int = 4, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 2), nn.GELU(),
            nn.Linear(d_model * 2, d_model), nn.Dropout(dropout),
        )
        self.adapt_gate = nn.Linear(d_model, 1)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        qkv = self.qkv(self.norm1(x))
        qkv = qkv.reshape(B, T, 3, self.n_heads, self.head_dim)
        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        scale = math.sqrt(self.head_dim)
        attn = torch.matmul(q, k.transpose(-2, -1)) / scale
        mask = torch.triu(torch.ones(T, T, device=x.device) * float('-inf'), diagonal=1)
        attn = attn + mask.unsqueeze(0)
        attn = F.softmax(attn, dim=-1)
        out = torch.matmul(attn, v).transpose(1, 2).reshape(B, T, D)
        out = self.proj(out)
        x = x + self.dropout(out)
        gate = torch.sigmoid(self.adapt_gate(self.norm2(x)))
        x = x + gate * self.dropout(self.ffn(self.norm2(x)))
        return x

# =============================================================================
# 5. FSI_ECHO MODEL
# =============================================================================
class FSIEchoModel(nn.Module):
    def __init__(self, vocab_size: int = 4096, d_model: int = 192,
                 n_swarm_layers: int = 3, n_assembly_layers: int = 3,
                 n_nanobots: int = 512):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.morph_embed = MorphEmbedding(vocab_size, d_model)
        self.swarm_layers = nn.ModuleList([
            NanobotSwarm(d_model, n_nanobots) for _ in range(n_swarm_layers)
        ])
        self.assembly_layers = nn.ModuleList([
            AssemblyBlock(d_model) for _ in range(n_assembly_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
        self.verify = nn.Sequential(
            nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1),
        )
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.morph_embed.base_embed.weight
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, mean=0.0, std=0.02)
                if hasattr(m, 'bias') and m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None) -> Dict:
        x = self.morph_embed(tokens)
        for layer in self.swarm_layers:
            x = layer(x)
        for layer in self.assembly_layers:
            x = layer(x)
        x = self.norm(x)
        logits = self.lm_head(x)
        confidence = torch.sigmoid(self.verify(x)).squeeze(-1)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1), ignore_index=0)
        return {'logits': logits, 'confidence': confidence, 'loss': loss}

    @torch.no_grad()
    def generate(self, tokenizer, prompt: str, max_tokens: int = 512,
                 temperature: float = 0.3, top_k: int = 5, top_p: float = 0.9) -> Dict:
        self.eval()
        device = next(self.parameters()).device
        toks = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device)
        generated_ids = []
        confs = []
        for _ in range(min(max_tokens, 2048 - toks.shape[1])):
            out = self.forward(toks)
            logits = out['logits'][0, -1, :] / max(temperature, 0.01)
            if top_k > 0:
                vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < vals[-1]] = float('-inf')
            if top_p < 1.0:
                sorted_lg, sorted_idx = torch.sort(logits, descending=True)
                cum = torch.cumsum(F.softmax(sorted_lg, dim=-1), dim=-1)
                rm = cum > top_p
                rm[1:] = rm[:-1].clone()
                rm[0] = False
                logits[sorted_idx[rm]] = float('-inf')
            logits = torch.nan_to_num(logits, nan=-100.0, posinf=100.0, neginf=-100.0)
            logits[0] = float('-inf')
            for rid in range(300, logits.size(-1)):
                logits[rid] = float('-inf')
            if (logits > float(-1e9)).sum() == 0:
                logits[tokenizer.eos_id] = 0.0
            nxt = logits.argmax().unsqueeze(0)
            confs.append(out['confidence'][0, -1].item())
            if nxt.item() == tokenizer.eos_id:
                break
            generated_ids.append(nxt.item())
            toks = torch.cat([toks, nxt.unsqueeze(0)], dim=1)
        generated = tokenizer.decode(generated_ids, skip_special=True)
        avg_conf = sum(confs) / max(len(confs), 1)
        return {'generated': generated, 'confidence': avg_conf, 'tokens': len(generated_ids)}

    def param_count(self) -> int:
        return sum(p.numel() for p in self.parameters())

# =============================================================================
# 6. TRAINER
# =============================================================================
class Trainer:
    def __init__(self, model: FSIEchoModel, tokenizer: CodeTokenizer, lr: float = 3e-4):
        self.model = model
        self.tokenizer = tokenizer
        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.95), weight_decay=0.1)
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=5000, eta_min=1e-5)

    def step(self, texts: List[str], batch_size: int = 2, device: str = 'cpu') -> float:
        self.model.train()
        batch = random.sample(texts, min(batch_size, len(texts)))
        encoded = []
        max_len = 0
        for t in batch:
            toks = self.tokenizer.encode(t)
            if 5 < len(toks) <= 2048:
                encoded.append(toks)
                max_len = max(max_len, len(toks))
        if not encoded:
            return 0.0
        padded = torch.zeros(len(encoded), max_len, dtype=torch.long)
        targets = torch.full((len(encoded), max_len), 0, dtype=torch.long)
        for i, toks in enumerate(encoded):
            padded[i, :len(toks)] = torch.tensor(toks)
            targets[i, :len(toks)-1] = torch.tensor(toks[1:])
            targets[i, len(toks)-1] = 1
        padded, targets = padded.to(device), targets.to(device)
        out = self.model(padded, targets)
        self.optimizer.zero_grad()
        out['loss'].backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.optimizer.step()
        self.scheduler.step()
        return out['loss'].item()

    def save(self, path: str):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save({
            'model': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
            'config': {'vocab_size': self.model.vocab_size, 'd_model': self.model.d_model,
                       'n_swarm': len(self.model.swarm_layers), 'n_assembly': len(self.model.assembly_layers),
                       'n_nanobots': self.model.swarm_layers[0].n_nanobots if self.model.swarm_layers else 512},
        }, path)

    @classmethod
    def load(cls, path: str, device: str = 'cpu') -> 'Trainer':
        data = torch.load(path, map_location=device, weights_only=True)
        cfg = data.get('config', {})
        model = FSIEchoModel(
            vocab_size=cfg.get('vocab_size', 4096), d_model=cfg.get('d_model', 192),
            n_swarm_layers=cfg.get('n_swarm', 3), n_assembly_layers=cfg.get('n_assembly', 3),
            n_nanobots=cfg.get('n_nanobots', 512),
        )
        model.load_state_dict(data['model'])
        tokenizer = CodeTokenizer(vocab_size=cfg.get('vocab_size', 4096))
        t = cls(model, tokenizer)
        if 'optimizer' in data:
            t.optimizer.load_state_dict(data['optimizer'])
        model.to(device)
        return t

# =============================================================================
# 7. CLOSED-LOOP DEBUG
# =============================================================================
class CodeVerifier:
    @staticmethod
    def check_syntax(code: str) -> Tuple[bool, str]:
        try:
            compile(code, '<debug>', 'exec')
            return True, ""
        except SyntaxError as e:
            return False, f"Line {e.lineno}: {e.msg}"
    @staticmethod
    def find_issues(code: str) -> List[str]:
        issues = []
        for i, line in enumerate(code.split('\n'), 1):
            s = line.strip()
            if s == 'except:':
                issues.append(f"L{i}: Bare except — specify exception")
        return issues
    def verify(self, code: str) -> Dict:
        ok, err = self.check_syntax(code)
        issues = self.find_issues(code)
        return {'valid': ok and not issues, 'syntax_ok': ok, 'error': err, 'issues': issues, 'code': code}

class ClosedLoopDebugger:
    def __init__(self, model: FSIEchoModel, tokenizer: CodeTokenizer, max_iters: int = 3):
        self.model = model
        self.tokenizer = tokenizer
        self.max_iters = max_iters
        self.verifier = CodeVerifier()
    def debug(self, code: str, requirement: str = "", max_iterations: int = None) -> Dict:
        iters = max_iterations or self.max_iters
        result = self._extract_code(code)
        if not result:
            return {'code': code, 'error': 'Could not parse code', 'iterations': 0, 'confidence': 0.0}
        buggy = result
        prompt = f"Fix this code:\n```python\n{buggy}\n```\nFixed:\n```python\n"
        best_code, best_v = buggy, self.verifier.verify(buggy)
        for i in range(iters):
            gen = self.model.generate(self.tokenizer, prompt, max_tokens=256, temperature=0.3, top_k=5)
            fixed = self._extract_code(gen['generated'])
            if not fixed:
                prompt += "\n```python\n"
                continue
            v = self.verifier.verify(fixed)
            if v['valid']:
                return {'code': fixed, 'verification': v, 'iterations': i+1, 'confidence': gen['confidence']}
            if len(v['issues']) < len(best_v['issues']):
                best_code, best_v = fixed, v
            if v['issues']:
                prompt += f"\nIssues: {'; '.join(v['issues'])}\nFixed:\n```python\n"
            elif not v['syntax_ok']:
                prompt += f"\n{v['error']}\nFixed:\n```python\n"
            else:
                break
        return {'code': best_code, 'verification': best_v, 'iterations': iters, 'confidence': 0.0}
    def _extract_code(self, text: str) -> str:
        m = re.search(r'```(?:python)?\n(.*?)```', text, re.DOTALL)
        if m: return m.group(1).strip()
        lines = text.split('\n')
        code = []
        in_code = False
        for line in lines:
            if line.startswith('```'): in_code = not in_code; continue
            if in_code: code.append(line)
        return '\n'.join(code).strip() if code else ''

# =============================================================================
# 8. GGUF EXPORT
# =============================================================================
def export_gguf(model: FSIEchoModel, tokenizer: CodeTokenizer, path: str):
    sd = model.state_dict()
    keys = sorted(sd.keys())
    meta = {
        'general.name': 'FSI_ECHO', 'general.architecture': 'fsi_echo',
        'general.description': 'Morphing Code Swarm',
        'general.file_type': 0, 'general.vocab_size': model.vocab_size,
        'general.context_length': 2048, 'general.parameter_count': model.param_count(),
        'fsi_echo.block_count': len(model.swarm_layers) + len(model.assembly_layers),
        'fsi_echo.embedding_length': model.d_model,
        'fsi_echo.feed_forward_length': model.d_model * 2,
        'fsi_echo.attention.head_count': 4,
        'fsi_echo.nanobot_count': model.swarm_layers[0].n_nanobots if model.swarm_layers else 512,
    }
    with open(path, 'wb') as f:
        f.write(b'GGUF' + struct.pack('<I', 3) + struct.pack('<Q', len(keys)) + struct.pack('<Q', len(meta)))
        for k, v in meta.items():
            f.write(struct.pack('<I', len(k)) + k.encode())
            if isinstance(v, str):
                f.write(struct.pack('<I', 8) + struct.pack('<Q', len(v)) + v.encode())
            elif isinstance(v, int):
                f.write(struct.pack('<I', 4) + struct.pack('<I', v))
            elif isinstance(v, float):
                f.write(struct.pack('<I', 6) + struct.pack('<f', v))
        offset = 0
        for name in keys:
            nb = len(name)
            f.write(struct.pack('<Q', nb) + name.encode() +
                    struct.pack('<I', len(sd[name].shape)) +
                    b''.join(struct.pack('<Q', d) for d in sd[name].shape) +
                    struct.pack('<I', 0) + struct.pack('<Q', offset))
            offset += sd[name].numel() * 4
        for name in keys:
            f.write(sd[name].float().numpy().tobytes())

# =============================================================================
# 9. MAIN
# =============================================================================
if __name__ == '__main__':
    import argparse
    p = argparse.ArgumentParser(description='FSI_ECHO')
    p.add_argument('--train', action='store_true')
    p.add_argument('--gen', type=str, help='Generate from prompt')
    p.add_argument('--gguf', type=str, help='Export to GGUF')
    p.add_argument('--load', type=str, help='Load checkpoint')
    p.add_argument('--steps', type=int, default=5000)
    args = p.parse_args()
    device = 'cpu'
    if args.load and os.path.exists(args.load):
        trainer = Trainer.load(args.load, device)
    else:
        model = FSIEchoModel()
        model.to(device)
        tok = CodeTokenizer()
        print(f"New model: {model.param_count():,} params")
        trainer = Trainer(model, tok)
    if args.gen:
        r = trainer.model.generate(trainer.tokenizer, args.gen, max_tokens=256)
        print(r['generated'])
    if args.gguf:
        export_gguf(trainer.model, trainer.tokenizer, args.gguf)