import json
import os
import random
import time
from collections import defaultdict

os.environ["CUDA_VISIBLE_DEVICES"] = "3"

import jax
import numpy as np
import tqdm
import wandb
from absl import app, flags
from agents import agents
from ml_collections import config_flags
from utils.datasets import Dataset, GCDataset, HGCDataset, GCOTADataset
from utils.env_utils import make_env_and_datasets, sample_fp_states
from utils.evaluation import evaluate
from utils.flax_utils import restore_agent, save_agent
from utils.log_utils import CsvLogger, get_exp_name, get_flag_dict, get_wandb_video, setup_wandb
from dataclasses import replace

FLAGS = flags.FLAGS

flags.DEFINE_string('run_group', 'Debug', 'Run group.')
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_string('env_name', 'antmaze-large-navigate-v0', 'Environment (dataset) name.')
flags.DEFINE_string('save_dir', 'exp/', 'Save directory.')
flags.DEFINE_string('project', 'OGBench', 'Wandb project')
flags.DEFINE_string('restore_path', None, 'Restore path.')
flags.DEFINE_integer('restore_epoch', None, 'Restore epoch.')

flags.DEFINE_integer('train_steps', 1000000, 'Number of training steps.')
flags.DEFINE_integer('log_interval', 5000, 'Logging interval.')
flags.DEFINE_integer('eval_interval', 100000, 'Evaluation interval.')
flags.DEFINE_integer('save_interval', 1000000, 'Saving interval.')
flags.DEFINE_integer('batch_size', 256, 'Batch size.')

flags.DEFINE_integer('eval_tasks', None, 'Number of tasks to evaluate (None for all).')
flags.DEFINE_integer('eval_episodes', 50, 'Number of episodes for each task.')
flags.DEFINE_float('eval_temperature', 0, 'Actor temperature for evaluation.')
flags.DEFINE_float('eval_gaussian', None, 'Action Gaussian noise for evaluation.')
flags.DEFINE_integer('video_episodes', 1, 'Number of video episodes for each task.')
flags.DEFINE_integer('video_frame_skip', 3, 'Frame skip for videos.')
flags.DEFINE_integer('eval_on_cpu', 1, 'Whether to evaluate on CPU.')
config_flags.DEFINE_config_file('agent', 'agents/gciql.py', lock_config=False)

#arguments related to reward shaping:
flags.DEFINE_integer('reward_shaping', 0, 'Whether to sample future states and append to dataset')
flags.DEFINE_integer('append_fp_states', 0, 'Whether to sample future states and append to dataset')
flags.DEFINE_string('future_restore_path', None, 'Restore path for future agent.')
flags.DEFINE_integer('fp_restore_epoch', None, 'Restore epoch for future agent.')
config_flags.DEFINE_config_file('rs_agent', 'agents/mc_occ.py', lock_config=False)


config_flags.DEFINE_config_file('rew_agent', 'agents/rew.py', lock_config=False)
flags.DEFINE_string('rew_restore_path', None, 'Restore path for future agent.')
flags.DEFINE_string('rew_restore_epoch', None, 'Restore path for future agent.')
flags.DEFINE_bool('use_rew', False, "use reward model to infer reward")
flags.DEFINE_bool('ota', False, "Use OTA critic")
flags.DEFINE_bool('state_action', False, "Use state-action occupancy and reward models")

def main(_):
    # Set up logger.
    exp_name = get_exp_name(FLAGS.seed)
    setup_wandb(project=FLAGS.project, group=FLAGS.run_group, name=exp_name)

    FLAGS.save_dir = os.path.join(FLAGS.save_dir, wandb.run.project, FLAGS.run_group, exp_name)
    os.makedirs(FLAGS.save_dir, exist_ok=True)
    flag_dict = get_flag_dict()
    with open(os.path.join(FLAGS.save_dir, 'flags.json'), 'w') as f:
        json.dump(flag_dict, f)

    # Set up environment and dataset.
    config = FLAGS.agent
    rs_config = FLAGS.rs_agent
    rew_config = FLAGS.rew_agent
    print("frame stack:", config['frame_stack'])
    
    config['reward_shaping'] = rs_config['reward_shaping'] = bool(FLAGS.reward_shaping)
    config['append_fp_states'] = rs_config['append_fp_states'] = bool(FLAGS.append_fp_states)
    config['use_rew'] = rew_config['use_rew'] = FLAGS.use_rew
    config['state_action'] = rew_config['state_action'] = rs_config['state_action'] = FLAGS.state_action
    env, train_dataset, val_dataset = make_env_and_datasets(FLAGS.env_name, frame_stack=config['frame_stack'])
    
    #add extra function to get the future states
    dataset_class = {
        'GCDataset': GCDataset,
        'GCOTADataset': GCOTADataset,
        'HGCDataset': HGCDataset
    }[config['dataset_class']]
    train_dataset = dataset_class(Dataset.create(**train_dataset), config)
    if val_dataset is not None:
        val_dataset = dataset_class(Dataset.create(**val_dataset), config)

    if FLAGS.ota:
        value_subgoal_steps = config['abstraction']
    else:
        value_subgoal_steps = None
    # Initialize agent.
    random.seed(FLAGS.seed)
    np.random.seed(FLAGS.seed)

    example_batch = train_dataset.sample(1, value_subgoal_steps = value_subgoal_steps)
    if config['discrete']:
        # Fill with the maximum action to let the agent know the action space size.
        example_batch['actions'] = np.full_like(example_batch['actions'], env.action_space.n - 1)


    #create agents.
    agent_class = agents[config['agent_name']]
    agent = agent_class.create(
        FLAGS.seed,
        example_batch['observations'],
        example_batch['actions'],
        config,
    )
    config['batch_size'] = FLAGS.batch_size
    if FLAGS.reward_shaping:
            agent_class = agents[rs_config['agent_name']]
            rs_config['future_prediction'] = True
            future_agent = agent_class.create(
                FLAGS.seed,
                example_batch['observations'],
                example_batch['actions'],
                rs_config,
            )
            if FLAGS.use_rew:
                agent_class = agents[rew_config['agent_name']]
                rew_agent = agent_class.create(
                    FLAGS.seed,
                    example_batch['observations'],
                    example_batch['actions'],
                    rew_config,
                )
            agent = replace(agent, reward_fn=rew_agent)
        
    # Restore agents.
    if FLAGS.restore_path is not None:
        agent = restore_agent(agent, FLAGS.restore_path, FLAGS.restore_epoch)

    if FLAGS.reward_shaping:
        if FLAGS.future_restore_path is not None:
            future_agent = restore_agent(future_agent, FLAGS.future_restore_path, FLAGS.fp_restore_epoch)   
        if FLAGS.use_rew:
            rew_agent = restore_agent(rew_agent, FLAGS.rew_restore_path, FLAGS.rew_restore_epoch)
        else:
            rew_agent = None
        
    print("Agents restored. Ready to train.")
    # Train agent.
    train_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'train.csv'))
    eval_logger = CsvLogger(os.path.join(FLAGS.save_dir, 'eval.csv'))
    first_time = time.time()
    last_time = time.time()
    for i in tqdm.tqdm(range(1, FLAGS.train_steps + 1), smoothing=0.1, dynamic_ncols=True):
        # Update agent.
        if FLAGS.reward_shaping:
            batch = train_dataset.sample(config['batch_size'], fwd_model = future_agent, rew_model = rew_agent, \
                                         value_subgoal_steps = value_subgoal_steps)
        else:
            batch = train_dataset.sample(config['batch_size'], value_subgoal_steps = value_subgoal_steps)
        agent, update_info = agent.update(batch)

        # Log metrics.
        if i % FLAGS.log_interval == 0:
            train_metrics = {f'training/{k}': v for k, v in update_info.items()}
            if val_dataset is not None:
                if FLAGS.reward_shaping:
                    val_batch = val_dataset.sample(config['batch_size'], fwd_model = future_agent, rew_model = rew_agent, \
                                                   value_subgoal_steps = value_subgoal_steps)
                else:
                    val_batch = val_dataset.sample(config['batch_size'], value_subgoal_steps = value_subgoal_steps)
                _, val_info = agent.total_loss(val_batch, grad_params=None)
                train_metrics.update({f'validation/{k}': v for k, v in val_info.items()})
            train_metrics['time/epoch_time'] = (time.time() - last_time) / FLAGS.log_interval
            train_metrics['time/total_time'] = time.time() - first_time
            last_time = time.time()
            wandb.log(train_metrics, step=i)
            train_logger.log(train_metrics, step=i)

        # Evaluate agent.
        if i == 1 or i % FLAGS.eval_interval == 0:
            if FLAGS.eval_on_cpu:
                eval_agent = jax.device_put(agent, device=jax.devices('cpu')[0])
            else:
                eval_agent = agent
            renders = []
            eval_metrics = {}
            overall_metrics = defaultdict(list)
            task_infos = env.unwrapped.task_infos if hasattr(env.unwrapped, 'task_infos') else env.task_infos
            num_tasks = FLAGS.eval_tasks if FLAGS.eval_tasks is not None else len(task_infos)
            for task_id in tqdm.trange(1, num_tasks + 1):
                task_name = task_infos[task_id - 1]['task_name']
                eval_info, trajs, cur_renders = evaluate(
                    agent=eval_agent,
                    env=env,
                    task_id=task_id,
                    config=config,
                    num_eval_episodes=FLAGS.eval_episodes,
                    num_video_episodes=FLAGS.video_episodes,
                    video_frame_skip=FLAGS.video_frame_skip,
                    eval_temperature=FLAGS.eval_temperature,
                    eval_gaussian=FLAGS.eval_gaussian,
                )
                renders.extend(cur_renders)
                metric_names = ['success']
                eval_metrics.update(
                    {f'evaluation/{task_name}_{k}': v for k, v in eval_info.items() if k in metric_names}
                )
                for k, v in eval_info.items():
                    if k in metric_names:
                        overall_metrics[k].append(v)
            for k, v in overall_metrics.items():
                eval_metrics[f'evaluation/overall_{k}'] = np.mean(v)

            if FLAGS.video_episodes > 0:
                video = get_wandb_video(renders=renders, n_cols=num_tasks)
                eval_metrics['video'] = video

            wandb.log(eval_metrics, step=i)
            eval_logger.log(eval_metrics, step=i)

        # Save agent.
        #if i % FLAGS.save_interval == 0:
        #    save_agent(agent, FLAGS.save_dir, i)

    train_logger.close()
    eval_logger.close()


if __name__ == '__main__':
    app.run(main)
