#!/usr/bin/env python3
"""
 CRNN Caligrafía OPTIMIZADO para Mac Intel i9 + AMD GPU
 Entrenamiento 98% accuracy con GPU AMD Radeon Pro
 predict(imagen_path) → {"caligrafia": 2, "confidence": 0.97}
 Integrado Flask API + Optimizaciones específicas Intel Mac
"""

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from PIL import Image
import os
import numpy as np
import warnings
warnings.filterwarnings('ignore')
import sys
from pathlib import Path
import shutil

import platform

IS_APPLE_SILICON = platform.machine() == 'arm64'
print(f"🏗️ Arquitectura: {platform.machine()}")

def get_device():
    return torch.device('cpu')  # Railway = CPU only

device = get_device()
print(f" CUDA: {torch.cuda.is_available()}")
print(f" MPS API: {torch.backends.mps.is_available()} (IGNORADO en Intel)")
print(f" Usando: {device}")


# ──────────────────────────────────────────────
#  ResizeWithPad: escala manteniendo proporción
#  y rellena con blanco hasta TARGET_H x TARGET_W
# ──────────────────────────────────────────────
class ResizeWithPad:
    """
    Redimensiona preservando la proporción original:
      1. Escala hasta que alto == target_h
      2. Si ancho < target_w → padding blanco a ambos lados (centrado)
      3. Si ancho > target_w → recorte centrado
    Resultado siempre: target_w × target_h px
    """
    def __init__(self, target_h=64, target_w=256, fill=255):
        self.target_h = target_h
        self.target_w = target_w
        self.fill     = fill

    def __call__(self, img):
        # Normaliza cualquier modo a RGB con fondo blanco
        if img.mode == 'RGBA':
            fondo = Image.new('RGB', img.size, (self.fill, self.fill, self.fill))
            fondo.paste(img, mask=img.split()[3])
            img = fondo
        else:
            img = img.convert('RGB')

        w, h = img.size
        scale = self.target_h / h
        new_w = int(w * scale)
        img   = img.resize((new_w, self.target_h), Image.LANCZOS)

        if new_w < self.target_w:
            pad_left = (self.target_w - new_w) // 2
            new_img  = Image.new('RGB', (self.target_w, self.target_h),
                                 (self.fill, self.fill, self.fill))
            new_img.paste(img, (pad_left, 0))
            img = new_img
        elif new_w > self.target_w:
            left = (new_w - self.target_w) // 2
            img  = img.crop((left, 0, left + self.target_w, self.target_h))

        return img


# ──────────────────────────────────────────────
#  CRNN + Attention
# ──────────────────────────────────────────────
class CaligrafiaCRNN(nn.Module):
    def __init__(self, num_classes=7):  # ← consistente con el .pth entrenado
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, 3, padding=1), nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
            nn.Dropout2d(0.25)
        )

        self.lstm = nn.LSTM(4096, 256, bidirectional=True, batch_first=True, dropout=0.3)
        self.attention  = nn.MultiheadAttention(512, 8, batch_first=True)
        self.classifier = nn.Sequential(
            nn.Linear(512, 256), nn.ReLU(), nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    def forward(self, x):
        features = self.cnn(x)
        B, C, H, W = features.shape
        features = features.permute(0, 3, 1, 2).reshape(B, W, C * H)

        lstm_out, _ = self.lstm(features)
        attn_out, _ = self.attention(lstm_out, lstm_out, lstm_out)
        pooled = attn_out.mean(dim=1)

        return self.classifier(pooled)


# ──────────────────────────────────────────────
#  PREDICTOR
# ──────────────────────────────────────────────
class CaligrafiaPredictor:
    def __init__(self, model_path="best_caligrafia_model.pth", num_classes=7):
        self.device = get_device()
        self.model  = CaligrafiaCRNN(num_classes).to(self.device)
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        self.model.eval()

        self.transform = transforms.Compose([
            transforms.Lambda(lambda img: ResizeWithPad(64, 256, fill=255)(img)),
            transforms.Grayscale(num_output_channels=1),
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,))
        ])

        self.clases = ['Impresa', 'Italica_cursiva', 'No_valido',
                       'Procesal_encadenada', 'Procesal_simple',
                       'Pseudorredonda', 'Redonda']
        print(f"✅ Clases cargadas: {self.clases}")

    def predict_image(self, image):
        """Predicción desde objeto PIL Image (para Flask)"""
        # Normaliza modo antes de transformar — evita crash con P, RGBA, CMYK
        if image.mode not in ('RGB', 'RGBA'):
            image = image.convert('RGB')

        image = self.transform(image).unsqueeze(0).to(self.device)
        with torch.no_grad():
            logits     = self.model(image)
            probs      = torch.softmax(logits, dim=1)
            pred_class = probs.argmax(dim=1).cpu().item()
            confidence = probs[0, pred_class].cpu().item()
        return {
            "caligrafia":    self.clases[pred_class],
            "confidence":    float(confidence),
            "probs":         probs[0].cpu().tolist(),
            "pred_class_id": int(pred_class)
        }

    def predict(self, image_path):
        """Predicción desde ruta de archivo"""
        image = Image.open(image_path)
        # Normaliza modo antes de transformar
        if image.mode not in ('RGB', 'RGBA'):
            image = image.convert('RGB')

        image = self.transform(image).unsqueeze(0).to(self.device, non_blocking=True)
        with torch.no_grad():
            logits     = self.model(image)
            probs      = torch.softmax(logits, dim=1)
            pred_class = probs.argmax(dim=1).cpu().item()
            confidence = probs[0, pred_class].cpu().item()
        return {
            "caligrafia":    self.clases[pred_class],
            "confidence":    float(confidence),
            "probs":         probs[0].cpu().tolist(),
            "pred_class_id": int(pred_class)
        }