import os
import numpy as np
from functools import partial
from easydict import EasyDict
import copy
import time
from tensorboardX import SummaryWriter

from core.envs import SimpleCarlaEnv, CarlaEnvWrapper
from core.utils.others.tcp_helper import parse_carla_tcp
from core.eval import SerialEvaluator, SingleCarlaEvaluator, CarlaBenchmarkEvaluator
from ding.envs import SyncSubprocessEnvManager, BaseEnvManager
from ding.policy import SACPolicy
from ding.worker import BaseLearner, SampleCollector, NaiveReplayBuffer
from ding.utils import set_pkg_seed

from demo.simple_rl.env_wrapper import DiscreteBenchmarkEnvWrapper, ContinuousBenchmarkEnvWrapper, ContinuousRgbBenchmarkEnvWrapper, ContinuousRgbCarlaEnvWrapper
from core.utils.data_utils.bev_utils import unpack_birdview
from core.utils.others.ding_utils import compile_config
from demo.simple_rl.model import CSACRGBRLModel

train_config = dict(
    exp_name='sac2_bev32_buf2e5_lr1e4_bs128_ns3000_update4_train_ft',
    env=dict(
        collector_env_num=7,
        evaluator_env_num=1,
        simulator=dict(
            town='Town01',
            disable_two_wheels=True,
            verbose=False,
            waypoint_num=32,
            planner=dict(
                type='behavior',
                resolution=0.1,
                min_distance = 3
            ),
            obs=(
                dict(
                    name='rgb',
                    type='rgb',
                    size=[320, 180],
                    position=[2.0,0,1.4],
                    rotataion=[0,0,0]
                ),
            )
        ),
        col_is_failure=True,
        stuck_is_failure=True,
        off_road_is_failure=True,
        off_route_is_failure=True,
        ignore_light=True,
        finish_reward=300,
        visualize=dict(
            type='rgb',
            outputs=['video'],
            show_text=True,
            save_dir='/home/qhzhang/video'
        ),
        manager=dict(
            collect=dict(
                auto_reset=True,
                shared_memory=False,
                context='spawn',
                max_retry=3,
                reset_timeout=1000,
                step_timeout=1000
            ),
            eval=dict(
                shared_memory=False,
                auto_reset=False,
                context='spawn',
                max_retry=3,
                reset_timeout=1000,
                step_timeout=1000
                )
        ),
        wrapper=dict(
            # Collect and eval suites for training
            collect=dict(suite='NoCrashTown01-v1', ),
            eval=dict(suite='NoCrashTown01-v1', ),
        ),
    ),
    server=[
        dict(carla_host='localhost', carla_ports=[9000, 9016, 2]),
    ],
    policy=dict(
        cuda=True,
        model=dict(
            action_shape=2,
            twin_critic=True,
            task_pretrained=False,
            img_pretrained=True,
            fix_perception=False,
            normalization=None,
        ),
        learn=dict(
            batch_size=256,
            learning_rate_q=1e-4,
            learning_rate_policy=1e-4,
            learning_rate_value=1e-4,
            learning_rate_alpha=1e-4,
            weight_decay=0.0001,
            learner=dict(
                hook=dict(
                    load_ckpt_before_run='',
                ),
            ),
        ),
        collect=dict(
            n_sample=3000,
            noise_sigma=0.1,
            collector=dict(
                collect_print_freq=1000,
                deepcopy_obs=True,
                transform_obs=True,
            ),
        ),
        eval=dict(
            evaluator=dict(
                suite='NoCrashTown01-v1',
                episodes_per_suite=50,
                eval_freq=5000,
                n_episode=3,
                stop_rate=0.7,
                transform_obs=True,
                render=True
            ),
        ),
        other=dict(
            replay_buffer=dict(
                replay_buffer_size=400000,
                replay_buffer_start_size=10000,
                max_use=16,
                monitor=dict(
                    sampled_data_attr=dict(
                        print_freq=10000,  # times
                    ),
                    periodic_thruput=dict(
                        seconds=120000000,
                    ),
                ),
            ),
        ),
    ),
)

main_config = EasyDict(train_config)


def wrapped_env(env_cfg, wrapper_cfg, host, port, tm_port=None):
    return ContinuousRgbBenchmarkEnvWrapper(SimpleCarlaEnv(env_cfg, host, port, tm_port), wrapper_cfg)

def single_wrapped_env_for_benchmark(env_cfg, host, port, tm_port=None):
    return ContinuousRgbCarlaEnvWrapper(SimpleCarlaEnv(env_cfg, host, port, tm_port))

def main(cfg, seed=0):
    cfg = compile_config(
        cfg,
        SyncSubprocessEnvManager,
        SACPolicy,
        BaseLearner,
        SampleCollector,
        buffer=NaiveReplayBuffer,
    )
    tcp_list = parse_carla_tcp(cfg.server)
    collector_env_num, evaluator_env_num = cfg.env.collector_env_num, cfg.env.evaluator_env_num
    assert len(tcp_list) >= collector_env_num + evaluator_env_num, \
        "Carla server not enough! Need {} servers but only found {}.".format(
            collector_env_num + evaluator_env_num, len(tcp_list)
    )

    collector_env = SyncSubprocessEnvManager(
        env_fn=[partial(wrapped_env, cfg.env, cfg.env.wrapper.collect, *tcp_list[i]) for i in range(collector_env_num)],
        cfg=cfg.env.manager.collect,
    )
     #evaluate_env = BaseEnvManager(
     #    env_fn=[partial(wrapped_env, cfg.env, cfg.env.wrapper.eval, *tcp_list[collector_env_num + i]) for i in range(evaluator_env_num)],
     #    cfg=cfg.env.manager.eval,
     #)

    vis_eval_single_env = wrapped_env(cfg.env, cfg.env.wrapper.eval, *tcp_list[collector_env_num])
     
    # Uncomment this to add save replay when evaluation
    # evaluate_env.enable_save_replay(cfg.env.replay_path)
    collector_env.seed(seed)
    vis_eval_single_env.seed(seed)
    # evaluate_env.seed(seed)
    set_pkg_seed(seed)

    model = CSACRGBRLModel(**cfg.policy.model)
    policy = SACPolicy(cfg.policy, model=model)

    tb_logger = SummaryWriter('./log/{}/'.format(cfg.exp_name))
    learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
    collector = SampleCollector(cfg.policy.collect.collector, collector_env, policy.collect_mode, tb_logger, exp_name=cfg.exp_name)
    # evaluator = SerialEvaluator(cfg.policy.eval.evaluator, evaluate_env, policy.eval_mode, tb_logger, exp_name=cfg.exp_name)
    replay_buffer = NaiveReplayBuffer(cfg.policy.other.replay_buffer, tb_logger, exp_name=cfg.exp_name)

    wandb_notes='task pretrain='+str(cfg.policy.model.task_pretrained) + '_img_pretrained='+str(cfg.policy.model.img_pretrained)+ '_fix-weight=' + str(cfg.policy.model.fix_perception)
    vis_evaluator = SingleCarlaEvaluator(wandb_notes, cfg.policy.eval.evaluator, vis_eval_single_env, policy.eval_mode)

    learner.call_hook('before_run')

    new_data = collector.collect(n_sample=10000, train_iter=learner.train_iter)
    replay_buffer.push(new_data, cur_collector_envstep=collector.envstep)

    while True:
        print('eval')
        vis_evaluator.eval()
        # if evaluator.should_eval(learner.train_iter):
        #     stop, rate = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
        #     if stop:
        #         break
        # Sampling data from environments
        new_data = collector.collect(train_iter=learner.train_iter)
        update_per_collect = len(new_data) // cfg.policy.learn.batch_size * 4
        replay_buffer.push(new_data, cur_collector_envstep=collector.envstep)
        # Training
        for i in range(update_per_collect):
            train_data = replay_buffer.sample(cfg.policy.learn.batch_size, learner.train_iter)
            if train_data is not None:
                train_data = copy.deepcopy(train_data)
                unpack_birdview(train_data)
                learner.train(train_data, collector.envstep)

    learner.call_hook('after_run')

    collector.close()
    evaluator.close()
    learner.close()
    replay_buffer.close()


if __name__ == '__main__':
    main(main_config)
