import argparse
import json
import torch
import os
import pandas as pd
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from genre_policy import GenreDecisionPolicy

def main():
    parser = argparse.ArgumentParser(description="Batch infer genre from JSONL or CSV")
    parser.add_argument("--input", required=True, type=str)
    parser.add_argument("--output", required=True, type=str)
    parser.add_argument("--batch-size", default=256, type=int)
    parser.add_argument("--model-dir", default=".", type=str)
    parser.add_argument("--device", default="cuda", type=str)
    args = parser.parse_args()

    device = torch.device(args.device if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(args.model_dir)
    model.to(device)
    model.eval()

    policy_path = os.path.join(args.model_dir, "decision_policy.json")
    policy = GenreDecisionPolicy(policy_path)
    id2label = model.config.id2label

    # Load inputs
    if args.input.endswith(".csv"):
        df = pd.read_csv(args.input)
        records = df.to_dict(orient="records")
    else:
        records = []
        with open(args.input, "r") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))

    results = []
    
    with open(args.output, "w") as f_out:
        for i in tqdm(range(0, len(records), args.batch_size)):
            batch = records[i:i + args.batch_size]
            
            texts = []
            for item in batch:
                t = item.get("title", "")
                a = item.get("author", "Unknown")
                texts.append(f"Title: {t}\nAuthor: {a}")
                
            inputs = tokenizer(texts, return_tensors="pt", truncation=True, padding=True, max_length=128)
            inputs = {k: v.to(device) for k, v in inputs.items()}
            
            with torch.no_grad():
                outputs = model(**inputs)
                probs = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().numpy()
                
            for j, prob in enumerate(probs):
                top_indices = prob.argsort()[::-1]
                top_k = [[id2label[idx], float(prob[idx])] for idx in top_indices]
                
                top1_label, top1_confidence = top_k[0]
                top2_label, top2_confidence = top_k[1]
                margin = top1_confidence - top2_confidence
                
                decision_out = policy.decide(top_k)
                
                res = {
                    "title": batch[j].get("title", ""),
                    "author": batch[j].get("author", "Unknown"),
                    "top1_label": top1_label,
                    "top1_confidence": round(top1_confidence, 4),
                    "top2_label": top2_label,
                    "top2_confidence": round(top2_confidence, 4),
                    "margin": round(margin, 4),
                    "top_k": [[l, round(p, 4)] for l, p in top_k],
                    "decision": decision_out["decision"],
                    "final_label": decision_out["final_label"],
                    "ambiguity_bucket": decision_out["ambiguity_bucket"],
                    "reason": decision_out["reason"],
                    "model_version": "book_genre_v5_title_author"
                }
                
                if "book_id" in batch[j]:
                    res["book_id"] = batch[j]["book_id"]
                if "path" in batch[j]:
                    res["path"] = batch[j]["path"]
                    
                f_out.write(json.dumps(res) + "\n")
                
    print(f"Batch inference complete. Output saved to {args.output}")

if __name__ == "__main__":
    main()