import pyCreeper.crData;
import pyCreeper.crGraphs;

from helpers import dataFormatHelpers;
from helpers import helpers;
import analysis;
import constants;


def visualiseICRData(icrData_, graphFilePath_, dataDimension0Labels_, dataDimension1Labels_, markers_, lineStyles_):
    """
    :param icrData: as returned by analysis.getICRData
    :param graphFilePath: full path to the output graphs file, using *GRAPH_TYPE* in it as a placeholder for different graph types
    :param dataDimension0Labels: list of labels for dimension 0 of icrData, e.g. a list of controller names
    :param dataDimension1Labels: list of labels for dimension 1 of icrData, e.g. a list of environment names
    :param markers_:
    :param lineStyles_:
    """

    #-- the graphs will be put into a subfolder
    subFolder = "icrf";
    folderPath = graphFilePath_[:(graphFilePath_.rfind("/")+1)] + subFolder;
    helpers.createFolder(folderPath)

    graphFileName = graphFilePath_[(graphFilePath_.rfind("/")+1):];
    graphFilePath_ = folderPath + "/" + graphFileName;
    #--

    infoGainRateData = [];
    uncertaintyCost = [];
    displacementCost = [];
    misinformationCost = [];
    uncertaintyCostPerMinute = [];
    displacementCostPerMinute = [];
    misinformationCostPerMinute = [];
    for a in range(len(dataDimension0Labels_)):
        infoGainRateData.append([]);
        uncertaintyCost.append([]);
        displacementCost.append([]);
        misinformationCost.append([]);
        uncertaintyCostPerMinute.append([]);
        displacementCostPerMinute.append([]);
        misinformationCostPerMinute.append([]);
        for b in range(len(dataDimension1Labels_)):
            infoGainRateData[a].append(icrData_[a][b][0]);
            uncertaintyCost[a].append(icrData_[a][b][1]);
            displacementCost[a].append(icrData_[a][b][2]);
            misinformationCost[a].append(icrData_[a][b][4]);
            uncertaintyCostPerMinute[a].append(icrData_[a][b][5]);
            displacementCostPerMinute[a].append(icrData_[a][b][6]);
            misinformationCostPerMinute[a].append(icrData_[a][b][7]);

    xAxisGroupSize = 3;
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        xAxisGroupSize = 4;

    #-- info gain rate
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;

    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_infoGainRate_Ti{}".format(analysis.INFO_GAIN_COMPRESSION_TIME_INTERVAL))
    pyCreeper.crGraphs.createLinePlot(infoGainRateData, markers_=markers_, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize, lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="Information gain rate", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    #-- uncertainty cost
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_uncertaintyCost");
    pyCreeper.crGraphs.createLinePlot(uncertaintyCost, markers_=markers_, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="Total $C_U$", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);


    #-- displacement cost
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_displacementCost");
    pyCreeper.crGraphs.createLinePlot(displacementCost, markers_=markers_, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="Total $C_D$", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    #-- misinformation cost
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_misinformationCost");
    pyCreeper.crGraphs.createLinePlot(misinformationCost, markers_=markers_, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="Total $C_M$", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    #-- uncertainty cost per minute
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        yMin = pyCreeper.crGraphs.INVALID_VALUE;
        yMax = pyCreeper.crGraphs.INVALID_VALUE;
        yTickStep = 0;
    else:
        yMin = 0;
        yMax = 3000;
        yTickStep = 500;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_uncertaintyCostPM");
    pyCreeper.crGraphs.createLinePlot(uncertaintyCostPerMinute, markers_=markers_, yTicksStep_=yTickStep, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="$C_U$ per minute", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    #-- displacement cost per minute
    yMin = 0;
    yMax = 5000;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_displacementCostPM");
    pyCreeper.crGraphs.createLinePlot(displacementCostPerMinute, markers_=markers_, doWilcoxon_=True, yTicksStep_=500, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="$C_D$ per minute", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);

    #-- misinformation cost per minute
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","icrf_misinformationCostPM");
    pyCreeper.crGraphs.createLinePlot(misinformationCostPerMinute, markers_=markers_, doWilcoxon_=True, lineWidth_=2, xAxisGroupSize_=xAxisGroupSize,lineStyles_=lineStyles_, yMin_=yMin,yMax_=yMax, yLabel_="Total $C_M$ per minute", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath, size_=(15,5), gridType_=pyCreeper.crGraphs.GRID_TYPE.MAJOR, markerSize_=10);


def visualiseICRDataOverTime(icrDataOverTime_, graphFilePath_="", maxTime_ = 0, timeBinSize_=20):
    """
    Create a visualisation of icr data over time, for a particular scenario.
    :param icrDataOverTime: either icrData for 1 run or multiple runs (array of icrData arrays)
    :param graphFilePath_: full path to the output graphs file, using *GRAPH_TYPE* in it as a placeholder for different graph types
    :param maxTime_: Maxium time to show on graph. The rest of the data will be chopped off. Set to 0 to show all data
    """

    #-- the graphs will be put into a subfolder
    subFolder = "icrfOverTime";
    folderPath = graphFilePath_[:(graphFilePath_.rfind("/")+1)] + subFolder;
    helpers.createFolder(folderPath)

    graphFileName = graphFilePath_[(graphFilePath_.rfind("/")+1):];
    graphFilePath_ = folderPath + "/" + graphFileName;
    #--

    #print(icrDataOverTime_);
    #print(len(icrDataOverTime_[0][0]));

    visData = [];
    if (pyCreeper.crData.getNumberOfListDimensions(icrDataOverTime_) == 2):
        visData.append(icrDataOverTime_);
    else:
        visData = icrDataOverTime_;

    #-- chop off the end of data if requested
    if (maxTime_ > 0):
        for x in range(len(visData)):
            for y in range(len(visData[x])):
                visData[x][y] = visData[x][y][:maxTime_];


    #-- compress data
    potentialRewardData = [];
    actualRewardData = [];
    infoGainData = [];
    infoGainSocialData = [];
    cuData = [];
    cdData = [];
    cmData = [];

    #print(visData[x][3])
    #for i in range(len(visData[x][3])):
    #    if (visData[x][3][i] > 0):
    #        print("non-zero info gain at t={} value={}".format(i, visData[x][3][i]));

    for x in range(len(visData)):
        potentialRewardData.append(pyCreeper.crData.compressList(visData[x][2],timeBinSize_));
        actualRewardData.append(pyCreeper.crData.compressList(visData[x][0],timeBinSize_));
        infoGainData.append(pyCreeper.crData.compressList(visData[x][3],timeBinSize_,useAverages_=False));
        infoGainSocialData.append(pyCreeper.crData.compressList(visData[x][4],timeBinSize_,useAverages_=False));
        cuData.append(pyCreeper.crData.compressList(visData[x][5],timeBinSize_));
        cdData.append(pyCreeper.crData.compressList(visData[x][6],timeBinSize_));
        cmData.append(pyCreeper.crData.compressList(visData[x][8],timeBinSize_));


    potentialRewardData = pyCreeper.crData.getListByFlippingColumnsAndRows(potentialRewardData);
    actualRewardData = pyCreeper.crData.getListByFlippingColumnsAndRows(actualRewardData);
    infoGainData = pyCreeper.crData.getListByFlippingColumnsAndRows(infoGainData);
    infoGainSocialData = pyCreeper.crData.getListByFlippingColumnsAndRows(infoGainSocialData);
    cuData = pyCreeper.crData.getListByFlippingColumnsAndRows(cuData);
    cdData = pyCreeper.crData.getListByFlippingColumnsAndRows(cdData);
    cmData = pyCreeper.crData.getListByFlippingColumnsAndRows(cmData);


    #-- create plot labels
    timeLabels = [];
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        xLabel = "Time [hours]";
        timeLabelInterval = 0.5;
        for i in range(len(potentialRewardData)):
            time = (i*timeBinSize_) / 3600 # converting seconds to hours;
            if (time % timeLabelInterval == 0):
                timeLabels.append("{}".format(time));
            else:
                timeLabels.append("")
    else:
        xLabel = "Time [minutes]";
        timeLabelInterval = 1;
        if (len(visData[0][0]) > 1000):
            timeLabelInterval = 2;

        for i in range(len(potentialRewardData)):
            time = (i*timeBinSize_) / 60 # converting seconds to minutes;
            if (time % timeLabelInterval == 0):
                timeLabels.append("{}".format(time));
            else:
                timeLabels.append("")


    yLabel = "Value";

    #-- info gain over time
    legendLabels = ["$\Delta$I'", "$\Delta$I' recruitment only"];
    pyCreeper.crGraphs.createLinePlot([infoGainData, infoGainSocialData],
        yLabel_=yLabel, yMin_=0, yMax_=1.0,xLabel_=xLabel,
        showAverages_=True, xTickLabels_=timeLabels, markerSize_=0, yTicksStep_=0.1, lineWidth_=4, colors_=['m','g'], lineStyles_=['-','--'], size_=(15,5),legendLabels_=legendLabels, numOfLegendColumns_=1, filePath_=graphFilePath_.replace("*GRAPH_TYPE*","infoGain"));

    #-- costs over time
    legendLabels = ['Actual reward ($R$)',"Uncertainty cost ($C_U$)","Displacement cost ($C_D$)", "Misinformation cost ($C_M$)"];

    pyCreeper.crGraphs.createLinePlot([actualRewardData, cuData, cdData,cmData],
        yLabel_=yLabel, yMin_=0, yMax_=dataFormatHelpers.getMaxRewardForExperiment(graphFilePath_),xLabel_=xLabel, size_=(12,4), yTicksStep_=20,
        showAverages_=True, xTickLabels_=timeLabels, markerSize_=0, lineWidth_=4, colors_=['b','r','g','m'], legendLabels_=legendLabels, numOfLegendColumns_=1, filePath_=graphFilePath_.replace("*GRAPH_TYPE*","costs"));



def visualiseICRSecondaryData(icrSecondaryData_, graphFilePath_, dataDimension0Labels_, dataDimension1Labels_, doWilcoxon_=True, pyCreeperStyle_ = None):
    """
    :param icrData: as returned by analysis.getICRSecondaryData
    :param graphFilePath: full path to the output graphs file, using *GRAPH_TYPE* in it as a placeholder for different graph types
    :param dataDimension0Labels: list of labels for dimension 0 of icrData, e.g. a list of controller names
    :param dataDimension1Labels: list of labels for dimension 1 of icrData, e.g. a list of environment names
    :param markers_:
    :param lineStyles_:
    """

    if (pyCreeperStyle_ == None):
        pyCreeperStyle_ = pyCreeper.crGraphStyle.crGraphStyle();

    pyCreeperStyle_.gridType = pyCreeper.crGraphStyle.GRID_TYPE.MAJOR;
    pyCreeperStyle_.figureSize = (15,5);
    pyCreeper.crGraphs.setStyle(pyCreeperStyle_);

    #-- the graphs will be put into a subfolder
    subFolder = "icrfSecondary";
    folderPath = graphFilePath_[:(graphFilePath_.rfind("/")+1)] + subFolder;
    helpers.createFolder(folderPath)

    graphFileName = graphFilePath_[(graphFilePath_.rfind("/")+1):];
    graphFilePath_ = folderPath + "/" + graphFileName;
    #--

    #print(icrSecondaryData_);
    unemploymentRates = [];
    unevenEmploymentRates = [];
    foragingTripTravelTimes = [];
    worksiteDiscoveryTimes = [];

    for a in range(len(dataDimension0Labels_)):
        unemploymentRates.append([]);
        unevenEmploymentRates.append([]);
        foragingTripTravelTimes.append([]);
        worksiteDiscoveryTimes.append([]);
        for b in range(len(dataDimension1Labels_)):
            unemploymentRates[a].append([]);
            unevenEmploymentRates[a].append([]);
            foragingTripTravelTimes[a].append(icrSecondaryData_[a][b][2])
            worksiteDiscoveryTimes[a].append([]);
            for r in range(len(icrSecondaryData_[a][b][0])):
                unemploymentRates[a][b].append((int)(icrSecondaryData_[a][b][0][r] * 100)) # convert to percentage
                unevenEmploymentRates[a][b].append(icrSecondaryData_[a][b][1][r] * 100) # convert to percentage
                worksiteDiscoveryTimes[a][b].append(icrSecondaryData_[a][b][3][r]/60.0) # convert to minutes


    #print(foragingTripTravelTimes);
    xAxisGroupSize = 3;
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        xAxisGroupSize = 4;

    #-- unemployment rate
    yMin = pyCreeper.crGraphs.INVALID_VALUE;
    yMax = pyCreeper.crGraphs.INVALID_VALUE;

    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","unemploymentRate")
    pyCreeper.crGraphs.createLinePlot(unemploymentRates, doWilcoxon_=doWilcoxon_, xAxisGroupSize_=xAxisGroupSize, yMin_=yMin,yMax_=yMax, useBoxPlotPadding_=True, yLabel_="Average unemploy-\nment rate [%]", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath);

    #-- uneven employment rate
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        yMax = pyCreeper.crGraphs.INVALID_VALUE;
    else:
        yMax = 70;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","unevenEmploymentRate");
    pyCreeper.crGraphs.createLinePlot(unevenEmploymentRates, doWilcoxon_=doWilcoxon_, xAxisGroupSize_=xAxisGroupSize, yMin_=yMin,yMax_=yMax, useBoxPlotPadding_=True, yLabel_="Average uneven \nemployment rate", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath);

    #-- foraging trip travel times
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        yMax = pyCreeper.crGraphs.INVALID_VALUE;
    else:
        yMax = 40;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","foragingTripTravelTimes");
    pyCreeper.crGraphs.createLinePlot(foragingTripTravelTimes, doWilcoxon_=doWilcoxon_, xAxisGroupSize_=xAxisGroupSize, yMin_=yMin,yMax_=yMax, useBoxPlotPadding_=True, yLabel_="Average foraging trip \n travel time [secs]", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath);

    #-- average worksite discovery time
    if (constants.USING_FOOTBOTS_SIMULATION_RESULTS):
        yMax = pyCreeper.crGraphs.INVALID_VALUE;
    else:
        yMax = pyCreeper.crGraphs.INVALID_VALUE;
    pyCreeperStyle_.yOffset = 0.021;
    plotFilePath = graphFilePath_.replace("*GRAPH_TYPE*","worksiteDiscoveryTime");
    pyCreeper.crGraphs.createLinePlot(worksiteDiscoveryTimes, doWilcoxon_=doWilcoxon_, xAxisGroupSize_=xAxisGroupSize, yMin_=yMin,yMax_=yMax, useBoxPlotPadding_=True, yLabel_="Average worksite \n discovery time [mins]", xTickLabels_=dataDimension1Labels_, legendLabels_=dataDimension0Labels_, showBoxPlots_=True, filePath_=plotFilePath);



def getPlotStylesForControllers(controllersArray_, noizeValsArray_):
    markers = [];
    lineStyles = [];
    for i in range(len(controllersArray_)):
        controllerId = dataFormatHelpers.getControllerPathName(controllersArray_[i])
        if (controllerId == 'S' or controllerId == "solitary"):
            if ("gpev2000" in noizeValsArray_[i]):
                markers.append('bo-');
            else:
                markers.append('bo-');
            lineStyles.append('-');
        elif (controllerId == 'S_OLDmov'):
            markers.append('bo--');
            lineStyles.append('--');
        elif (controllerId == 'LB_commD12500' or controllerId == "localBroadcaster"):
            if ("gpev2000" in noizeValsArray_[i]):
                markers.append('rs-');
            else:
                markers.append('rs-');

            lineStyles.append('-');
    return (markers, lineStyles);



