import pyCreeper;
pyCreeper.crGraphs.DEBUG_LEVEL = 0;

import analysis;
import visualisation;

import constants;
constants.USING_FOOTBOTS_SIMULATION_RESULTS = True;

from helpers import dataFormatHelpers, footbotAnalysisHelpers;

#-- setup icrfAnalyser
import icrfAnalyser;
icrfAnalyser.EVENTS_RECORD_TIME_COLUMN = 0;
icrfAnalyser.EVENTS_RECORD_EVENT_TYPE_COLUMN = 1;
icrfAnalyser.EVENTS_RECORD_ROBOT_ID_COLUMN = 2;
icrfAnalyser.EVENTS_RECORD_SITE_ID_COLUMN = 3;
icrfAnalyser.EVENTS_RECORD_SITE_VALUE_COLUMN = 4;
icrfAnalyser.EVENTS_RECORD_SITE_VOLUME_COLUMN = 5;
icrfAnalyser.EVENTS_RECORD_RECRUITER_ID_COLUMN = 7;
icrfAnalyser.EVENTS_RECORD_REWARD_RECEIVED_COLUMN = -1;

icrfAnalyser.GET_ROBOT_WORKSITE_STATE_DATA_FUNCTION = footbotAnalysisHelpers.getRobotWorksiteStateDataForFootbots;

analysis.INFO_GAIN_COMPRESSION_TIME_INTERVAL = 60;
analysis.RECALC_CLEANED_DATA = False;
#--



nrVals = [10];

experimentNamePrefix = "footBotsArgos"
controllers = ['Solitary','Local broadcaster'];

envs = ['Heap1-8','Heap2-8','Heap4-8','Scatter25-8','Heap1-12','Heap2-12','Heap4-12','Scatter25-12'];

controllers = ['Local broadcaster'];
envs = ['Heap2-12'];
#envs = ['Heap1-12','Heap2-12','Heap4-12','Scatter25-12'];

numRuns = 1;

DO_COMPLETION_TIME = False;
DO_ICR_ANALYSIS = False;
DO_ICR_ANALYSIS_OVER_TIME = True;
DO_ICR_SECONDARY_ANALYSIS = False;

for nr in range(len(nrVals)):
    nrVal = nrVals[nr];

    results_completionTime = [[[] for j in range(len(envs))] for k in range(len(controllers))];
    results_icrData = [[[] for j in range(len(envs))] for k in range(len(controllers))];
    results_icrDataOverTime = [[[] for j in range(len(envs))] for k in range(len(controllers))];
    results_icrSecondaryData = [[[] for j in range(len(envs))] for k in range(len(controllers))];

    for c in range(len(controllers)):
        controllerVal = controllers[c];
        print(controllerVal);

        for env in range(len(envs)):
            envName = envs[env];
            print(envName);
            icrData = [];

            dataPathPrefix = "A_{}/COLLECT_RT1-1_{}_{}_NR{}".format(dataFormatHelpers.getControllerPathName(controllerVal), dataFormatHelpers.getEnvironmentPathName(envName), dataFormatHelpers.getControllerPathName(controllerVal), nrVal);
            if (DO_COMPLETION_TIME):
                results_completionTime[c][env] = analysis.getCompletionTimes(dataPathPrefix, numRuns, convertToHours_=True);

            if (DO_ICR_ANALYSIS):
                results_icrData[c][env] = analysis.getICRData(dataPathPrefix, numRuns, nrVal, dataFormatHelpers.getMaxTimeForExperiment(experimentNamePrefix, nrVal), dataFormatHelpers.getRewardGainRateForExperiment(experimentNamePrefix + dataPathPrefix))

            if (DO_ICR_ANALYSIS_OVER_TIME):
                results_icrDataOverTime[c][env] = analysis.getICRDataOverTime(dataPathPrefix, numRuns, nrVal, dataFormatHelpers.getMaxTimeForExperiment(experimentNamePrefix, nrVal), dataFormatHelpers.getRewardGainRateForExperiment(experimentNamePrefix + dataPathPrefix));

            if (DO_ICR_SECONDARY_ANALYSIS):
                results_icrSecondaryData[c][env] = analysis.getICRSecondaryData(dataPathPrefix, numRuns, nrVal, dataFormatHelpers.getMaxTimeForExperiment(experimentNamePrefix, nrVal));

    markers, lineStyles = visualisation.getPlotStylesForControllers(controllers,[experimentNamePrefix for x in range(len(controllers))]);

    if (DO_COMPLETION_TIME):
        plotFilePath="{}/{}/completionTime_NR{}".format(constants.PLOTS_PATH,experimentNamePrefix,nrVal);
        plotFilePath += ".png";
        yMin = 0;
        yMax = 5;
        pyCreeper.crGraphs.createLinePlot(results_completionTime, markers_=markers, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=4, lineStyles_=lineStyles, yTicksStep_=1, yMin_=yMin,yMax_=yMax, yLabel_="Completion time [hours]", xTickLabels_=envs, legendLabels_=controllers, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    if (DO_ICR_ANALYSIS):
        plotFilePath="{}/{}/*GRAPH_TYPE*_NR{}.png".format(constants.PLOTS_PATH,experimentNamePrefix,nrVal,);
        visualisation.visualiseICRData(results_icrData,plotFilePath,controllers,envs,markers,lineStyles);



    if (DO_ICR_ANALYSIS_OVER_TIME):
        for c in range(len(controllers)):
            controllerVal = controllers[c];
            for env in range(len(envs)):
                envName = envs[env];
                plotFilePath="{}/{}/overTime_*GRAPH_TYPE*_NR{}_{}_{}".format(constants.PLOTS_PATH,experimentNamePrefix,nrVal,envName,dataFormatHelpers.getControllerPathName(controllerVal));
                if (constants.USING_SIMULATION_RESULTS):
                    plotFilePath+="_sim";
                plotFilePath+="_{}runs.png".format(numRuns);
                visualisation.visualiseICRDataOverTime(results_icrDataOverTime[c][env],plotFilePath,maxTime_=4*3600,timeBinSize_=180);

    if (DO_ICR_SECONDARY_ANALYSIS):
        plotFilePath="{}/{}/*GRAPH_TYPE*_NR{}.png".format(constants.PLOTS_PATH,experimentNamePrefix,nrVal);
        visualisation.visualiseICRSecondaryData(results_icrSecondaryData,plotFilePath,controllers,envs,markers,lineStyles);