#!/usr/bin/env python

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

# beta-1/sim-20180511-163319
# beta-10/sim-20180512-105824
# beta-0-1/sim-20180512-105719
# loadfile = 'beta_experiment/beta-0-1/sim-20180512-105719'
os.chdir('D:\Masterarbeit_ausgelagert')
loadfile = 'sim-20191114-000009_server'
plot_var = 'v' #variable of ising object to plot, write as string
# iter = 2000
iter_list = np.arange(0, 2000, 1)
numAgents = 150
saveFigBool = True
energy_model = True

def upper_tri_masking(A):
    m = A.shape[0]
    r = np.arange(m)
    mask = r[:, None] < r
    return A[mask]

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

# filename = 'save/' + loadfile + '/isings/gen[' + iter + ']-isings.pickle'

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

# isings = pickle.load(open(filename, 'rb'))

FOOD = np.zeros((len(iter_list), numAgents))
for ii, iter in enumerate(iter_list):
    filename = 'save/' + loadfile + '/isings/gen[' + str(iter) + ']-isings.pickle'

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

    isings = pickle.load(open(filename, 'rb'))

    food = []

    for i, I in enumerate(isings):
        exec('to_plot = I.' + plot_var)
        food.append(to_plot)

    # food = np.divide(food, 6)
    FOOD[ii, :] = food

# TODO: fix font glitches?!
# plt.rc('text', usetex=True)
# font = {'family': 'serif', 'size': 28, 'serif': ['computer modern roman']}
# plt.rc('font', **font)
plt.rc('legend', **{'fontsize': 20})

fig, ax = plt.subplots(1, 1, figsize=(19, 10))
fig.text(0.51, 0.035, r'$Generation$', ha='center', fontsize=20)
fig.text(0.07, 0.5, plot_var, va='center', rotation='vertical', fontsize=20)
title = 'Food consumed per organism'
fig.suptitle(title)

## SCATTER
for i in range(0, numAgents):
    ax.scatter(iter_list, FOOD[:, i], color=[0, 0, 0], alpha=0.05, s=10)

## MEAN + STD
# muF = np.mean(FOOD, axis=1)
# ax.plot(iter_list, muF, color=[0, 0, 0])
#
# maxF = np.max(FOOD, axis=1)
# minF = np.min(FOOD, axis=1)
# ax.fill_between(iter_list, maxF, minF,
#                 color=[0.9, 0.9, 0.9])
#
# sigmaF = FOOD.std(axis=1)
# ax.fill_between(iter_list, muF + sigmaF, muF - sigmaF,
#                 color=[0.6, 0.6, 0.6]
#                 )


#yticks = np.arange(0, 100, 20)
#ax.set_yticks(yticks)


# xticks = [0.1, 0.5, 1, 2, 4, 10, 50, 100, 200, 500, 1000, 2000]
# ax.set_xscale("log", nonposx='clip')
# ax.set_xticks(xticks)
# ax.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())


savefolder = folder + '/figs/' + plot_var + '/'
savefilename = savefolder + 'fitness_gen_' + str(iter_list[0]) + '-' + str(iter_list[-1]) + '.png'
if not path.exists(savefolder):
    makedirs(savefolder)

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

    savemsg = 'Saving ' + savefilename
    print(savemsg)

plt.show()
