Area_Cover_along_probability

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


plot_title = "Area Covered 225 Robots Finite Food"

#list of string
data_to_import = ['total_area_covered',
                'n_bots']

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

#nitialises an empty dictionary
data = {}
data_to_plot = {}

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

path = []

nbots_array = ['1', '25', '81', '169', '225']
levyP_array = ['10', '20', '30', '40', '50', '60', '70', '80', '90', '100']

testNumber = 'simulation_results'
cwd = '/Users/ammacpro1/simulation_results/self_energy_robot/' + testNumber

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

UTprefix = 'Levywalklike_P'

#UTfolder = ['No_Trophallaxis','Trophallaxis_generous', 'Th50','Th70','Th90',]
UTfolder = [UTprefix + p for p in levyP_array]

trophFolder = 'FiniteFood'

for f in UTfolder:
    path.append(cwd + '/' + f + '/' + trophFolder + '/Bot*/*Summary.csv')

def analyseFolders(folder, curr_p):
    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)

    swarmSize = sorted(list(set(data["n_bots"])), key=int)

    for SSize in swarmSize:
        for _data in [key for key, value in data.items()]:
            data_to_plot[_data + "_levyP" + levyP_array[curr_p] + "_bots" + SSize] = \
                    [float(i)
                     for i, j in zip(
                        data[_data],
                        data["n_bots"])
                     if j == SSize]

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

def plotData(b, drawError):
    #print(data_to_plot)

    data_mean = []
    data_sd = []

    for PSize in levyP_array:
        c = colors[curr_b]
        #w = widths[curr_b]
        w = 2
        vals = []
        sd = []
        vals = np.mean(data_to_plot["total_area_covered_levyP" + PSize + "_bots" + b])
        if (drawError):
            sd = np.std(data_to_plot["total_area_covered_levyP" + PSize + "_bots" + b])
        else:
            sd = 0
        data_mean.append(vals)
        data_sd.append(sd)

    levyP_bins = [int(x) for x in levyP_array]
    ax.errorbar(levyP_bins,
                data_mean,
                data_sd,
                #sd, fmt = 'o',
                capsize = 2,
                linestyle = LineStyles[curr_b % len(LineStyles)],
                #linestyle = "None",
                linewidth = w,
                marker = Markers[curr_b % len(Markers)],
                #marker='o',
                color = c,
                #label = "Number of food patches = " + str(PSize))
                label = nbots_array[curr_b] + " Robots")

    ax.legend(frameon=False, loc="best", prop={'size': 10})
    ax.set_xlabel("% Probability")
    #if (drawY == 1):
    ax.set_ylabel("Average Area Covered")

    ax.set_title(plot_title)

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


curr_p = 0   #curr_path = current path
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
curr_b = 0  #current bots group
for b in nbots_array:
    plotData(b, drawError)
    curr_b += 1

plt.show()