# ---------- veryfy.py ----------
# (주식회사엠케이디 | MKD)
# Compare base vs fine-tuned model generations — Fully compatible with your mkd_core.py

import os
import torch
from mkd_core import MKDCore, MyTokenizer

# ============================================================
# CONFIGURATION
# ============================================================
BASE_DIR = "/home/mkd-s003/AI/finetune"
TOKENIZER_JSON = os.path.join(BASE_DIR, "prod_tokenizer_v2.json")
BASE_MODEL = os.path.join(BASE_DIR, "mkdcore_v2.pt")

# Automatically find the latest fine-tuned model
def find_latest_finetuned():
    files = [
        f for f in os.listdir(BASE_DIR)
        if f.endswith(".pt") and ("finetuned" in f or "segment_finetuned" in f)
    ]
    if not files:
        return None
    files.sort(key=lambda n: os.path.getmtime(os.path.join(BASE_DIR, n)))
    return os.path.join(BASE_DIR, files[-1])

LATEST_MODEL = find_latest_finetuned()
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ============================================================
# TOKENIZER
# ============================================================
tok = MyTokenizer(TOKENIZER_JSON)
vocab_size = tok.vocab_size

# ============================================================
# TEXT GENERATION
# ============================================================
@torch.no_grad()
def generate_text(model, prompt, max_new_tokens=60, temperature=0.8):
    model.eval()

    # ✅ Safe tokenizer handling (works with your MyTokenizer)
    try:
        ids = tok.encode(prompt)
        if not isinstance(ids, list):
            ids = ids.ids
    except Exception as e:
        print(f"⚠️ Encoding error: {e}")
        return ""

    input_ids = torch.tensor(ids, dtype=torch.long, device=DEVICE).unsqueeze(0)

    for _ in range(max_new_tokens):
        with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
            logits = model(input_ids)
            probs = torch.softmax(logits[0, -1] / temperature, dim=-1)
            next_token = torch.multinomial(probs, 1)
        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)

    # ✅ Safe decode
    try:
        return tok.decode(input_ids[0].tolist())
    except Exception:
        return "⚠️ Decode failed"

# ============================================================
# MODEL LOADING
# ============================================================
def load_model(path):
    print(f"🔹 Loading model → {path}")
    model = MKDCore(vocab_size).to(DEVICE)
    state = torch.load(path, map_location=DEVICE)
    model.load_state_dict(state, strict=False)
    return model

base_model = load_model(BASE_MODEL)
latest_model = load_model(LATEST_MODEL) if LATEST_MODEL else None

# ============================================================
# TEST PROMPTS
# ============================================================
prompts = [
    "Artificial Intelligence will",
    "The future of digital technology",
    "AI in healthcare and robotics",
    "How climate change impacts innovation",
]

# ============================================================
# COMPARISON
# ============================================================
print("\n================ COMPARISON START ================\n")

for prompt in prompts:
    print(f"🧩 PROMPT: {prompt}\n")

    base_output = generate_text(base_model, prompt)
    print(f"--- BASE MODEL OUTPUT ---\n{base_output}\n")

    if latest_model:
        ft_output = generate_text(latest_model, prompt)
        print(f"--- FINE-TUNED MODEL OUTPUT ---\n{ft_output}\n")

    print("=" * 80 + "\n")

print("✅ Comparison complete. Observe how the fine-tuned model evolves with each segment.")