"""
Inference code for Code Analyzer Model
This file enables the "Use this model" button on Hugging Face.
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch


def load_model_and_tokenizer(model_name: str):
    """Load model and tokenizer"""
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
        trust_remote_code=True
    )
    return model, tokenizer


def build_input(task, code):
    """Build input in the same format as used during training"""
    parts = []
    if task.strip():
        parts.append(f"Задача:\n{task.strip()}")
    if code.strip():
        parts.append(f"Решение (код):\n```python\n{code.strip()}\n```")
    return "\n\n".join(parts)


def generate_response(
    model,
    tokenizer,
    task: str,
    code: str,
    max_new_tokens: int = 1024,
    temperature: float = 0.7,
    top_p: float = 0.8,
    top_k: int = 20,
    repetition_penalty: float = 1.05,
):
    """Generate analysis response for task and student code"""
    # Format input in training style
    input_text = build_input(task, code)
    prompt = f"{input_text}\n\nОтвет:\n"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Extract only the answer part
    if "Ответ:" in response:
        response = response.split("Ответ:")[-1].strip()
    
    return response


if __name__ == "__main__":
    # Example usage
    import json
    
    model_name = "Vilyam888/Code_analyze.1.0"
    
    print("Loading model...")
    model, tokenizer = load_model_and_tokenizer(model_name)
    
    # Example: task and student code
    task = "Напишите функцию, которая принимает список чисел и возвращает сумму всех элементов."
    code = """def sum_list(numbers):
    total = 0
    for num in numbers:
        total += num
    return total"""
    
    print(f"\nЗадача: {task}")
    print(f"\nКод студента:\n{code}\n")
    print("Generating analysis...")
    
    response = generate_response(model, tokenizer, task, code)
    
    # Try to parse as JSON
    try:
        result = json.loads(response)
        print(f"\nРезультат анализа (JSON):")
        print(json.dumps(result, ensure_ascii=False, indent=2))
    except json.JSONDecodeError:
        print(f"\nРезультат анализа:\n{response}")