#!/usr/bin/env python3
"""Continue training from saved model with augmented data."""
import json
import os
import sys
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
from datasets import Dataset

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import RAW_DIR, MODELS_DIR, MAX_SEQ_LEN, LABEL_MAP, BATCH_SIZE, LEARNING_RATE

FINAL_MODEL_DIR = os.path.join(MODELS_DIR, "final_model")

def load_all_jsonl():
    examples = []
    for fname in os.listdir(RAW_DIR):
        if not fname.endswith(".jsonl") or "extra" in fname:
            continue
        fpath = os.path.join(RAW_DIR, fname)
        with open(fpath) as f:
            for line in f:
                line = line.strip()
                if not line: continue
                obj = json.loads(line)
                text = obj.get("text", "").strip()
                label_str = obj.get("label", "")
                if text and label_str in LABEL_MAP:
                    examples.append({"text": text, "label": LABEL_MAP[label_str]})
    return examples

def main():
    print("Loading saved model...")
    tokenizer = AutoTokenizer.from_pretrained(FINAL_MODEL_DIR)
    model = AutoModelForSequenceClassification.from_pretrained(FINAL_MODEL_DIR)
    
    print("Loading augmented dataset...")
    examples = load_all_jsonl()
    print(f"Total examples: {len(examples)}")
    
    # Check label balance
    labels = [ex["label"] for ex in examples]
    print(f"  GENERIC: {labels.count(0)}")
    print(f"  SEMANTIC: {labels.count(1)}")
    
    dataset = Dataset.from_list(examples)
    splits = dataset.train_test_split(test_size=0.15, seed=42)
    
    def tokenize(examples):
        return tokenizer(examples["text"], padding="max_length", 
                        truncation=True, max_length=MAX_SEQ_LEN)
    
    tokenized = splits.map(tokenize, batched=True)
    tokenized = tokenized.remove_columns(["text"])
    tokenized = tokenized.rename_column("label", "labels")
    tokenized.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
    
    print("\nContinuing training for 1 epoch...")
    training_args = TrainingArguments(
        output_dir=os.path.join(MODELS_DIR, "checkpoints_v2"),
        eval_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=200,
        logging_steps=25,
        learning_rate=5e-6,  # Lower LR for continued training
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE * 2,
        num_train_epochs=1,
        weight_decay=0.01,
        warmup_ratio=0.05,
        fp16=torch.cuda.is_available(),
        save_total_limit=1,
        load_best_model_at_end=True,
        metric_for_best_model="accuracy",
        greater_is_better=True,
        report_to="none",
        seed=42,
        dataloader_num_workers=2,
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized["train"],
        eval_dataset=tokenized["test"],
        compute_metrics=lambda p: (
            {"accuracy": (p.predictions.argmax(-1) == p.label_ids).mean()}
        ),
    )
    
    trainer.train()
    
    # Evaluate on specific problem cases
    print("\nEvaluating problem cases:")
    model.eval()
    problem_queries = [
        "I love spicy food",
        "my name is John",
        "मेरा नाम रवि है",
        "नमस्ते",
        "hello",
        "मुझे कॉफी पसंद है",
        "I work as a software engineer",
        "my favorite color is blue",
    ]
    
    for query in problem_queries:
        inputs = tokenizer(query, return_tensors="pt", padding="max_length",
                          truncation=True, max_length=MAX_SEQ_LEN)
        with torch.no_grad():
            logits = model(**inputs).logits
        probs = torch.nn.functional.softmax(logits, dim=-1).numpy()[0]
        pred = "SEMANTIC" if probs[1] > probs[0] else "GENERIC"
        print(f"  [{pred:8s}] gen={probs[0]:.3f} sem={probs[1]:.3f}  \"{query}\"")
    
    # Save model again
    trainer.save_model(FINAL_MODEL_DIR)
    tokenizer.save_pretrained(FINAL_MODEL_DIR)
    print(f"\nModel saved to {FINAL_MODEL_DIR}")

if __name__ == "__main__":
    main()