#!/usr/bin/env python3
"""CPU LoRA SFT for identity calibration on AksaraLLM/Kiel-Pro-0.5B-v3.

- Loads base in bf16.
- LoRA r=8, alpha=16, on q_proj, k_proj, v_proj, o_proj.
- 50 identity prompts, 3 epochs, batch 1 with grad-accum 4.
- Uses Qwen2 ChatML template.
"""
import os, sys, json, math, time, shutil
from pathlib import Path
import torch
from torch.utils.data import Dataset
from transformers import (
    AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments,
    DataCollatorForLanguageModeling, set_seed,
)
from peft import LoraConfig, get_peft_model, TaskType

sys.path.insert(0, str(Path(__file__).parent))
from identity_data import get_dataset

BASE = 'AksaraLLM/Kiel-Pro-0.5B-v3'
OUT = Path('/home/ubuntu/aksara_audit/sft/kiel-pro-0.5b-v3-identity')
OUT.mkdir(parents=True, exist_ok=True)

SEED = 7
set_seed(SEED)

# Qwen2 ChatML template pieces
SYS = "Kamu adalah Kiel-Pro, model bahasa Indonesia dari proyek AksaraLLM."


def format_chatml(system: str, user: str, assistant: str) -> str:
    return (
        f"<|im_start|>system\n{system}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        f"<|im_start|>assistant\n{assistant}<|im_end|>\n"
    )


class IdentityDS(Dataset):
    def __init__(self, tok, max_len=256):
        data = get_dataset()
        self.items = []
        for q, a in data:
            text = format_chatml(SYS, q, a)
            enc = tok(text, truncation=True, max_length=max_len, padding='max_length')
            ids = enc['input_ids']
            att = enc['attention_mask']
            # Mask loss on prompt portion — we only want to teach the assistant response.
            prompt = (
                f"<|im_start|>system\n{SYS}<|im_end|>\n"
                f"<|im_start|>user\n{q}<|im_end|>\n"
                f"<|im_start|>assistant\n"
            )
            prompt_len = len(tok(prompt, truncation=True, max_length=max_len)['input_ids'])
            labels = list(ids)
            for i in range(min(prompt_len, len(labels))):
                labels[i] = -100
            # also mask pad
            for i, m in enumerate(att):
                if m == 0:
                    labels[i] = -100
            self.items.append({
                'input_ids': torch.tensor(ids, dtype=torch.long),
                'attention_mask': torch.tensor(att, dtype=torch.long),
                'labels': torch.tensor(labels, dtype=torch.long),
            })

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        return self.items[idx]


def main():
    print('loading tokenizer...', flush=True)
    tok = AutoTokenizer.from_pretrained(BASE, token=os.environ['HF_TOKEN'])
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    print('loading model (fp32, CPU; AVX2 only here)...', flush=True)
    t0 = time.time()
    model = AutoModelForCausalLM.from_pretrained(
        BASE,
        dtype=torch.float32,
        low_cpu_mem_usage=True,
        token=os.environ['HF_TOKEN'],
    )
    model.config.use_cache = False
    print(f'  loaded in {time.time()-t0:.1f}s, {sum(p.numel() for p in model.parameters())/1e6:.1f}M params', flush=True)

    lora_cfg = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
        target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj'],
        bias='none',
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()

    # Small max_len because identity answers are short.
    ds = IdentityDS(tok, max_len=128)
    print(f'dataset: {len(ds)} examples', flush=True)

    args = TrainingArguments(
        output_dir=str(OUT / 'runs'),
        per_device_train_batch_size=1,
        gradient_accumulation_steps=2,
        num_train_epochs=2,
        learning_rate=2e-4,
        warmup_ratio=0.1,
        lr_scheduler_type='cosine',
        logging_steps=1,
        save_steps=500,
        save_total_limit=1,
        report_to='none',
        bf16=False,   # CPU bf16 autocast via model dtype
        fp16=False,
        remove_unused_columns=False,
        seed=SEED,
        dataloader_num_workers=0,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ds,
        processing_class=tok,
    )
    print('starting LoRA training on CPU...', flush=True)
    t0 = time.time()
    trainer.train()
    print(f'done in {(time.time()-t0)/60:.1f} min', flush=True)

    # Save LoRA
    lora_path = OUT / 'lora'
    model.save_pretrained(str(lora_path))
    tok.save_pretrained(str(lora_path))
    print(f'saved LoRA to {lora_path}', flush=True)

    # Merge + save full model
    print('merging LoRA into base...', flush=True)
    merged = model.merge_and_unload()
    merged_path = OUT / 'merged'
    merged.save_pretrained(str(merged_path), safe_serialization=True)
    tok.save_pretrained(str(merged_path))
    print(f'saved merged model to {merged_path}', flush=True)


if __name__ == '__main__':
    main()