from fastapi import FastAPI, File, UploadFile, HTTPException
from PIL import Image
import numpy as np
import pickle
from io import BytesIO
import math

app = FastAPI()

# Cargar el modelo SOM previamente entrenado
with open("som.pkl", "rb") as tf:
    som = pickle.load(tf)

M = 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.]
])

def predict_fingerprint(image):
    try:
        processed_image = preprocess_image(image)
        winner = som.winner(processed_image)
        fingerprint_type = get_fingerprint_type(winner)
        return fingerprint_type
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def preprocess_image(image):
    # Guardar la imagen en formato TIFF
    image.resize((256,256))
    image.save("temp_image.tiff")
    processed_image = representativo("temp_image.tiff")
    processed_image_resized = processed_image.reshape(1, -1)
    return processed_image_resized

def get_fingerprint_type(winner):
    labels = {0: "LL", 1: "RL", 2: "WH", 3: "AR"}
    fingerprint_type = labels[int(M[winner[0], winner[1]])]
    return fingerprint_type

@app.post("/predict/")
async def predict_fingerprint_api(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        image = Image.open(BytesIO(contents)).convert('L')
        print(f"IMAGEN\n\n{list(image.getdata())}")
        fingerprint_type = predict_fingerprint(image)
        return {"prediction": fingerprint_type}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def sobel(I):
    m, n = I.shape
    Gx = np.zeros([m-2, n-2], np.float32)
    Gy = np.zeros([m-2, n-2], 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-2):
        for i in range(1, n-2):
            Gx[j-1, i-1] = np.sum(I[j-1:j+2, i-1:i+2] * gx)
            Gy[j-1, i-1] = 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.zeros([m+2*(d//2), n+2*(d//2)], np.float32)
    salida = np.zeros([m, n], np.float32)
    temp[1:m+1, 1:n+1] = G
    for i in range(1, m):
        for j in range(1, n):
            A = np.sort(temp[i-1:i+2, j-1:j+2].reshape(-1))
            salida[i-1, j-1] = A[d+1]
    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:j+1] * Gy[i*w:(i+1)*w, j:j+1])
            XX = np.sum(Gx[i*w:(i+1)*w, j:j+1]**2 - Gy[i*w:(i+1)*w, j:j+1]**2)
            mOrientaciones[i, j] = (0.5 * math.atan2(YY, XX) + math.pi/2.0) * (180.0 / math.pi)
    return mOrientaciones

def representativo(archivo):
    im = Image.open(archivo)
    m, n = im.size
    imarray = np.array(im, np.float32)
    patron = imarray[1:m-1, 1:n-1]
    EE = orientacion(patron, 14)
    return np.asarray(EE).reshape(-1)