import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers, models
from flask import Flask, request, jsonify, render_template
from PIL import Image
import io
from datasets import load_dataset

# ------------------------------
# Configuración
# ------------------------------
MODEL_PATH = "modelo_sexo.keras"  # Cambiado para evitar conflictos con el viejo
IMG_SIZE = (64, 64)

def build_model():
    model = models.Sequential([
        layers.Conv2D(32, (3, 3), activation='relu', input_shape=(64, 64, 3)),
        layers.MaxPooling2D((2, 2)),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.Flatten(),
        layers.Dense(64, activation='relu'),
        layers.Dense(1, activation='sigmoid')   # 1 = mujer, 0 = hombre
    ])
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model

def train_and_save_model():
    print("Cargando dataset FairFace...")
    ds = load_dataset("HuggingFaceM4/FairFace", "0.25")
    
    # Usamos todo el dataset de entrenamiento disponible
    ds_train = ds['train']
    
    images = []
    labels = []
    print("Procesando imágenes para clasificación de género...")
    for item in ds_train:
        img = item['image'].convert('RGB').resize(IMG_SIZE)
        images.append(np.array(img) / 255.0)
        # En FairFace: 0 corresponde a 'Male' (Hombre) y 1 a 'Female' (Mujer)
        labels.append(1.0 if item['gender'] == 1 else 0.0)
    
    X = np.array(images)
    y = np.array(labels)
    
    print(f"Datos cargados: {len(X)} imágenes (mujeres: {sum(y)}, hombres: {len(y)-sum(y)})")
    
    # Entrenar modelo
    model = build_model()
    print("Entrenando clasificador de género...")
    model.fit(X, y, epochs=7, batch_size=64, validation_split=0.2, verbose=1)
    
    # Guardar el nuevo modelo
    model.save(MODEL_PATH)
    print(f"Modelo guardado como {MODEL_PATH}")
    return model

# Cargar modelo o entrenar si no existe
def load_or_train_model():
    if os.path.exists(MODEL_PATH):
        try:
            model = tf.keras.models.load_model(MODEL_PATH)
            print("✅ Modelo de género cargado desde disco correctamente")
            return model
        except Exception as e:
            print(f"⚠️ Error al cargar el modelo: {e}")
            print("Entrenando uno nuevo...")
            return train_and_save_model()
    else:
        print("Modelo no encontrado. Entrenando desde cero...")
        return train_and_save_model()

# Inicializar el modelo global
model = load_or_train_model()

# ------------------------------
# Aplicación Flask
# ------------------------------
app = Flask(__name__)

def predecir_sexo(imagen_bytes):
    img = Image.open(io.BytesIO(imagen_bytes)).convert('RGB')
    img = img.resize(IMG_SIZE)
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    prediccion = model.predict(img_array, verbose=0)[0][0]
    return "mujer" if prediccion >= 0.5 else "hombre"

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({'error': 'No se encontró ninguna imagen'}), 400
    file = request.files['image']
    if file.filename == '':
        return jsonify({'error': 'Nombre de archivo vacío'}), 400
    try:
        imagen_bytes = file.read()
        
        # 1. Obtenemos el valor numérico puro de la predicción
        img = Image.open(io.BytesIO(imagen_bytes)).convert('RGB')
        img = img.resize(IMG_SIZE)
        img_array = np.array(img) / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        
        # Guardamos la predicción como un float para poder operar con ella
        prediccion_raw = float(model.predict(img_array, verbose=0)[0][0])
        
        # 2. Calculamos la clase y el porcentaje de confianza
        if prediccion_raw >= 0.5:
            genero = "mujer"
            # Si es mujer, la confianza es directamente el valor predicho (ej: 0.90 -> 90%)
            confianza = prediccion_raw * 100 
        else:
            genero = "hombre"
            # Si es hombre, la confianza es lo que le falta para 1 (ej: 0.10 predicho -> 90% seguro de que es hombre)
            confianza = (1.0 - prediccion_raw) * 100
            
        # 3. Devolvemos el JSON con ambos datos, formateando la confianza a 2 decimales
        return jsonify({
            'genero': genero,
            'confianza': f"{confianza:.2f}%"
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 7860))
    app.run(host='0.0.0.0', port=port)