import os, sys, argparse, math as _math
import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

HERE = os.path.dirname(os.path.abspath(__file__))
PROJ = os.path.dirname(HERE)


# ─── V2 model (backward compat: GGUF + QLoRA checkpoints) ───────────────

class ALiBiAttention(torch.nn.Module):
    def __init__(self, hidden: int, num_heads: int, num_kv_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = hidden // num_heads
        self.num_groups = num_heads // num_kv_heads
        self.q_proj = torch.nn.Linear(hidden, hidden, bias=False)
        self.k_proj = torch.nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
        self.v_proj = torch.nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
        self.o_proj = torch.nn.Linear(hidden, hidden, bias=False)

    @staticmethod
    def _get_alibi_slopes(num_heads):
        n = 2 ** _math.ceil(_math.log2(num_heads))
        return torch.tensor([2.0 ** (-(i + 1)) for i in range(n)][:num_heads])

    def forward(self, x):
        B, T, C = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        k = k.repeat_interleave(self.num_groups, dim=1)
        v = v.repeat_interleave(self.num_groups, dim=1)
        scores = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
        slopes = self._get_alibi_slopes(self.num_heads).to(x.device, x.dtype)
        pos = torch.arange(T, device=x.device)
        alibi = (pos.view(1, T) - pos.view(T, 1)).abs().neg().unsqueeze(0).unsqueeze(0)
        scores = scores + alibi * slopes.view(-1, 1, 1)
        causal = torch.triu(torch.full((T, T), float('-inf'), device=x.device, dtype=x.dtype), diagonal=1)
        scores = scores + causal
        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(x.dtype)
        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.o_proj(out)


class RMSNormV2(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = torch.nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


class SwiGLU(torch.nn.Module):
    def __init__(self, hidden: int, intermediate: int):
        super().__init__()
        self.gate = torch.nn.Linear(hidden, intermediate, bias=False)
        self.up = torch.nn.Linear(hidden, intermediate, bias=False)
        self.down = torch.nn.Linear(intermediate, hidden, bias=False)

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


class TransformerBlockV2(torch.nn.Module):
    def __init__(self, hidden: int, intermediate: int, num_heads: int, num_kv_heads: int):
        super().__init__()
        self.ln1 = RMSNormV2(hidden)
        self.attn = ALiBiAttention(hidden, num_heads, num_kv_heads)
        self.ln2 = RMSNormV2(hidden)
        self.mlp = SwiGLU(hidden, intermediate)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))
        return x


class TinyModelV2(torch.nn.Module):
    def __init__(self, vocab_size=1757, hidden=128, intermediate=640,
                 num_layers=3, num_heads=8, num_kv_heads=4, max_seq_len=2048,
                 tie_weights=True):
        super().__init__()
        self.max_seq_len = max_seq_len
        self.token_embed = torch.nn.Embedding(vocab_size, hidden)
        self.blocks = torch.nn.ModuleList([
            TransformerBlockV2(hidden, intermediate, num_heads, num_kv_heads)
            for _ in range(num_layers)
        ])
        self.ln_f = RMSNormV2(hidden)
        self.lm_head = torch.nn.Linear(hidden, vocab_size, bias=False)
        if tie_weights:
            self.lm_head.weight = self.token_embed.weight

    def forward(self, input_ids):
        x = self.token_embed(input_ids)
        for block in self.blocks:
            x = block(x)
        x = self.ln_f(x)
        return self.lm_head(x)

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=128, temperature=0.7, top_p=0.9, stream_callback=None):
        self.eval()
        for _ in range(max_new_tokens):
            if input_ids.size(1) > self.max_seq_len:
                input_ids = input_ids[:, -self.max_seq_len:]
            logits = self.forward(input_ids)
            logits = logits[:, -1, :]
            if temperature > 0:
                logits = logits / temperature
            if top_p < 1.0:
                sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
                cum_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
                cutoff = cum_probs > top_p
                cutoff[..., 1:] = cutoff[..., :-1].clone()
                cutoff[..., 0] = False
                logits[~cutoff] = float('-inf')
            probs = F.softmax(logits, dim=-1)
            if temperature > 0:
                next_token = torch.multinomial(probs, 1)
            else:
                next_token = probs.argmax(dim=-1, keepdim=True)
            input_ids = torch.cat([input_ids, next_token], dim=1)
            if stream_callback:
                stream_callback(next_token.item())
            if next_token.item() == 2:
                break
        return input_ids


# ─── QLoRA (NF4 + LoRA) for V2 ─────────────────────────────────────────

NF4_LEVELS = torch.tensor([
    -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453,
    -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
    0.07958029955625534, 0.16093020141124725, 0.24611230194568634,
    0.33791524171829224, 0.44070982933044434, 0.5626170039176941,
    0.7229568362236023, 1.0,
])

def unpack_nf4(qweight, shape):
    n = qweight.numel() * 2
    lo = (qweight & 0x0F).view(-1)
    hi = ((qweight >> 4) & 0x0F).view(-1)
    indices = torch.stack([lo, hi], dim=1).reshape(n)
    return indices[:shape[0] * shape[1]].reshape(shape)

def pack_nf4(indices):
    n = indices.numel()
    even = indices[0::2]
    odd = indices[1::2]
    packed = (odd << 4) | even
    return packed.to(torch.uint8)


class QLoRALinear(torch.nn.Module):
    def __init__(self, in_features: int, out_features: int, r: int = 8, alpha: float = 16, dropout: float = 0.0):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.r = r
        self.scaling = alpha / r
        self.dropout = torch.nn.Dropout(dropout)
        self.register_buffer("qweight", torch.zeros((in_features * out_features + 1) // 2, dtype=torch.uint8))
        self.register_buffer("scales", torch.zeros(out_features))
        self.register_buffer("bias", torch.zeros(out_features))
        self.lora_A = torch.nn.Parameter(torch.zeros(r, in_features))
        self.lora_B = torch.nn.Parameter(torch.zeros(out_features, r))
        torch.nn.init.kaiming_uniform_(self.lora_A, a=_math.sqrt(5))

    def dequantize(self):
        indices = unpack_nf4(self.qweight, (self.out_features, self.in_features))
        return NF4_LEVELS.to(self.qweight.device)[indices.long()] * self.scales[:, None]

    def quantize_from(self, weight, bias=None):
        w = weight.float()
        scales = w.abs().max(dim=1, keepdim=True).values.clamp(min=1e-8)
        scaled = (w / scales).clamp(-1, 1)
        idx = (scaled[:, :, None] - NF4_LEVELS[None, None, :].to(w.device)).abs().argmin(dim=-1)
        self.qweight.copy_(pack_nf4(idx.reshape(-1)))
        self.scales.copy_(scales.squeeze(1))
        if bias is not None:
            self.bias.copy_(bias.float())

    def forward(self, x):
        base = F.linear(x, self.dequantize(), self.bias)
        return base + self.dropout(x) @ self.lora_A.T @ self.lora_B.T * self.scaling


def apply_qlora_v2(model, target_modules=None, r=8, alpha=16, dropout=0.0):
    if target_modules is None:
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate", "up", "down"]
    qlora_params = 0
    for name, module in model.named_modules():
        if not isinstance(module, torch.nn.Linear):
            continue
        key = name.split(".")[-1]
        if key not in target_modules:
            continue
        parent = model
        parts = name.split(".")
        for p in parts[:-1]:
            parent = getattr(parent, p)
        qlora = QLoRALinear(module.in_features, module.out_features, r=r, alpha=alpha, dropout=dropout)
        qlora.quantize_from(module.weight, module.bias)
        setattr(parent, parts[-1], qlora)
        qlora_params += 2 * r * module.in_features + module.out_features * r
    n = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f" QLoRA applied: {qlora_params:,} LoRA params | trainable: {n:,}")


# ─── V3 model (from model_tiny) ─────────────────────────────────────────

def _load_v3_model():
    from scripts.model_tiny import TinyModel as TinyModelV3
    return TinyModelV3


# ─── Tokenizer ──────────────────────────────────────────────────────────

def load_tokenizer(ckpt_type=None):
    from tokenizers import Tokenizer as Tk
    paths = [os.path.join(PROJ, "tokenizer", "tokenizer.json"),
             os.path.join(PROJ, "tokenizer.json")]
    if ckpt_type and ckpt_type.startswith("v2"):
        v2_path = os.path.join(os.path.dirname(PROJ), "lumia-v1", "tokenizer", "tokenizer.json")
        if os.path.exists(v2_path):
            paths.insert(0, v2_path)
    for p in paths:
        if os.path.exists(p):
            tok = Tk.from_file(p)
            print(f" Tokenizer: {p} ({tok.get_vocab_size()} vocab)")
            return tok
    print(" Tokenizer not found")
    sys.exit(1)


def build_chat(system, user, tokenizer):
    parts = [f"<|system|>\n{system}", f"<|user|>\n{user}", "<|assistant|>\n"]
    text = "\n".join(parts)
    return tokenizer.encode(text).ids


# ─── Detect checkpoint type ─────────────────────────────────────────────

def _detect_ckpt_type(state):
    keys = list(state.keys())
    if any("rpw" in k or "vcr" in k for k in keys):
        return "v3"
    if any("qweight" in k for k in keys):
        return "v2_qlora"
    if any("gate" in k for k in keys) or any("mlp.gate" in k for k in keys):
        return "v2_fp32"
    if any(k.startswith("blocks.") for k in keys) and any("ln1.weight" in k for k in keys):
        blk_keys = [k.split(".")[1] for k in keys if k.startswith("blocks.")]
        max_blk = max(int(b) for b in blk_keys) if blk_keys else 0
        if max_blk >= 3:
            return "v3"
    return "v2_fp32"


# ─── Loaders ────────────────────────────────────────────────────────────

def load_gguf(gguf_path):
    print(f" Loading GGUF: {gguf_path}")
    import gguf as _gguf
    reader = _gguf.GGUFReader(gguf_path)

    model = TinyModelV2(vocab_size=1757, hidden=128, intermediate=640,
                        num_layers=3, num_heads=8, num_kv_heads=4,
                        max_seq_len=2048, tie_weights=True)
    model.eval()

    state = {}
    for tensor in reader.tensors:
        name = tensor.name
        data = torch.from_numpy(tensor.data.copy())
        if name == "token_embd.weight":
            state["token_embed.weight"] = data
        elif name == "output_norm.weight":
            state["ln_f.weight"] = data
        elif name == "output.weight":
            state["lm_head.weight"] = data
        elif name.startswith("blk."):
            parts = name.split(".")
            blk = int(parts[1])
            sub = parts[2]
            if sub == "attn_norm":
                state[f"blocks.{blk}.ln1.{parts[3]}"] = data
            elif sub == "ffn_norm":
                state[f"blocks.{blk}.ln2.{parts[3]}"] = data
            elif sub == "attn_q":
                state[f"blocks.{blk}.attn.q_proj.weight"] = data
            elif sub == "attn_k":
                state[f"blocks.{blk}.attn.k_proj.weight"] = data
            elif sub == "attn_v":
                state[f"blocks.{blk}.attn.v_proj.weight"] = data
            elif sub == "attn_output":
                state[f"blocks.{blk}.attn.o_proj.weight"] = data
            elif sub == "ffn_gate":
                state[f"blocks.{blk}.mlp.gate.weight"] = data
            elif sub == "ffn_up":
                state[f"blocks.{blk}.mlp.up.weight"] = data
            elif sub == "ffn_down":
                state[f"blocks.{blk}.mlp.down.weight"] = data

    model.load_state_dict(state, strict=False)
    print(f" Loaded: {len(state)} tensors")
    return model


def load_checkpoint(ckpt_path):
    print(f" Loading checkpoint: {ckpt_path}")
    state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
    ckpt_type = _detect_ckpt_type(state)
    print(f" Detected: {ckpt_type}")

    if ckpt_type.startswith("v2"):
        has_qlora = ckpt_type == "v2_qlora"
        model = TinyModelV2(vocab_size=1757, hidden=128, intermediate=640,
                            num_layers=3, num_heads=8, num_kv_heads=4,
                            max_seq_len=2048, tie_weights=True)
        if has_qlora:
            apply_qlora_v2(model)
        model.load_state_dict(state, strict=False)
        model.eval()
        return model, ckpt_type

    TV3 = _load_v3_model()
    model = TV3()
    model.load_state_dict(state)
    model.eval()
    return model, ckpt_type


# ─── Main ───────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Infer Lumia V2 (GGUF/QLoRA) or V3 (best.pt)")
    parser.add_argument("--gguf", default=None, help="V2 GGUF path")
    parser.add_argument("--checkpoint", default="best.pt", help="Checkpoint path (default: best.pt)")
    parser.add_argument("--prompt", default="What is 2+2?")
    parser.add_argument("--system", default="You are a helpful AI assistant.")
    parser.add_argument("--max-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.7)
    parser.add_argument("--top-p", type=float, default=0.9)
    args = parser.parse_args()

    if args.gguf and args.checkpoint:
        print("Specify only one of --gguf or --checkpoint")
        sys.exit(1)

    model = None
    ckpt_type = None
    if args.gguf:
        if not os.path.exists(args.gguf):
            print(f"GGUF not found: {args.gguf}")
            sys.exit(1)
        model = load_gguf(args.gguf)
        ckpt_type = "gguf"
    elif args.checkpoint:
        if not os.path.exists(args.checkpoint):
            print(f"Checkpoint not found: {args.checkpoint}")
            sys.exit(1)
        model, ckpt_type = load_checkpoint(args.checkpoint)
    else:
        print("Specify --gguf or --checkpoint")
        sys.exit(1)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    print(f" Device: {device}")

    tok = load_tokenizer(ckpt_type)
    input_ids = build_chat(args.system, args.prompt, tok)
    input_ids = torch.tensor([input_ids], device=device)

    print("\n" + tok.decode(input_ids[0].tolist()), end="", flush=True)

    def stream(id_):
        print(tok.decode([id_]), end="", flush=True)

    model.generate(
        input_ids,
        max_new_tokens=args.max_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        stream_callback=stream,
    )
    print()


if __name__ == "__main__":
    main()