import mainaux.PlotCompilation as PlotCompilation
import mainaux.PlotHelpers as PlotHelpers
import mainaux.SimHelpers as SimHelpers
from state.Proteins import Proteins
import numpy as np

# Only use this file for running simulations where param. values
# remain constant across all simulations

# Controls the # of simulations that will run
# If NUM_OF_SIMULATIONS > 1, the average of the simulation data
# will be plotted (+std. dev, range, all data, or None)
NUM_OF_SIMULATIONS = 1

NUM_OF_TIMESTEPS = 2400

# Parameter that determines the increment of timestep you care about, starting from t = 0
# Default is 1
SAMPLING_RATE = 1

# Determines what keys will be plotted
LIST_OF_KEY_NAMES = ['multi_splice_transcript_nuc', 'multi_splice_transcript_cyt', 'total_proteins', 'proteins_nuc', 'proteins_cyt', 'proteins_virion', 'proteins_mem']

# Export raw data from simulations instead of plotting the data
# Histograms are still plotted (if any)
EXPORT_RAW_DATA_NO_PLOT = False

# Group set of rows in a csv file by their row number in the
# specified key(s) instead of by their simulation number
# These group of rows are separated by newlines.
# Setting is only relevant if EXPORT_RAW_DATA_NO_PLOT is True
GROUP_BY_ROW_NUM = True

# Determines how all the simulation data will be displayed
#   options: "standard deviation", "range", "all data", None
#   Note average is always plotted for >1 sim.
# Setting value is only relevant if EXPORT_RAW_DATA_NO_PLOT is False
TYPE_OF_PLOT = "standard deviation"

# Use the default plot on log axes setting of each plot
# If False, will turn on nonlog y-axes for ALL plots
# Setting is only relevant if EXPORT_RAW_DATA_NO_PLOT is False
USE_DEFAULT_LOG_SETTING = True

# Set to True if want to save states of simulation
SAVE_STATE = False

# List of timesteps to save state at
# Must hold true: each element must have a value smaller than NUM_OF_TIMESTEPS
SAVE_STATE_TIMESTEP_LIST = [1500, 2300]

# DO NOT TOUCH ANYTHING BELOW THIS LINE UNLESS YOU'VE READ THE DOCS
if SAVE_STATE:
    batch_label = SimHelpers.initialize_label("samplestates")
else:
    batch_label = ""

# master_tracking_dict is a dictionary that maps key to
# a list of matrices, where each key is a key in variable_tracking_dict
# and the value is a list of matrices. Each matrix corresponds to a
# particular simulation
# COMMENT OUT THIS LINE IF YOU WANT TO RETAIN PREVIOUS master_tracking_dict
if NUM_OF_SIMULATIONS == 1:
    master_tracking_dict, list_of_plotting_keys, hist_dict = SimHelpers.initialize_runsim_dict(LIST_OF_KEY_NAMES, NUM_OF_SIMULATIONS, NUM_OF_TIMESTEPS, SAMPLING_RATE, save_state_mode=SAVE_STATE, timestep_list=SAVE_STATE_TIMESTEP_LIST, batch_label=batch_label, rtn_hist_dict=True)

    # add histogram plots here
    # see function prototype in PlotHelpers for more info on usage
    PlotHelpers.plot_histogram(hist_dict, ['num_of_successful_Env_t_of_diff_progeny', 'num_of_unsuccessful_Env_t_of_diff_progeny'], 'Frequency of virons', 'Number of Env Trimers', 'Frequency', ['r', 'k'], bins=np.linspace(0, 25, 25))
    PlotHelpers.plot_histogram(hist_dict, ['num_of_' + protein_name + '_of_diff_progeny' for protein_name in ['Gag', 'Vif', 'GagProPol', 'Vpr', 'Nef', 'successful_Env_t']], 'Frequency of different proteins in particular progeny', 'Number of Specific Protein', 'Frequency', ['r', 'k', 'b', 'g', 'c', 'y'], bins=np.linspace(0, 25, 25))    
else:
    master_tracking_dict, list_of_plotting_keys = SimHelpers.initialize_runsim_dict(LIST_OF_KEY_NAMES, NUM_OF_SIMULATIONS, NUM_OF_TIMESTEPS, SAMPLING_RATE, save_state_mode=SAVE_STATE, timestep_list=SAVE_STATE_TIMESTEP_LIST, batch_label=batch_label, rtn_hist_dict=False)

# for key in LIST_OF_KEY_NAMES:
#     PlotCompilation.generate_plotting_key(master_tracking_dict, key)

if EXPORT_RAW_DATA_NO_PLOT:
    output_data_label = SimHelpers.initialize_label("outputdata")
    SimHelpers.create_output_data_info_file(output_data_label, GROUP_BY_ROW_NUM, NUM_OF_SIMULATIONS)
    for key_name in LIST_OF_KEY_NAMES:
        SimHelpers.write_key(master_tracking_dict, key_name, output_data_label, GROUP_BY_ROW_NUM)

# NOTE: Can only have one plot_tracking_dictionary line enabled at any time if planning to run the entire script
# Use the plots generated by this line to plot data from the master_tracking_dict
# This will give ideas on how to filter the master_tracking_dict using generate_subset_dict
# Comment this line and uncomment plot_tracking_dictionary a couple of lines below to enable
# plotting of the subset tracking dictionary
# OPTIONS HERE ARE None (for average only), "all data", "range", "standard deviation"
record = SimHelpers.plot_tracking_dictionary(master_tracking_dict, NUM_OF_TIMESTEPS, list_of_plotting_keys, TYPE_OF_PLOT, log_setting=USE_DEFAULT_LOG_SETTING, sampling_rate=SAMPLING_RATE, export_raw_data_no_plot=EXPORT_RAW_DATA_NO_PLOT)

# NOTE: Can only have one plot_tracking_dictionary line enabled at any time if planning to run the entire script
# Comment this line and uncomment plot_tracking_dictionary a couple of lines above to enable
# plotting of the master tracking dictionary (instead of the subset tracking dict below)
# Can call subset_tracking_dict multiple times with input_dict=subset_tracking_dict
# to filter subset_tracking_dict multiple times
#subset_tracking_dict = SimHelpers.generate_subset_dict(input_dict=master_tracking_dict, key='proteins_nuc', row_index=Proteins.index['Rev'], timestep=70, value=0, operator=">=")
#record = SimHelpers.plot_tracking_dictionary(subset_tracking_dict, NUM_OF_TIMESTEPS, list_of_plotting_keys, 'standard deviation', True)
