from typing import Dict, Tuple
import numpy as np
import streamlit as st
import scipy as sp

import matplotlib

import matplotlib.pyplot as plt
import seaborn as sns

from util import generate_random_discrete_distribution
from util import compute_progression_of_nash_during_training


def exploration_view(results_dir):
    name = 'Evolution of policy sampling distribution VS Nash averaging'
    st.write(f'# {name}')

    selfplay_choice = st.sidebar.radio('Select Self-Play algorithm',
                                       ('Naive SP', 'Full History SP', 'Iterated Nash Response'))


    st.write('## Distance between Policy sampling distribution and max-ent Nash \
                 during training?')

    min_checkpoint, max_checkpoint, step_checkpoint = 100, 1000, 100
    range_checkpoint = range(min_checkpoint, max_checkpoint + 1, step_checkpoint)
    checkpoint = st.sidebar.slider('Choose training iteration', min_checkpoint, max_checkpoint, step=step_checkpoint)

    distributions = compute_distributions(range_checkpoint)
    # import ipdb; ipdb.set_trace()
    plot_distances_between_nash_and_policy_sampling_distribution(distributions)

    st.write('## Maxent Nash VS policy sampling distribution')
    plot_selected_policy_sampling_distribution_and_maxent_nash(maxent_nash=distributions[checkpoint][0],
                                                               policy_sampling_distribution=distributions[checkpoint][1],
                                                               checkpoint=checkpoint)

    st.write('## Evolution of support under maxent-nash for policies in menagerie')


    # NOTE: For the real experiment these supports should come from meta-games
    #       generated by head to head matches using the menagerie at {checkpoint}
    progression_nash = compute_progression_of_nash_during_training(range_checkpoint)

    plot_evolution_of_support_for_menagerie(progression_nash, max_x=range_checkpoint[-1])


def plot_selected_policy_sampling_distribution_and_maxent_nash(maxent_nash,
                                                               policy_sampling_distribution,
                                                               checkpoint):
    width = 0.7
    indices = np.arange(len(maxent_nash.ravel()))

    fig, ax = plt.subplots(1,1)
    ax.bar(indices, maxent_nash.ravel(), width=width,
            color='b', label='max-ent Nash')
    ax.bar(indices, policy_sampling_distribution.ravel(),
            width=0.8*width, color='r', alpha=0.5, label='Policy Sampling Distribution')

    plt.title(f'Policy sampling distriubtion and maxent-Nash Equilibrium at checkpoint: {checkpoint}')
    plt.xlabel('Agent ID')
    plt.ylabel('Support')
    plt.xticks(indices)
    plt.legend()
    st.pyplot()
    plt.close()


def plot_distances_between_nash_and_policy_sampling_distribution(distributions):
    x_ticks, wass_distances, kl_distances = generate_plot_elements_evolution_distance_plot(distributions)

    f, ax = plt.subplots(1, 1)
    ax = sns.pointplot(x=x_ticks, y=wass_distances, label='Wasserstein distance', ax=ax)
    ax = sns.pointplot(x=x_ticks, y=kl_distances, color='orange', label='KL divergence', dodge=True, ax=ax)
    f.legend(labels=['Wasserstein distance','KL divergence'])
    plt.xlabel('training iteration')
    plt.ylabel('Distance')

    st.pyplot()
    plt.close()


def plot_evolution_of_support_for_menagerie(progression_nash, max_x):
    fig, axes = plt.subplots(len(progression_nash), 1, figsize=(7,12))
    checkpoints = progression_nash.columns
    for i, first_checkpoint_appearence in zip(range(len(progression_nash)), progression_nash.columns):
        evolution_support_policy_i = progression_nash.iloc[i]
        single_menagerie_policy_support_evolution(checkpoints=checkpoints[i:],
                                                  supports_nash=list(evolution_support_policy_i.array)[i:],
                                                  max_x=max_x,
                                                  first_checkpoint_appearence=first_checkpoint_appearence,
                                                  axis=axes[i],
                                                  title=f'ID: {i}')
    plt.subplots_adjust(hspace=0.3)
    plt.xlabel('Training episode')
    plt.ylabel('Support under Nash')
    st.pyplot()


def single_menagerie_policy_support_evolution(checkpoints, supports_nash, max_x,
                                              first_checkpoint_appearence, axis, title):
    axis.vlines(x=first_checkpoint_appearence, ymin=0, ymax=1, linestyles='dashed', color='black', linewidth=2, alpha=0.5)
    axis.plot(checkpoints, supports_nash, 'b-o')
    axis.set_ylim(0, 1.1) # We are working with probability distributions, the extra 0.1 is to make room
    axis.set_xlim(0, max_x * 1.05)
    axis.set_xscale('linear')
    axis.set_ylabel(title)
    

# TODO: pls find better name.
def generate_plot_elements_evolution_distance_plot(distributions: Dict[int, Tuple[np.array, np.array]]):
    symmetric_kl = lambda p, q: (np.sum(sp.special.rel_entr(p, q)) + np.sum(sp.special.rel_entr(q, p))) / 2.0

    x_ticks, wass_distances, kl_distances = [], [], []
    for i, (maxent_nash, policy_sampling_supports) in distributions.items():
        x_ticks.append(i)
        wass_distance = compute_distances(maxent_nash.ravel(), policy_sampling_supports.ravel(),
                                          distance_function=sp.stats.wasserstein_distance)
        kl_distance = compute_distances(maxent_nash.ravel(), policy_sampling_supports,
                                        distance_function=symmetric_kl)
        wass_distances.append(wass_distance)
        kl_distances.append(kl_distance)
    return x_ticks, wass_distances, kl_distances


def compute_distances(dist1, dist2, distance_function):
    return distance_function(dist1, dist2)


@st.cache
def compute_distributions(range_checkpoint):
    distributions = {checkpoint: (generate_random_discrete_distribution(size+1),
                                  generate_random_discrete_distribution(size+1))
		     for size, checkpoint in enumerate(range_checkpoint)}
    return distributions
