template_header = r"""
\begin{table}[H]
\centering
\caption{All final values}
\label{tab:Results}
\begin{tabular}{llllll}
\hline
\textbf{Level} & \textbf{Techn.} & \textbf{Reward} & \textbf{S.dev.} & \textbf{Ep. len.} & \textbf{S. dev.} \\
\hline
"""

template_header = r"""
\begin{table}[H]
\centering
\caption{All final values}
\label{tab:Results}
\begin{tabular}{llllll}
\hline
\textbf{Level} & \textbf{Techn.} & \textbf{Reward} & \textbf{S.dev.} \\
\hline
"""

template_footer = r"""
\end{tabular}
\end{table}
"""

template = r"""
LEVEL_NAME & \textbf{TRPO}  & TRPO_0   &  TRPO_1  & TRPO_2  & TRPO_3  \\
LEVEL_NAME & \textbf{PPO}   &  PPO_0   &   PPO_1  &  PPO_2  &  PPO_3  \\
LEVEL_NAME & \textbf{A2C}   &  A2C_0   &   A2C_1  &  A2C_2  &  A2C_3  \\
LEVEL_NAME & \textbf{DQN}   &  DQN_0   &   DQN_1  &  DQN_2  &  DQN_3  \\
LEVEL_NAME & \textbf{NACE}  & NACE_0   &  NACE_1  & NACE_2  & NACE_3 \\
"""

template = r"""
LEVEL_NAME & \textbf{TRPO}  & TRPO_0   &  TRPO_1  \\
LEVEL_NAME & \textbf{PPO}   &  PPO_0   &   PPO_1  \\
LEVEL_NAME & \textbf{A2C}   &  A2C_0   &   A2C_1  \\
LEVEL_NAME & \textbf{DQN}   &  DQN_0   &   DQN_1  \\
LEVEL_NAME & \textbf{IMPALA}   &  IMPALA_0   &  IMPALA_1    \\
LEVEL_NAME & \textbf{COUNT}   &  COUNT_0   &   COUNT_1    \\
LEVEL_NAME & \textbf{RIDE}   &  RIDE_0   &   RIDE_1    \\
LEVEL_NAME & \textbf{CURIOSITY}   &  CURIOSITY_0   &  CURIOSITY_1    \\
LEVEL_NAME & \textbf{RND}   &  RND_0   &   RND_1    \\
LEVEL_NAME & \textbf{AMIGO}   &  AMIGO_0   &   AMIGO_1   \\
LEVEL_NAME & \textbf{NACE}  & NACE_0   &  NACE_1  
"""

import sys
import ast
if "NACEOnly" in sys.argv:
    template = r"""
LEVEL_NAME & NACE_0   &  NACE_1 \\
"""
    template_header = r"""
\begin{table}[H]
\centering
\caption{All final values}
\label{tab:Results1K}
\begin{tabular}{llllll}
\hline
\textbf{Level} & \textbf{Reward} & \textbf{S.dev.} \\
\hline
"""

if len(sys.argv) >= 2 and sys.argv[1] == "t=1000":
    template_header = template_header.replace("final values", "values after 1K steps")



import glob
import os
import sys

def load_files(folder_path):
    # Construct the search pattern
    search_pattern = os.path.join(folder_path, "*_*.csv")
    # List all files matching the pattern
    files = glob.glob(search_pattern)
    # Filter files to ensure METHOD part doesn't contain underscores
    filtered_files = [(os.path.basename(f), f) for f in files if 'MiniGrid' in os.path.basename(f) or "BabyAI" in os.path.basename(f)]
    return filtered_files

LevelResults = dict([])
def EnterContent(file):
    filename, fullpath = file
    with open(fullpath) as f:
        Level, Method = filename.split("_")
        Method = Method.split(".csv")[0]
        ResultLines = f.read().split("\n")
        if len(sys.argv) >= 2 and sys.argv[1].startswith("t="):
            T = sys.argv[1].split("t=")[1]
            for Line in ResultLines:
                if Line.startswith(T + ","):
                    Result = Line
                    break
        else:
            Result = ResultLines[-2] #last in this case 
        if Level not in LevelResults:
            LevelResults[Level] = []
        LevelResults[Level] += [(Method, [float(x) for x in Result.split(",")])]
        #print(Level, Method, lines[-2]) #-2 since last is newline


# Example usage
folder_path = 'SavedCSV/'  # Replace with your folder path
files = load_files(folder_path)

for file in files:
    EnterContent(file)

TABLE = template_header

ExcludedLevels = ["MiniGrid-Empty-5x5-v0", "MiniGrid-MultiRoom-N6-v0", "MiniGrid-FourRooms-v0", "MiniGrid-UnlockPickup-v0", "MiniGrid-BlockedUnlockPickup-v0"]
for Level in sorted(LevelResults.keys()):
    if Level in ExcludedLevels:
        continue
    table = template.replace("LEVEL_NAME", Level)
    #print(Level, LevelResults[Level])
    BestResultIndex = 0
    BestResultValue = float("-inf")
    for i, (Technique, Results) in enumerate(LevelResults[Level]):
        Result = Results[1]
        if Result > BestResultValue:
            BestResultValue = Result
            BestResultIndex = i
    """for j, (Technique, Results) in enumerate(LevelResults[Level]):
        BoldValue = j == BestResultIndex
        #print(Level, Technique, Results)
        for i,X in enumerate(Results[1:]): #first value is just timestep, but irrelevant as we take the last for the table
            #print(Level, Technique, i, X)
            #print("REPLACE:", Technique + "_" + str(i), str(X))
            Y = f"{X:.3f}"
            BOLD = True
            table = table.replace(Technique + "_" + str(i), r"\textbf{" + str(Y) + "}" if BoldValue and BOLD and "NACEOnly" not in sys.argv else str(Y))"""
    for j, (Technique, Results) in enumerate(LevelResults[Level]):
        BoldValue = False
        #BoldValue = j == BestResultIndex
        #print(Level, Technique, Results)
        for i,X in enumerate(Results[1:]): #first value is just timestep, but irrelevant as we take the last for the table
            #print(Level, Technique, i, X)
            #print("REPLACE:", Technique + "_" + str(i), str(X))
            Y = f"{X:.3f}"
            if i == 0:
                BestResultValueSTR = f"{BestResultValue:.3f}"
                BoldValue = BestResultValueSTR == Y
            if i == 1 and Technique in ["IMPALA", "COUNT", "RIDE", "CURIOSITY", "RND", "AMIGO"]:
                Y = "-"
                KEY = Technique + "__" + Level
                with open("stddevs_dict", "r") as f:
                    text = f.read()
                val = ast.literal_eval(text)[KEY]
                Y = f"{val:.3f}"
            BOLD = True
            table = table.replace(Technique + "_" + str(i), r"\textbf{" + str(Y) + "}" if BoldValue and BOLD else str(Y))
    table += r"\hline"
    for technique in (["NACE", "DQN", "AC2", "PPO", "TRPO", "IMPALA", "COUNT", "RIDE", "CURIOSITY", "RND", "AMIGO"] if "NACEOnly" not in sys.argv else ["NACE"]):
        for index in range(4):
            table = table.replace(technique + "_" + str(index), "")
    TABLE += table

TABLE+= template_footer
print(TABLE)

