"""
Export the trained model to HuggingFace-compatible format.

Creates:
  - model.safetensors (weights)
  - config.json (architecture config)
  - generation_config.json
  - tokenizer.json, tokenizer_config.json, special_tokens_map.json
"""

import os
import sys
import json
import torch
from collections import OrderedDict
from safetensors.torch import save_file

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model.config import ModelConfig
from model.transformer import Transformer
from model.data import get_tokenizer

CHECKPOINT = "/jfs/deepak-kumar/checkpoints_dpo/dpo_final.pt"
OUTPUT_DIR = "/home/jovyan/training/hf_model"

os.makedirs(OUTPUT_DIR, exist_ok=True)

print("=" * 60)
print("  EXPORTING MODEL TO HUGGING FACE FORMAT")
print("=" * 60)

# --- 1. Load model ---
print("\n[1/4] Loading checkpoint...")
tokenizer = get_tokenizer()
special_tokens = ["<|user|>", "<|assistant|>", "<|end|>"]
vocab = tokenizer.get_vocab()
new_tokens = [t for t in special_tokens if t not in vocab]
if new_tokens:
    tokenizer.add_tokens(new_tokens, special_tokens=True)

model_config = ModelConfig()
model_config.vocab_size = len(tokenizer)

model = Transformer(model_config)
ckpt = torch.load(CHECKPOINT, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model"])
step = ckpt.get("step", 0)
del ckpt
print(f"  Loaded DPO model (step {step}, vocab {model_config.vocab_size})")

# --- 2. Convert state dict keys to HF-style naming ---
print("\n[2/4] Converting weights to safetensors...")

state_dict = model.state_dict()
hf_state = OrderedDict()

KEY_MAP = {
    "tok_embeddings.weight": "model.embed_tokens.weight",
    "norm.weight": "model.norm.weight",
    "output.weight": "lm_head.weight",
}

for key, tensor in state_dict.items():
    if key in KEY_MAP:
        hf_state[KEY_MAP[key]] = tensor
        continue

    if key.startswith("layers."):
        parts = key.split(".")
        layer_idx = parts[1]
        rest = ".".join(parts[2:])

        layer_map = {
            "attention_norm.weight": f"model.layers.{layer_idx}.input_layernorm.weight",
            "ffn_norm.weight": f"model.layers.{layer_idx}.post_attention_layernorm.weight",
            "attention.wq.weight": f"model.layers.{layer_idx}.self_attn.q_proj.weight",
            "attention.wk.weight": f"model.layers.{layer_idx}.self_attn.k_proj.weight",
            "attention.wv.weight": f"model.layers.{layer_idx}.self_attn.v_proj.weight",
            "attention.wo.weight": f"model.layers.{layer_idx}.self_attn.o_proj.weight",
            "ffn.w_gate.weight": f"model.layers.{layer_idx}.mlp.gate_proj.weight",
            "ffn.w_up.weight": f"model.layers.{layer_idx}.mlp.up_proj.weight",
            "ffn.w_down.weight": f"model.layers.{layer_idx}.mlp.down_proj.weight",
        }

        if rest in layer_map:
            hf_state[layer_map[rest]] = tensor
        else:
            print(f"  WARNING: unmapped key {key}")
            hf_state[key] = tensor
    elif key == "freqs_cis":
        continue
    else:
        print(f"  WARNING: unmapped key {key}")
        hf_state[key] = tensor

# Convert all to bfloat16 for storage
for k in hf_state:
    if hf_state[k].dtype == torch.float32:
        hf_state[k] = hf_state[k].to(torch.bfloat16)

safetensors_path = os.path.join(OUTPUT_DIR, "model.safetensors")
save_file(hf_state, safetensors_path)
size_gb = os.path.getsize(safetensors_path) / 1e9
print(f"  Saved {len(hf_state)} tensors -> {safetensors_path} ({size_gb:.2f} GB)")

# --- 3. Write config files ---
print("\n[3/4] Writing config files...")

config_json = {
    "architectures": ["LlamaForCausalLM"],
    "model_type": "llama",
    "vocab_size": model_config.vocab_size,
    "hidden_size": model_config.hidden_dim,
    "intermediate_size": model_config.intermediate_dim,
    "num_hidden_layers": model_config.num_layers,
    "num_attention_heads": model_config.num_attention_heads,
    "num_key_value_heads": model_config.num_kv_heads,
    "max_position_embeddings": model_config.max_seq_len,
    "rope_theta": model_config.rope_theta,
    "rms_norm_eps": model_config.rms_norm_eps,
    "hidden_act": "silu",
    "initializer_range": 0.02,
    "tie_word_embeddings": False,
    "torch_dtype": "bfloat16",
    "transformers_version": "4.40.0",
    "use_cache": True,
    "bos_token_id": tokenizer.bos_token_id,
    "eos_token_id": tokenizer.eos_token_id,
    "pad_token_id": tokenizer.pad_token_id,
}

with open(os.path.join(OUTPUT_DIR, "config.json"), "w") as f:
    json.dump(config_json, f, indent=2)
print("  config.json")

gen_config = {
    "bos_token_id": tokenizer.bos_token_id,
    "eos_token_id": tokenizer.eos_token_id,
    "pad_token_id": tokenizer.pad_token_id,
    "do_sample": True,
    "temperature": 0.7,
    "top_k": 50,
    "top_p": 0.9,
    "repetition_penalty": 1.15,
    "max_new_tokens": 512,
    "transformers_version": "4.40.0",
}

with open(os.path.join(OUTPUT_DIR, "generation_config.json"), "w") as f:
    json.dump(gen_config, f, indent=2)
print("  generation_config.json")

# --- 4. Export tokenizer ---
print("\n[4/4] Exporting tokenizer...")
tokenizer.save_pretrained(OUTPUT_DIR)
print("  Tokenizer files saved")

print("\n" + "=" * 60)
print("  EXPORT COMPLETE -> " + OUTPUT_DIR)
print("=" * 60)
print("\nFiles:")
for f in sorted(os.listdir(OUTPUT_DIR)):
    size = os.path.getsize(os.path.join(OUTPUT_DIR, f))
    if size > 1e6:
        print(f"  {f:40s} {size/1e6:.1f} MB")
    else:
        print(f"  {f:40s} {size/1e3:.1f} KB")