import functools
import numpy as np

from core.tf_config import *
from utility.utils import Every, TempStore
from utility.graph import video_summary
from utility.timer import Timer
from utility.run import Runner, evaluate, RunMode
from utility import pkg
from env.func import create_env
from replay.func import create_replay
from core.dataset import create_dataset


def train(agent, env, eval_env, replay):
    collect_fn = pkg.import_module('agent', algo=agent.name).collect
    collect = functools.partial(collect_fn, replay)
    
    em = pkg.import_module(env.name.split("_")[0], pkg='env')
    info_func = em.info_func if hasattr(em, 'info_func') else None

    env_step = agent.get_env_step()
    runner = Runner(env, agent, step=env_step,
        run_mode=RunMode.TRAJ, info_func=info_func)
    agent.TRAIN_PERIOD = env.max_episode_steps
    while not replay.good_to_learn():
        env_step = runner.run(step_fn=collect)
        replay.finish_episodes()

    to_eval = Every(agent.EVAL_PERIOD)
    to_record = Every(agent.LOG_PERIOD, agent.LOG_PERIOD)
    to_eval = Every(agent.EVAL_PERIOD)
    to_record = Every(agent.EVAL_PERIOD*10)
    rt = Timer('run')
    tt = Timer('train')
    # et = Timer('eval')
    lt = Timer('log')
    print('Training starts...')
    while env_step <= int(agent.MAX_STEPS):
        with rt:
            env_step = runner.run(step_fn=collect)
        assert np.all(runner.env_output.reset), runner.env_output.reset
        replay.finish_episodes()

        with tt:
            agent.train_record(env_step)

        # if to_eval(env_step):
        #     with TempStore(agent.get_states, agent.reset_states):
        #         with et:
        #             record = agent.RECORD and to_record(env_step)
        #             eval_score, eval_epslen, video = evaluate(
        #                 eval_env, agent, n=agent.N_EVAL_EPISODES, 
        #                 record=agent.RECORD, size=(64, 64))
        #             if record:
        #                 video_summary(f'{agent.name}/sim', video, step=env_step)
        #             agent.store(
        #                 eval_score=eval_score, 
        #                 eval_epslen=eval_epslen)

        if to_record(env_step):
            with lt:
                fps = rt.average() * agent.TRAIN_PERIOD
                tps = tt.average() * agent.N_UPDATES
                
                agent.store(
                    env_step=agent.get_env_step(),
                    train_step=agent.get_train_step(),
                    fps=fps, 
                    tps=tps,
                )
                agent.store(**{
                    'train_step': agent.get_train_step(),
                    'time/run': rt.total(), 
                    'time/train': tt.total(),
                    # 'time/eval': et.total(),
                    'time/log': lt.total(),
                    'time/run_mean': rt.average(), 
                    'time/train_mean': tt.average(),
                    # 'time/eval_mean': et.average(),
                    'time/log_mean': lt.average(),
                })
                agent.record(env_step)
                agent.save()

def main(env_config, model_config, agent_config, replay_config):
    silence_tf_logs()
    configure_gpu()
    configure_precision(agent_config.get('precision', 32))

    use_ray = env_config.get('n_workers', 1) > 1
    if use_ray:
        import ray
        from utility.ray_setup import sigint_shutdown_ray
        ray.init()
        sigint_shutdown_ray()

    env = create_env(env_config)
    # eval_env_config = env_config.copy()
    # eval_env_config['n_workers'] = 1
    # eval_env_config['n_envs'] = 1
    # reward_key = [k for k in eval_env_config.keys() if 'reward' in k]
    # [eval_env_config.pop(k) for k in reward_key]
    # eval_env = create_env(eval_env_config, force_envvec=True)
    eval_env = None

    agent_config['N_UPDATES'] *= env_config['n_workers'] * env_config['n_envs']
    create_model, Agent = pkg.import_agent(config=agent_config)
    models = create_model(model_config, env)

    n_workers = env_config.get('n_workers', 1)
    n_envs = env_config.get('n_envs', 1)
    replay_config['n_envs'] = n_workers * n_envs
    replay_config['seqlen'] = env.max_episode_steps
    if getattr(models, 'state_keys', ()):
        replay_config['state_keys'] = list(models.state_keys)
    replay = create_replay(replay_config)
    replay.load_data()

    am = pkg.import_module('agent', config=agent_config)
    data_format = am.get_data_format(
        env=env, replay_config=replay_config, 
        agent_config=agent_config, model=models)
    dataset = create_dataset(replay, env, data_format=data_format)
    
    agent = Agent(
        config=agent_config, 
        models=models, 
        dataset=dataset, 
        env=env)
    
    agent.save_config(dict(
        env=env_config,
        model=model_config,
        agent=agent_config,
        replay=replay_config
    ))

    train(agent, env, eval_env, replay)

    if use_ray:
        ray.shutdown()
