import cv2
import numpy as np
import pickle
import math
from fastapi import FastAPI, File, UploadFile, HTTPException
from PIL import Image
from io import BytesIO
from pydantic import BaseModel
from typing import List
from scipy.ndimage import median_filter
from scipy.signal import convolve2d
from minisom import MiniSom

def load_model():
    with open('somlucuma.pkl', 'rb') as fid:
        som = pickle.load(fid)
    MM = np.loadtxt('matrizMM.txt', delimiter=" ")
    return som, MM

def sobel(patron):
    gx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float32)
    gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=np.float32)

    Gx = convolve2d(patron, gx, mode='valid')
    Gy = convolve2d(patron, gy, mode='valid')

    return Gx, Gy

def medfilt2(G, d=3):
    return median_filter(G, size=d)

def orientacion(patron, w):
    Gx, Gy = sobel(patron)
    Gx = medfilt2(Gx)
    Gy = medfilt2(Gy)

    m, n = Gx.shape
    mOrientaciones = np.zeros((m // w, n // w), dtype=np.float32)

    for i in range(m // w):
        for j in range(n // w):
            Gx_patch = Gx[i*w:(i+1)*w, j*w:(j+1)*w]
            Gy_patch = Gy[i*w:(i+1)*w, j*w:(j+1)*w]

            YY = np.sum(2 * Gx_patch * Gy_patch)
            XX = np.sum(Gx_patch**2 - Gy_patch**2)

            mOrientaciones[i, j] = (0.5 * np.arctan2(YY, XX) + np.pi / 2.0) * (18.0 / np.pi)

    return mOrientaciones

def representativo(imarray):
    imarray = np.squeeze(imarray)
    m, n = imarray.shape
    patron = imarray[1:m-1, 1:n-1]
    EE = orientacion(patron, 14)
    return np.asarray(EE).reshape(-1)

class InputData(BaseModel):
    data: List[float]

app = FastAPI()

som, MM = load_model()

@app.post("/predict/")
async def predict(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(BytesIO(contents)).convert('L')
        image = np.asarray(image)
        if image.shape != (256, 256):
            raise ValueError("La imagen debe ser de tamaño 256x256.")
        image = image.reshape(256, 256, 1)
        print(f"Imagen convertida a matriz: {image.shape}")
        
        representative_data = representativo(image)
        print(f"Datos representativos de la imagen: {representative_data.shape}")
        representative_data = representative_data.reshape(1, -1)
        
        w = som.winner(representative_data)
        print(f"Índice ganador del SOM: {w}")
        prediction = MM[w]
        print(f"Predicción: {prediction}")
        
        return {"prediction": prediction}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))