__all__ = ["plot_vae"]

from pathlib import Path
from typing import Optional

import matplotlib
import matplotlib.cm as cm
import matplotlib.figure
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import torch


def plot_vae(
    path: Path,
    savepath: Path,
    filename: str,
    title: str,
    num_input: int,
    num_hidden: int,
    num_latent: int,
    plot_edges=True,
    input_sample: Optional[torch.Tensor] = None,
    output_sample: Optional[torch.Tensor] = None,
    mu: Optional[torch.Tensor] = None,
    logvar: Optional[torch.Tensor] = None,
) -> matplotlib.figure.Figure:
    """
    This function is aimed to visualize MOVE's architecture.

    Args:
        path: path where the trained model with defined weights is to be found
        filename: name of the model
        title: title of the figure
        num_input: number of input nodes
        num_hidden: number of output nodes
        num_latent: number of latent nodes
        plot_edges: plot edges, i.e. connections with assigned weights between nodes
        input_sample: array with input values to fill with a mapped color value
        output_sample : "      " output " "
        mu: "                  " mean (latent) "     "
        logvar: "              " log variance  "     "

    Returns:
        figure

    Notes:
        k: input node index
        j: hidden node index
        i: latent node index
    """
    model_weights = torch.load(path / filename)
    G = nx.Graph()

    # Position of the layers:
    layer_distance = 10
    node_distance = 550
    latent_node_distance = 550
    latent_sep = 5 * latent_node_distance

    # Adding nodes to the graph ##############################
    # Bias nodes
    G.add_node(
        "input_bias",
        pos=(-6 * layer_distance, -3 * node_distance - num_input * node_distance / 2),
        color=0.0,
    )
    G.add_node(
        "mu_bias",
        pos=(
            -3 * layer_distance,
            (num_hidden + 3) * node_distance
            - num_hidden * node_distance / 2
            + latent_sep / 2,
        ),
        color=0.0,
    )
    G.add_node(
        "var_bias",
        pos=(
            -3 * layer_distance,
            -3 * node_distance - num_hidden * node_distance / 2 - latent_sep / 2,
        ),
        color=0.0,
    )
    G.add_node(
        "sam_bias",
        pos=(
            0.5 * layer_distance,
            -3 * latent_node_distance - num_latent * latent_node_distance / 2,
        ),
        color=0.0,
    )
    G.add_node(
        "out_bias",
        pos=(3 * layer_distance, -3 * node_distance - num_hidden * node_distance / 2),
        color=0.0,
    )

    # Actual nodes
    for k in range(num_input):
        G.add_node(
            f"input_{k}",
            pos=(
                -6 * layer_distance,
                k * node_distance - num_input * node_distance / 2,
            ),
            color=[input_sample[k] if input_sample is not None else 0.0][0],
        )
        G.add_node(
            f"output_{k}",
            pos=(6 * layer_distance, k * node_distance - num_input * node_distance / 2),
            color=[output_sample[k] if output_sample is not None else 0.0][0],
        )
    for j in range(num_hidden):
        G.add_node(
            f"encoder_hidden_{j}",
            pos=(
                -3 * layer_distance,
                j * node_distance - num_hidden * node_distance / 2,
            ),
            color=0.0,
        )
        G.add_node(
            f"decoder_hidden_{j}",
            pos=(
                3 * layer_distance,
                j * node_distance - num_hidden * node_distance / 2,
            ),
            color=0.0,
        )
    for i in range(num_latent):
        G.add_node(
            f"mu_{i}",
            pos=(0 * layer_distance, i * latent_node_distance + latent_sep / 2),
            color=[mu[i] if mu is not None else 0.0][0],
        )
        G.add_node(
            f"var_{i}",
            pos=(0 * layer_distance, -i * latent_node_distance - latent_sep / 2),
            color=[np.exp(logvar[i] / 2) if logvar is not None else 0.0][0],
        )
        G.add_node(
            f"sam_{i}",
            pos=(
                0.5 * layer_distance,
                i * latent_node_distance - num_latent * latent_node_distance / 2,
            ),
            color=0.0,
        )

    # Adding weights to the graph #########################

    if plot_edges:
        for layer, values in model_weights.items():
            if layer == "encoderlayers.0.weight":
                for k in range(values.shape[1]):  # input
                    for j in range(values.shape[0]):  # encoder_hidden
                        G.add_edge(
                            f"input_{k}",
                            f"encoder_hidden_{j}",
                            weight=values.numpy()[j, k],
                        )

            elif layer == "encoderlayers.0.bias":
                for j in range(values.shape[0]):  # encoder_hidden
                    G.add_edge(
                        "input_bias", f"encoder_hidden_{j}", weight=values.numpy()[j]
                    )

            elif layer == "mu.weight":
                for j in range(values.shape[1]):  # encoder hidden
                    for i in range(values.shape[0]):  # mu
                        G.add_edge(
                            f"encoder_hidden_{j}",
                            f"mu_{i}",
                            weight=values.numpy()[i, j],
                        )

            elif layer == "mu.bias":
                for i in range(values.shape[0]):  # encoder_hidden
                    G.add_edge("mu_bias", f"mu_{i}", weight=values.numpy()[i])

            elif layer == "var.weight":
                for j in range(values.shape[1]):  # encoder hidden
                    for i in range(values.shape[0]):  # var
                        G.add_edge(
                            f"encoder_hidden_{j}",
                            f"var_{i}",
                            weight=values.numpy()[i, j],
                        )

            elif layer == "var.bias":
                for i in range(values.shape[0]):  # encoder_hidden
                    G.add_edge("var_bias", f"var_{i}", weight=values.numpy()[i])

            # Sampled layer from mu and var:
            elif layer == "decoderlayers.0.weight":
                for i in range(values.shape[1]):  # sampled latent
                    for j in range(values.shape[0]):  # decoder_hidden
                        G.add_edge(
                            f"sam_{i}",
                            f"decoder_hidden_{j}",
                            weight=values.numpy()[j, i],
                        )

            # Sampled layer from mu and var:
            elif layer == "decoderlayers.0.bias":
                for j in range(values.shape[0]):  # decoder_hidden
                    G.add_edge(
                        "sam_bias", f"decoder_hidden_{j}", weight=values.numpy()[j]
                    )

            elif layer == "out.weight":
                for j in range(values.shape[1]):  # decoder_hidden
                    for k in range(values.shape[0]):  # output
                        G.add_edge(
                            f"output_{k}",
                            f"decoder_hidden_{j}",
                            weight=values.numpy()[k, j],
                        )

            elif layer == "out.bias":
                for k in range(values.shape[0]):  # output
                    G.add_edge("out_bias", f"output_{k}", weight=values.numpy()[k])

    fig = plt.figure(figsize=(60, 60))
    pos = nx.get_node_attributes(G, "pos")
    color = list(nx.get_node_attributes(G, "color").values())
    edge_color = list(nx.get_edge_attributes(G, "weight").values())
    edge_width = list(nx.get_edge_attributes(G, "weight").values())

    edge_cmap = matplotlib.colormaps["seismic"]
    node_cmap = matplotlib.colormaps["seismic"]

    abs_max = np.max([abs(np.min(color)), abs(np.max(color))])
    abs_max_edge = np.max([abs(np.min(edge_color)), abs(np.max(edge_color))])

    _ = cm.ScalarMappable(
        cmap=node_cmap, norm=matplotlib.colors.Normalize(vmin=-abs_max, vmax=abs_max)
    )
    sm_edge = cm.ScalarMappable(
        cmap=edge_cmap,
        norm=matplotlib.colors.Normalize(vmin=-abs_max_edge, vmax=abs_max_edge),
    )

    nx.draw(
        G,
        pos=pos,
        with_labels=True,
        node_size=100,
        node_color=color,
        edge_color=edge_color,
        width=edge_width,
        font_color="black",
        font_size=10,
        edge_cmap=edge_cmap,
        cmap=node_cmap,
        vmin=-abs_max,
        vmax=abs_max,
    )

    # plt.colorbar(sm_node, label="Node value", shrink = .2)
    plt.colorbar(sm_edge, label="Edge value", shrink=0.2)
    plt.tight_layout()
    fig.savefig(savepath / f"{title}.png", format="png", dpi=200)
    return fig