import os
import torch
from PIL import Image

from model import get_model
from dataset import decode_prediction, get_val_transforms, ALPHABET
from word_detector import split_into_words


class OCR:
    def __init__(
        self,
        checkpoint_path: str = "best_model.pth",  # По умолчанию смотрим в корень для Hugging Face
        device: str = None,
    ):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        # Если переданный путь не существует, но файл есть в корне — используем корень
        if not os.path.exists(checkpoint_path) and os.path.exists("best_model.pth"):
            checkpoint_path = "best_model.pth"
        elif not os.path.exists(checkpoint_path) and os.path.exists("../models/best_model.pth"):
            checkpoint_path = "../models/best_model.pth"

        num_classes = len(ALPHABET) + 1
        self.model = get_model(num_classes=num_classes, device=self.device)

        print(f"[OCR] загружаем веса: {checkpoint_path}")
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        
        # Поддержка загрузки как чистого state_dict, так и сохраненного через чекпоинт словаря
        if "model_state_dict" in checkpoint:
            self.model.load_state_dict(checkpoint["model_state_dict"])
        else:
            self.model.load_state_dict(checkpoint)
            
        self.model.eval()

        print(f"[OCR] точность модели при обучении: {checkpoint.get('accuracy', 'неизвестно') if isinstance(checkpoint, dict) else 'неизвестно'}")

        self.transform = get_val_transforms(img_height=32, img_width=128)

    @torch.no_grad()
    def _recognize_single_word(self, image: Image.Image) -> str:
        image_tensor = self.transform(image.convert("RGB")).unsqueeze(0).to(self.device)

        logits = self.model(image_tensor)           # [T, B, C]
        preds = logits.argmax(dim=2).permute(1, 0)   # [B, T]

        text = decode_prediction(preds[0].tolist())
        return text

    def recognize(self, image_path: str) -> str:
        image = Image.open(image_path)
        return self._recognize_single_word(image)

    def recognize_line(self, image_path: str, **detector_kwargs) -> str:
        image = Image.open(image_path)
        word_images = split_into_words(image, **detector_kwargs)

        if not word_images:
            return self._recognize_single_word(image)

        words = [self._recognize_single_word(w) for w in word_images]
        return " ".join(words)

    @torch.no_grad()
    def recognize_batch(self, image_paths: list[str]) -> list[str]:
        images = []
        for path in image_paths:
            image = Image.open(path).convert("RGB")
            images.append(self.transform(image))

        batch = torch.stack(images).to(self.device)
        logits = self.model(batch)
        preds = logits.argmax(dim=2).permute(1, 0)

        results = [decode_prediction(seq.tolist()) for seq in preds]
        return results


if __name__ == "__main__":
    # Для локальных тестов проверяет обе папки
    default_path = "best_model.pth" if os.path.exists("best_model.pth") else "../models/best_model.pth"
    ocr = OCR(checkpoint_path=default_path)

    test_image = "test_word.jpg"

    try:
        result = ocr.recognize_line(test_image)
        print(f"\nРаспознанный текст: '{result}'")
    except FileNotFoundError:
        print(f"\nКартинка '{test_image}' не найдена.")
        print("укажи нормальный путь.")