import sys
from pathlib import Path
from matplotlib import pyplot as plt
import itertools as it

try:
    # noinspection PyUnboundLocalVariable
    if __package__ is None:
        DIR = Path(__file__).resolve().parent
        sys.path.insert(0, str(DIR.parent))
        __package__ = DIR.name
    else:
        DIR = None
except NameError:
    DIR = None
    pass

import simplejson
from stable_baselines3.common.vec_env import SubprocVecEnv

from environments import helpers as h
from environments.factory.factory_dirt import DirtFactory
from environments.factory.dirt_util import DirtProperties
from environments.factory.combined_factories import DirtItemFactory
from environments.factory.factory_item import ItemFactory
from environments.factory.additional.item.item_util import ItemProperties
from environments.logging.envmonitor import EnvMonitor
from environments.utility_classes import MovementProperties, ObservationProperties, AgentRenderOptions
import pickle
from plotting.compare_runs import compare_seed_runs, compare_model_runs
import pandas as pd
import seaborn as sns

import multiprocessing as mp

"""
In this studie, we want to explore the macro behaviour of multi agents which are trained on the same task, 
but never saw each other in training.
Those agents learned 


We start with training a single policy on a single task (dirt cleanup / item pickup).
Then multiple agent equipped with the same policy are deployed in the same environment.

There are further distinctions to be made:

1. No Observation - ['no_obs']:
- Agent do not see each other but their consequences of their combined actions
- Agents can collide

2. Observation in seperate slice - [['seperate_0'], ['seperate_1'], ['seperate_N']]:
- Agents see other entitys on a seperate slice
- This slice has been filled with $0 | 1 | \mathbb{N}(0, 1)$
-- Depending ob the fill value, agents will react diffently
   -> TODO: Test this! 

3. Observation in level slice - ['in_lvl_obs']:
- This tells the agent to treat other agents as obstacle. 
- However, the state space is altered since moving obstacles are not part the original agent observation. 
- We are out of distribution.

4. Obseration (similiar to camera read out) ['in_lvl_0.5', 'in_lvl_n']
- This tells the agent to treat other agents as obstacle, but "sees" them encoded as a different value. 
- However, the state space is altered since moving obstacles are not part the original agent observation. 
- We are out of distribution.
"""

n_agents = 4
ood_monitor_file = f'e_1_{n_agents}_agents'
baseline_monitor_file = 'e_1_baseline'


def policy_model_kwargs():
    return dict()  # gae_lambda=0.25, n_steps=16, max_grad_norm=0.25, use_rms_prop=True)


def dqn_model_kwargs():
    return dict(buffer_size=50000,
                learning_starts=64,
                batch_size=64,
                target_update_interval=5000,
                exploration_fraction=0.25,
                exploration_final_eps=0.025
                )


def encapsule_env_factory(env_fctry, env_kwrgs):

    def _init():
        with env_fctry(**env_kwrgs) as init_env:
            return init_env

    return _init


def load_model_run_baseline(seed_path, env_to_run):
    # retrieve model class
    model_cls = next(val for key, val in h.MODEL_MAP.items() if key in seed_path.parent.name)
    # Load both agents
    model = model_cls.load(seed_path / 'model.zip', device='cpu')
    # Load old env kwargs
    with next(seed_path.glob('*.json')).open('r') as f:
        env_kwargs = simplejson.load(f)
        env_kwargs.update(done_at_collision=True)
    # Init Env
    with env_to_run(**env_kwargs) as env_factory:
        monitored_env_factory = EnvMonitor(env_factory)

        # Evaluation Loop for i in range(n Episodes)
        for episode in range(100):
            env_state = monitored_env_factory.reset()
            rew, done_bool = 0, False
            while not done_bool:
                action = model.predict(env_state, deterministic=True)[0]
                env_state, step_r, done_bool, info_obj = monitored_env_factory.step(action)
                rew += step_r
                if done_bool:
                    break
            print(f'Factory run {episode} done, reward is:\n    {rew}')
        monitored_env_factory.save_run(filepath=seed_path / f'{baseline_monitor_file}.pick')


def load_model_run_study(seed_path, env_to_run, additional_kwargs_dict):
    global model_cls
    # retrieve model class
    model_cls = next(val for key, val in h.MODEL_MAP.items() if key in seed_path.parent.name)
    # Load both agents
    models = [model_cls.load(seed_path / 'model.zip', device='cpu') for _ in range(n_agents)]
    # Load old env kwargs
    with next(seed_path.glob('*.json')).open('r') as f:
        env_kwargs = simplejson.load(f)
        env_kwargs.update(
            n_agents=n_agents,
            done_at_collision=True,
            **additional_kwargs_dict.get('post_training_kwargs', {}))
    # Init Env
    with env_to_run(**env_kwargs) as env_factory:
        monitored_factory_env = EnvMonitor(env_factory)
        # Evaluation Loop for i in range(n Episodes)
        for episode in range(50):
            env_state = monitored_factory_env.reset()
            rew, done_bool = 0, False
            while not done_bool:
                try:
                    actions = [model.predict(env_state[model_idx], deterministic=True)[0]
                               for model_idx, model in enumerate(models)]
                except ValueError as e:
                    print(e)
                    print('Env_Kwargs are:\n')
                    print(env_kwargs)
                    print('Path is:\n')
                    print(seed_path)
                    exit()
                env_state, step_r, done_bool, info_obj = monitored_factory_env.step(actions)
                rew += step_r
                if done_bool:
                    break
            print(f'Factory run {episode} done, reward is:\n    {rew}')
        monitored_factory_env.save_run(filepath=seed_path / f'{ood_monitor_file}.pick')
    # Eval monitor outputs are automatically stored by the monitor object
    del models, env_kwargs, env_factory
    import gc
    gc.collect()


def start_mp_study_run(envs_map, policies_path):
    paths = list(y for y in policies_path.iterdir() if y.is_dir() and not (y / f'{ood_monitor_file}.pick').exists())
    if paths:
        with mp.get_context("spawn").Pool(mp.cpu_count()) as pool:
            print("Starting MP with: ", pool._processes, " Processes")
            _ = pool.starmap(load_model_run_study,
                             it.product(paths,
                                        (envs_map[policies_path.parent.name][0],),
                                        (observation_modes[policies_path.parent.parent.name],))
                             )


def start_mp_baseline_run(envs_map, policies_path):
    paths = list(y for y in policies_path.iterdir() if y.is_dir() and
                 not (y / f'{baseline_monitor_file}.pick').exists())
    if paths:
        with mp.get_context("spawn").Pool(mp.cpu_count()) as pool:
            print("Starting MP with: ", pool._processes, " Processes")
            _ = pool.starmap(load_model_run_baseline,
                             it.product(paths,
                                        (envs_map[policies_path.parent.name][0],))
                             )


if __name__ == '__main__':
    # What to do:
    train = True
    baseline_run = True
    ood_run = True
    plotting = True

    train_steps = 1e6
    n_seeds = 3
    frames_to_stack = 3

    # Define a global studi save path
    start_time = 'new_reward'  # int(time.time())
    study_root_path = Path(__file__).parent.parent / 'study_out' / f'{Path(__file__).stem}_{start_time}'

    # Define Global Env Parameters
    # Define properties object parameters
    obs_props = ObservationProperties(render_agents=AgentRenderOptions.NOT,
                                      omit_agent_self=True,
                                      additional_agent_placeholder=None,
                                      frames_to_stack=frames_to_stack,
                                      pomdp_r=2
                                      )
    move_props = MovementProperties(allow_diagonal_movement=True,
                                    allow_square_movement=True,
                                    allow_no_op=False)
    dirt_props = DirtProperties(initial_dirt_ratio=0.35, initial_dirt_spawn_r_var=0.1,
                                clean_amount=0.34,
                                max_spawn_amount=0.1, max_global_amount=20,
                                max_local_amount=1, spawn_frequency=0, max_spawn_ratio=0.05,
                                dirt_smear_amount=0.0)
    item_props = ItemProperties(n_items=10,
                                spawn_frequency=30, n_drop_off_locations=2,
                                max_agent_inventory_capacity=15)
    factory_kwargs = dict(n_agents=1, max_steps=400, parse_doors=True,
                          level_name='rooms', doors_have_area=True,
                          verbose=False,
                          mv_prop=move_props,
                          obs_prop=obs_props,
                          done_at_collision=True
                          )

    # Bundle both environments with global kwargs and parameters
    env_map = {}
    env_map.update({'dirt': (DirtFactory, dict(dirt_prop=dirt_props,
                                               **factory_kwargs.copy()))})
    if False:
        env_map.update({'item': (ItemFactory, dict(item_prop=item_props,
                                                   **factory_kwargs.copy()))})
        env_map.update({'itemdirt': (DirtItemFactory, dict(dirt_prop=dirt_props, item_prop=item_props,
                                                           **factory_kwargs.copy()))})
    env_names = list(env_map.keys())

    # Define parameter versions according with #1,2[1,0,N],3
    observation_modes = {}
    if False:
        observation_modes.update({
            'seperate_1': dict(
                post_training_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.COMBINED,
                    additional_agent_placeholder=None,
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                ),
                additional_env_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.NOT,
                    additional_agent_placeholder=1,
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                )
            )})
        observation_modes.update({
            'seperate_0': dict(
                post_training_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.COMBINED,
                    additional_agent_placeholder=None,
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                ),
                additional_env_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.NOT,
                    additional_agent_placeholder=0,
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                )
            )})
        observation_modes.update({
            'seperate_N': dict(
                post_training_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.COMBINED,
                    additional_agent_placeholder=None,
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                ),
                additional_env_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.NOT,
                    additional_agent_placeholder='N',
                    omit_agent_self=True,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                )
            )})
        observation_modes.update({
            'in_lvl_obs': dict(
                post_training_kwargs=
                dict(obs_prop=ObservationProperties(
                    render_agents=AgentRenderOptions.LEVEL,
                    omit_agent_self=True,
                    additional_agent_placeholder=None,
                    frames_to_stack=frames_to_stack,
                    pomdp_r=2)
                )
            )})
    observation_modes.update({
        #  No further adjustment needed
        'no_obs': dict(
            post_training_kwargs=
            dict(obs_prop=ObservationProperties(
                render_agents=AgentRenderOptions.NOT,
                additional_agent_placeholder=None,
                omit_agent_self=True,
                frames_to_stack=frames_to_stack,
                pomdp_r=2)
            )
        )
    })

    # Train starts here ############################################################
    # Build Major Loop  parameters, parameter versions, Env Classes and models
    if train:
        for obs_mode in observation_modes.keys():
            for env_name in env_names:
                for model_cls in [h.MODEL_MAP['A2C']]:
                    # Create an identifier, which is unique for every combination and easy to read in filesystem
                    identifier = f'{model_cls.__name__}_{start_time}'
                    # Train each combination per seed
                    combination_path = study_root_path / obs_mode / env_name / identifier
                    env_class, env_kwargs = env_map[env_name]
                    env_kwargs = env_kwargs.copy()
                    # Retrieve and set the observation mode specific env parameters
                    additional_kwargs = observation_modes.get(obs_mode, {}).get("additional_env_kwargs", {})
                    env_kwargs.update(additional_kwargs)
                    for seed in range(n_seeds):
                        env_kwargs.update(env_seed=seed)
                        # Output folder
                        seed_path = combination_path / f'{str(seed)}_{identifier}'
                        if (seed_path / 'monitor.pick').exists():
                            continue
                        seed_path.mkdir(parents=True, exist_ok=True)

                        # Env Init & Model kwargs definition
                        if model_cls.__name__ in ["PPO", "A2C"]:
                            # env_factory = env_class(**env_kwargs)

                            env_factory = SubprocVecEnv([encapsule_env_factory(env_class, env_kwargs)
                                                         for _ in range(6)], start_method="spawn")
                            model_kwargs = policy_model_kwargs()

                        elif model_cls.__name__ in ["RegDQN", "DQN", "QRDQN"]:
                            with env_class(**env_kwargs) as env_factory:
                                model_kwargs = dqn_model_kwargs()

                        else:
                            raise NameError(f'The model "{model_cls.__name__}" has the wrong name.')

                        param_path = seed_path / f'env_params.json'
                        try:
                            env_factory.env_method('save_params', param_path)
                        except AttributeError:
                            env_factory.save_params(param_path)

                        # EnvMonitor Init
                        callbacks = [EnvMonitor(env_factory)]

                        # Model Init
                        model = model_cls("MlpPolicy", env_factory,
                                          verbose=1, seed=seed, device='cpu',
                                          **model_kwargs)

                        # Model train
                        model.learn(total_timesteps=int(train_steps), callback=callbacks)

                        # Model save
                        save_path = seed_path / f'model.zip'
                        model.save(save_path)

                        # Monitor Save
                        callbacks[0].save_run(seed_path / 'monitor.pick')

                        # Better be save then sorry: Clean up!
                        del env_factory, model
                        import gc
                        gc.collect()

                    # Compare performance runs, for each seed within a model
                    try:
                        compare_seed_runs(combination_path, use_tex=False)
                    except ValueError:
                        pass
                    # Better be save then sorry: Clean up!
                    try:
                        del env_kwargs
                        del model_kwargs
                        import gc
                        gc.collect()
                    except NameError:
                        pass

                # Compare performance runs, for each model
                # FIXME: Check THIS!!!!
                try:
                    compare_model_runs(study_root_path / obs_mode / env_name, f'{start_time}', 'step_reward',
                                       use_tex=False)
                except ValueError:
                    pass
                pass
            pass
        pass
    pass
    # Train ends here ############################################################

    # Evaluation starts here #####################################################
    # First Iterate over every model and monitor "as trained"
    if baseline_run:
        print('Start Baseline Tracking')
        for obs_mode in observation_modes:
            obs_mode_path = next(x for x in study_root_path.iterdir() if x.is_dir() and x.name == obs_mode)
            # For trained policy in study_root_path / identifier
            for env_path in [x for x in obs_mode_path.iterdir() if x.is_dir()]:
                for policy_path in [x for x in env_path.iterdir() if x. is_dir()]:
                    # Iteration
                    start_mp_baseline_run(env_map, policy_path)

                    # for policy_path in (y for y in policy_path.iterdir() if y.is_dir()):
                    #    load_model_run_baseline(policy_path)
        print('Baseline Tracking done')

    # Then iterate over every model and monitor "ood behavior" - "is it ood?"
    if ood_run:
        print('Start OOD Tracking')
        for obs_mode in observation_modes:
            obs_mode_path = next(x for x in study_root_path.iterdir() if x.is_dir() and x.name == obs_mode)
            # For trained policy in study_root_path / identifier
            for env_path in [x for x in obs_mode_path.iterdir() if x.is_dir()]:
                for policy_path in [x for x in env_path.iterdir() if x. is_dir()]:
                    # FIXME: Pick random seed or iterate over available seeds
                    # First seed path version
                    # policy_path = next((y for y in policy_path.iterdir() if y.is_dir()))
                    # Iteration
                    start_mp_study_run(env_map, policy_path)
                    #for policy_path in (y for y in policy_path.iterdir() if y.is_dir()):
                    #    load_model_run_study(policy_path, env_map[env_path.name][0], observation_modes[obs_mode])
        print('OOD Tracking Done')

    # Plotting
    if plotting:
        # TODO: Plotting
        print('Start Plotting')
        df_list = list()
        for observation_folder in (x for x in study_root_path.iterdir() if x.is_dir()):

            for env_folder in (x for x in observation_folder.iterdir() if x.is_dir()):
                for model_folder in (x for x in env_folder.iterdir() if x.is_dir()):
                    # Gather per seed results in this list

                    for seed_folder in (x for x in model_folder.iterdir() if x.is_dir()):
                        for monitor_file in [f'{baseline_monitor_file}.pick', f'{ood_monitor_file}.pick']:

                            with (seed_folder / monitor_file).open('rb') as f:
                                monitor_df = pickle.load(f)

                            monitor_df = monitor_df.fillna(0)
                            monitor_df['seed'] = int(seed_folder.name.split('_')[0])
                            monitor_df['monitor'] = monitor_file.split('.')[0]
                            monitor_df['monitor'] = monitor_df['monitor'].astype(str)
                            monitor_df['env'] = env_folder.name

                            monitor_df['obs_mode'] = observation_folder.name
                            monitor_df['obs_mode'] = monitor_df['obs_mode'].astype(str)
                            monitor_df['model'] = model_folder.name.split('_')[0]

                            df_list.append(monitor_df)

        id_cols = ['monitor', 'env', 'obs_mode', 'model']

        df = pd.concat(df_list, ignore_index=True)
        df = df.fillna(0)
        for env_name in env_names:
            for id_col in id_cols:
                df[id_col] = df[id_col].astype(str)

                if True:
                    # df['fail_sum'] = df.loc[:, df.columns.str.contains("failed")].sum(1)
                    df['pick_up'] = df.loc[:, df.columns.str.contains("]_item_pickup")].sum(1)
                    df['drop_off'] = df.loc[:, df.columns.str.contains("]_item_dropoff")].sum(1)
                    df['failed_item_action'] = df.loc[:, df.columns.str.contains("]_failed_item_action")].sum(1)
                    df['failed_cleanup'] = df.loc[:, df.columns.str.contains("]_failed_dirt_cleanup")].sum(1)
                    df['coll_lvl'] = df.loc[:, df.columns.str.contains("]_vs_LEVEL")].sum(1)
                    df['coll_agent'] = df.loc[:, df.columns.str.contains("]_vs_Agent")].sum(1) / 2
                    # df['`collis`ions'] = df['coll_lvl'] + df['coll_agent']

                    value_vars = ['pick_up', 'drop_off', 'failed_item_action', 'failed_cleanup',
                                  'coll_lvl', 'coll_agent', 'dirt_cleaned']

                df_grouped = df.groupby(id_cols + ['seed']
                                        # 'sum' if "agent" in key else 'mean'
                                        ).agg({key: 'sum' for key in df.columns
                                               if key not in (id_cols + ['seed'])})
                df_melted = df_grouped.reset_index().melt(id_vars=id_cols,
                                                          value_vars=value_vars,  # 'step_reward',
                                                          var_name="Measurement",
                                                          value_name="Score")
                # df_melted["Measurements"] = df_melted["Measurement"] + " " + df_melted["monitor"]

                # Plotting
                fig, ax = plt.subplots(figsize=(11.7, 8.27))

                c = sns.catplot(data=df_melted[df_melted['env'] == env_name],
                                x='Measurement', hue='monitor', row='model', col='obs_mode', y='Score',
                                sharey=True, kind="box", height=4, aspect=.7, legend_out=False, legend=False,
                                showfliers=False)
                c.set_xticklabels(rotation=65, horizontalalignment='right')
                # c.fig.subplots_adjust(top=0.9)  # adjust the Figure in rp
                c.fig.suptitle(f"Cat plot for {env_name}")
                plt.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
                plt.tight_layout()
                plt.savefig(study_root_path / f'results_{n_agents}_agents_{env_name}.png')
        pass