
import icrfHelper;
import pyCreeper.crFiles;
import pyCreeper.crData;

import copy;
from helpers import dataFormatHelpers;
import constants;

from helpers import footbotAnalysisHelpers;


EVENTS_RECORD_TIME_COLUMN = 0;
EVENTS_RECORD_EVENT_TYPE_COLUMN = 1;
EVENTS_RECORD_ROBOT_ID_COLUMN = 2;
EVENTS_RECORD_SITE_ID_COLUMN = 3;
EVENTS_RECORD_SITE_VALUE_COLUMN = -1;
EVENTS_RECORD_SITE_VOLUME_COLUMN = 6;
EVENTS_RECORD_RECRUITER_ID_COLUMN = 7;
EVENTS_RECORD_REWARD_RECEIVED_COLUMN = 7;

EVENTS_RECORD_HAS_ZERO_BASED_SITE_NUMBERING = True;




def getRobotWorksiteStateData(filePath_, maxTime_, sitesProduceReward_):

    infoEventsData = pyCreeper.crFiles.fileToArray(filePath_,"\t", False);

    infoEventTimes = pyCreeper.crData.getListColumnAsArray(infoEventsData, EVENTS_RECORD_TIME_COLUMN);

    sitesCompleted = [[] for i in range(maxTime_)];
    sitesSubscribedRobots = [[] for i in range(maxTime_)]
    sitesSubscribedRobotsByRecruitment = [[] for i in range(maxTime_)]

    sitesReachedRobots = [[] for i in range(maxTime_)]
    sitesLadenRobots = [[] for i in range(maxTime_)]
    sitesEarningRobots = [[] for i in range(maxTime_)]
    sitesRobotsToAbandon = [[] for i in range(maxTime_)]

    sitesInitialVolumes = [];
    sitesInitialValues = [];

    numOfSites = [0 for i in range(maxTime_)];
    totalEnvReward = [0 for i in range(maxTime_)];
    actualReward = [0 for i in range(maxTime_)];

    robotIds = []; # an array of known robot ids - robots might be numbered arbitrarily
    robotCurrentTaskId = []; # an array, indexed by robotIndex from robotIds[] with site ids where each robot works
    brokenRobotIds = [[] for i in range(maxTime_)] # an array of robots that broke and should be ignored


    #--
    #-- resolve each event to sort robots into various arrays, then calculate various costs for this time step
    #--
    for t in range(maxTime_):


        #print("-------------------- t={}".format(t));

        #-- maintain actual received reward from previous time step
        if (t > 0):

            numOfSites[t] = numOfSites[t-1];
            totalEnvReward[t] = totalEnvReward[t-1];
            actualReward[t] = actualReward[t-1]

            sitesCompleted[t] = copy.deepcopy(sitesCompleted[t-1]);
            sitesSubscribedRobots[t] = copy.deepcopy(sitesSubscribedRobots[t-1]);
            sitesSubscribedRobotsByRecruitment[t] = copy.deepcopy(sitesSubscribedRobotsByRecruitment[t-1]);
            sitesReachedRobots[t] = copy.deepcopy(sitesReachedRobots[t-1]);
            sitesLadenRobots[t] = copy.deepcopy(sitesLadenRobots[t-1]);
            sitesEarningRobots[t] = copy.deepcopy(sitesEarningRobots[t-1]);
            sitesRobotsToAbandon[t] = copy.deepcopy(sitesRobotsToAbandon[t-1]);
            brokenRobotIds[t] = copy.deepcopy(brokenRobotIds[t-1]);

        if (t in infoEventTimes):
            #print("-- t={}".format(t))
            #-- there was an event at this time
            for loadingTimeDataIndex, timeValue in enumerate(infoEventTimes):

                if (timeValue == t):
                    time, eventType, siteId, robotId = icrfHelper.getInfoEventDataFromRow(infoEventsData[loadingTimeDataIndex]);

                    if (False):
                        printEvents=True;
                        printEventDataRows=True;
                        printTaskInfo = True;
                    else:
                        printEvents=False;
                        printEventDataRows=False;
                        printTaskInfo = False;

                    #-- maintain a set of known robot ids
                    pyCreeper.crData.addUniqueElementToList(robotId,robotIds);
                    robotArrayIndex = robotIds.index(robotId);
                    #print("Known robot ids: {}".format(robotIds));

                    #-- maintain an array of site ids where each robot works
                    if (len(robotCurrentTaskId) < len(robotIds)):
                        robotCurrentTaskId.append(-1);

                    #-- sometimes, siteId needs to be adjusted because it is incorrect in the data. Figure our the siteId out based on subscriptions
                    if ((eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_RECEIVED) or
                            eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_STARTED) or
                            eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_FINISHED) or
                            eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_BROKE))):
                        siteId = robotCurrentTaskId[robotArrayIndex];
                        #print("   changed site id to {}".format(siteId));


                    #--
                    #-- resolve the current event
                    #--
                    if (printEvents):
                        print("-- resolving event {} siteId {} robotId {}".format(eventType, siteId, robotId));
                        if (robotId in brokenRobotIds[t]):
                            print ("      ignoring due to robot broken");


                    if (siteId >= 0 and eventType != dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_SWITCHED) and
                            (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_FIXED) or not(robotId in brokenRobotIds[t]))):

                        if (printEventDataRows):
                            print("t={} {} robot {} (index {}) site {} ".format(timeValue,eventType,robotId,robotArrayIndex,siteId));

                        #-- addition of a new site
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_ADDED)):
                            siteValue = 1.0;
                            if (EVENTS_RECORD_SITE_VALUE_COLUMN >= 0):
                                siteValue = infoEventsData[loadingTimeDataIndex][EVENTS_RECORD_SITE_VALUE_COLUMN];
                            siteVolume = int(round(infoEventsData[loadingTimeDataIndex][EVENTS_RECORD_SITE_VOLUME_COLUMN]));

                            numOfSites[t] += 1;

                            siteReward = siteValue * siteVolume;


                            #-- maintain various lists associated with sites
                            sitesSubscribedRobots[t].append([]);
                            sitesSubscribedRobotsByRecruitment[t].append([]);
                            sitesReachedRobots[t].append([]);
                            sitesEarningRobots[t].append([]);
                            sitesLadenRobots[t].append([]);
                            sitesRobotsToAbandon[t].append([]);

                            sitesInitialVolumes.append(siteVolume);
                            sitesInitialValues.append(siteValue);

                            totalEnvReward[t] += siteReward;

                            if (printEvents):
                                print("t=" + str(t) + " ==========================[]====== NEW SITE vol=" + str(
                                     siteVolume) + " val " + str(siteValue) + "  Nt=" + str(numOfSites[t]) + " TOTAL env R " + str(totalEnvReward[t]));

                        #-- site deletion
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_COMPLETED) or eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_DESTROYED)):
                            if (siteId not in sitesCompleted[t]):  # make sure this only happens once, either when site is completed by robots or destroyed by the environment

                                sitesCompleted[t].append(siteId);
                                siteReward = sitesInitialValues[siteId] * sitesInitialVolumes[siteId];
                                totalEnvReward[t] -= siteReward;

                                if (printEvents):
                                    print("t=" + str(t) + " ============================[]===== SITE DONE " + str(siteId) + " TOTAL env R " + str(totalEnvReward[t]) + " completed sites:" + str(sitesCompleted[t]));

                            #-- if robot stay at sites to get reward, clear various arrays after site deleted, assuming that not all events would be correctly recorded
                            if (sitesProduceReward_):
                                for robotId in sitesReachedRobots[t][siteId]:
                                    if (robotId in sitesSubscribedRobots[t][siteId]):
                                        sitesSubscribedRobots[t][siteId].remove(robotId);
                                    if (robotId in sitesSubscribedRobotsByRecruitment[t][siteId]):
                                        sitesSubscribedRobotsByRecruitment[t][siteId].remove(robotId);
                                    if (robotId in sitesRobotsToAbandon[t][siteId]):
                                        sitesRobotsToAbandon[t][siteId].remove(robotId);

                                    #pyCreeper.crData.removeElementFromList(siteId, robotsTaskSubscriptions[robotArrayIndex])

                                    if (printEvents):
                                        print("     automatically unsibscribed robot " + str(robotId) + " from task " + str(siteId));
                                sitesReachedRobots[t][siteId] = [];
                                sitesEarningRobots[t][siteId] = [];
                                sitesLadenRobots[t][siteId] = [];

                        #-- for any events apart from site creation, the site value should be read from the info that was recorded when site was created
                        if (siteId < len(sitesInitialValues)):
                            siteValue = sitesInitialValues[siteId];
                        else:
                            print ("!! SITE VALUE NOT RECORDED for siteId{}".format(siteId));
                            print ("    t={} {} robot {} ".format(timeValue, eventType, robotId));


                        #-- keep track of robots that know about a site
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_SCOUTED) or eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_RECRUITED)):
                             debugMessage = "t=" + str(t) + " robot " + str(robotId) + " subscribed to site " + str(siteId) + "  ";
                             pyCreeper.crData.addUniqueElementToList(robotId, sitesSubscribedRobots[t][siteId], debugMessage, printEvents);
                             robotCurrentTaskId[robotArrayIndex] = siteId;

                             if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_RECRUITED)):
                                 debugMessage = "t=" + str(t) + " robot " + str(robotId) + " recruited to site " + str(siteId) + "  ";
                                 pyCreeper.crData.addUniqueElementToList(robotId, sitesSubscribedRobotsByRecruitment[t][siteId],debugMessage, printEvents)


                        #-- keep track of robots that reached the site
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_REACHED)):
                            pyCreeper.crData.addUniqueElementToList(robotId, sitesReachedRobots[t][siteId]);
                            if (sitesProduceReward_):
                                pyCreeper.crData.addUniqueElementToList(robotId, sitesLadenRobots[t][siteId]);

                        #-- keep track of robots that left the site laden with resource
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_PAUSED)):
                            pyCreeper.crData.addUniqueElementToList(robotId,sitesLadenRobots[t][siteId]);
                            pyCreeper.crData.removeElementFromList(robotId, sitesReachedRobots[t][siteId]);

                        #-- keep track of robots that started doing a task
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_STARTED)):

                            pyCreeper.crData.addUniqueElementToList(robotId, sitesReachedRobots[t][siteId]);
                            pyCreeper.crData.addUniqueElementToList(robotId, sitesLadenRobots[t][siteId]);

                            if (printEvents):
                                print("t=" + str(t) + " robot " + str(robotId) + " started site " + str(siteId) + "  ");

                        #-- robots that are earning reward
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_STARTED)):

                            outputStr = "t=" + str(t) + " rew started for robot " + str(robotId) + " site " + str(siteId);
                            pyCreeper.crData.addUniqueElementToList(robotId, sitesLadenRobots[t][siteId]);
                            pyCreeper.crData.addUniqueElementToList(robotId, sitesEarningRobots[t][siteId], outputStr, printEvents);


                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_FINISHED)):
                            debugMessage = "t=" + str(t) + " robot " + str(robotId) + " reward finished from site " + str(siteId) + "   " + str(sitesEarningRobots[t][siteId]);
                            pyCreeper.crData.removeElementFromList(robotId, sitesEarningRobots[t][siteId], debugMessage,printEvents);
                            pyCreeper.crData.removeElementFromList(robotId, sitesLadenRobots[t][siteId]);

                            #-- unsubscribe the robot if it was set to abandon the worksite after it unloaded cart
                            if (robotId in sitesRobotsToAbandon[t][siteId]):
                                pyCreeper.crData.removeElementFromList(robotId, sitesSubscribedRobots[t][siteId], debugMessage,printEvents);
                                pyCreeper.crData.removeElementFromList(robotId, sitesSubscribedRobotsByRecruitment[t][siteId]);
                                pyCreeper.crData.removeElementFromList(robotId, sitesReachedRobots[t][siteId]);

                                if (printEvents):
                                     print("t=" + str(t) + " robot " + str(robotId) + " unsubscribed from task " + str(siteId) + " AFTER EMPTYING CART ");


                        #-- reward received
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_RECEIVED)):
                            #print("adding to reward: %".format(infoEventsData[loadingTimeDataIndex][EVENTS_RECORD_REWARD_RECEIVED_COLUMN]));
                            actualReward[t] += infoEventsData[loadingTimeDataIndex][EVENTS_RECORD_REWARD_RECEIVED_COLUMN]/10000.0;

                        #-- abandonment
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_ABANDONED) or eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_SITE_MISSING)):
                            if (robotId in sitesSubscribedRobots[t][siteId]):
                                #-- for the sake of costs, calculation, robot should only unsubscribe when it receives reward from worksite that it wants to abandon
                                if not (robotId in sitesLadenRobots[t][siteId]):
                                    debugMessage = "t=" + str(t) + " robot " + str(robotId) + " unsubscribed from site " + str(siteId);
                                    pyCreeper.crData.removeElementFromList(robotId, sitesSubscribedRobots[t][siteId], debugMessage,printEvents);
                                    pyCreeper.crData.removeElementFromList(robotId, sitesSubscribedRobotsByRecruitment[t][siteId]);
                                    pyCreeper.crData.removeElementFromList(robotId, sitesReachedRobots[t][siteId]);

                                    #pyCreeper.crData.removeElementFromList(siteId, robotsTaskSubscriptions[robotArrayIndex])

                                else:
                                    # -- robot still needs to return the load, remember it should be counted as abandoned after it's unloaded
                                    debugMessage = "t=" + str(t) + " robot " + str(robotId) + " WILL unsubscribe from task " + str(siteId);
                                    pyCreeper.crData.addUniqueElementToList(robotId, sitesRobotsToAbandon[t][siteId], debugMessage,printEvents)

                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_BROKE)):
                            pyCreeper.crData.addUniqueElementToList(robotId, brokenRobotIds[t]);
                            pyCreeper.crData.removeElementFromList(robotId,sitesSubscribedRobots[t][siteId],"",printEvents);
                            pyCreeper.crData.removeElementFromList(robotId,sitesSubscribedRobotsByRecruitment[t][siteId]);
                            pyCreeper.crData.removeElementFromList(robotId,sitesReachedRobots[t][siteId]);
                            pyCreeper.crData.removeElementFromList(robotId,sitesEarningRobots[t][siteId]);
                            pyCreeper.crData.removeElementFromList(robotId,sitesLadenRobots[t][siteId]);

                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_ROBOT_FIXED)):
                            pyCreeper.crData.removeElementFromList(robotId, brokenRobotIds[t]);

                    else:
                        if (eventType == dataFormatHelpers.getEventDescriptionForType(constants.E_REWARD_RECEIVED)):
                            errorStr = "  \n !!! asked to ignore a reward received event is this correct?"
                            errorStr += "  \n event row:{}".format(infoEventsData[loadingTimeDataIndex]);
                            errorStr += "  \n broken robots:{}".format(brokenRobotIds[t]);
                            errorStr += "  \n in file: {}".format(filePath_);
                            raise ValueError(errorStr)


            if (printTaskInfo):
                print("   COMPLETED TASKS: " + str(sorted(sitesCompleted[t])));
                print("   SUB:             " + str(icrfHelper.getTasksInfoForPrinting(sitesSubscribedRobots[t])));
                print("   REACHED:         " + str(icrfHelper.getTasksInfoForPrinting(sitesReachedRobots[t])));
                print("   EARN :           " + str(icrfHelper.getTasksInfoForPrinting(sitesEarningRobots[t])));
                print("   IN CART :        " + str(icrfHelper.getTasksInfoForPrinting(sitesLadenRobots[t])));
                print("   reward :        " + str(actualReward[t]));


        else:
            # -- there was no info event for this time
            pass


    return([
        sitesInitialVolumes, sitesInitialValues,                        #0-1
        numOfSites, totalEnvReward, actualReward,                       #2-4
        sitesCompleted,                                                 #5
        sitesSubscribedRobots, sitesSubscribedRobotsByRecruitment,      #6-7
        sitesReachedRobots, sitesLadenRobots, sitesEarningRobots,        #8-10
        brokenRobotIds,robotIds                                          #11-12

    ])



GET_ROBOT_WORKSITE_STATE_DATA_FUNCTION = getRobotWorksiteStateData;

def getRunICRDataOverTime(filePath_, numOfRobots_, maxTime_, rewardGainRate_, sitesProduceReward_=False, expectedFinalReward_=-1):

    """

    :param runDataPathPrefix_:
    :param numOfRobots_:
    :param maxTime_:
    :param rewardGainRate_: how much reward can be loaded from worksite per second
    :param sitesProduceReward_:
    :param expectedFinalReward_: (optional, default = -1). If >=0, the total reward obtained at the end of the run will be checked and error thrown if the expected reward hasn't been reached
    :return: arrays of data, where the 0th dimension is different types of data and 1st dimension is values over time. The different returned types of data are:
        actualRewardData, expectedRewardData, potentialRewardData,
        infoGainData, infoGainSocialData,
        uncertaintyCostData, displacementCostData, displacementCostCoeffData, misinformationCostData,
        brokenRobotIds,robotIds
    """

    expectedRewardData = [0 for i in range(maxTime_)]; # projected expected reward gain based on number of subscribed robots
    potentialRewardData = [0 for i in range(maxTime_)]; # expected reward + reward that unsubscribed robots could get from active worksites.
                                                        # this is effectively total reward in the environment at the beginning - a multiply of the number of robots subscribed to depleted sites


    infoGainData = [0 for i in range(maxTime_)]; # normalised information gain
    infoGainSocialData = [0 for i in range(maxTime_)]; # normalised information gain only resulting from recruitment

    uncertaintyCostDecreaseData = [0 for i in range(maxTime_)]; # amount by which C_U decreases for a given time step
    uncertaintyCostData = [0 for i in range(maxTime_)]; # C_U
    displacementCostData = [0 for i in range(maxTime_)]; # C_D
    displacementCostCoeffData = [0 for i in range(maxTime_)];
    misinformationCostData = [0 for i in range(maxTime_)]; # C_M

    numOfRewardReceivedEvents = 0; # used for data checks only


    print(">>> Calculating ICR data for {}...".format(filePath_));

    #-- process the info events to obtain states of robots over time
    (sitesInitialVolumes, sitesInitialValues,                        #0-1
        numOfSites, totalEnvReward, actualRewardData,                       #2-4
        sitesCompleted,                                                 #5
        sitesSubscribedRobots, sitesSubscribedRobotsByRecruitment,      #6-7
        sitesReachedRobots, sitesLadenRobots, sitesEarningRobots,       #8-10
        brokenRobotIds,robotIds) = GET_ROBOT_WORKSITE_STATE_DATA_FUNCTION(filePath_,maxTime_,sitesProduceReward_);

    #--
    for t in range(maxTime_):
        printCosts = False;
        printInfoGain = False;
        printTaskInfo = False;
        printRewards = False;

        if (False):
            printCosts = True;
            printInfoGain = True;
            printTaskInfo = True;
            printRewards = True;

        #--
        #-- finished resolving all events for this time frame. Calculate information flow and costs
        #--
        #                           R_T'                 N_W'           N_R
        rewardPerRobotPerSite = (totalEnvReward[0] / ( numOfSites[t] * numOfRobots_));  # adjustemnt for different types of environments
        rewardGainRate = rewardGainRate_;

        uncertaintyCostData[t] = 0
        uncertaintyCostDecreaseData[t] = 0;
        displacementCostData[t] = 0;
        misinformationCostData[t] = 0;

        misinformationCostPaidForLastSite = 0;  # amount of C_M paid when all worksites have been depleted. Used in the checks.

        for siteId in range(numOfSites[t]):
            siteValue = sitesInitialValues[siteId];

            #--
            #-- information gain
            #--
            if (t > 0):
                if not (siteId in sitesCompleted[t]):
                    subscribedRobotsDiff = len(sitesSubscribedRobots[t][siteId]) - len(sitesSubscribedRobots[t-1][siteId]);
                    infoGainData[t] += subscribedRobotsDiff / float(numOfRobots_);

                    subscribedRobotsByRecruitmentDiff = len(sitesSubscribedRobotsByRecruitment[t][siteId]) - len(sitesSubscribedRobotsByRecruitment[t-1][siteId]);
                    infoGainSocialData[t] += subscribedRobotsByRecruitmentDiff / float(numOfRobots_);

                    if (printInfoGain):
                        print("~~ t={} site {} adding to info gain {} (subscribed robot diff={})".format(t,siteId,subscribedRobotsDiff / float(numOfRobots_), subscribedRobotsDiff))


            #--
            #-- expected reward
            #--
            expectedRewardData[t] += len(sitesSubscribedRobots[t][siteId]) * siteValue * rewardPerRobotPerSite;


            #--
            #-- potential reward:
            #--
            potentialRewardData[t] += len(sitesSubscribedRobots[t][siteId]) * siteValue * rewardPerRobotPerSite; # the part from expected reward
            outpPotentialReward = "     potential rew from site {} part 1={}".format(siteId,len(sitesSubscribedRobots[t][siteId]) * siteValue * rewardPerRobotPerSite)
            if not(siteId in sitesCompleted[t]):
                potentialRewardData[t] += (numOfRobots_ - len(sitesSubscribedRobots[t][siteId])) * siteValue * rewardPerRobotPerSite # if active worksites, add reward that could be gained by unemployed robots
                outpPotentialReward += " part 2={}".format((numOfRobots_ - len(sitesSubscribedRobots[t][siteId])) * siteValue * rewardPerRobotPerSite)

            #--
            #-- actual reward:
            #--
            #actualRewardData[t] += len(sitesEarningRobots[t][siteId]) * siteValue * rewardGainRate;

            #--
            #-- uncertainty cost
            #--
            if (not (siteId in sitesCompleted[t])):
                uncertaintyCostData[t]+= (numOfRobots_ * rewardPerRobotPerSite) - len(sitesSubscribedRobots[t][siteId]) * rewardPerRobotPerSite
                uncertaintyCostDecreaseData[t] += len(sitesSubscribedRobots[t][siteId]) * rewardPerRobotPerSite
            else:
                if (len(sitesLadenRobots[t][siteId]) > 0):
                    #print("t={} task {} removed U for {} robots (before U={}) ".format(t, siteId, len(sitesLadenRobots[t][siteId]), uncertaintyCostData[t]))
                    uncertaintyCostDecreaseData[t] += len(sitesLadenRobots[t][siteId]) * rewardPerRobotPerSite; # TODO: wrong - this appear in the paper calculation?

            #--
            #-- displacement cost
            #--
            if (not (siteId in sitesCompleted[t])):
                displacementCostData[t] += (len(sitesSubscribedRobots[t][siteId]) - len(sitesEarningRobots[t][siteId])) * rewardPerRobotPerSite
                #print("t={} task {} ADDed MC for {} robots ".format(t, siteId, (len(sitesSubscribedRobots[t][siteId]) - len(sitesEarningRobots[t][siteId])) ))
            else:
                displacementCostData[t] += (len(sitesLadenRobots[t][siteId]) - len(sitesEarningRobots[t][siteId])) * rewardPerRobotPerSite
                #print("t={} task {} ADDed MC for {} robots ".format(t, siteId, (len(sitesLadenRobot[t]s[siteId]) - len(sitesEarningRobots[t][siteId])) ))

            #--
            #-- misinformation cost
            #--
            if ((siteId in sitesCompleted[t])):
                costVal = (len(sitesSubscribedRobots[t][siteId]) - len(sitesLadenRobots[t][siteId])) * rewardPerRobotPerSite
                if (len(sitesCompleted[t]) < numOfSites[t]):
                    misinformationCostData[t] += costVal
                else:
                    misinformationCostPaidForLastSite += costVal;

        #-- finished looping through site ids


        #-- displacement cost coefficient
        if (uncertaintyCostDecreaseData[t] > 0):
            displacementCostCoeffData[t] = displacementCostData[t] / uncertaintyCostDecreaseData[t];


        #--
        #-- print stuff
        #--
        if (printCosts):
            if (t > 0 and (infoGainData[t] != infoGainData[t-1] or uncertaintyCostData[t] != uncertaintyCostData[t - 1] or displacementCostData[t] != displacementCostData[t - 1] or misinformationCostData[t] != misinformationCostData[t - 1])):
                print("   t={} C_U={} C_D={} C_M={} dC_U=-{} d={}  r={}    infoGain'={} ".format(
                         t, uncertaintyCostData[t], displacementCostData[t], misinformationCostData[t], uncertaintyCostDecreaseData[t], displacementCostCoeffData[t],
                         rewardPerRobotPerSite, infoGainData[t]));
                print(outpPotentialReward)

                if (printTaskInfo):
                    print("   COMPLETED TASKS: " + str(sorted(sitesCompleted[t])));
                    print("   SUB:             " + str(icrfHelper.getTasksInfoForPrinting(sitesSubscribedRobots[t])));
                    print("   REACHED:         " + str(icrfHelper.getTasksInfoForPrinting(sitesReachedRobots[t])));
                    #print("PROC: " + str(tasksProcessingRobots));
                    #print("READY TO EARN : " + str(tasksReadyToGainRobots));
                    print("   EARN :           " + str(icrfHelper.getTasksInfoForPrinting(sitesEarningRobots[t])));
                    print("   IN CART :        " + str(icrfHelper.getTasksInfoForPrinting(sitesLadenRobots[t])));

        if (printRewards):
            if (t > 0 and (actualRewardData[t-1] < totalEnvReward[t])):
                print(".. t={} R'={} R*={}    R={} earning robots:{}".format(t, expectedRewardData[t], potentialRewardData[t], actualRewardData[t], sitesEarningRobots[t]));

        #--
        #-- do checks
        #--

        #-- various equations
        if (t > 0):
            dR = actualRewardData[t] - actualRewardData[t-1];

            #-- robot cannot be subscribed to 2 worksites at once
            for robotId in range(numOfRobots_):
                numOfSubscriptions = 0;
                for siteId in range(numOfSites[t]):
                    if (robotId in sitesSubscribedRobots[t][siteId]):
                        numOfSubscriptions +=1
                        if (numOfSubscriptions > 1):
                            errorStr = "Robot {} subscribed in > 1 worksites)".format(robotId)
                            errorStr += ("\nCOMPLETED TASKS: " + str(sorted(sitesCompleted[t])));
                            errorStr += ("\nSUB:             " + str(icrfHelper.getTasksInfoForPrinting(sitesSubscribedRobots[t])));
                            errorStr += ("\nREACHED:         " + str(icrfHelper.getTasksInfoForPrinting(sitesReachedRobots[t])));
                            errorStr += ("\nEARN :           " + str(icrfHelper.getTasksInfoForPrinting(sitesEarningRobots[t])));
                            errorStr += ("\nIN CART :        " + str(icrfHelper.getTasksInfoForPrinting(sitesLadenRobots[t])));
                            errorStr += "   in file: {}".format(filePath_);
                            raise ValueError(errorStr)

            #-- C_D + C_M = R' - (r/ro)*dR
            eqLeftSide = displacementCostData[t] + misinformationCostData[t] + misinformationCostPaidForLastSite;
            eqRightSide = expectedRewardData[t] - (rewardPerRobotPerSite / rewardGainRate)*dR;
            if (abs(eqLeftSide - eqRightSide) > 0.0001):
                errorStr = "!!!!!!!!! t={} C_D + C_M != R' - (r/ro)*dR : costsAddedUp={}  C_M for last site={},   R'-(r/ro)*dR={}   (r/ro)={}    dR={} ".format(
                    t, eqLeftSide, misinformationCostPaidForLastSite, eqRightSide, (rewardPerRobotPerSite / rewardGainRate), dR);
                errorStr += "   in file: {}".format(filePath_);
                raise ValueError(errorStr)


            #-- C_U + C_D + C_M = R* - (r/ro)*dR
            eqLeftSide = uncertaintyCostData[t] + displacementCostData[t] + misinformationCostData[t] + misinformationCostPaidForLastSite;
            eqRightSide = potentialRewardData[t] - (rewardPerRobotPerSite / rewardGainRate)*dR;
            if (abs(eqLeftSide - eqRightSide) > 0.0001):
                errorStr = "!!!!!!!!! t={} C_U + C_D + C_M != R* - (r/ro)*dR : costsAddedUp={}  C_M for last site={},   R*'-(r/ro)*dR={}   (r/ro)={}    dR={} ".format(
                    t, eqLeftSide, misinformationCostPaidForLastSite, eqRightSide, (rewardPerRobotPerSite / rewardGainRate), dR);
                errorStr += "   in file: {}".format(filePath_);
                raise ValueError(errorStr)

            #-- m >= 0
            if (displacementCostCoeffData[t] < 0):
                errorStr = "!!!!!!!!! t={} m < 0)".format(t);
                errorStr += "\nC_U=" + str(uncertaintyCostData[t]) + " C_D=" + str(displacementCostData[t]) + " C_M=" + str(misinformationCostData[t]) + " dR=" + str(dR)
                errorStr += "   in file: {}".format(filePath_);
                raise ValueError(errorStr)

            #-- C_M >= 0
            if (misinformationCostData[t] < 0):
                errorStr = "!!!!!!!!! t={} C_M < 0 ({})".format(t, misinformationCostData[t]);
                errorStr += "   in file: {}".format(filePath_);
                raise ValueError(errorStr)

        #-- finished resolving one time frame

    if (expectedFinalReward_ >=0 and actualRewardData[maxTime_-1] < expectedFinalReward_):
        errorStr = " !!! total reward gained R={}, which is less than expected ({}) - maybe the maximum run time needs to be extended?".format(actualRewardData[maxTime_-1], expectedFinalReward_);
        errorStr += "    num of reward received events {}".format(numOfRewardReceivedEvents);
        errorStr += "   in file: {}".format(filePath_);
        raise ValueError(errorStr)

    return [actualRewardData, expectedRewardData, potentialRewardData,
        infoGainData, infoGainSocialData,
        uncertaintyCostData, displacementCostData, displacementCostCoeffData, misinformationCostData
        ]

#===================== SPECIAL CALCULATIONS =====================


def calculateInformationGainRate(infoGainDataOverTime_, infoGainCompressionTimeInterval_):

    infoGainDataOverTimeCompressed = pyCreeper.crData.compressList(infoGainDataOverTime_,infoGainCompressionTimeInterval_, useAverages_=False);

    #print(infoGainDataOverTime_);
    #print(infoGainDataOverTimeCompressed);

    positiveRegionStartTime = -1;
    positiveRegionStartVal = -1;
    regionInfoGainRates = [];
    positiveRegionTotalInfoGain = 0;

    for timeBin in range(0,len(infoGainDataOverTimeCompressed)):
        currentInfoGain = infoGainDataOverTimeCompressed[timeBin];
        time = timeBin * infoGainCompressionTimeInterval_;

        #-- positive region boundaries
        if (currentInfoGain > 0):
            if (positiveRegionStartTime < 0):
                positiveRegionStartTime = time;
                positiveRegionStartVal = currentInfoGain # infoGainDataOverTimeCompressed[timeBin-1][0];
                #print("Positive region started at " + str(time) + " with info gain " + str(positiveRegionStartVal))
                positiveRegionTotalInfoGain = 0;

        #-- max value of a current positive region
        if (positiveRegionStartTime >= 0):
            positiveRegionTotalInfoGain += currentInfoGain;

        if (currentInfoGain <= 0 and positiveRegionStartTime >= 0):
            endBin = timeBin-1;
            endTime = endBin * infoGainCompressionTimeInterval_;
            #print("-- Positive region end at " + str(endTime))

            #-- calculate the info gain rate in this region
            dx = endTime - positiveRegionStartTime
            duration = dx + infoGainCompressionTimeInterval_ # if start and end at the same time, duration of that positive region is 1 time bin interval
            #regionInfoGainRate = positiveRegionTotalInfoGain / (duration / compressionTimeInterval);
            regionInfoGainRate = positiveRegionTotalInfoGain / duration ;

            #print("   Start: [{},{}] duration: {}, total info gain={}, region info gain rate={} ".format(positiveRegionStartTime,positiveRegionStartVal, duration, positiveRegionTotalInfoGain, regionInfoGainRate))
            regionInfoGainRates.append(regionInfoGainRate);
            positiveRegionStartTime = -1;


    #-- get max info speed for this run
    if (len(regionInfoGainRates) > 0):
        #print("Returning {}".format(max(regionInfoGainRates)))
        return (max(regionInfoGainRates));
    else:
        return 0;



