"""
Plot mutual information of olfactory systems as a function of the 
stimulus magnitude. This uses data generated by CS-variability-adaptation/
entropy_run and aggregated by CS-variability-adaptation/
aggregate_entropy_objects.


Created by Nirag Kadakia at 12:40 04-28-2018
This work is licensed under the 
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 
International License. 
To view a copy of this license, visit 
http://creativecommons.org/licenses/by-nc-sa/4.0/.
"""


import scipy as sp
import sys
import matplotlib as mpl
import matplotlib.pyplot as plt
sys.path.append('../../shared_src')
from save_load_figure_data import save_fig
from plot_formats import fig_MI_trace

# The location of the source code for CS-variability-adaptation is listed
# in the ../../shared_src/local_methods file within src_dir()
from local_methods import src_dir
sys.path.append(src_dir())
from utils import get_flag
from load_specs import read_specs_file
from load_data import load_aggregated_temporal_objects


def plot_MI(data_flag, Kk_idxs=[4, 4], Tt_idxs=[0, 10, 20, 50, 
			80, 100, 140, 195]):
	"""
	Kk_idxs is a list of which of the iterated variables to plot
	(inner loop) for each data_flag (outer loop).
	"""
	
	colors = plt.cm.viridis(sp.linspace(0.1, 0.9, len(Tt_idxs)))
	
	list_dict = read_specs_file(data_flag)
	iter_vars = list_dict['iter_vars']
	iter_vars_dims = []
	for iter_var in list_dict['iter_vars']:
		iter_vars_dims.append(len(list_dict['iter_vars'][iter_var]))	
	
	assert len(iter_vars) == 3, "Need 3 iter_vars"
	iter_var_names = ['signal_trace_multiplier', 'Kk_1', 'Kk_2']
	for iName, name in enumerate(iter_var_names):
		assert list(iter_vars.keys())[iName] == name, "%sth variable "\
			"must have name %s" % (iName, name)
	mu_Ss0_vals = iter_vars['signal_trace_multiplier']
	Kk_1_vals = iter_vars['Kk_1']
	Kk_2_vals = iter_vars['Kk_2']
	x = mu_Ss0_vals
	
	print ('Loading object list...'),
	CS_object_array = load_aggregated_temporal_objects(data_flag)
	print ('...loaded.')
	
	for iTt, Tt_idx in enumerate(Tt_idxs):
	
		data = CS_object_array['entropy'][Tt_idx, :, Kk_idxs[0], Kk_idxs[1], :]

		fig = fig_MI_trace()
		for iM in range(list_dict['params']['Mm']):
			plt.plot(x, data[...,iM], lw=0.3, alpha=0.5,
						color=colors[iTt], linestyle='-')
		plt.plot(x, sp.average(data, axis=-1), lw=2,
					color=colors[iTt], linestyle='-')
		plt.ylim(0, 7.01)
		plt.xlim(1, 10**4)
	
		save_fig('MI_trace_%s_%s' % (Kk_idxs, Tt_idx), subdir=data_flag, 
									tight_layout=False)
		
				
if __name__ == '__main__':
	data_flag = get_flag()
	plot_MI(data_flag)