# ---------- verify.py ----------
# (주식회사엠케이디 | MKD)
# Compare base, previous fine-tuned, and latest fine-tuned model generations
# Automatically logs comparison results into /logs/

import os
import torch
from datetime import datetime
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")
LOG_DIR = os.path.join(BASE_DIR, "logs")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

os.makedirs(LOG_DIR, exist_ok=True)

# Create timestamped log file
timestamp = datetime.now().strftime("%Y%m%d_%H%M")
LOG_FILE = os.path.join(LOG_DIR, f"compare_results_{timestamp}.txt")

def write_log(line=""):
    print(line)
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(line + "\n")

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

# ============================================================
# MODEL FINDING
# ============================================================
def find_finetuned_models():
    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, None
    files.sort(key=lambda n: os.path.getmtime(os.path.join(BASE_DIR, n)))
    latest = os.path.join(BASE_DIR, files[-1])
    previous = os.path.join(BASE_DIR, files[-2]) if len(files) > 1 else None
    return previous, latest

PREV_MODEL, LATEST_MODEL = find_finetuned_models()

# ============================================================
# TEXT GENERATION
# ============================================================
@torch.no_grad()
def generate_text(model, prompt, max_new_tokens=60, temperature=0.8):
    model.eval()
    try:
        ids = tok.encode(prompt)
        if not isinstance(ids, list):
            ids = ids.ids
    except Exception as e:
        return f"Encoding error: {e}"
    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)
    try:
        return tok.decode(input_ids[0].tolist())
    except Exception:
        return "Decode failed."

# ============================================================
# MODEL LOADING
# ============================================================
def load_model(path):
    write_log(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)
prev_model = load_model(PREV_MODEL) if PREV_MODEL else None
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
# ============================================================
write_log("=" * 90)
write_log("MKD MODEL COMPARISON REPORT")
write_log(f"Timestamp: {timestamp}")
write_log("=" * 90)
write_log(f"Base model: {BASE_MODEL}")
if PREV_MODEL:
    write_log(f"Previous fine-tuned model: {PREV_MODEL}")
else:
    write_log("Previous fine-tuned model: None found")
if LATEST_MODEL:
    write_log(f"Latest fine-tuned model: {LATEST_MODEL}")
else:
    write_log("Latest fine-tuned model: None found")
write_log("=" * 90 + "\n")

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

    base_output = generate_text(base_model, prompt)
    write_log(f"--- BASE MODEL OUTPUT ({os.path.basename(BASE_MODEL)}) ---\n{base_output}\n")

    if prev_model:
        prev_output = generate_text(prev_model, prompt)
        write_log(f"--- PREVIOUS FINE-TUNED MODEL OUTPUT ({os.path.basename(PREV_MODEL)}) ---\n{prev_output}\n")

    if latest_model:
        latest_output = generate_text(latest_model, prompt)
        write_log(f"--- LATEST FINE-TUNED MODEL OUTPUT ({os.path.basename(LATEST_MODEL)}) ---\n{latest_output}\n")

    write_log("=" * 90 + "\n")

write_log("Comparison complete. Observe how outputs evolve segment by segment.")
write_log(f"Results saved to: {LOG_FILE}")