"""Private phrase classification; inputs are not persisted by this application."""
import json
import os
from functools import lru_cache
from pathlib import Path

import gradio as gr
from huggingface_hub import snapshot_download
import numpy as np
import torch
import spaces
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from inference import calibrate, rank_average

CONFIG = json.loads((Path(__file__).parent / 'deployment_config.json').read_text())
LABELS = CONFIG['label2id']
ID2LABEL = {int(index): label for label, index in LABELS.items()}
ZERO_GPU = os.getenv('MS_ZERO_GPU') == '1'
DEVICE = 'cuda' if ZERO_GPU or torch.cuda.is_available() else 'cpu'
RECORDS = {record['name']: record for record in CONFIG['models']}
CHOICES = CONFIG['choices']


@lru_cache(maxsize=len(RECORDS) if ZERO_GPU else 1)
def load_model(name):
    record = RECORDS[name]
    token = os.environ.get('HF_TOKEN')
    print('Loading classifier:', name, flush=True)
    model_path = snapshot_download(
        record['repo_id'], revision=record.get('revision'), token=token, max_workers=2,
        allow_patterns=['*.safetensors', 'config.json', 'tokenizer*.json',
                        'special_tokens_map.json', 'added_tokens.json', 'vocab.*',
                        'merges.txt', '*.model'],
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        model_path, local_files_only=True, trust_remote_code=False,
        use_safetensors=True,
    ).to(DEVICE).eval()
    if {label: int(index) for label, index in model.config.label2id.items()} != LABELS:
        raise RuntimeError('Model label mapping differs from the deployment configuration')
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True, trust_remote_code=False)
    print('Classifier ready:', name, flush=True)
    return model, tokenizer


def probabilities(name, text):
    model, tokenizer = load_model(name)
    inputs = tokenizer(text, return_tensors='pt', truncation=True,
                       max_length=int(RECORDS[name]['max_length'])).to(DEVICE)
    with torch.inference_mode():
        return torch.softmax(model(**inputs).logits, dim=-1).cpu().numpy()


@spaces.GPU(duration=60)
def classify(text, choice):
    if not isinstance(text, str) or not text.strip():
        raise gr.Error('Enter a phrase first.')
    if len(text) > 20000:
        raise gr.Error('Enter a shorter phrase (at most 20,000 characters).')
    if choice not in CHOICES:
        raise gr.Error('Select one of the available classifiers.')
    selected = CHOICES[choice]
    if selected == '__ensemble__':
        values = rank_average([probabilities(name, text) for name in CONFIG['model_order']])
        values = calibrate(values, CONFIG['calibration']['temperature'])
    else:
        values = probabilities(selected, text)
    row = values[0]
    index = int(np.argmax(row))
    scores = {ID2LABEL[i]: float(value) for i, value in enumerate(row)}
    return ID2LABEL[index], scores


if ZERO_GPU:
    # ZeroGPU requires CUDA placement at startup, outside the GPU function.
    for model_name in RECORDS:
        load_model(model_name)


with gr.Blocks(title='MS life-event classifier', analytics_enabled=False) as demo:
    gr.Markdown('# MS life-event classifier\nPaste one phrase and choose a classifier. Each prediction returns one event category.')
    phrase = gr.Textbox(label='Phrase', lines=4)
    choice = gr.Dropdown(list(CHOICES), value=list(CHOICES)[0], label='Classifier')
    button = gr.Button('Classify', variant='primary')
    label = gr.Textbox(label='Assigned category', interactive=False)
    scores = gr.Label(label='Top category scores', num_top_classes=5)
    gr.Markdown('Scores are diagnostic; they do not guarantee correctness. Long inputs are truncated to each model’s training token limit.')
    button.click(classify, [phrase, choice], [label, scores], api_name='predict')

demo.queue(default_concurrency_limit=1, max_size=16)
if __name__ == '__main__':
    demo.launch()