off_plot

from matplotlib import pyplot as plt
import matplotlib.cm as cm
import numpy as np
import os
from glob import glob
import csv
from collections import OrderedDict
import ast
import re
import fnmatch



#list of string
data_to_import = ['total_steps_taken']
#data_to_import = ['total_area_covered']

# ->
# 0 don't draw error bars
# 1 draw error bars
drawError = 1

path = []

eff_array = ['0','40', '50','60','70', '80', '90', '100']

testNumber = 'FiniteFood'
cwd = '/Users/ammacpro1/AnacondaProjects/Trophallaxis_Model_for_journal_ALR_2500_maxfood100_checkboardmap1_100runs/Simulation_results/' + testNumber

##cwd = os.getcwd()
##cwd = cwd + "/" + testNumber

UTfolder_reference =  ['Random_No_Trophallaxis']  #'Levyrandom_No_Trophallaxis'

UTprefix = ['Random_Trophallaxis_Generous_eff',
            'Random_Trophallaxis_UT70LT20average_eff',
            'Levywalklike_Trophallaxis_Generous_eff',
            'Levywalklike_Trophallaxis_UT70LT20average_eff',
            #'Levyrandom_Trophallaxis_Generous_eff',
            #'Levyrandom_Trophallaxis_UT70LT20average_eff',
           ]
UTprefixLegend = ['Random_Trophallaxis_Generous',
            'Random_Trophallaxis_UT70LT20average',
            'Levywalklike_Trophallaxis_Generous',
            'Levywalklike_Trophallaxis_UT70LT20average',
            #'Levyrandom_Trophallaxis_Generous',
            #'Levyrandom_Trophallaxis_UT70LT20average',
           ]

trophFolder = 'Bots225'
plot_title = "Total Step Taken_" +str(trophFolder) + '_' +str(testNumber)
#UTfolder = ['No_Trophallaxis','Trophallaxis_generous', 'Th50','Th70','Th90',]
UTfolderP = []
for p in UTprefix:
    for n in range(len(eff_array)-1):
        UTfolderP.append(p + eff_array[n+1])
UTfolder = UTfolder_reference
for pf in UTfolderP:
    UTfolder.append(pf)

UTfolderP2 = UTfolder_reference
for pf in UTprefixLegend:
    UTfolderP2.append(pf)

for f in UTfolder:
    path.append(cwd + '/' + f + '/' + trophFolder + '/Bot*/*Summary.csv')
data_to_plot = {}
def analyseFolders(folder, curr_p):
    #nitialises an empty dictionary
    data = {}
    #data_to_plot = {}

    for _data in data_to_import:
        data[_data] = []

    files = glob(folder)
    for file in files:
        with open(file, 'rt') as f:
            reader = csv.reader(f, delimiter=',')
            for row in reader:
                for key in data_to_import:
                    if row[0] in [key]:
                        data[key].append(row[1])

    sorted_data = dict.fromkeys(data.keys(),[])

    keys = [key for key, value in data.items()]
    #print(keys)
    #print('curr_p', curr_p)
    for _data in [key for key, value in data.items()]:
        #data_to_plot[_data + "_f" + UTfolderP2[curr_pref] + "_eff" + eff_array[curr_p]] = \
        data_to_plot[_data + "_eff" + eff_array[curr_p]] = \
                                                        [float(i) for i in data[_data]]
    keys_plot = [key for key, value in data_to_plot.items()]
    #print('keys_plot', keys_plot)

fig, ax = plt.subplots()
colors = cm.rainbow(np.linspace(0.4, 1, len(UTfolderP2)))
widths = np.linspace(2, 0.5, len(UTfolderP2))
#Markers = ['+','+','o','o','o','^','o','p','s']
#LineStyles = ['--','--','-','-','-','-.',':','--']
Markers = ['+','+','+','o','o','s','s','^','s','s','s','o','^','o','p','s']
LineStyles = ['-','-','--','-','--','-','--','-','--',':',':','-.','-.','-','-.',':','--']

def plotData():
    #print(data_to_plot)
    #print('len:', len(data_to_plot))
    print('curr_pref:', curr_pref)

    data_mean = []
    data_sd = []

    if curr_pref < len(UTfolder_reference):
        eff_cycle = eff_array[curr_p]
    else:
        eff_cycle = eff_array[1:]
        #eff_cycle = eff_array[len(UTfolder_reference):]
    print('eff_array:',eff_array)
    print('eff_cycle:',eff_cycle)
    for effSize in eff_cycle:
        print('effSize:',effSize)
        c = colors[curr_pref]
        #w = widths[curr_b]
        w = 2
        vals = []
        sd = []
        #tmpArray = data_to_plot["total_area_covered_eff" + effSize]
        tmpArray = data_to_plot["total_steps_taken_eff" + effSize]
        vals = np.mean(tmpArray)
        if (drawError):
            sd = np.std(tmpArray)
        else:
            sd = 0
        if effSize == '0':
            for i in range(len(eff_array)-1):
                data_mean.append(vals)
                data_sd.append(sd)
        else:
            data_mean.append(vals)
            data_sd.append(sd)

    eff_bins = [int(x) for x in eff_array]
    eff_bins = eff_bins[1:]
    ax.errorbar(eff_bins,
                data_mean,
                data_sd,
                #sd, fmt = 'o',
                capsize = 2,
                linestyle = LineStyles[curr_pref % len(LineStyles)],
                #linestyle = "None",
                linewidth = w,
                marker = Markers[curr_pref % len(Markers)],
                #marker='o',
                color = c,
                #label = "Number of food patches = " + str(effSize))
                label = UTfolderP2[curr_pref])

    ax.legend(frameon=False, loc="best", prop={'size': 10})
    ax.set_xlabel("Efficiency of Energy transfer (%)")
    ax.set_ylabel("Average Step Taken")

    ax.set_title(plot_title)

    #plotname = path + '/' + plot_title + "_" + _data
    #fig.savefig(plotname + '.pdf', orientation='landscape')
    #fig.savefig(plotname + '.png', orientation='landscape')

    plt.show(block=False)

########## main

curr_p = 0   #curr_path = current path
curr_pref = 0

for f in range(len(UTfolder_reference)):
    data_to_plot = {}
    analyseFolders(path[curr_p], curr_p)
    plotData()
    curr_p += 1
    curr_pref += 1


analyseFolders(path[curr_p], curr_p)
plotData()

curr_pref += 1
for pp in UTprefix:
    UTfolderP = []
    for n in range(len(eff_array)-1):
        UTfolderP.append(pp + eff_array[n+1])
    path = []
    for f in UTfolderP:
        path.append(cwd + '/' + f + '/' + trophFolder + '/Bot*/*Summary.csv')
    #print('path', path)
    #curr_p = 1
    for p in path: #p is the full path depending on which f is currently being used, f is a part of full path
        analyseFolders(p, curr_p)
        curr_p += 1
    plotData()
    curr_pref += 1


plt.show()