# from neuroxlink.src.neuroxlink import NeuroxLink
from neuroxlink import NeuroxLink

import sys
import os
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.io as pio
import numpy as np
pio.renderers.default = "plotly_mimetype"

nlx = NeuroxLink(cdn_url="https://cdn.neurolibre.org")

# Define and create a list of DOIs and figures
doi1 = "10.55458/neurolibre.alberta1"
doi2 = "10.55458/neurolibre.alberta2"
doi3 = "10.55458/neurolibre.alberta3"
doi4 = "10.55458/neurolibre.alberta4"
doi5 = "10.55458/neurolibre.alberta5"
dois = [doi1, doi2, doi3, doi4, doi5]
figures = ['fig1', 'fig2', 'fig3']
nlx.import_papers(dois)

paper = nlx.papers[0]
print("\n------------> AUTHORS AND AFFILIATIONS")
print(paper.get_authors(), "\n------------> DEPENDENCIES")
print(paper.get_dependencies(), "\n------------> HEADINGS")
print(paper.headings)

paper.inspect_plotly_figures()

fig = paper.create_plotly_object_from('fig0')
fig.show()

nlx.papers[doi2].create_plotly_object_from('fig0').show()

nlx.papers[doi3].create_plotly_object_from('fig0').show()

combined_df = pd.DataFrame()

for doi in dois:
    data = {}
    
    # Loop through each figure for the current DOI
    # Get the data for the current DOI and figure
    fig = "fig0"
    fig_data = nlx.papers[doi].get_plotly_data(fig, select_trace_type="scatter", select_trace_mode="markers")

    # Add data for each figure, ensure the column name is unique
    data['y'] = fig_data['y']
    data['x'] = fig_data['x']  # Use the figure-specific column name
        
    # Convert the collected data into a DataFrame
    df = pd.DataFrame(data)
    
    # Add the DOI to each row in the DataFrame
    df['doi'] = doi
    
    # Append to the combined DataFrame
    combined_df = pd.concat([combined_df, df], ignore_index=True)

combined_df.head()

fig = px.scatter(combined_df,x="x",y="y",color="doi",template="plotly_dark")
fig.update_layout(
    width=800,
    height=800,
    title = "Putting together AI from 5 articles",
    xaxis=dict(scaleanchor="y", scaleratio=1,constrain="domain"),  # Link x and y axis scaling
    yaxis=dict(scaleanchor="x", scaleratio=1,constrain="domain"),
    margin=dict(l=0, r=0, t=50, b=0),
        legend=dict(
        orientation='v',  # Vertical orientation
        yanchor='middle',  # Centered vertically
        xanchor='right',   # Right-aligned
        x=1.4,             # Position just outside the right side
        y=0.5              # Centered on the y-axis
    )
)
fig.update_traces(marker=dict(size=12, opacity=0.9))
fig.show()

# Define and create a list of DOIs and figures
doi1 = "10.55458/neurolibre.alberta1"
doi2 = "10.55458/neurolibre.alberta2"
doi3 = "10.55458/neurolibre.alberta3"
doi4 = "10.55458/neurolibre.alberta4"
doi5 = "10.55458/neurolibre.alberta5"
dois = [doi1, doi2, doi3, doi4, doi5]
figures = ['fig1', 'fig2', 'fig3']

nlx.import_papers([doi1,doi2,doi3,doi4,doi5])

nlx.papers[doi2].inspect_plotly_figures()

nlx.papers["10.55458/neurolibre.alberta2"].create_plotly_object_from('fig1')
#nlx.papers["10.55458/neurolibre.alberta2"].create_plotly_object_from('fig1', select_trace_type="histogram")
#nlx.papers["10.55458/neurolibre.alberta2"].create_plotly_object_from('fig1', select_trace_type="box")
#nlx.papers["10.55458/neurolibre.alberta2"].create_plotly_object_from('fig1', select_trace_type="scatter")
#nlx.papers["10.55458/neurolibre.alberta2"].create_plotly_object_from('fig1', select_trace_type="scatter",select_trace_mode="markers")

data = []

for doi in dois:
    fig_data = []
    for fig in figures:
        print(f"Fetching data: {fig} from {doi}")
        fig_data.append(nlx.papers[doi].create_plotly_object_from(fig,select_trace_type="scatter",select_trace_mode="markers"))
    data.append(fig_data)
print("done")    

fig = make_subplots(rows=5, cols=3, subplot_titles=("ABV", "ABD", "ABE") * 5)

for row in range(5):
    fig.update_yaxes(title_text=f"Score", row=row+1, col=1)
    fig.update_yaxes(title_text=f"Paper {row+1}", row=row+1, col=3, side="right")
    for col in range(3):
        subplot_figure = data[row][col]
        for trace in subplot_figure.data:
            fig.add_trace(trace, row=row + 1, col=col + 1)
            
# Update layout
fig.update_layout(title_text="Mosaic Plot of 15 figures from 5 articles",
                  height=800, width=800,
                  showlegend=True,
                  template="plotly_dark")

# Show the figure
fig.show()

column_names = {'fig1': 'vol', 'fig2': 'dens', 'fig3': 'eff'}

# Initialize an empty DataFrame to store the combined data
combined_df = pd.DataFrame()

# Loop through each DOI and figure
for doi in dois:
    data = {}
    for fig in figures:
        # Get the data for the current DOI and figure
        fig_data = nlx.papers[doi].get_plotly_data(fig,select_trace_type="scatter",select_trace_mode="markers")
        
        #data['score'] = fig_data['x']
        data[column_names[fig]] = fig_data['y']
        
    # Create a DataFrame for the current DOI with DOI as a column
    df = pd.DataFrame(data)
    df['DOI'] = doi  # Add the DOI to each row
    
    # Append to the combined DataFrame
    combined_df = pd.concat([combined_df, df], ignore_index=True)

print(combined_df)

doi_names = {'10.55458/neurolibre.alberta1': 'Hippocampus', '10.55458/neurolibre.alberta2': 'Amygdala', '10.55458/neurolibre.alberta3': 'Temporal',"10.55458/neurolibre.alberta4":"Parietal","10.55458/neurolibre.alberta5":"Prefrontal"}

# Initialize an empty DataFrame for combining all DOIs' data
combined_df = pd.DataFrame()

# Loop through each DOI
for doi in dois:
    data = {}
    
    # Loop through each figure for the current DOI
    for fig in figures:
        # Get the data for the current DOI and figure
        fig_data = nlx.papers[doi].get_plotly_data(fig, select_trace_type="scatter", select_trace_mode="markers")
        
        # Add data for each figure, ensure the column name is unique
        data['score'] = fig_data['y']
        data[column_names[fig]] = fig_data['x']  # Use the figure-specific column name
        
    # Convert the collected data into a DataFrame
    df = pd.DataFrame(data)
    
    # Add the DOI to each row in the DataFrame
    df['Region'] = doi_names[doi]
    
    # Append to the combined DataFrame
    combined_df = pd.concat([combined_df, df], ignore_index=True)

# Print or check the combined DataFrame
combined_df.head()

# numeric_cols = combined_df.select_dtypes(include=[np.number]).columns
# df_zscore[numeric_cols] = (combined_df[numeric_cols] - combined_df[numeric_cols].mean())/combined_df[numeric_cols].std()
# df_zscore["DOI"] = combined_df["DOI"]

fig = px.scatter_matrix(combined_df,
    dimensions=["vol", "dens", "eff","score"],
    color="Region", template="plotly_dark")
fig.update_traces(diagonal_visible=False)
fig.update_layout(margin=dict(l=0,r=0,t=0,b=0))
fig.show()

features = ['score', 'vol', 'dens', 'eff']
X = combined_df[features].values
X_mean = np.mean(X, axis=0)
X_std = X - X_mean
cov_matrix = np.cov(X_std.T)
eigenvalues, eigenvectors = np.linalg.eig(cov_matrix)
sorted_indices = np.argsort(eigenvalues)[::-1]
sorted_eigenvalues = eigenvalues[sorted_indices]
sorted_eigenvectors = eigenvectors[:, sorted_indices]

top_2_eigenvectors = sorted_eigenvectors[:, :2]
X_pca = X_std.dot(top_2_eigenvectors)

pca_df = pd.DataFrame(data=X_pca, columns=['PCA1', 'PCA2'])
pca_df['Region'] = combined_df['Region']

fig = px.scatter(pca_df, x='PCA1', y='PCA2', color='Region', title='PCA of Regions based on Features', template="plotly_dark")
fig.update_traces(marker=dict(size=15, opacity=0.7))
fig.update_layout(margin=dict(l=0,r=0,t=0,b=0))
fig.show()

# If you don't believe me that the variability is coming from volume and score...

# Aliens are soo linear, omg. Just stick with the ones with bigger heads :) 

fig = px.scatter(combined_df,x="vol",y="score",color="Region",template="plotly_dark")
fig.update_traces(marker=dict(size=15, opacity=0.7))
fig.update_layout(margin=dict(l=0,r=0,t=0,b=0))
fig.show()

score_corr_pearson = combined_df.corr(numeric_only=True, method="pearson")
score_corr_kendall = combined_df.corr(numeric_only=True, method="kendall")
score_corr_spearman = combined_df.corr(numeric_only=True, method="spearman")

# Create a list of the correlation matrices and corresponding labels
corr_matrices = [score_corr_pearson, score_corr_kendall, score_corr_spearman]
methods = ['Pearson', 'Kendall', 'Spearman']

# Initialize the first heatmap (Pearson by default)
fig = px.imshow(corr_matrices[0],
                color_continuous_scale='Viridis',
                zmin=-0.5, zmax=0.9, 
                title=f'{methods[0]} Correlation Matrix', template="ggplot2")

# Update the layout to add the slider
fig.update_layout(
    updatemenus=[
        dict(
            type="buttons",
            direction="down",
            buttons=[
                dict(
                    args=[{"z": [corr_matrices[0].values]}],
                    label="Pearson",
                    method="restyle"
                ),
                dict(
                    args=[{"z": [corr_matrices[1].values]}],
                    label="Kendall",
                    method="restyle"
                ),
                dict(
                    args=[{"z": [corr_matrices[2].values]}],
                    label="Spearman",
                    method="restyle"
                )
            ]
        )
    ]
)

fig.update_layout(margin=dict(l=0,r=0,t=0,b=0))
fig.show()