# %% [markdown]
# # Imports
import os
import pickle
import warnings
from operator import itemgetter
from pathlib import Path
from timeit import default_timer as timer

import colorcet as cc
import matplotlib.colors as mplc
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import pandas as pd
import seaborn as sns
from bokeh.embed import file_html
from bokeh.io import output_file, output_notebook, show
from bokeh.layouts import gridplot
from bokeh.models import ColumnDataSource, FactorRange, Legend, Span, PreText, Circle
from bokeh.palettes import Spectral4, all_palettes
from bokeh.plotting import curdoc, figure, output_file, show
from bokeh.resources import CDN
from bokeh.sampledata.stocks import AAPL, GOOG, IBM, MSFT
from joblib import Parallel, delayed
from matplotlib.cm import ScalarMappable
from mpl_toolkits.axes_grid1 import make_axes_locatable
from sklearn.metrics import adjusted_rand_score
from sklearn.model_selection import ParameterGrid
from sklearn.neighbors import NearestNeighbors

from graspy.cluster import AutoGMMCluster, GaussianCluster
from graspy.embed import AdjacencySpectralEmbed, LaplacianSpectralEmbed
from graspy.plot import degreeplot, edgeplot, gridplot, heatmap, pairplot
from graspy.utils import symmetrize
from src.cluster import DivisiveCluster
from src.data import load_everything, load_metagraph, load_networkx
from src.embed import ase, lse, preprocess_graph
from src.graph import MetaGraph
from src.hierarchy import signal_flow
from src.io import savefig, saveobj, saveskels
from src.utils import (
    get_blockmodel_df,
    get_sbm_prob,
    invert_permutation,
    meta_to_array,
    savefig,
)
from src.visualization import (
    bartreeplot,
    get_color_dict,
    get_colors,
    remove_spines,
    sankey,
    screeplot,
)

from bokeh.layouts import column, row
from bokeh.models import Select
from bokeh.palettes import Spectral5
from bokeh.plotting import curdoc, figure
from scipy.linalg import orthogonal_procrustes


def pair_augment(mg):
    pair_df = pd.read_csv(
        "maggot_models/data/raw/Maggot-Brain-Connectome/pairs/bp-pairs-2020-01-13_continuedAdditions.csv"
    )

    skeleton_labels = mg.meta.index.values

    # extract valid node pairings
    left_nodes = pair_df["leftid"].values
    right_nodes = pair_df["rightid"].values

    left_right_pairs = list(zip(left_nodes, right_nodes))

    left_nodes_unique, left_nodes_counts = np.unique(left_nodes, return_counts=True)
    left_duplicate_inds = np.where(left_nodes_counts >= 2)[0]
    left_duplicate_nodes = left_nodes_unique[left_duplicate_inds]

    right_nodes_unique, right_nodes_counts = np.unique(right_nodes, return_counts=True)
    right_duplicate_inds = np.where(right_nodes_counts >= 2)[0]
    right_duplicate_nodes = right_nodes_unique[right_duplicate_inds]

    left_nodes = []
    right_nodes = []
    for left, right in left_right_pairs:
        if left not in left_duplicate_nodes and right not in right_duplicate_nodes:
            if left in skeleton_labels and right in skeleton_labels:
                left_nodes.append(left)
                right_nodes.append(right)

    pair_nodelist = np.concatenate((left_nodes, right_nodes))
    not_paired = np.setdiff1d(skeleton_labels, pair_nodelist)
    sorted_nodelist = np.concatenate((pair_nodelist, not_paired))

    # sort the graph and metadata according to this
    sort_map = dict(zip(sorted_nodelist, range(len(sorted_nodelist))))
    inv_perm_inds = np.array(itemgetter(*skeleton_labels)(sort_map))
    perm_inds = invert_permutation(inv_perm_inds)

    mg.reindex(perm_inds)

    side_labels = mg["Hemisphere"]
    side_labels = side_labels.astype("<U2")
    for i, l in enumerate(side_labels):
        if mg.meta.index.values[i] in not_paired:
            side_labels[i] = "U" + l
    mg["Hemisphere"] = side_labels
    n_pairs = len(left_nodes)
    return mg, n_pairs


def max_symmetrize(mg, n_pairs):
    """ assumes that mg is sorted
    
    Parameters
    ----------
    mg : [type]
        [description]
    n_pairs : [type]
        [description]
    """
    adj = mg.adj
    left_left_adj = adj[:n_pairs, :n_pairs]
    left_right_adj = adj[:n_pairs, n_pairs : 2 * n_pairs]
    right_right_adj = adj[n_pairs : 2 * n_pairs, n_pairs : 2 * n_pairs]
    right_left_adj = adj[n_pairs : 2 * n_pairs, :n_pairs]

    # max, average gives similar results
    sym_ipsi_adj = np.maximum(left_left_adj, right_right_adj)
    sym_contra_adj = np.maximum(left_right_adj, right_left_adj)

    sym_adj = adj.copy()
    sym_adj[:n_pairs, :n_pairs] = sym_ipsi_adj
    sym_adj[n_pairs : 2 * n_pairs, n_pairs : 2 * n_pairs] = sym_ipsi_adj
    sym_adj[:n_pairs, n_pairs : 2 * n_pairs] = sym_contra_adj
    sym_adj[n_pairs : 2 * n_pairs, :n_pairs] = sym_contra_adj

    sym_mg = MetaGraph(sym_adj, mg.meta)  # did not change indices order so this ok

    return sym_mg


FNAME = os.path.basename(__file__)[:-3]
print(FNAME)
BRAIN_VERSION = "2020-01-21"


# %% [markdown]
# #
thresh = 2
mg = load_metagraph("Gad", BRAIN_VERSION)
mg.adj[mg.adj < thresh] = 0
mg, n_pairs = pair_augment(mg)

left_pair_ids = mg["Pair ID"][:n_pairs]
right_pair_ids = mg["Pair ID"][n_pairs : 2 * n_pairs]
print((left_pair_ids == right_pair_ids).all())

sym_mg = max_symmetrize(mg, n_pairs)
left_pair_ids = sym_mg["Pair ID"][:n_pairs]
right_pair_ids = sym_mg["Pair ID"][n_pairs : 2 * n_pairs]
print((left_pair_ids == right_pair_ids).all())

sym_mg.make_lcc()
n_pairs = sym_mg.meta["Pair ID"].nunique() - 1
left_pair_ids = sym_mg["Pair ID"][:n_pairs]
right_pair_ids = sym_mg["Pair ID"][n_pairs : 2 * n_pairs]
print((left_pair_ids == right_pair_ids).all())

uni_pair, counts = np.unique(sym_mg["Pair ID"], return_counts=True)
print(np.min(counts))
# %% [markdown]
# #
left_pair_ids = sym_mg["Pair ID"][:n_pairs]
right_pair_ids = sym_mg["Pair ID"][n_pairs : 2 * n_pairs]

latent = lse(sym_mg.adj, n_components=None)
left_latent = latent[:n_pairs, :]
right_latent = latent[n_pairs : 2 * n_pairs, :]
R, scalar = orthogonal_procrustes(left_latent, right_latent)

n_components = latent.shape[1]
class_labels = sym_mg["Class 1"]
n_unique = len(np.unique(class_labels))
sym_mg.meta["Original index"] = range(len(sym_mg.meta))
left_df = sym_mg.meta[sym_mg.meta["Hemisphere"] == "L"]
left_inds = left_df["Original index"].values
left_latent = latent[left_inds, :]
left_latent = left_latent @ R
latent[left_inds, :] = left_latent

latent_cols = [f"dim {i}" for i in range(latent.shape[1])]
latent_df = pd.DataFrame(data=latent, index=sym_mg.meta.index, columns=latent_cols)

df = pd.concat((sym_mg.meta, latent_df), axis=1)
df.index.name = "Skeleton ID"

x = "dim 0"
y = "dim 1"

color_by = "Class 1"
uni_color_classes = np.unique(df[color_by])
colormap = dict(zip(uni_color_classes, cc.glasbey_light))
colors = np.array(itemgetter(*df[color_by])(colormap))
df["colors"] = colors
source = ColumnDataSource(df)

TOOLS = "pan,reset,zoom_in,zoom_out,box_zoom,lasso_select,xbox_select"

scatter_tooltips = [
    ("ID", "@{Skeleton ID}"),
    ("Class 1", "@{Class 1}"),
    ("Class 2", "@{Class 2}"),
    ("Side", "@Hemisphere"),
    ("Name", "@neuron_name"),
    ("Lineage", "@lineage"),
]

stats = PreText(text="HERE", width=500)


def create_figure():

    scatter_fig = figure(
        tools=TOOLS,
        plot_width=900,
        plot_height=800,
        min_border=10,
        min_border_left=50,
        toolbar_location="above",
        x_axis_location=None,
        y_axis_location=None,
        tooltips=scatter_tooltips,
    )
    scatter_fig.xgrid.visible = False
    scatter_fig.ygrid.visible = False

    legend_it = []
    for c in uni_color_classes:
        temp_df = df[df[color_by] == c]
        temp_source = ColumnDataSource(temp_df)
        out = scatter_fig.circle(
            x=x.value,
            y=y.value,
            source=temp_source,
            alpha=0.6,
            line_color=None,
            fill_color="colors",
            # legend_group=color_by,
        )
        legend_it.append((c, [out]))
        # selected = Circle(fill_alpha=1)
        # nonselected = Circle(fill_alpha=0.2)
        # out.selection_glyph = selected
        # out.nonselection_glyph = nonselected

    out = scatter_fig.circle(
        x=x.value,
        y=y.value,
        source=source,
        alpha=0.1,
        line_color=None,
        fill_color="colors",
        # legend_group=color_by,
    )

    legend = Legend(items=legend_it[: len(legend_it) // 2], location=(0, 0))
    legend.click_policy = "hide"
    legend.label_text_font_size = "5pt"
    legend.orientation = "horizontal"
    legend.padding = 0
    legend.margin = 0
    legend2 = Legend(items=legend_it[len(legend_it) // 2 :], location=(0, 0))
    legend2.click_policy = "hide"
    legend2.label_text_font_size = "6pt"
    legend2.orientation = "horizontal"
    legend2.padding = 0
    legend2.margin = 0
    scatter_fig.add_layout(legend, "below")
    scatter_fig.add_layout(legend2, "below")
    return scatter_fig


def update(attr, old, new):
    layout.children[0] = create_figure()


cols = ["Skeleton ID", "Merge Class", "Hemisphere", "Pair ID", "lineage"]


def update_stats(data):
    stats.text = data[cols].to_string(index=False)


def selection_change(attr, old, new):
    selected = source.selected.indices
    data = source.to_df()
    if selected:
        update_stats(data.iloc[selected])


x = Select(title="X-Axis", value="dim 0", options=latent_cols)
x.on_change("value", update)

y = Select(title="Y-Axis", value="dim 1", options=latent_cols)
y.on_change("value", update)


source.selected.on_change("indices", selection_change)


controls = column(x, y, width=200)
left_col = column(create_figure())
right_col = column(row(controls), row(stats))
# bottom_row = row(controls, stats)
layout = row(left_col, right_col)

curdoc().add_root(layout)
curdoc().title = BRAIN_VERSION
show(layout)
# html = file_html(scatter_fig, CDN, "mb_adj")
# with open("full_new_adj.html", "w") as f:
#     f.write(html)


# %%
