# interface.py

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PIL import Image
import io

from bioprocess_model import BioprocessModel
from decorators import gpu_decorator  # Asegúrate de que la ruta es correcta

# Nuevas importaciones para Yi-Coder
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import json

# Inicialización del dispositivo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model_path = "Qwen/Qwen2.5-32B-Instruct" #"01-ai/Yi-Coder-9B-Chat"

# Carga del tokenizer y modelo
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path).to(device).eval()

def parse_bounds(bounds_str, num_params):
    try:
        # Reemplazar 'inf' por 'np.inf' si el usuario lo escribió así
        bounds_str = bounds_str.replace('inf', 'np.inf')
        # Evaluar la cadena de límites
        bounds = eval(f"[{bounds_str}]")
        if len(bounds) != num_params:
            raise ValueError("Número de límites no coincide con el número de parámetros.")
        lower_bounds = [b[0] for b in bounds]
        upper_bounds = [b[1] for b in bounds]
        return lower_bounds, upper_bounds
    except Exception as e:
        print(f"Error al parsear los límites: {e}. Usando límites por defecto.")
        lower_bounds = [-np.inf] * num_params
        upper_bounds = [np.inf] * num_params
        return lower_bounds, upper_bounds

def generate_analysis(prompt, max_length=100):  # Reducido a 100
    """
    Genera un análisis utilizando el modelo Yi-Coder-9B-Chat.
    """
    try:
        # Tokenizar el prompt y mover los tensores al dispositivo correcto
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,  # Limitar la generación a 100 tokens
                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id
            )
        
        analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return analysis
    except Exception as e:
        print(f"Error al generar el análisis con Yi-Coder: {e}. Usando análisis por defecto.")
        return "Análisis generado por el modelo de lenguaje."

@gpu_decorator(duration=100)  # Reducido de 600 a 100 segundos
def process_and_plot(
    file,
    biomass_eq1, biomass_eq2, biomass_eq3,
    biomass_param1, biomass_param2, biomass_param3,
    biomass_bound1, biomass_bound2, biomass_bound3,
    substrate_eq1, substrate_eq2, substrate_eq3,
    substrate_param1, substrate_param2, substrate_param3,
    substrate_bound1, substrate_bound2, substrate_bound3,
    product_eq1, product_eq2, product_eq3,
    product_param1, product_param2, product_param3,
    product_bound1, product_bound2, product_bound3,
    legend_position,
    show_legend,
    show_params,
    biomass_eq_count,
    substrate_eq_count,
    product_eq_count
):
    # Leer el archivo Excel
    df = pd.read_excel(file.name)
    
    # Verificar que las columnas necesarias estén presentes
    expected_columns = ['Tiempo', 'Biomasa', 'Sustrato', 'Producto']
    for col in expected_columns:
        if col not in df.columns:
            raise KeyError(f"La columna esperada '{col}' no se encuentra en el archivo Excel.")

    # Asignar los datos desde las columnas
    time = df['Tiempo'].values
    biomass_data = df['Biomasa'].values
    substrate_data = df['Sustrato'].values
    product_data = df['Producto'].values

    # Convierte los contadores a enteros
    biomass_eq_count = int(biomass_eq_count)
    substrate_eq_count = int(substrate_eq_count)
    product_eq_count = int(product_eq_count)

    # Recolecta las ecuaciones, parámetros y límites según los contadores
    biomass_eqs = [biomass_eq1, biomass_eq2, biomass_eq3][:biomass_eq_count]
    biomass_params = [biomass_param1, biomass_param2, biomass_param3][:biomass_eq_count]
    biomass_bounds = [biomass_bound1, biomass_bound2, biomass_bound3][:biomass_eq_count]

    substrate_eqs = [substrate_eq1, substrate_eq2, substrate_eq3][:substrate_eq_count]
    substrate_params = [substrate_param1, substrate_param2, substrate_param3][:substrate_eq_count]
    substrate_bounds = [substrate_bound1, substrate_bound2, substrate_bound3][:substrate_eq_count]

    product_eqs = [product_eq1, product_eq2, product_eq3][:product_eq_count]
    product_params = [product_param1, product_param2, product_param3][:product_eq_count]
    product_bounds = [product_bound1, product_bound2, product_bound3][:product_eq_count]

    biomass_results = []
    substrate_results = []
    product_results = []

    # Inicializar el modelo principal
    main_model = BioprocessModel()

    # Ajusta los modelos de Biomasa
    for i in range(len(biomass_eqs)):
        equation = biomass_eqs[i]
        params_str = biomass_params[i]
        bounds_str = biomass_bounds[i]

        try:
            main_model.set_model_biomass(equation, params_str)
        except ValueError as ve:
            raise ValueError(f"Error en la configuración del modelo de biomasa {i+1}: {ve}")

        params = [param.strip() for param in params_str.split(',')]
        lower_bounds, upper_bounds = parse_bounds(bounds_str, len(params))

        try:
            y_pred = main_model.fit_model(
                'biomass', time, biomass_data,
                bounds=(lower_bounds, upper_bounds)
            )
            biomass_results.append({
                'y_pred': y_pred.tolist(),  # Convertir a lista para serialización
                'equation': equation,
                'params': main_model.params['biomass']
            })
        except Exception as e:
            raise RuntimeError(f"Error al ajustar el modelo de biomasa {i+1}: {e}")

    # Ajusta los modelos de Sustrato
    for i in range(len(substrate_eqs)):
        equation = substrate_eqs[i]
        params_str = substrate_params[i]
        bounds_str = substrate_bounds[i]

        try:
            main_model.set_model_substrate(equation, params_str)
        except ValueError as ve:
            raise ValueError(f"Error en la configuración del modelo de sustrato {i+1}: {ve}")

        params = [param.strip() for param in params_str.split(',')]
        lower_bounds, upper_bounds = parse_bounds(bounds_str, len(params))

        try:
            y_pred = main_model.fit_model(
                'substrate', time, substrate_data,
                bounds=(lower_bounds, upper_bounds)
            )
            substrate_results.append({
                'y_pred': y_pred.tolist(),  # Convertir a lista para serialización
                'equation': equation,
                'params': main_model.params['substrate']
            })
        except Exception as e:
            raise RuntimeError(f"Error al ajustar el modelo de sustrato {i+1}: {e}")

    # Ajusta los modelos de Producto
    for i in range(len(product_eqs)):
        equation = product_eqs[i]
        params_str = product_params[i]
        bounds_str = product_bounds[i]

        try:
            main_model.set_model_product(equation, params_str)
        except ValueError as ve:
            raise ValueError(f"Error en la configuración del modelo de producto {i+1}: {ve}")

        params = [param.strip() for param in params_str.split(',')]
        lower_bounds, upper_bounds = parse_bounds(bounds_str, len(params))

        try:
            y_pred = main_model.fit_model(
                'product', time, product_data,
                bounds=(lower_bounds, upper_bounds)
            )
            product_results.append({
                'y_pred': y_pred.tolist(),  # Convertir a lista para serialización
                'equation': equation,
                'params': main_model.params['product']
            })
        except Exception as e:
            raise RuntimeError(f"Error al ajustar el modelo de producto {i+1}: {e}")

    # Genera las gráficas
    fig, axs = plt.subplots(3, 1, figsize=(10, 15))

    # Gráfica de Biomasa
    axs[0].plot(time, biomass_data, 'o', label='Datos de Biomasa')
    for i, result in enumerate(biomass_results):
        axs[0].plot(time, result['y_pred'], '-', label=f'Modelo de Biomasa {i+1}')
    axs[0].set_xlabel('Tiempo')
    axs[0].set_ylabel('Biomasa')
    if show_legend:
        axs[0].legend(loc=legend_position)

    # Gráfica de Sustrato
    axs[1].plot(time, substrate_data, 'o', label='Datos de Sustrato')
    for i, result in enumerate(substrate_results):
        axs[1].plot(time, result['y_pred'], '-', label=f'Modelo de Sustrato {i+1}')
    axs[1].set_xlabel('Tiempo')
    axs[1].set_ylabel('Sustrato')
    if show_legend:
        axs[1].legend(loc=legend_position)

    # Gráfica de Producto
    axs[2].plot(time, product_data, 'o', label='Datos de Producto')
    for i, result in enumerate(product_results):
        axs[2].plot(time, result['y_pred'], '-', label=f'Modelo de Producto {i+1}')
    axs[2].set_xlabel('Tiempo')
    axs[2].set_ylabel('Producto')
    if show_legend:
        axs[2].legend(loc=legend_position)

    plt.tight_layout()
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    image = Image.open(buf)

    # Construcción del prompt para Yi-Coder
    prompt = f"""
Analiza estos resultados de cinética y parámetros de bioprocesos, como experto:

Biomasa:
{json.dumps(biomass_results, indent=2)}

Sustrato:
{json.dumps(substrate_results, indent=2)}

Producto:
{json.dumps(product_results, indent=2)}

Para cada cinética (Biomasa, Sustrato, Producto):
1. Evalúa el modelo (ajuste, parámetros)
2. Identifica problemas específicos
3. Sugiere mejoras concretas

Concluye con un veredicto general sobre la calidad del modelado.
"""

    # Generar el análisis utilizando Yi-Coder
    analysis = generate_analysis(prompt, max_length=550)  # Reducido a 100

    return image, analysis