import gradio as gr
import numpy as np
import matplotlib as mpl
mpl.use('agg')
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from umap import UMAP
import plotly.express as px
import pandas as pd

class recon_encoder(nn.Module):

    def __init__(self, latent_size, nconv=16, pool=4, drop=0.05):
        super(recon_encoder, self).__init__()


        self.encoder = nn.Sequential( # Appears sequential has similar functionality as TF avoiding need for separate model definition and activ
          nn.Conv2d(in_channels=1, out_channels=nconv, kernel_size=3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv, nconv, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          nn.Conv2d(nconv, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),          
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          nn.Conv2d(nconv*2, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),     
          nn.Dropout(drop),     
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          #nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),     
          #nn.Dropout(drop),     
          #nn.ReLU(),
          #nn.MaxPool2d((pool,pool)),
        )


        self.bottleneck = nn.Sequential(
          # FC layer at bottleneck -- dropout might not make sense here
          nn.Flatten(),
          nn.Linear(1024, latent_size),
          #nn.Dropout(drop),
          nn.ReLU(),
#          nn.Linear(latent_size, 1024),
#          #nn.Dropout(drop),
#          nn.ReLU(),
#          nn.Unflatten(1,(64,4,4))# 0 is batch dimension
          )


        self.decoder1 = nn.Sequential(

          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*4, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),
            
          #nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*2, 1, 3, stride=1, padding=(1,1)), #Output conv layer has 2 for mu and sigma
          nn.Sigmoid() #Amplitude mode
          )
    

    def forward(self,x):
        with torch.cuda.amp.autocast():
            x1 = self.encoder(x)
            x1 = self.bottleneck(x1)
            #print(x1.shape)
            return x1


    #Helper function to calculate size of flattened array from conv layer shapes    
    def calc_fc_shape(self):
        x0 = torch.zeros([256,256]).unsqueeze(0)
        x0 = self.encoder(x0)

        self.conv_bock_output_shape = x0.shape
        #print ("Output of conv block shape is", self.conv_bock_output_shape)
        self.flattened_size = x0.flatten().shape[0]
        #print ("Flattened layer size is", self.flattened_size)
        return self.flattened_size

class recon_model(nn.Module):

    def __init__(self, latent_size, nconv=16, pool=4, drop=0.05):
        super(recon_model, self).__init__()


        self.encoder = nn.Sequential( # Appears sequential has similar functionality as TF avoiding need for separate model definition and activ
          nn.Conv2d(in_channels=1, out_channels=nconv, kernel_size=3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv, nconv, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          nn.Conv2d(nconv, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),          
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          nn.Conv2d(nconv*2, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),     
          nn.Dropout(drop),     
          nn.ReLU(),
          nn.MaxPool2d((pool,pool)),

          #nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),     
          #nn.Dropout(drop),     
          #nn.ReLU(),
          #nn.MaxPool2d((pool,pool)),
        )


        self.bottleneck = nn.Sequential(
          # FC layer at bottleneck -- dropout might not make sense here
          nn.Flatten(),
          nn.Linear(1024, latent_size),
          #nn.Dropout(drop),
          nn.ReLU(),
          nn.Linear(latent_size, 1024),
          #nn.Dropout(drop),
          nn.ReLU(),
          nn.Unflatten(1,(64,4,4))# 0 is batch dimension
          )


        self.decoder1 = nn.Sequential(

          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*4, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          nn.Dropout(drop),
          nn.ReLU(),
          nn.Upsample(scale_factor=pool, mode='bilinear'),
            
          #nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)),
          #nn.Dropout(drop),
          #nn.ReLU(),
          #nn.Upsample(scale_factor=pool, mode='bilinear'),

          nn.Conv2d(nconv*2, 1, 3, stride=1, padding=(1,1)), #Output conv layer has 2 for mu and sigma
          nn.Sigmoid() #Amplitude mode
          )
    

    def forward(self,x):
        with torch.cuda.amp.autocast():
            x1 = self.encoder(x)
            x1 = self.bottleneck(x1)
            #print(x1.shape)
            return self.decoder1(x1)


    #Helper function to calculate size of flattened array from conv layer shapes    
    def calc_fc_shape(self):
        x0 = torch.zeros([256,256]).unsqueeze(0)
        x0 = self.encoder(x0)

        self.conv_bock_output_shape = x0.shape
        #print ("Output of conv block shape is", self.conv_bock_output_shape)
        self.flattened_size = x0.flatten().shape[0]
        #print ("Flattened layer size is", self.flattened_size)
        return self.flattened_size

full_model = torch.load('betst_model_100x_0064.pth',map_location=torch.device('cpu'))
encoder_model = recon_encoder(latent_size=64)
encoder_state_dict = encoder_model.state_dict()

checkpoint = torch.load('betst_model_100x_0064_statedict.pth',map_location=torch.device('cpu'))
pretrained_dict = {k: v for k, v in checkpoint.items() if k in encoder_state_dict}

encoder_model.load_state_dict(pretrained_dict)
#
#all_data = np.load('E031_256.npy').astype(np.float32)
#all_data = all_data.reshape(-1,1,256,256)
#dataloader = DataLoader(all_data,batch_size=32,shuffle=False)

def load_data(file):
    all_data = np.load(file.name).astype(np.float32)
    all_data = all_data.reshape(-1,1,256,256)
    dataloader = DataLoader(all_data,batch_size=32,shuffle=False)
    return all_data, dataloader, 'upload complete: {}'.format(all_data.shape)



def show_image(selection, all_data):
    fig1, ax1 = plt.subplots()
    ax1.imshow(all_data[selection][0],plt.cm.inferno,origin='lower')
    ax1.axis('off')
    fig1.tight_layout()
    
    fig2, ax2 = plt.subplots()
    prediction = full_model(torch.tensor(all_data[selection].reshape(-1,1,256,256))).detach().cpu().numpy()
    ax2.imshow(prediction[0,0],plt.cm.inferno,origin='lower')
    ax2.axis('off')
    fig2.tight_layout()
    
    return fig1, fig2

def encode_data(dataloader):
    preds_full = []
    preds_enc = []
    

    for i, images in enumerate(dataloader):
        if i > 5:
            break
        pred_full = full_model(images)
        pred_enc = encoder_model(images)
        for j in range(images.shape[0]):
            preds_full.append(pred_full[j].detach().cpu().numpy())
            preds_enc.append(pred_enc[j].detach().cpu().numpy())

    
    processed_images = np.array(preds_full).squeeze()
    encoded_images = np.array(preds_enc)
    message = 'finished'
    
    return message, processed_images, encoded_images

def print_state(state):
    return state.shape

def latent_vis(encoded_data,decomp_method,clustering_method,cluster_number,all_data):
    if decomp_method == 'PCA':
        pca = PCA(n_components=2)
        decomp = pca.fit_transform(encoded_data)
    elif decomp_method == 'tSNE':
        tsne = TSNE(n_components=2)
        decomp = tsne.fit_transform(encoded_data)
    elif decomp_method == 'UMAP':
        reducer = UMAP()
        decomp = reducer.fit_transform(encoded_data)
        
    if clustering_method == 'KMeans':
        kmeans = KMeans(n_clusters=int(cluster_number))
        cluster_labels = kmeans.fit_predict(encoded_data)
        
    
    df = pd.DataFrame(decomp,columns=['x','y'])
    df['cluster'] = cluster_labels
    df['value'] = np.ones_like(cluster_labels) * np.arange(len(decomp))
    
    fig = px.scatter(df,x='x',y='y',color='cluster',color_continuous_scale='viridis',hover_name='value',hover_data={'x': False,
                                                                                                                    'y': False,
                                                                                                                    'cluster': False,
                                                                                                                   'value': False})
#    fig = px.scatter(x=decomp[:,0],y=decomp[:,1],color=clusters,hover_data=np.arange(len(decomp)))
    fig.update_layout(clickmode='event+select')
    fig.update_traces(marker=dict(size=12),
                      selector=dict(mode='markers'))
    
    fig1 = plt.figure(figsize=(20,5))
    n_rows = 1
    n_cols = int(cluster_number)
    colors = plt.cm.viridis(np.linspace(0,1,len(np.unique(cluster_labels))))
    
    for i in np.unique(cluster_labels):
        ind = np.where(cluster_labels[:] == i)[0]
        #ax.scatter(decomp[cluster_labels[:] == i,0],decomp[cluster_labels[:] == i,1],color=colors[i],label='class {}'.format(i))
        
        r = np.random.choice(ind)
        ax1 = fig1.add_subplot(n_rows,n_cols,i+1)
        ax1.imshow(all_data[r][0],plt.cm.inferno,origin='lower')
        ax1.set_title('Class {}: {}'.format(i,len(ind)),color=colors[i],fontsize=20,weight='bold')
        
    #ax.legend()
    
    #fig.tight_layout()
    fig1.tight_layout()
    return decomp, cluster_labels, fig, fig1

def interactive_vis(decomp,clusters,images):
    df = pd.DataFrame(decomp,columns=['x','y'])
    df['cluster'] = clusters
    df['value'] = np.ones_like(clusters) * np.arange(len(decomp))
    df['im'] = images
    
    fig = px.scatter(df,x='x',y='y',color='cluster',custom_data='im',color_continuous_scale='viridis',hover_name='value',hover_data={'x': False,
                                                                                                                    'y': False,
                                                                                                                    'cluster': False,
                                                                                                                   'value': False})
#    fig = px.scatter(x=decomp[:,0],y=decomp[:,1],color=clusters,hover_data=np.arange(len(decomp)))
    fig.update_layout(clickmode='event+select')
    fig.update_traces(marker=dict(size=20),
                      selector=dict(mode='markers'))
    
    return fig

def neighbor_vis(decomp,neighbor_index,n_neighbors,all_data):
    neighbor_index = int(neighbor_index)
    
    d = np.sqrt((decomp[:,0] - decomp[neighbor_index,0]) ** 2 + (decomp[:,1] - decomp[neighbor_index,1]) ** 2)
    ar = np.argsort(d)
    
    n_rows = int(np.ceil(n_neighbors/5))
    n_cols = 5
    fig = plt.figure(figsize=(20,5*n_rows))
    
    n = 1
    ax = fig.add_subplot(n_rows,n_cols,n)
    ax.imshow(all_data[neighbor_index][0],plt.cm.inferno,origin='lower')
    ax.set_title('{}'.format(neighbor_index),fontsize=20,weight='bold')
    ax.axis('off')
    n += 1
    
    neighbors = ar[1:1+n_neighbors-1]
    
    for i in neighbors:
        ax = fig.add_subplot(n_rows,n_cols,n)
        ax.imshow(all_data[i][0],plt.cm.inferno,origin='lower')
        ax.set_title('{}'.format(i),fontsize=20)
        ax.axis('off')
        n += 1
        
    return fig
    

intro_text1 = '# AI-NERD: Artificial Intelligence for Non-Equilibrium Relaxation Dynamics'
intro_text2 = 'AI-NERD is a platform for applying unsupervised image classification to X-ray Photon Corrleation Spectroscopy (XPCS) data. Here, we demonstrate how raw experimental data can be automatically processed and clustered, and how latent space analysis can be used to understand the physics of relaxing systems without any background information or assumptions.<br><br>Please see out [preprint](https://arxiv.org/abs/2212.03984) for more information.<br><br>'
l = 900
with gr.Blocks() as demo:
    gr.Markdown(intro_text1)
    gr.Markdown(intro_text2)
    
    gr.Markdown('### Evaluation of Training Results')
    gr.Markdown('Use the dropdown menu below to select a sample image.  The frame on the left will show the raw C2 data, and the frame on the right will show the neural network output.  After sampling individual images, click _Process All Images_ to run the entire dataset through the Autoencoder')
    
    with gr.Row():
        file_path = gr.File()
        with gr.Column():
            upload_status = gr.Textbox(label='file upload status')
            file_upload = gr.Button(value='load data')     
            
    all_data = gr.State()
    dataloader = gr.State()
    file_upload.click(load_data,file_path,[all_data,dataloader,upload_status])
    
    selection = gr.Dropdown(list(np.arange(2000)),value=200,label='select sample image')
    with gr.Row():
        output_image_1 = gr.Plot(label='input C2 data')
        output_image_2 = gr.Plot(label='Autoencoder Reproduction')
    
    selection.change(show_image,[selection, all_data],[output_image_1,output_image_2])
    
    with gr.Row():
        process_all = gr.Button(value='Process All Images')
        status = gr.Textbox(label='batch processing status')

    
    proc_im = gr.State()
    enc_im = gr.State()
    process_all.click(encode_data,inputs=[dataloader],outputs=[status,proc_im,enc_im],show_progress=True,status_tracker=None)
    
#    check_type = gr.Button(value='check state info')
#    check_stat = gr.Textbox()
#    check_type.click(print_state,inputs=proc_im,outputs=check_stat)
    gr.Markdown('<br><br>') 
    gr.Markdown('### Latent Space Visualization')
    gr.Markdown('Select the decomposition and clustering method for latent space visualization')
    with gr.Row():
        with gr.Column():
            decomp_method = gr.Dropdown(choices=['PCA','tSNE','UMAP'],label='select decomposition method',value='UMAP')
            with gr.Row():
                clustering_method = gr.Dropdown(choices=['KMeans','Agglomerative','DBSCAN'],label='select clusterting algorithm',value='KMeans')
                cluster_number = gr.Number(label='input number of clusters',value=5)
                
            process_vis = gr.Button(value='Visualize Latent Space')
        latent_scatter = gr.Plot()
    latent_sample = gr.Plot()
        
    save_decomp_coords = gr.State()
    save_cluster_labels = gr.State()
    process_vis.click(latent_vis,[enc_im,decomp_method,clustering_method,cluster_number,all_data],[save_decomp_coords,save_cluster_labels,latent_scatter,latent_sample])
    
    gr.Markdown('<br><br><br>')
    gr.Markdown('### Visualize Nearest Neighbors')
    gr.Markdown('Hover over data points in the scatter plot above, to identify the index of points of interest.  Enter the desired index in the box below, and click _Visualize Neighbors_.')
    
    with gr.Row():
        with gr.Column():
            neighbor_index = gr.Number(label='input point index',value=110)
            n_neighbors = gr.Slider(label='select number of neighbors to view',minimum=5,maximum=10,value=5,step=1)
            
        neighbor_button = gr.Button(value='Visualize Neighbors')
    
    neighbor_plot = gr.Plot()
    neighbor_button.click(neighbor_vis,[save_decomp_coords,neighbor_index,n_neighbors,all_data],neighbor_plot)
    #neighbor_button.click(interactive_vis,[save_decomp_coords,save_cluster_labels,proc_im],interactive_plot)
    
    
    
demo.launch()