import os
import datetime
import gym
import json

from stable_baselines3 import PPO
import optuna

def check_json_exists(filename):
    try:
        open(filename, 'r')
    except:
        main = {
            "env" : "RoadRunner-v0",
            "algo" : "PPO",
            "policy" : "MlpPolicy",
            "counter" : 0,
        }
        with open(filename, "w") as f:
            json.dump(main, f)

check_json_exists("previous_request.json")

class MegaTrainer:
    def __init__(self, env_name, algo):
        with open("previous_request.json", "r") as f:
            self.data = json.loads(f.read())

        self.env_name = env_name
        self.algorithm = algo
        self.counter = self.data["counter"] + 1

        self.env = gym.make(self.env_name)
        self.env.reset()

        self.data["counter"] += 1
        with open("previous_request.json", "w") as f:
            json.dump(self.data, f)
        
    def make_directories(self, models_dir, logdir):
        if(not os.path.exists(models_dir)):
            os.makedirs(models_dir)

        if(not os.path.exists(logdir)):
            os.makedirs(logdir)

    def learn(self, timestep, iterations, hyps, directory, trial):
        model_dir = f"model/{directory}/{self.counter}"
        logdir = f"logs/{directory}"
        self.make_directories(model_dir, logdir)

        model = PPO(**hyps)

        for i in range(1, iterations+1):
            if(trial.should_prune()):
                raise optuna.TrialPruned()
            dt = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f")
            model.learn(total_timesteps=timestep, reset_num_timesteps=False, tb_log_name=str(self.counter))
            model.save(model_dir+f"/{dt}_{timestep*i}")

        return self.counter





if(__name__=="__main__"):
    load_from_json("previous_request.json")