"""Minimální torch/transformers starter — inference + trénovací smyčka (CPU-safe)."""

from __future__ import annotations

import os
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

HF_TOKEN = os.environ.get("HF_TOKEN")  # výhradně z env
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_ID = "gpt2"  # malý model, funguje na CPU


def load_model(model_id: str = MODEL_ID):
    tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
    model = AutoModelForCausalLM.from_pretrained(model_id, token=HF_TOKEN)
    model.to(DEVICE)
    model.eval()
    return tokenizer, model


def infer(tokenizer, model, prompt: str, max_new_tokens: int = 50) -> str:
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.7,
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


def mini_training_loop() -> None:
    """Ukázková trénovací smyčka na syntetických datech (CPU-safe)."""
    model = torch.nn.Linear(10, 1).to(DEVICE)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    criterion = torch.nn.MSELoss()

    for step in range(5):
        x = torch.randn(4, 10, device=DEVICE)
        y = torch.randn(4, 1, device=DEVICE)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        print(f"Krok {step + 1}: loss = {loss.item():.4f}")


def main() -> None:
    print(f"Zařízení: {DEVICE}")
    if DEVICE == "cpu":
        print("Poznámka: Běží se na CPU. Pro Wan2.1/Dia je vyžadován GPU.")

    tokenizer, model = load_model()
    result = infer(tokenizer, model, "Dnes je hezký den,", max_new_tokens=20)
    print(f"Inference: {result}")

    mini_training_loop()


if __name__ == "__main__":
    main()