#!/usr/bin/env python3
# train_lm_adamw_8bit.py
#
# Classic Hugging Face Trainer script for causal LM fine-tuning on:
#   ./train.txt and ./valid.txt
#
# Extra:
# - Replaces literal "\\n" with real newline "\n" inside each sample.
# - Ensures each text sample ends with literal "</s>".
# - Masks loss so model learns ONLY after the first "\n".
# - Prints “what the model sees” before training.
# - Uses bitsandbytes AdamW 8-bit optimizer:
#       optim="adamw_bnb_8bit"

import os
import json
import random
import importlib.util

import torch

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    set_seed,
)


# ----------------------------
# Paths / knobs
# ----------------------------

TOKENIZER_PATH = "./checkpoint-55000"
POLICY_CKPT = "./checkpoint-55000"

TRAIN_TXT = "./train.txt"
VALID_TXT = "./valid.txt"

OUT_DIR = "./runs/lm_finetune_adamw8bit"
LOG_DIR = os.path.join(OUT_DIR, "tb")

SEED = 42
MAX_LENGTH = 512

EOS_STR = "</s>"

# Preview controls
PREVIEW_TEXT_SAMPLES = 8
PREVIEW_TOKEN_TRUNC = 256
PREVIEW_BATCH_ITEMS = 2
WRITE_PREVIEW_JSONL = True

# Optimizer
USE_ADAMW_8BIT = True
OPTIM_NAME = "adamw_bnb_8bit"

# Training knobs
MAX_STEPS = 50000
PER_DEVICE_TRAIN_BATCH_SIZE = 16
PER_DEVICE_EVAL_BATCH_SIZE = 64
GRADIENT_ACCUMULATION_STEPS = 1

LEARNING_RATE = 5e-5
WEIGHT_DECAY = 0.0
WARMUP_RATIO = 0.01
LR_SCHEDULER_TYPE = "cosine"

DATALOADER_NUM_WORKERS = 16


def check_bitsandbytes_available():
    if not USE_ADAMW_8BIT:
        return

    if importlib.util.find_spec("bitsandbytes") is None:
        raise ImportError(
            "USE_ADAMW_8BIT=True but bitsandbytes is not installed.\n"
            "Install it with something like:\n"
            "    pip install bitsandbytes\n"
            "or, in conda/mamba environments:\n"
            "    pip install bitsandbytes\n"
        )

    print("bitsandbytes detected. Using optimizer:", OPTIM_NAME)


def preprocess_text(text: str) -> str:
    """
    1) Convert literal '\\n' into a real newline '\n' inside the sample.
    2) Append literal EOS_STR if missing.
    """
    t = text.rstrip("\n")
    t = t.replace("\\n", "\n")
    if not t.rstrip().endswith(EOS_STR):
        t = t + EOS_STR
    return t


def safe_makedirs(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def tokens_preview(tokenizer, ids, max_tokens=80):
    toks = tokenizer.convert_ids_to_tokens(ids)
    if len(toks) > max_tokens:
        return toks[:max_tokens] + ["…"]
    return toks


def dump_jsonl(path: str, rows):
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


class CausalLMMaskedCollator:
    """
    Pads input_ids with tokenizer.pad_token_id and pads labels with -100.

    Expects each feature to contain:
      - input_ids: List[int]
      - labels:    List[int]
    """

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def __call__(self, features):
        batch = self.tokenizer.pad(
            {"input_ids": [f["input_ids"] for f in features]},
            padding=True,
            return_tensors="pt",
        )

        max_len = batch["input_ids"].shape[1]

        padded_labels = []
        for f in features:
            lab = f["labels"]
            if len(lab) < max_len:
                lab = lab + [-100] * (max_len - len(lab))
            else:
                lab = lab[:max_len]
            padded_labels.append(lab)

        batch["labels"] = torch.tensor(padded_labels, dtype=torch.long)
        return batch


def main():
    check_bitsandbytes_available()

    set_seed(SEED)
    torch.backends.cuda.matmul.allow_tf32 = True

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print("CUDA available:", torch.cuda.is_available())
    if torch.cuda.is_available():
        print("GPU:", torch.cuda.get_device_name(0))
    print("Optimizer:", OPTIM_NAME if USE_ADAMW_8BIT else "default")

    # ----------------------------
    # Load tokenizer/model
    # ----------------------------
    tokenizer = AutoTokenizer.from_pretrained(
        TOKENIZER_PATH,
        use_fast=True,
        trust_remote_code=True,
        return_token_type_ids=False,
    )

    if tokenizer.pad_token is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "<|pad|>"})

    model = AutoModelForCausalLM.from_pretrained(
        POLICY_CKPT,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else None,
        device_map=None,
        trust_remote_code=True,
    )

    if len(tokenizer) != model.get_input_embeddings().num_embeddings:
        model.resize_token_embeddings(len(tokenizer))

    model.config.use_cache = False
    model.to(device)

    # ----------------------------
    # Load datasets
    # ----------------------------
    raw = load_dataset(
        "text",
        data_files={"train": TRAIN_TXT, "validation": VALID_TXT},
    )

    raw = raw.filter(lambda x: len(x["text"].strip()) > 0)

    def preprocess_map(batch):
        processed = [preprocess_text(t) for t in batch["text"]]
        return {"text": processed}

    raw = raw.map(
        preprocess_map,
        batched=True,
        desc="Preprocess: literal \\n -> newline + ensure </s>",
    )

    # ----------------------------
    # Tokenize + build labels mask
    # ----------------------------
    def tokenize_and_mask_batch(batch):
        texts = batch["text"]

        enc = tokenizer(
            texts,
            truncation=True,
            max_length=MAX_LENGTH,
            padding=False,
        )

        prompt_texts = []
        for t in texts:
            if "\n" in t:
                head, _tail = t.split("\n", 1)
                prompt_texts.append(head + "\n")
            else:
                prompt_texts.append(t)

        enc_prompt = tokenizer(
            prompt_texts,
            truncation=True,
            max_length=MAX_LENGTH,
            padding=False,
        )

        labels = []
        prompt_lens = []

        for ids, p_ids in zip(enc["input_ids"], enc_prompt["input_ids"]):
            p_len = min(len(ids), len(p_ids))
            prompt_lens.append(p_len)
            labels.append([-100] * p_len + ids[p_len:])

        enc["labels"] = labels
        enc["prompt_len"] = prompt_lens
        return enc

    tok = raw.map(
        tokenize_and_mask_batch,
        batched=True,
        remove_columns=[],
        desc="Tokenizing + masking loss before first newline",
    )

    drop_cols = []
    if "text" in tok["train"].column_names:
        drop_cols.append("text")
    if "prompt_len" in tok["train"].column_names:
        drop_cols.append("prompt_len")

    train_ds = tok["train"].remove_columns(drop_cols) if drop_cols else tok["train"]
    valid_ds = tok["validation"].remove_columns(drop_cols) if drop_cols else tok["validation"]

    # ----------------------------
    # Preview: sample-level
    # ----------------------------
    safe_makedirs(OUT_DIR)

    train_len = len(tok["train"])
    k = min(PREVIEW_TEXT_SAMPLES, train_len)

    random.seed(SEED)
    preview_indices = random.sample(range(train_len), k=k) if train_len > 0 else []

    preview_rows = []

    print("\n=== PREVIEW: samples after preprocess + tokenization + label masking ===")

    for idx in preview_indices:
        ex = tok["train"][idx]

        text = ex["text"]
        input_ids = ex["input_ids"]
        labels = ex["labels"]
        prompt_len = ex["prompt_len"]

        ids_prompt = input_ids[:prompt_len]
        ids_target = input_ids[prompt_len:]
        label_ids = [x for x in labels if x != -100]

        row = {
            "split": "train",
            "index": idx,
            "text": text,
            "prompt_len_tokens": int(prompt_len),
            "len_ids": len(input_ids),
            "input_ids_head": input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)],
            "tokens_head": tokens_preview(
                tokenizer,
                input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)],
                PREVIEW_TOKEN_TRUNC,
            ),
            "decoded_full": tokenizer.decode(input_ids, skip_special_tokens=False),
            "decoded_prompt": tokenizer.decode(ids_prompt, skip_special_tokens=False),
            "decoded_target": tokenizer.decode(ids_target, skip_special_tokens=False),
            "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False),
        }

        preview_rows.append(row)

        print(f"\n--- sample idx={idx} ---")
        print("TEXT postprocessed:")
        print(text)
        print(f"len(input_ids) = {len(input_ids)} | prompt_len_tokens = {prompt_len}")
        print("TOKENS head:")
        print(row["tokens_head"])
        print("DECODED PROMPT masked from loss:")
        print(row["decoded_prompt"])
        print("DECODED TARGET learned / loss applies here:")
        print(row["decoded_target"])

    if WRITE_PREVIEW_JSONL:
        dump_path = os.path.join(OUT_DIR, "preview_samples.jsonl")
        dump_jsonl(dump_path, preview_rows)
        print(f"\nWrote sample previews to: {dump_path}")

    # ----------------------------
    # Collator + Trainer
    # ----------------------------
    data_collator = CausalLMMaskedCollator(tokenizer=tokenizer)

    args = TrainingArguments(
        output_dir=OUT_DIR,
        overwrite_output_dir=True,

        max_steps=MAX_STEPS,

        per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH_SIZE,
        per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,

        learning_rate=LEARNING_RATE,
        weight_decay=WEIGHT_DECAY,
        warmup_ratio=WARMUP_RATIO,
        lr_scheduler_type=LR_SCHEDULER_TYPE,

        logging_dir=LOG_DIR,
        logging_steps=1,
        report_to=["tensorboard"],

        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=2,

        bf16=torch.cuda.is_available(),
        fp16=False,

        dataloader_num_workers=DATALOADER_NUM_WORKERS,
        remove_unused_columns=False,

        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,

        # Main modification:
        # bitsandbytes AdamW 8-bit optimizer.
        optim=OPTIM_NAME if USE_ADAMW_8BIT else "adamw_torch",
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_ds,
        eval_dataset=valid_ds,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # ----------------------------
    # Preview: collated batch
    # ----------------------------
    print("\n=== PREVIEW: a collated training batch after padding + masked labels ===")

    n_take = min(2, len(train_ds))
    features = [train_ds[i] for i in range(n_take)]
    first_batch = data_collator(features)

    input_ids_b = first_batch["input_ids"].detach().cpu()
    labels_b = first_batch["labels"].detach().cpu()
    attn_b = first_batch["attention_mask"].detach().cpu()

    print("Batch keys:", list(first_batch.keys()))
    print("input_ids shape:", tuple(input_ids_b.shape))
    print("labels shape:", tuple(labels_b.shape))
    print("attention_mask shape:", tuple(attn_b.shape))

    batch_preview_rows = []
    n_show = min(PREVIEW_BATCH_ITEMS, input_ids_b.shape[0])

    for i in range(n_show):
        ids = input_ids_b[i].tolist()
        labels = labels_b[i].tolist()
        label_ids = [x for x in labels if x != -100]

        row = {
            "batch_item": i,
            "input_len": int(attn_b[i].sum().item()),
            "input_ids_head": ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)],
            "tokens_head": tokens_preview(
                tokenizer,
                ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)],
                PREVIEW_TOKEN_TRUNC,
            ),
            "decoded_input": tokenizer.decode(ids, skip_special_tokens=False),
            "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False),
        }

        batch_preview_rows.append(row)

        print(f"\n--- batch item {i} ---")
        print("TOKENS head:")
        print(row["tokens_head"])
        print("DECODED INPUT, specials kept:")
        print(row["decoded_input"])
        print("DECODED LABELS, only learned tokens:")
        print(row["decoded_labels"])

    if WRITE_PREVIEW_JSONL:
        dump_path = os.path.join(OUT_DIR, "preview_first_batch.jsonl")
        dump_jsonl(dump_path, batch_preview_rows)
        print(f"\nWrote batch preview to: {dump_path}")

    # ----------------------------
    # Train / save
    # ----------------------------
    train_result = trainer.train()

    final_dir = os.path.join(OUT_DIR, "final")
    trainer.save_model(final_dir)
    tokenizer.save_pretrained(final_dir)

    metrics = train_result.metrics
    metrics["train_samples"] = len(train_ds)

    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    trainer.save_state()

    eval_metrics = trainer.evaluate()
    eval_metrics["eval_samples"] = len(valid_ds)

    trainer.log_metrics("eval", eval_metrics)
    trainer.save_metrics("eval", eval_metrics)

    print("\nDone.")
    print(f"Final model saved to: {final_dir}")
    print(f"TensorBoard: tensorboard --logdir {OUT_DIR}")


if __name__ == "__main__":
    main()