import math
mvanilla="IMPALA"
mcount="COUNT"
mride="RIDE"
mcuriosity = "CURIOSITY"
mrnd = "RND"
mamigo = "AMIGO"

M={ 1: ("MiniGrid-Unlock-v0", mvanilla),
    2: ("MiniGrid-Unlock-v0", mcount),
    3: ("MiniGrid-Unlock-v0", mride),
    4: ("MiniGrid-DoorKey-8x8-v0", mcount),
    5: ("MiniGrid-Unlock-v0", mcuriosity),
    6: ("MiniGrid-Unlock-v0", mrnd),
    7: ("MiniGrid-DoorKey-8x8-v0", mride),
    8: ("MiniGrid-DoorKey-8x8-v0", mvanilla),
    9: ("MiniGrid-LavaGapS7-v0", mamigo),
   10: ("MiniGrid-Unlock-v0", mamigo),
   11: ("MiniGrid-DoorKey-8x8-v0", mamigo),
   12: ("MiniGrid-Empty-16x16-v0", mamigo),
   13: ("MiniGrid-LavaGapS7-v0", mamigo),
   14: ("MiniGrid-SimpleCrossingS11N5-v0", mamigo),
   15: ("MiniGrid-DistShift2-v0", mamigo),
   16: ("MiniGrid-DoorKey-8x8-v0", mrnd),
   17: ("MiniGrid-DoorKey-8x8-v0", mcuriosity),
   18: ("MiniGrid-Empty-16x16-v0", mvanilla),
   19: ("MiniGrid-Empty-16x16-v0", mcount),
   20: ("MiniGrid-Empty-16x16-v0", mcuriosity),
   21: ("MiniGrid-Empty-16x16-v0", mrnd),
   22: ("MiniGrid-Empty-16x16-v0", mride),
   23: ("MiniGrid-DistShift2-v0", mcount),
   24: ("MiniGrid-DistShift2-v0", mcuriosity),
   25: ("MiniGrid-DistShift2-v0", mride),
   26: ("MiniGrid-DistShift2-v0", mvanilla),
   27: ("MiniGrid-DistShift2-v0", mrnd),
   28: ("MiniGrid-LavaGapS7-v0", mride),
   29: ("MiniGrid-LavaGapS7-v0", mcount),
   30: ("MiniGrid-LavaGapS7-v0", mvanilla),
   31: ("MiniGrid-LavaGapS7-v0", mrnd),
   32: ("MiniGrid-SimpleCrossingS11N5-v0", mrnd),
   33: ("MiniGrid-SimpleCrossingS11N5-v0", mvanilla),
   34: ("MiniGrid-SimpleCrossingS11N5-v0", mride),
   35: ("MiniGrid-SimpleCrossingS11N5-v0", mcuriosity),
   36: ("MiniGrid-SimpleCrossingS11N5-v0", mcount),
   37: ("MiniGrid-LavaGapS7-v0", mcuriosity)}
   
for i in range(1,max(M)+1):
    if i==15:
        continue
    fname = f"./OtherLogs/logs{i}.csv"
    #Timesteps,Mean Episode Reward,Standard Deviation of Episode Reward,Mean Episode Length,Standard Deviation of Episode Length
    # _tick,_time,frames,total_episodes,episode_wins,interactions,visited_states,mean_episode_length,mean_episode_return,total_loss,pg_loss,baseline_loss,entropy_loss,state_counts,change_counts,state_counts_std,change_counts_std
    level, technique = M[i]
    ReferenceReadLine = dict([]) #for each k store the last readline with frames <= k
    MEAN_EPISODE_RETURN_INDEX = -1
    MEAN_EPISODE_LENGTH_INDEX = -1
    TOTAL_EPISODES_INDEX = -1
    with open(fname) as fileread:
        alllines = fileread.read().split("\n")
        readlines = alllines[1:]
        for i,segment in enumerate(alllines[0].split(",")):
            if segment.strip() == "mean_episode_return":
                MEAN_EPISODE_RETURN_INDEX = i
            elif segment.strip() == "mean_episode_length":
                MEAN_EPISODE_LENGTH_INDEX = i
            elif segment.strip() == "total_episodes":
                TOTAL_EPISODES_INDEX = i
        timesteps = 100
        lastreadline = ""
        for readline in readlines:
            if readline.strip() == "":
                continue
            frame = int(readline.split(",")[2])
            while timesteps < frame:
                ReferenceReadLine[timesteps] = lastreadline
                timesteps += 100
                #print(timesteps, lastreadline)
            lastreadline = readline

        #print(ReferenceReadLine[3100]); exit(0)
        #mean_episode_reward = 0.0
        #sdev_episode_reward = 0.0
        #mean_episode_length = 0.0
        #sdev_episode_length = 0.0
        
        with open(f"./SavedCSV/{level}_{technique}.csv", "w") as filewrite:
            filewrite.write("Timesteps,Mean Episode Reward,Standard Deviation of Episode Reward,Mean Episode Length,Standard Deviation of Episode Length\n")
            timesteps = 100
            
            #print(readlinefields, readlinefields[8])
            #mean_episode_length
            mean_episode_reward = 0.0
            sdev_episode_reward = 0.0
            mean_episode_length = 0.0
            sdev_episode_length = 0.0
            READ = False
            while timesteps < 10000000:
                #mean_episode_reward = 0.0
                #sdev_episode_reward = 0.0
                #mean_episode_length = 0.0
                #sdev_episode_length = 0.0
                if timesteps in ReferenceReadLine:
                    readlinefields = ReferenceReadLine[timesteps].split(",")
                    if readlinefields != [""]: # or True:
                        if readlinefields[TOTAL_EPISODES_INDEX] != "" and float(readlinefields[TOTAL_EPISODES_INDEX]) >= 15: #not reasonable to take averages of too low sample cnt
                            if readlinefields[MEAN_EPISODE_RETURN_INDEX] != "":
                                v = float(readlinefields[MEAN_EPISODE_RETURN_INDEX])
                                if not math.isnan(v):
                                    mean_episode_reward = v
                                READ = True
                            if readlinefields[MEAN_EPISODE_LENGTH_INDEX] != "":
                                v = float(readlinefields[MEAN_EPISODE_LENGTH_INDEX])
                                if not math.isnan(v):
                                    mean_episode_length = v
                                READ = True
                #print(timesteps, mean_episode_reward)
                if READ:
                    filewrite.write(f"{timesteps}, {mean_episode_reward}, {sdev_episode_reward}, {mean_episode_length}, {sdev_episode_length}\n")
                timesteps += 100
