# ---------- fine_tune_v7_auto_segment.py ----------
# (주식회사엠케이디 | MKD)
# MKDCore Fine-Tuning — Auto-Segments from PKL, Safe Stop/Resume, GPU Monitor, Checkpoints

import os
import sys
import time
import gc
import pickle
import signal
import random
import datetime
import subprocess
import threading

import torch
from torch import nn, optim
from tqdm import tqdm
import psutil

# Your model/tokenizer wrapper
from mkd_core import MKDCore, MyTokenizer, SEQ_LEN


# ============================================================
# CONFIGURATION
# ============================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASE_DIR = "/home/mkd-s003/AI/finetune"

# Required files
TOKENIZER_JSON = os.path.join(BASE_DIR, "prod_tokenizer_v2.json")
TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl")
VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl")
MODEL_BASE = os.path.join(BASE_DIR, "mkdcore_v2.pt")

# State and outputs
STATE_PATH = os.path.join(BASE_DIR, "mkdcore_state_finetune.pt")
CKPT_DIR = os.path.join(BASE_DIR, "checkpoints_ft")
LOG_REASON = os.path.join(BASE_DIR, "reason_of_interruption.log")
PROGRESS_FILE = os.path.join(BASE_DIR, "segment_progress.pkl")   # remembers how many tokens consumed so far

os.makedirs(CKPT_DIR, exist_ok=True)
torch.cuda.empty_cache()

# Segment sizes
MAX_TRAIN_TOKENS_SEG = 20_000_000   # per run (segment)
MAX_VAL_TOKENS_SEG   = 5_000_000    # per run (segment)

# Training loop
EPOCHS = 16
STEPS_PER_EPOCH = 20_000
LR = 2e-5
LR_DECAY = 0.9
GRAD_CLIP = 1.0
SAVE_EVERY = 5_000
PRINT_EVERY = 1_000
RESTART_DELAY = 60

# Data streaming diagnostics
CHUNK_PRINT_EVERY = 1000  # print every N pickled sequences while loading


# ============================================================
# MONITORING
# ============================================================
def system_status():
    cpu = psutil.cpu_percent(interval=0.1)
    mem = psutil.virtual_memory().percent
    if not torch.cuda.is_available():
        return f"CPU {cpu:.0f}% | RAM {mem:.1f}%"
    mem_alloc = torch.cuda.memory_allocated(0) / 1e9
    total_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
    try:
        util = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"]
        ).decode().strip().split("\n")[0]
    except Exception:
        util = "?"
    return f"GPU {mem_alloc:.2f}/{total_mem:.0f} GB | Util {util}% | CPU {cpu:.0f}% | RAM {mem:.1f}%"


def log_reason(reason: str):
    with open(LOG_REASON, "a", encoding="utf-8") as f:
        f.write(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {reason}\n")


# ============================================================
# SEGMENT PROGRESS
# ============================================================
def load_segment_progress():
    """
    Returns dict with consumed token offsets:
      {"train_pos": int, "val_pos": int}
    If missing, returns zeros.
    """
    if not os.path.exists(PROGRESS_FILE):
        return {"train_pos": 0, "val_pos": 0}
    try:
        with open(PROGRESS_FILE, "rb") as f:
            data = pickle.load(f)
        return {
            "train_pos": int(data.get("train_pos", 0)),
            "val_pos": int(data.get("val_pos", 0)),
        }
    except Exception:
        return {"train_pos": 0, "val_pos": 0}


def save_segment_progress(train_pos: int, val_pos: int):
    with open(PROGRESS_FILE, "wb") as f:
        pickle.dump({"train_pos": int(train_pos), "val_pos": int(val_pos)}, f)


# ============================================================
# STREAM-LOAD A TOKEN SEGMENT FROM A PKL
# ============================================================
def load_tokens_segment(pkl_path: str, start_offset_tokens: int, max_tokens: int, label: str):
    """
    Reads pickled sequences from pkl_path and returns exactly up to max_tokens
    starting at token index start_offset_tokens (token-level offset).
    Returns: tensor[long], new_offset_tokens (start_offset_tokens + loaded_count)
    """
    tokens = []
    loaded_tokens = 0
    consumed = 0            # how many tokens we have skipped to reach start_offset_tokens
    chunks = 0

    print(f"\nLoading {os.path.basename(pkl_path)} segment | start={start_offset_tokens:,} tokens | limit={max_tokens:,} tokens")
    with open(pkl_path, "rb") as f:
        while loaded_tokens < max_tokens:
            try:
                seq = pickle.load(f)
            except EOFError:
                break

            if not isinstance(seq, list) or len(seq) == 0:
                continue

            seq_len = len(seq)

            # Skip whole sequence if we haven't reached the start offset yet
            if consumed + seq_len <= start_offset_tokens:
                consumed += seq_len
                chunks += 1
                if chunks % CHUNK_PRINT_EVERY == 0:
                    print(f"  • {chunks:,} chunks scanned | skipped={consumed:,} | loaded={loaded_tokens:,} | {label}")
                continue

            # Determine starting index within this sequence
            in_seq_start = max(0, start_offset_tokens - consumed)
            take = seq_len - in_seq_start
            take = min(take, max_tokens - loaded_tokens)

            if take > 0:
                tokens.extend(seq[in_seq_start:in_seq_start + take])
                loaded_tokens += take

            consumed += seq_len
            chunks += 1

            if chunks % CHUNK_PRINT_EVERY == 0:
                print(f"  • {chunks:,} chunks | loaded={loaded_tokens:,} | RAM {psutil.virtual_memory().percent:.1f}% | {label}")

            if loaded_tokens >= max_tokens:
                break

    new_offset = start_offset_tokens + loaded_tokens
    print(f"Loaded {loaded_tokens:,} tokens from {os.path.basename(pkl_path)} | new_offset={new_offset:,} | {label}")
    if loaded_tokens == 0:
        print(f"Warning: no tokens loaded from {pkl_path} at start offset {start_offset_tokens:,}. Possibly at file end.")
    return torch.tensor(tokens, dtype=torch.long), new_offset


# ============================================================
# BATCHING
# ============================================================
def get_batch(source: torch.Tensor, seq_len: int):
    if len(source) <= seq_len + 1:
        raise ValueError("Dataset segment too small for the configured sequence length.")
    i = random.randint(0, len(source) - seq_len - 1)
    x = source[i:i + seq_len].unsqueeze(0).to(DEVICE)
    y = source[i + 1:i + seq_len + 1].unsqueeze(0).to(DEVICE)
    return x, y


@torch.no_grad()
def evaluate_loss(model, loss_fn, vocab_size, val_data):
    model.eval()
    total = 0.0
    iters = 25
    for _ in range(iters):
        x, y = get_batch(val_data, SEQ_LEN)
        with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
            logits = model(x)
            loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))
        total += loss.item()
    return total / iters


# ============================================================
# CTRL+C SAFE HANDLER
# ============================================================
def save_state_on_interrupt(model, opt, scaler):
    try:
        torch.save({
            "model": model.state_dict(),
            "optimizer": opt.state_dict(),
            "scaler": scaler.state_dict(),
            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }, STATE_PATH)
        print(f"\nProgress saved → {STATE_PATH}")
    except Exception as e:
        print(f"Save failed: {e}")


def ctrl_c_handler_factory(model, opt, scaler):
    def handle_ctrl_c(sig, frame):
        print("\n[CTRL+C] Training interrupted. Saving state...")
        save_state_on_interrupt(model, opt, scaler)
        print("[1] Save & Stop")
        print("[2] Save & Restart in 60s")
        print(f"(No input → auto restart in {RESTART_DELAY}s)")

        choice = {"val": None}
        answered = threading.Event()

        def _read_choice():
            try:
                choice["val"] = input("Choose (1/2): ").strip()
            except Exception:
                choice["val"] = None
            finally:
                answered.set()

        t = threading.Thread(target=_read_choice, daemon=True)
        t.start()
        for sec in range(RESTART_DELAY, 0, -1):
            if answered.is_set():
                break
            print(f"\rRestarting in {sec:02d}s | {system_status()}", end="", flush=True)
            time.sleep(1)
        print()

        if not answered.is_set() or choice["val"] == "2":
            log_reason("Manual or auto restart")
            print("Restarting now...")
            os.execv(sys.executable, [sys.executable] + sys.argv)
        else:
            log_reason("Stopped safely by user")
            print("Training stopped safely.")
            sys.exit(0)

    return handle_ctrl_c


# ============================================================
# MAIN
# ============================================================
def main():
    print(f"Device: {DEVICE}")

    # Load tokenizer and base model (or previous fine-tuned if you prefer)
    tok = MyTokenizer(TOKENIZER_JSON)
    vocab_size = tok.vocab_size
    print(f"Vocab size: {vocab_size}")

    model = MKDCore(vocab_size).to(DEVICE)
    state = torch.load(MODEL_BASE, map_location=DEVICE)
    model.load_state_dict(state, strict=False)
    print(f"Starting from: {MODEL_BASE}")
    print("Model loaded successfully.")

    opt = optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95))
    loss_fn = nn.CrossEntropyLoss()
    scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())

    # Attach Ctrl+C handler
    signal.signal(signal.SIGINT, ctrl_c_handler_factory(model, opt, scaler))

    # Load segment progress
    prog = load_segment_progress()
    start_train = prog["train_pos"]
    start_val = prog["val_pos"]

    # Stream current segment
    train_data, new_train_pos = load_tokens_segment(TRAIN_PKL, start_train, MAX_TRAIN_TOKENS_SEG, "train")
    val_data, new_val_pos = load_tokens_segment(VAL_PKL, start_val, MAX_VAL_TOKENS_SEG, "val")

    # If nothing could be loaded, exit cleanly
    if len(train_data) == 0 or len(val_data) == 0:
        print("No tokens available for this segment. Consider increasing PKL caches or resetting segment progress.")
        print(f"Current offsets → train_pos={start_train:,}, val_pos={start_val:,}")
        return

    print(f"Dataset segment ready: {len(train_data):,} train | {len(val_data):,} val tokens")
    # Save new progress immediately so next run advances even if interrupted later
    save_segment_progress(new_train_pos, new_val_pos)
    print(f"Advanced segment offsets saved → train_pos={new_train_pos:,} | val_pos={new_val_pos:,}")

    # Training
    model_save = os.path.join(
        BASE_DIR,
        f"mkdcore_v2_segment_finetuned_{datetime.datetime.now().strftime('%Y%m%d_%H%M')}.pt"
    )
    print(f"\nFine-tuning started → {model_save}\n")

    train_losses, val_losses = [], []

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0.0
        start_time = time.time()

        with tqdm(total=STEPS_PER_EPOCH, desc=f"Epoch {epoch + 1}/{EPOCHS}", ncols=120) as pbar:
            for step in range(STEPS_PER_EPOCH):
                x, y = get_batch(train_data, SEQ_LEN)

                with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
                    logits = model(x)
                    loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))

                scaler.scale(loss).backward()
                scaler.unscale_(opt)
                torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
                scaler.step(opt)
                scaler.update()
                opt.zero_grad(set_to_none=True)

                total_loss += loss.item()
                avg_loss = total_loss / (step + 1)
                pbar.set_postfix({"Loss": f"{avg_loss:.4f}", "GPU": system_status()})
                pbar.update(1)

                if (step + 1) % PRINT_EVERY == 0:
                    elapsed = (time.time() - start_time) / 60
                    print(f"Step {step + 1} | Loss {avg_loss:.4f} | {system_status()} | Elapsed {elapsed:.1f}m")

                if (step + 1) % SAVE_EVERY == 0:
                    ckpt = os.path.join(CKPT_DIR, f"mkdcore_seg_ckpt_e{epoch + 1}_s{step + 1}.pt")
                    torch.save(model.state_dict(), ckpt)
                    print(f"Checkpoint saved → {ckpt}")

                if (step + 1) % 2000 == 0:
                    gc.collect()
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

        val_loss = evaluate_loss(model, loss_fn, vocab_size, val_data)
        train_losses.append(total_loss / STEPS_PER_EPOCH)
        val_losses.append(val_loss)
        torch.save(model.state_dict(), model_save)
        print(f"Epoch {epoch + 1} | Train {train_losses[-1]:.4f} | Val {val_loss:.4f} | Saved → {model_save}")

        for g in opt.param_groups:
            g["lr"] *= LR_DECAY

    print("\nFine-tuning completed successfully.")
    print(f"Final model → {model_save}")
    print(f"Next run will start at offsets → train_pos={new_train_pos:,} | val_pos={new_val_pos:,}")
    print("You can re-run this same script to consume the next segment automatically.")


if __name__ == "__main__":
    main()