import numpy as np
import matplotlib.pyplot as plt
from scipy import linalg
from sklearn.datasets import make_sparse_spd_matrix
from sklearn.covariance import GraphicalLassoCV, 
_wolf

import gradio as gr

prng = np.random.RandomState(1)

def get_precision_matrix(alpha = 0.98, smallest_coef = 0.4, largest_coef = 0.7):
    
    prec = make_sparse_spd_matrix(
        n_features, alpha=alpha, smallest_coef=smallest_coef, largest_coef=largest_coef, random_state=prng
    )

    return prec

def get_covariance_matrix(precision_matrix):

    return linalg.inv(precision_matrix)

def scaled_covariance_matrix(precision_matrix):

    covariance_matrix = get_covariance_matrix(precision_matrix)
    d = np.sqrt(np.diag(covariance_matrix))
    scaled_covariance_matrix = covariance_matrix / d
    scaled_covariance_matrix /= d[:, np.newaxis]

    return scaled_covariance_matrix

def scaled_precision_matrix(precision_matrix):

    covariance_matrix = get_covariance_matrix(precision_matrix)
    d = np.sqrt(np.diag(covariance_matrix))
    scaled_precision_matrix = precision_matrix * d
    scaled_precision_matrix *= d[:, np.newaxis]

    return scaled_precision_matrix


def get_samples(n_features, n_samples, scaled_covariance_matrix):
    
    X = prng.multivariate_normal(np.zeros(n_features), cov, size=n_samples)
    X -= X.mean(axis=0)
    X /= X.std(axis=0)

    return X

def get_empirical_covariance(X, n_samples):

    return np.dot(X.T, X) / n_samples

def estimate_covariance_lasso(X):

    model = GraphicalLassoCV()
    model.fit(X)
    return model.covariance_

def estimate_precision_lasso(X):
    
    model = GraphicalLassoCV()
    model.fit(X)
    return model.precision_

def estimate_covariance_leidotwolf(X):
    
    lw_cov_, _ = ledoit_wolf(X)

    return lw_cov_

def estimate_precision_leidotwolf(leidot_cov):

    return linalg.inv(leidot_cov)


# main function that will be called in the block
def compute_and_plot(alpha = 0.98, smallest_coef = 0.4, largest_coef = 0.7,
                    n_features = 20, n_samples = 60, measure = None, model = None):

    prec = get_precision_matrix(alpha = alpha, smallest_coef = smallest_coef, largest_coef = largest_coef)
    prec = scaled_precision_matrix(prec)
    cov = scaled_covariance_matrix(prec)
    X = get_samples(n_features, n_samples, cov)
                        
    if measure == 'covariance':
        if model == 'empirical':
            emp_cov = get_empirical_covariance(X, n_samples)
            fig, ax = plt.subplots()
            ax.imshow(emp_cov, interpolation="nearest", cmap=plt.cm.RdBu_r)
        elif model == 'lasso':
            lasso_cov = estimate_covariance_lasso(X)
            fig, ax = plt.subplots()
            ax.imshow(lasso_cov, interpolation="nearest", cmap=plt.cm.RdBu_r)
        elif model == 'leidot-wolf':
            lw_cov = estimate_covariance_leidotwolf(X)
            fig, ax = plt.subplots()
            ax.imshow(lw_cov, interpolation="nearest", cmap=plt.cm.RdBu_r)
        else:
            print('invalid')
    # elif measure == 'precision':

    # else:
    #     # TO DO: add empty plot
    #     print('invalid')
    
    
    #lasso_prec = estimate_precision_lasso(X)
    
    #lw_prec = estimate_precision_leidotwolf(leidot_cov)

    return fig
    

    
def iter_grid(n_rows, n_cols):
    # create a grid using gradio Block
    for _ in range(n_rows):
        with gr.Row():
            for _ in range(n_cols):
                with gr.Column():
                    yield

title = "Sparse inverse covariance estimation"
with gr.Blocks(title=title) as demo:
    gr.Markdown(f"## {title}")
    gr.Markdown("Estimating covariance and sparse precision from a small number of samples using GraphicalLasso and Ledoit-Wolf algorithms.")
    n_samples = gr.Slider(minimum=20, maximum=100, step=5, 
    label = "Number of Samples")
    n_features = gr.Slider(minimum=10, maximum=100, step=5, 
    label = "Number of features")
    alpha = gr.Slider(minimum=0, maximum=1, step=0.1, 
    label = "sparsity coefficient (alpha)")
    smallest_coef = gr.Slider(minimum=0, maximum=1, step=0.1, 
    label = "minimum correlation value")
    largest_coef = gr.Slider(minimum=0, maximum=1, step=0.1, 
    label = "maximum correlation value")

    models = ['empirical', 'lasso', 'leidot-wolf']
    model_counter = 0
    for _ in iter_grid(1, 3):

        model = models[model_counter]
        plot = gr.Plot(label=input_model)
        n_samples.change(fn=compute_and_plot, inputs=[0.98, 0.4, 0.7, 20, 60, 'covariance', model], outputs=plot)
        n_features.change(fn=compute_and_plot, inputs=[0.98, 0.4, 0.7, 20, 60, 'covariance', model], outputs=plot)
        alpha.change(fn=compute_and_plot, inputs=[0.98, 0.4, 0.7, 20, 60, 'covariance', model], outputs=plot)
        smallest_coef.change(fn=compute_and_plot, inputs=[0.98, 0.4, 0.7, 20, 60, 'covariance', model], outputs=plot)
        largest_coef.change(fn=compute_and_plot, inputs=[0.98, 0.4, 0.7, 20, 60, 'covariance', model], outputs=plot)
    
    
demo.launch()