import argparse
import json
import torch
import os
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from genre_policy import GenreDecisionPolicy

def main():
    parser = argparse.ArgumentParser(description="Infer genre from Title and Author")
    parser.add_argument("--title", required=True, type=str)
    parser.add_argument("--author", default="Unknown", type=str)
    parser.add_argument("--model-dir", default=".", type=str)
    args = parser.parse_args()

    # Load Model and Tokenizer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(args.model_dir)
    model.to(device)
    model.eval()

    # Load Policy
    policy_path = os.path.join(args.model_dir, "decision_policy.json")
    policy = GenreDecisionPolicy(policy_path)

    # Input format MUST NOT contain breadcrumbs or categories
    input_text = f"Title: {args.title}\nAuthor: {args.author}"

    inputs = tokenizer(input_text, return_tensors="pt", truncation=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)[0].cpu().numpy()

    # Get labels
    id2label = model.config.id2label
    
    top_indices = probs.argsort()[::-1]
    top_k = [[id2label[i], float(probs[i])] for i 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)

    result = {
        "title": args.title,
        "author": args.author,
        "input_text": input_text,
        "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"
    }

    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()