# ============================================================
# AYURVEDA GEMMA-3 4B TRAINING — FINAL v2 (ChatGPT-reviewed fixes)
# For: A100 MIG 20GB | Linux | Unsloth
# Restart kernel before running Cell 1
# ============================================================

# ============================================================
# CELL 1: Environment & Config
# ============================================================
import os
import torch

# --- CRITICAL: Set cache BEFORE any unsloth/transformers import ---
CACHE = "/nlsasfs/home/aikosh/prod-aikosh35/.cache/unsloth"
os.makedirs(CACHE, exist_ok=True)
os.makedirs(f"{CACHE}/torch", exist_ok=True)
os.makedirs(f"{CACHE}/triton", exist_ok=True)

os.environ["UNSLOTH_COMPILED_CACHE"] = CACHE
os.environ["UNSLOTH_CACHE_DIR"] = CACHE
os.environ["TORCHINDUCTOR_CACHE_DIR"] = f"{CACHE}/torch"
os.environ["TRITON_CACHE_DIR"] = f"{CACHE}/triton"
os.environ["XDG_CACHE_HOME"] = CACHE
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# --- CONFIG ---
MODEL_NAME = "unsloth/gemma-3-4b-it"
DATASET_PATH = "/nlsasfs/home/aikosh/prod-aikosh35/gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl"
OUTPUT_DIR = "/nlsasfs/home/aikosh/prod-aikosh35/gemma3_4b_ayurveda_output"

MAX_SEQ_LENGTH = 1024
BATCH_SIZE = 1
GRAD_ACCUM = 32
LORA_R = 16
LORA_ALPHA = 16  # Changed from 32 to 16 (better for citation-heavy QA)
LR = 2e-4
EPOCHS = 1

assert torch.cuda.is_available(), "No GPU!"
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
print(f"Cache: {CACHE}")
print(f"Model: {MODEL_NAME}")
print("Config ready. Run Cell 2.")


# ============================================================
# CELL 2: Load Dataset
# ============================================================
from datasets import load_dataset
import gc

dataset = load_dataset("json", data_files=DATASET_PATH, split="train")
print(f"Loaded {len(dataset):,} conversations")

sample = dataset[0]
msgs = sample.get("messages", [])
print(f"Roles: {[m['role'] for m in msgs]}")
print(f"User preview: {msgs[1]['content'][:80]}...")

gc.collect()
print("Dataset ready. Run Cell 3.")


# ============================================================
# CELL 3: Load Model + Tokenizer
# ============================================================
from unsloth import FastLanguageModel
import gc

print(f"Loading {MODEL_NAME}...")

# Note: text_only=True skips vision components (saves ~2GB VRAM).
# If your Unsloth version doesn't support it, remove this parameter.
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LENGTH,
    dtype=None,
    load_in_4bit=True,
    text_only=True,  # Remove if Unsloth throws "unexpected keyword argument"
)

print(f"Model VRAM: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print(f"Chat template: {tokenizer.chat_template is not None}")

gc.collect()
torch.cuda.empty_cache()
print("Model ready. Run Cell 4.")


# ============================================================
# CELL 4: Format Dataset with Chat Template + Diagnostics
# ============================================================
# Convert messages -> text string using Gemma chat template
def format_chat_template(examples):
    texts = []
    for messages in examples["messages"]:
        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False  # False for training data
        )
        texts.append(text)
    return {"text": texts}

dataset = dataset.map(format_chat_template, batched=True, remove_columns=["messages"])

print(f"Formatted dataset: {len(dataset):,}")

# --- Dataset length diagnostics ---
lengths = [len(x["text"]) for x in dataset.select(range(min(1000, len(dataset))))]
print(f"Max text length: {max(lengths)}")
print(f"Avg text length: {sum(lengths)/len(lengths):.0f}")

# --- Filter extremely long conversations (prevents OOM) ---
original_count = len(dataset)
dataset = dataset.filter(lambda x: len(x["text"]) < 12000)
filtered_count = len(dataset)
print(f"Filtered: {original_count - filtered_count} conversations > 12000 chars removed")
print(f"Final dataset: {filtered_count:,}")

print(f"Sample length: {len(dataset[0]['text'])} chars")
print(f"Preview:\n{dataset[0]['text'][:500]}...")

gc.collect()
print("Dataset formatted. Run Cell 5.")


# ============================================================
# CELL 5: Apply LoRA Adapter (Guarded)
# ============================================================
from unsloth import FastLanguageModel
from peft import PeftModel
import gc

# Guard: prevents double-LoRA error if cell is re-run
if isinstance(model, PeftModel):
    print("LoRA already attached. Skipping.")
else:
    print("Attaching LoRA adapter...")
    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=3407,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
    )

model.print_trainable_parameters()
model.config.use_cache = False

gc.collect()
torch.cuda.empty_cache()
print(f"VRAM after LoRA: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print("LoRA ready. Run Cell 6.")


# ============================================================
# CELL 6: Training Config
# ============================================================
from trl import SFTConfig

total_steps = len(dataset) // (BATCH_SIZE * GRAD_ACCUM) * EPOCHS
print(f"Estimated steps: {total_steps:,}")

sft_args = SFTConfig(
    output_dir=OUTPUT_DIR,
    overwrite_output_dir=True,
    num_train_epochs=EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LR,
    lr_scheduler_type="cosine",
    warmup_steps=100,
    optim="adamw_8bit",
    weight_decay=0.01,
    max_grad_norm=0.3,
    bf16=True,
    fp16=False,
    logging_steps=25,
    save_steps=250,  # More frequent saves for 970k dataset
    save_total_limit=2,
    report_to="none",
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LENGTH,  # Changed from max_length to max_seq_length
    seed=3407,
)

print("SFTConfig ready. Run Cell 7.")


# ============================================================
# CELL 7: Initialize Trainer
# ============================================================
from trl import SFTTrainer
import gc

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    processing_class=tokenizer,
    args=sft_args,
    packing=False,  # Explicit: don't merge conversations into one sequence
)

print("Trainer initialized.")
print(f"VRAM before training: {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print("Run Cell 8 to start training.")


# ============================================================
# CELL 8: Start Training
# ============================================================
import gc

gc.collect()
torch.cuda.empty_cache()

print("=" * 60)
print("STARTING GEMMA-3 4B LoRA TRAINING")
print("=" * 60)
print(f"Dataset    : {len(dataset):,} conversations")
print(f"Epochs     : {EPOCHS}")
print(f"Batch      : {BATCH_SIZE} (accum: {GRAD_ACCUM}, effective: {BATCH_SIZE * GRAD_ACCUM})")
print(f"Seq Length : {MAX_SEQ_LENGTH}")
print(f"LoRA       : r={LORA_R}, alpha={LORA_ALPHA}")
print(f"LR         : {LR}")
print(f"VRAM       : {torch.cuda.memory_allocated() / 1024**3:.1f} GB")
print("=" * 60)

trainer.train()

print("=" * 60)
print("TRAINING COMPLETE")
print("=" * 60)
print("Run Cell 9 to save the adapter.")


# ============================================================
# CELL 9: Save Adapter
# ============================================================
import os

os.makedirs(OUTPUT_DIR, exist_ok=True)
adapter_path = os.path.join(OUTPUT_DIR, "adapter")

model.save_pretrained(adapter_path)
tokenizer.save_pretrained(adapter_path)

print(f"Adapter saved to: {adapter_path}")
print("Run Cell 10 for inference test.")


# ============================================================
# CELL 10: Quick Inference Test
# ============================================================
from unsloth import FastLanguageModel

FastLanguageModel.for_inference(model)

messages = [
    {"role": "user", "content": "What does Charaka say about the qualities of a good physician in Sutrasthana?"}
]

inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    return_tensors="pt",
    add_generation_prompt=True,  # True for inference
).to("cuda")

print("Generating...")
outputs = model.generate(
    inputs,
    max_new_tokens=512,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("\n" + "=" * 60)
print("RESPONSE")
print("=" * 60)
print(response)
print("=" * 60)


# ============================================================
# BONUS: Sanity Check Cell (Run before full training)
# ============================================================
# Test on 100 records first to verify everything works:
#
# small_dataset = dataset.select(range(100))
# trainer = SFTTrainer(
#     model=model,
#     train_dataset=small_dataset,
#     processing_class=tokenizer,
#     args=sft_args,
#     packing=False,
# )
# trainer.train()
# print("Sanity check passed! Now switch back to full dataset.")