# app.py

import gradio as gr
from utils import initialize_gmm, generate_grid, generate_contours, generate_intermediate_points, plot_samples_and_contours
import matplotlib.pyplot as plt
import torch
import numpy as np

def validate_inputs(mu_list, Sigma_list, pi_list):
    try:
        mu = eval(mu_list)
        Sigma = eval(Sigma_list)
        pi = eval(pi_list)
        
        if not (isinstance(mu, list) and all(isinstance(i, list) for i in mu)):
            return False, "Mu list is invalid."
        if not (isinstance(Sigma, list) and all(isinstance(i, list) for i in Sigma)):
            return False, "Sigma list is invalid."
        if not isinstance(pi, list):
            return False, "Pi list is invalid."
        
        if not torch.isclose(torch.tensor(pi).sum(), torch.tensor(1.0)):
            return False, "Mixture weights must sum to 1."
        
        return True, ""
    except Exception as e:
        return False, str(e)

def visualize_gmm(mu_list, Sigma_list, pi_list, dx, dtheta, T, N):
    is_valid, error_message = validate_inputs(mu_list, Sigma_list, pi_list)
    if not is_valid:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, f'Invalid input: {error_message}', horizontalalignment='center', verticalalignment='center')
        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_aspect('equal', adjustable='box')
        plt.close(fig)
        return fig, fig

    try:
        gmm = initialize_gmm(eval(mu_list), eval(Sigma_list), eval(pi_list))
        grid_points = generate_grid(dx)
        std_normal_contours = generate_contours(dtheta)
        gmm_samples = gmm.sample(500)
        normal_samples = torch.distributions.MultivariateNormal(torch.zeros(2), torch.eye(2)).sample((500,))
        (intermediate_points_gmm_to_normal, contour_intermediate_points_gmm_to_normal, grid_intermediate_points_gmm_to_normal,
            intermediate_points_normal_to_gmm, contour_intermediate_points_normal_to_gmm, grid_intermediate_points_normal_to_gmm) = \
                    generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, normal_samples, T, N)
        
        final_frame_gmm_to_normal = intermediate_points_gmm_to_normal.cpu().detach().numpy()
        final_frame_normal_to_gmm = intermediate_points_normal_to_gmm.cpu().detach().numpy()

        fig1, ax1 = plot_samples_and_contours(final_frame_gmm_to_normal, contour_intermediate_points_gmm_to_normal.cpu().detach().numpy(), grid_intermediate_points_gmm_to_normal.cpu().detach().numpy(), "GMM to Normal Final Frame")
        fig2, ax2 = plot_samples_and_contours(final_frame_normal_to_gmm, contour_intermediate_points_normal_to_gmm.cpu().detach().numpy(), grid_intermediate_points_normal_to_gmm.cpu().detach().numpy(), "Normal to GMM Final Frame")
        
        return fig1, fig2
    except Exception as e:
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, f'Error during visualization: {str(e)}', horizontalalignment='center', verticalalignment='center')
        ax.set_xlim(-5, 5)
        ax.set_ylim(-5, 5)
        ax.set_aspect('equal', adjustable='box')
        plt.close(fig)
        return fig, fig

demo = gr.Interface(
    fn=visualize_gmm,
    inputs=[
        gr.Textbox(label="Mu List", value="[[2, 1], [-1, -2], [3, -2]]", placeholder="Enter means as a list of lists, e.g., [[0,0], [1,1]]"),
        gr.Textbox(label="Sigma List", value="[[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]], [[0.05, 0.0], [0.0, 0.05]]]", placeholder="Enter covariances as a list of lists, e.g., [[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]]]"),
        gr.Textbox(label="Pi List", value="[0.05, 0.8, 0.15]", placeholder="Enter weights as a list, e.g., [0.5, 0.5]"),
        gr.Slider(minimum=0.01, maximum=1.0, label="dx", value=0.1),
        gr.Slider(minimum=2*np.pi/3600, maximum=2*np.pi/36, label="dtheta", value=2*np.pi/360),
        gr.Slider(minimum=1, maximum=100, label="T", value=10),
        gr.Slider(minimum=1, maximum=500, label="N", value=100)
    ],
    outputs=[
        gr.Plot(label="GMM to Normal Flow Final Frame"),
        gr.Plot(label="Normal to GMM Flow Final Frame")
    ],
    live=True
)

demo.launch()