import os.path
import sys

from Utils import environments
import gym
import d4rl
import time
import datetime
import torch
import itertools
import numpy as np
from loguru import logger

import argparse

from diffusion_predictor.Predictor_config import DMBP_config, update_DMBP_config
from diffusion_predictor.Predictor_net import Diffusion_instance
from diffusion_predictor.render_img import MuJoCoRenderer
from Utils.Batch_Buffer import batch_buffer
from Utils.seed import setup_seed, seed_env

def train_DMBP(args):
    env_name, dataset_name, seed = args.env_name, args.dataset, args.seed

    config = update_DMBP_config(env_name, DMBP_config, args)

    device = torch.device(args.device)
    setup_seed(seed)
    
    Buffer = batch_buffer(config, env_name, dataset_name, device, buffer_mode='normal', buffer_normalization=False)
    AGD = Diffusion_instance(Buffer.obs_dim, Buffer.act_dim, device, config, "naive")

    if args.checkpoint_ith is None:
        Training_steps_per_epoch, A_Epoches, N_Epoches = config["steps_per_epoch"], config["ambient_epochs"], config["naive_epochs"]

        for i in range(A_Epoches + N_Epoches):
            if i < A_Epoches:
                # Watch for this change
                AGD.train(Buffer, Training_steps_per_epoch, config['batch_size'], None, 1.0, "ambient")
            else:
                # Watch for this change
                AGD.train(Buffer, Training_steps_per_epoch, config['batch_size'], None, 0.5, "naive")
            Buffer.simple_detect(AGD, AGD)

            if i % 5 == 4:
                AGD.save_model(f"models_cp/{env_name}_{dataset_name}_epoch_{i}.pkl")
    else:
        AGD.load_model(f"models_cp/{env_name}_{dataset_name}_epoch_{args.checkpoint_ith}.pkl")
    
    loops, steps, ith = config["detect_denoise_loops"], config["detect_denoise_steps"], args.checkpoint_ith
    ith = config["ambient_epochs"] + config["naive_epochs"] - 1 if ith == None else ith

    Buffer.get_dataset(AGD, AGD, f"datasets_cp/{env_name}_{dataset_name}_loops_{loops}_steps_{steps}_ith_{ith}", 0.7)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--env_name', default='hopper', type=str)
    parser.add_argument('--dataset', default='medium-replay', type=str)
    parser.add_argument('--device', default="cuda:0", type=str)

    parser.add_argument('--load_model_path', default=None, type=str)
    parser.add_argument('--dataset_path', default=None, type=str)

    parser.add_argument('--seed', default=0, type=int)
    
    parser.add_argument('--checkpoint_ith', default=None, type=int)
    parser.add_argument('--denoisng_loops', default=None, type=int)
    parser.add_argument('--denoisng_steps', default=None, type=int)

    args = parser.parse_args()

    train_DMBP(args)