"""
Simple prediction demo: load a fine-tuned HF model and classify cleaned files.

Usage (bash):
python scripts/predict_demo.py --model_dir "./darkbert_finetuned_final" --input data/cleaned/cleaned_extracted_text.jsonl --out predictions.jsonl

"""
import argparse
import json
import os
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F


def predict_text_single(model, tokenizer, text, device='cpu', max_length=512, use_prompt=False, format_prompt_fn=None):
    inp = format_prompt_fn(text) if use_prompt and format_prompt_fn is not None else text
    enc = tokenizer(inp, truncation=True, padding=True, return_tensors='pt', max_length=max_length).to(device)
    with torch.no_grad():
        outputs = model(**enc)
        probs = F.softmax(outputs.logits, dim=-1).cpu().numpy()[0]
        pred = int(np.argmax(probs))
        return pred, float(probs[pred]), probs.tolist()


def read_jsonl(path):
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            yield json.loads(line)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_dir', required=True)
    parser.add_argument('--input', required=True)
    parser.add_argument('--out', default='predictions.jsonl')
    parser.add_argument('--device', default='cpu')
    parser.add_argument('--use_prompt', action='store_true')
    args = parser.parse_args()

    predict_file(args.model_dir, args.input, args.out, device=args.device, use_prompt=args.use_prompt)


def predict_file(model_dir: str, input_path: str, out_path: str, device: str = 'cpu', use_prompt: bool = False):
    """Programmatic API: load model from `model_dir`, predict on `input_path` JSONL, write predictions to `out_path`.
    Returns the number of predictions written.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(device)

    # try to import format_prompt function from notebook code if available
    try:
        from finetune_prompt_template import format_prompt_for_training as fmt
    except Exception:
        fmt = None

    count = 0
    os.makedirs(os.path.dirname(out_path) or '.', exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as out_f:
        for obj in read_jsonl(input_path):
            text = obj.get('text', '')
            pred, conf, probs = predict_text_single(model, tokenizer, text, device=device, use_prompt=use_prompt, format_prompt_fn=fmt)
            out = {
                'text': text,
                'pred_id': pred,
                'confidence': conf,
                'probs': probs
            }
            if 'label' in obj:
                out['label'] = obj['label']
            out_f.write(json.dumps(out, ensure_ascii=False) + '\n')
            count += 1

    print(f"Predictions saved to {out_path} ({count} items)")
    return count


if __name__ == '__main__':
    main()