#!/usr/bin/env python3
"""Predict whether Lean source code compiles with a trained classifier."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from typing import Any, Sequence

DEFAULT_CACHE_DIR = None
DEFAULT_CLASSIFIER_OUTPUT = 'EPFL-LARA/lean-encoder-modernbert-large'

def encode_head_tail(
    tokenizer: Any,
    texts: Sequence[str],
    *,
    max_length: int,
    head_fraction: float = 0.25,
) -> dict[str, torch.Tensor]:
    """Tokenize, retaining the beginning and end when a source is too long."""
    if not 0.0 < head_fraction < 1.0:
        raise ValueError("head_fraction must be between 0 and 1")
    payload_budget = max_length - tokenizer.num_special_tokens_to_add(pair=False)
    if payload_budget <= 0:
        raise ValueError("max_length is smaller than the tokenizer's special tokens")
    head_tokens = max(1, round(payload_budget * head_fraction))
    tail_tokens = payload_budget - head_tokens

    # A few compiler_data rows are over a million characters. Bound tokenizer
    # memory with a generous character pre-slice, then perform the exact slice
    # in token space below.
    prepared_texts: list[str] = []
    head_character_budget = max(2_048, head_tokens * 8)
    tail_character_budget = max(2_048, tail_tokens * 8)
    for text in texts:
        if len(text) > head_character_budget + tail_character_budget:
            text = (
                text[:head_character_budget]
                + "\n-- middle omitted for classifier context window --\n"
                + text[-tail_character_budget:]
            )
        prepared_texts.append(text)

    tokenized = tokenizer(
        prepared_texts,
        add_special_tokens=False,
        truncation=False,
        return_attention_mask=False,
        verbose=False,
    )["input_ids"]

    features: list[dict[str, list[int]]] = []
    for input_ids in tokenized:
        if len(input_ids) > payload_budget:
            input_ids = input_ids[:head_tokens] + input_ids[-tail_tokens:]
        with_special_tokens = tokenizer.build_inputs_with_special_tokens(input_ids)
        features.append(
            {
                "input_ids": with_special_tokens,
                "attention_mask": [1] * len(with_special_tokens),
            }
        )
    return tokenizer.pad(
        features,
        padding=True,
        pad_to_multiple_of=8,
        return_tensors="pt",
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("path", nargs="?", help="Lean file; omit to read stdin")
    parser.add_argument("--model", default=DEFAULT_CLASSIFIER_OUTPUT)
    parser.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR)
    parser.add_argument("--max-length", type=int, default=8192)
    parser.add_argument("--head-fraction", type=float, default=0.25)
    args = parser.parse_args()

    source_code = (
        Path(args.path).read_text(encoding="utf-8") if args.path else sys.stdin.read()
    )
    if not source_code.strip():
        parser.error("Lean source is empty")

    tokenizer = AutoTokenizer.from_pretrained(args.model, cache_dir=args.cache_dir)
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model,
        cache_dir=args.cache_dir,
        dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        attn_implementation="sdpa",
    )
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device).eval()
    inputs = encode_head_tail(
        tokenizer,
        [source_code],
        max_length=args.max_length,
        head_fraction=args.head_fraction,
    )
    inputs = {key: value.to(device) for key, value in inputs.items()}
    with torch.inference_mode():
        probabilities = model(**inputs).logits.float().softmax(dim=-1)[0].cpu()
    valid_probability = float(probabilities[1])
    print(
        json.dumps(
            {
                "is_valid": valid_probability >= 0.5,
                "valid_probability": valid_probability,
                "invalid_probability": float(probabilities[0]),
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()