#!/usr/bin/env python

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
from scipy.special import comb
import pickle
from sklearn import manifold
from os import path, makedirs

# beta-1/sim-20180511-163319
# beta-10/sim-20180512-105824
# beta-0-1/sim-20180512-105719
loadfiles = ['beta_experiment/beta-1/sim-20180511-163319',
             'beta_experiment/beta-10/sim-20180512-105824',
             'beta_experiment/beta-0-1/sim-20180512-105719']
# iter_list = [0, 25, 50, 75, 200, 500, 1000]
iter_list = np.arange(0, 4000, 1)
# iter_list = np.arange(0, 2000, 6)
numAgents = 50
# ---- plot settings ------
saveFigBool = True
autoLoad = True
alpha = 0.3

def calc_sus(ising):
    # x = ising.position[0, :] + 4
    # y = ising.position[1, :] + 4

    x = np.diff(ising.position[0, :])
    y = np.diff(ising.position[1, :])

    T = float(len(x))

    xm = np.mean(x)
    ym = np.mean(y)

    x2m = np.mean(x ** 2)
    y2m = np.mean(y ** 2)

    Sx = (x2m - xm ** 2)
    Sy = (y2m - ym ** 2)

    S = Sx + Sy

    return S


for loadfile in loadfiles:
    folder = 'save/' + loadfile
    # fname = folder + \
    #         '/figs/tsne_gen-' + \
    #         str(iter_list[0]) + '-' + str(iter_list[1] - iter_list[0]) + '-' + str(iter_list[-1]) + \
    #         '.png'

    folder2 = folder + '/figs/S/'
    fname2 = folder2 + 'S-' + \
             str(iter_list[0]) + '-' + str(iter_list[1] - iter_list[0]) + '-' + str(iter_list[-1]) + \
             '.npz'

    # See if you've generated this tsne already (they take a long time)
    if path.isfile(fname2) and autoLoad:
        txt = 'Loading: ' + fname2
        print(txt)
        data = np.load(fname2)
        S = data['S']
    else:
        # CALCULUATE T-SNE
        # [ organism(gen),  edges]
        S = np.zeros( ( len(iter_list), numAgents ) )
        iter_label = []

        for iter in iter_list:
            filename = 'save/' + loadfile + '/isings/gen[' + str(iter) + ']-isings.pickle'

            startstr = 'Loading simulation:' + filename
            print(startstr)

            isings = pickle.load(open(filename, 'rb'))
            iter_label.append('Gen: ' + str(iter))

            for orgNum, I in enumerate(isings):
                S[iter, orgNum] = calc_sus(I)


        if not path.exists(folder2):
            makedirs(folder2)

        np.savez(fname2, S=S)


    # PLOT

    fig, ax = plt.subplots(1, 1, figsize=(18, 9))
    fig.text(0.51, 0.035, r'$Generation$', ha='center', fontsize=28)
    fig.text(0.07, 0.5, r'$\chi_{x,y}$', va='center', rotation='vertical', fontsize=28)
    title = r'Positional Susceptibility of Organisms'
    fig.suptitle(title)

    # for orgNum in range(0, numAgents):
    #     ax.scatter(iter_list, S[:, orgNum], color=[0, 0, 0], alpha=alpha)

    # for iter in iter_list:
    #     ax.scatter(iter, np.mean(S[iter, :]), color=[0, 0, 0], alpha=alpha)

    muF = np.mean(S, axis=1)
    ax.scatter(iter_list, np.log10(muF), color=[0, 0, 0], alpha=alpha)
    # ax.set_yscale("log")

    # maxF = np.max(S, axis=1)
    # minF = np.min(S, axis=1)
    # ax.fill_between(iter_list, maxF, minF,
    #                 color=[0.9, 0.9, 0.9])

    # sigmaF = S.std(axis=1)
    # ax.fill_between(iter_list, np.log10(muF + sigmaF, muF - sigmaF),
    #                 color=[0.6, 0.6, 0.6]
    #                 )


    savefilename = fname2[:-4] + '.png'
    if not path.exists(folder2):
        makedirs(folder2)

    if saveFigBool:
        plt.savefig(savefilename, bbox_inches='tight', dpi=150)
        # plt.close()

        savemsg = 'Saving ' + savefilename
        print(savemsg)

plt.show()