import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
import numpy as np

# Helper functions to visualize model outputs
def plot_validation_vs_training_accuracy(epoch, train_acc, val_acc_softmax, val_acc_bayes):
    r"""
    Inputs:
        - **epoch**: int number of epochs that was chosen during training
        - **train_acc**: list of float
        - **val_acc_softmax**: list of float 
        - **val_acc_bayes**: list of float 
    """
    x_range = np.linspace(0, epoch)

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=x_range,
            y=train_acc,
            mode='lines',
            name='training accuracy softmax'
        )
    )

    fig.add_trace(
        go.Scatter(
            x=x_range,
            y=val_acc_softmax,
            mode='lines',
            name='validation accuracy softmax'
        )
    )

    fig.add_trace(
        go.Scatter(
            x=x_range,
            y=val_acc_bayes,
            mode='lines',
            name='validation accuracy bayes'))

    fig.show()

def plot_rbfcenters_embeddings(code, labels, mu, gamma):
    r"""
    Inputs:
        - **code**: float values that are feed to the RBF layer, also known as embeddings or latent space
        - **labels**: float values of the predicted labels generated by the model
        - **mu**: float values of the centers for the RBF neurons
        - **gamma**: float values of gamma for the RBF neurons 
    """
    fig = go.Figure()
    fig.add_trace( 
        go.Scatter3d(
            x=code[0,0,:,:].flatten(),
            y=code[0,1,:,:].flatten(), 
            z=code[0,2,:,:].flatten(),
            name='class',
            mode='markers',
            marker=dict(
                size=0.9, 
                color=labels[0].flatten(),
                colorscale='jet'
            ) 
        )
    )
    
    fig.add_trace(
        go.Scatter3d(
            x=mu[:,0],
            y=mu[:,1], 
            z=mu[:,2],
            name='centers',
            mode='markers',
            marker=dict(
                #color=gamma,
                opacity=1,
                #colorscale='jet'
            ) 
        )
    )
    
    fig.update_layout(template="plotly_dark", title="Embeddings and Rbf nodes")
    
    fig.show()