#!/usr/bin/env python3
"""
TinyLlama GENOME₀ Loader

Load TinyLlama weights from local GGUF file (not HuggingFace cache).
This is the GENOME₀ path: weights come from local storage, not internet.

Usage:
    from helix_cdc.loaders import load_tinyllama_genome0

    model, tokenizer, metadata = load_tinyllama_genome0()
    # model is frozen, ready for HelixLobe attachment
"""

import hashlib
import re
from pathlib import Path
from typing import Dict, Any, Tuple, Optional

import numpy as np
import torch
import torch.nn as nn

# Default paths
DEFAULT_GGUF_PATH = "models/tinyllama-1.1b-fp32.gguf"
HF_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"


def gguf_name_to_hf_name(gguf_name: str) -> str:
    """
    Convert GGUF tensor name to HuggingFace LlamaForCausalLM name.

    GGUF format:
        token_embd.weight -> model.embed_tokens.weight
        blk.0.attn_q.weight -> model.layers.0.self_attn.q_proj.weight
        blk.0.attn_k.weight -> model.layers.0.self_attn.k_proj.weight
        blk.0.attn_v.weight -> model.layers.0.self_attn.v_proj.weight
        blk.0.attn_output.weight -> model.layers.0.self_attn.o_proj.weight
        blk.0.ffn_gate.weight -> model.layers.0.mlp.gate_proj.weight
        blk.0.ffn_up.weight -> model.layers.0.mlp.up_proj.weight
        blk.0.ffn_down.weight -> model.layers.0.mlp.down_proj.weight
        blk.0.attn_norm.weight -> model.layers.0.input_layernorm.weight
        blk.0.ffn_norm.weight -> model.layers.0.post_attention_layernorm.weight
        output.weight -> lm_head.weight
        output_norm.weight -> model.norm.weight
    """
    # Special cases
    if gguf_name == "token_embd.weight":
        return "model.embed_tokens.weight"
    if gguf_name == "output.weight":
        return "lm_head.weight"
    if gguf_name == "output_norm.weight":
        return "model.norm.weight"

    # Block patterns: blk.N.xxx.weight
    m = re.match(r"blk\.(\d+)\.(.+)\.weight", gguf_name)
    if m:
        layer_idx = m.group(1)
        suffix = m.group(2)

        # Attention projections
        if suffix == "attn_q":
            return f"model.layers.{layer_idx}.self_attn.q_proj.weight"
        if suffix == "attn_k":
            return f"model.layers.{layer_idx}.self_attn.k_proj.weight"
        if suffix == "attn_v":
            return f"model.layers.{layer_idx}.self_attn.v_proj.weight"
        if suffix == "attn_output":
            return f"model.layers.{layer_idx}.self_attn.o_proj.weight"

        # FFN projections
        if suffix == "ffn_gate":
            return f"model.layers.{layer_idx}.mlp.gate_proj.weight"
        if suffix == "ffn_up":
            return f"model.layers.{layer_idx}.mlp.up_proj.weight"
        if suffix == "ffn_down":
            return f"model.layers.{layer_idx}.mlp.down_proj.weight"

        # Norms
        if suffix == "attn_norm":
            return f"model.layers.{layer_idx}.input_layernorm.weight"
        if suffix == "ffn_norm":
            return f"model.layers.{layer_idx}.post_attention_layernorm.weight"

    # Fallback: return as-is (will show up as unmapped)
    return gguf_name


def load_gguf_tensors(gguf_path: str) -> Dict[str, np.ndarray]:
    """
    Load all tensors from GGUF file.

    Returns dict: {gguf_name: numpy_array}

    Note: The gguf library's tensor.data is a numpy memmap that is ALREADY
    reshaped correctly. For example, token_embd.weight is (32000, 2048) which
    matches HF's (vocab_size, hidden_dim). We should NOT reshape based on
    tensor.shape which reports the GGUF storage format (2048, 32000).
    """
    from gguf import GGUFReader

    reader = GGUFReader(gguf_path)
    tensors = {}

    for tensor_info in reader.tensors:
        name = tensor_info.name
        # tensor.data is already a properly shaped numpy memmap
        # Just copy to a regular numpy array
        arr = np.array(tensor_info.data, dtype=np.float32)
        tensors[name] = arr

    return tensors


def gguf_to_pytorch_state_dict(gguf_path: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, Any]]:
    """
    Convert GGUF file to PyTorch state_dict.

    Returns:
        state_dict: {hf_name: torch.Tensor}
        mapping_info: {gguf_name: hf_name, unmapped: [...]}
    """
    gguf_tensors = load_gguf_tensors(gguf_path)

    state_dict = {}
    mapping_info = {"mapped": {}, "unmapped": []}

    for gguf_name, arr in gguf_tensors.items():
        hf_name = gguf_name_to_hf_name(gguf_name)

        # The gguf library already returns tensors in the correct shape
        # (e.g., embeddings are (vocab_size, hidden_dim) which is HF format)
        # No transpose needed!
        tensor = torch.from_numpy(arr.copy())
        state_dict[hf_name] = tensor

        if hf_name != gguf_name:
            mapping_info["mapped"][gguf_name] = hf_name
        else:
            mapping_info["unmapped"].append(gguf_name)

    return state_dict, mapping_info


def load_tinyllama_genome0(
    gguf_path: Optional[str] = None,
    device: str = "cpu",
    dtype: torch.dtype = torch.float32,
    freeze: bool = True,
) -> Tuple[nn.Module, Any, Dict[str, Any]]:
    """
    Load TinyLlama from local GGUF file (GENOME₀ path).

    This loads weights from a local GGUF file instead of HuggingFace cache,
    proving the model can be loaded without internet/HF dependency.

    Args:
        gguf_path: Path to TinyLlama GGUF file (default: models/tinyllama-1.1b-fp32.gguf)
        device: Target device ("cpu" or "cuda")
        dtype: Target dtype (torch.float32 or torch.float16)
        freeze: Whether to freeze all parameters (default True for lobe experiments)

    Returns:
        model: LlamaForCausalLM with weights loaded from GGUF
        tokenizer: HF tokenizer (still from HF, but could be local)
        metadata: Dict with loading info, mapping stats, SHA256
    """
    from transformers import AutoConfig, AutoTokenizer, LlamaForCausalLM

    # Resolve path
    if gguf_path is None:
        # Try relative to helix-cdc root
        helix_root = Path(__file__).parent.parent.parent
        gguf_path = helix_root / DEFAULT_GGUF_PATH
    else:
        gguf_path = Path(gguf_path)

    if not gguf_path.exists():
        raise FileNotFoundError(
            f"GGUF file not found: {gguf_path}\n"
            f"Expected TinyLlama FP32 GGUF at: {DEFAULT_GGUF_PATH}"
        )

    print(f"[GENOME₀] Loading TinyLlama from local GGUF...")
    print(f"[GENOME₀] Path: {gguf_path}")
    print(f"[GENOME₀] Size: {gguf_path.stat().st_size / 1e9:.2f} GB")

    # Compute GGUF SHA256 (first 1MB for speed)
    with open(gguf_path, "rb") as f:
        gguf_hash = hashlib.sha256(f.read(1024 * 1024)).hexdigest()[:16]
    print(f"[GENOME₀] SHA256 (1MB): {gguf_hash}...")

    # Load tensors from GGUF
    print(f"[GENOME₀] Parsing GGUF tensors...")
    state_dict, mapping_info = gguf_to_pytorch_state_dict(str(gguf_path))
    print(f"[GENOME₀] Loaded {len(state_dict)} tensors")
    print(f"[GENOME₀] Mapped: {len(mapping_info['mapped'])}, Unmapped: {len(mapping_info['unmapped'])}")

    # Create model shell from HF config
    print(f"[GENOME₀] Creating model architecture from HF config...")
    config = AutoConfig.from_pretrained(HF_MODEL_NAME)
    model = LlamaForCausalLM(config)

    # Load state dict
    print(f"[GENOME₀] Loading weights into model...")
    missing, unexpected = model.load_state_dict(state_dict, strict=False)

    if missing:
        print(f"[GENOME₀] Missing keys: {len(missing)}")
        for k in missing[:5]:
            print(f"  - {k}")
        if len(missing) > 5:
            print(f"  ... and {len(missing) - 5} more")

    if unexpected:
        print(f"[GENOME₀] Unexpected keys: {len(unexpected)}")
        for k in unexpected[:5]:
            print(f"  - {k}")

    # Freeze if requested
    if freeze:
        print(f"[GENOME₀] Freezing all parameters...")
        for p in model.parameters():
            p.requires_grad = False

    # Move to device/dtype
    model = model.to(device=device, dtype=dtype)

    # Load tokenizer (still from HF for now)
    print(f"[GENOME₀] Loading tokenizer from HF...")
    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)

    # Build metadata
    total_params = sum(p.numel() for p in model.parameters())
    metadata = {
        "source": "GENOME0_GGUF",
        "gguf_path": str(gguf_path),
        "gguf_size_bytes": gguf_path.stat().st_size,
        "gguf_sha256_1mb": gguf_hash,
        "tensors_loaded": len(state_dict),
        "tensors_mapped": len(mapping_info["mapped"]),
        "tensors_unmapped": len(mapping_info["unmapped"]),
        "missing_keys": len(missing),
        "unexpected_keys": len(unexpected),
        "total_params": total_params,
        "frozen": freeze,
        "device": device,
        "dtype": str(dtype),
    }

    print(f"[GENOME₀] Model ready: {total_params:,} params on {device}")
    print(f"[GENOME₀] Source: LOCAL GGUF (not HuggingFace cache)")

    return model, tokenizer, metadata


if __name__ == "__main__":
    # Quick test
    print("=" * 70)
    print("GENOME₀ LOADER TEST")
    print("=" * 70)

    model, tokenizer, metadata = load_tinyllama_genome0()

    print()
    print("Metadata:")
    for k, v in metadata.items():
        print(f"  {k}: {v}")

    # Quick inference test
    print()
    print("Quick inference test...")
    inputs = tokenizer("Hello, my name is", return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_new_tokens=20,
            do_sample=False,
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Input: 'Hello, my name is'")
    print(f"Output: '{result}'")
