import pickle
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import math

class InputData(BaseModel):
    array: List[List[int]]

app = FastAPI()

# Cargar el modelo SOM
def load_model():
    with open('som.pkl', 'rb') as fid:
        som = pickle.load(fid)
    return som

def sobel(I):
    m, n = I.shape
    Gx = np.zeros([m, n], np.float32)
    Gy = np.zeros([m, n], np.float32)
    gx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    gy = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
    for j in range(1, m-1):
        for i in range(1, n-1):
            Gx[j, i] = np.sum(I[j-1:j+2, i-1:i+2] * gx)
            Gy[j, i] = np.sum(I[j-1:j+2, i-1:i+2] * gy)
    return Gx, Gy

def medfilt2(G, d=3):
    m, n = G.shape
    temp = np.pad(G, pad_width=d//2, mode='constant', constant_values=0)
    salida = np.zeros([m, n], np.float32)
    for i in range(m):
        for j in range(n):
            A = temp[i:i+d, j:j+d].flatten()
            salida[i, j] = np.median(A)
    return salida

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], np.float32)
    for i in range(m//w):
        for j in range(n//w):
            YY = np.sum(2 * Gx[i*w:(i+1)*w, j*w:(j+1)*w] * Gy[i*w:(i+1)*w, j*w:(j+1)*w])
            XX = np.sum(Gx[i*w:(i+1)*w, j*w:(j+1)*w]**2 - Gy[i*w:(i+1)*w, j*w:(j+1)*w]**2)
            mOrientaciones[i, j] = (0.5 * math.atan2(YY, XX) + math.pi / 2.0) * (180.0 / math.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)

som = load_model()

MM = np.array([
    [ 0., -1., -1., -1., -1.,  2., -1., -1., -1.,  3.],
    [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
    [-1., -1., -1.,  1., -1., -1., -1., -1., -1., -1.],
    [ 1., -1., -1., -1., -1., -1., -1., -1., -1.,  0.],
    [-1., -1., -1., -1.,  1., -1., -1., -1., -1., -1.],
    [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
    [ 3., -1., -1., -1., -1., -1., -1., -1., -1.,  3.],
    [-1., -1., -1.,  0., -1., -1.,  3., -1., -1., -1.],
    [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
    [ 2., -1., -1., -1.,  1., -1., -1., -1., -1.,  2.]
])

@app.post("/predict/")
async def predict(data: InputData):
    try:
        print(data.array)
        input_data = np.array(data.array).reshape(256, 256)
        representative_data = representativo(input_data)
        representative_data = representative_data.reshape(1, -1)
        
        w = som.winner(representative_data)
        prediction = MM[w]
        
        return {"prediction": prediction.tolist()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))