import random

import torch
from ray.rllib.agents import ppo, pg, dqn, a3c
from ray.rllib.agents.a3c.a3c_torch_policy import (
    A3CTorchPolicy,
    add_advantages,
)
from ray.rllib.agents.callbacks import DefaultCallbacks
from ray.rllib.agents.dqn.dqn_torch_policy import postprocess_nstep_and_prio
from ray.rllib.agents.pg.pg_torch_policy import post_process_advantages
from ray.rllib.agents.ppo.ppo_torch_policy import postprocess_ppo_gae
from ray.rllib.evaluation.sample_batch_builder import (
    MultiAgentSampleBatchBuilder,
)
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils import merge_dicts

from marltoolbox.algos.exploiters.influence_evader import (
    InfluenceEvaderTorchPolicy,
)
from marltoolbox.envs.matrix_sequential_social_dilemma import (
    IteratedPrisonersDilemma,
)
from marltoolbox.utils import postprocessing, miscellaneous

TEST_POLICIES = (
    (ppo.PPOTorchPolicy, postprocess_ppo_gae, ppo.DEFAULT_CONFIG),
    (dqn.DQNTorchPolicy, postprocess_nstep_and_prio, dqn.DEFAULT_CONFIG),
    (A3CTorchPolicy, add_advantages, a3c.DEFAULT_CONFIG),
    (pg.PGTorchPolicy, post_process_advantages, pg.DEFAULT_CONFIG),
)


def given_an_evader_policy():
    for policy_class, postprocessing_fn, default_config in TEST_POLICIES:
        print("policy_class", policy_class)
        coop_policy_class = policy_class.with_updates(
            postprocess_fn=miscellaneous.merge_policy_postprocessing_fn(
                postprocessing.welfares_postprocessing_fn(), postprocessing_fn
            )
        )

        if "target_network_update_freq" in default_config.keys():
            default_config["target_network_update_freq"] = 1

        config = merge_dicts(
            default_config,
            {
                "nested_policies": [
                    {
                        "Policy_class": coop_policy_class,
                        "config_update": {
                            postprocessing.ADD_UTILITARIAN_WELFARE: True
                        },
                    },
                    {"Policy_class": policy_class, "config_update": {}},
                ],
                "start_exploit_at_step_n": random.randint(1, 1000),
                "copy_weights_every_n_steps": random.randint(1, 1000),
                "welfare_key": postprocessing.WELFARE_UTILITARIAN,
            },
        )
        evader = InfluenceEvaderTorchPolicy(
            observation_space=IteratedPrisonersDilemma.OBSERVATION_SPACE,
            action_space=IteratedPrisonersDilemma.ACTION_SPACE,
            config=config,
        )
        yield evader, config["start_exploit_at_step_n"], config[
            "copy_weights_every_n_steps"
        ]


def when_reaching_the_start_of_exploitation(evader, start_exploit_at_step_n):
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    return evader


def get_fake_training_batch_for_ppo_in_ipd(policy):
    policy_id = "fake_player"
    players = {policy_id: policy}

    multi_agent_batch_builder = MultiAgentSampleBatchBuilder(
        policy_map={
            player_id: player for player_id, player in players.items()
        },
        clip_rewards=False,
        callbacks=DefaultCallbacks(),
    )

    n_steps_in_epi = 20
    for step_n in range(n_steps_in_epi):
        step_player_values = {
            SampleBatch.EPS_ID: 0,
            SampleBatch.OBS: 0,
            SampleBatch.NEXT_OBS: 0,
            SampleBatch.ACTIONS: 0,
            SampleBatch.REWARDS: random.randint(0, 10),
            SampleBatch.PREV_REWARDS: random.randint(0, 10),
            SampleBatch.VF_PREDS: random.randint(0, 10),
            SampleBatch.DONES: step_n == n_steps_in_epi - 1,
            SampleBatch.ACTION_DIST_INPUTS: [random.random(), random.random()],
            SampleBatch.ACTION_LOGP: random.random(),
        }
        multi_agent_batch_builder.add_values(
            agent_id=policy_id, policy_id=policy_id, **step_player_values
        )

    multiagent_batch = multi_agent_batch_builder.build_and_reset()
    return multiagent_batch.policy_batches[policy_id]


def assert_weights_of_both_policies_are_equal(evader):
    evading_weights, in_use_weights = get_models_weights(evader)
    for k, v in in_use_weights.items():
        assert torch.all(evading_weights[k] == v)


def get_models_weights(evader):
    evading_weights = {
        k: flatten_list_tensor(v)
        for k, v in evader.get_evading_weights().items()
    }
    in_use_weights = {
        k: flatten_list_tensor(v)
        for k, v in evader.get_in_use_weights().items()
    }
    return evading_weights, in_use_weights


def flatten_list_tensor(dict_tensor):
    list_tensor = [
        torch.reshape(w_tensor, shape=(-1,))
        for w_tensor in dict_tensor.values()
    ]
    return torch.cat(list_tensor, dim=0)


def when_before_the_start_of_exploitation(evader, start_exploit_at_step_n):
    n_steps_played = random.randint(0, start_exploit_at_step_n - 1)

    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def set_current_step_and_perform_one_optimization_step(
    evader, n_steps_played, samples=None
):
    evader.on_global_var_update({"timestep": n_steps_played})
    print("n_steps_played", n_steps_played)

    samples = (
        get_fake_training_batch_for_ppo_in_ipd(evader)
        if samples is None
        else samples
    )
    other_agent_batches = {
        "fake_opp": {1: get_fake_training_batch_for_ppo_in_ipd(evader)}
    }
    samples = evader.postprocess_trajectory(samples, other_agent_batches)
    evader.learn_on_batch(samples)
    if hasattr(evader, "update_target"):
        print("has update_target")
        evader.update_target()
    return evader


def assert_weights_of_both_policies_are_not_equal(evader):
    evading_weights, in_use_weights = get_models_weights(evader)
    for k, v in in_use_weights.items():
        print("k", k)
        assert torch.any(evading_weights[k] != v)


def when_every_n_step_after_start_of_exploitation(
    evader, start_exploit_at_step_n, copy_weights_every_n_steps, nth_repetition
):
    n_steps_played = (
        start_exploit_at_step_n + copy_weights_every_n_steps * nth_repetition
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def when_in_between_n_step_after_start_of_exploitation(
    evader, start_exploit_at_step_n, copy_weights_every_n_steps, nth_repetition
):
    play_intermediary_copies(
        evader,
        start_exploit_at_step_n,
        nth_repetition,
        copy_weights_every_n_steps,
    )

    additional_steps_be_be_in_between_two_copies = random.randint(
        1, copy_weights_every_n_steps - 1
    )
    n_steps_played = (
        start_exploit_at_step_n
        + copy_weights_every_n_steps * (nth_repetition - 1)
        + additional_steps_be_be_in_between_two_copies
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def play_intermediary_copies(
    evader, start_exploit_at_step_n, nth_repetition, copy_weights_every_n_steps
):
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    if (nth_repetition - 1) > 0:
        for i in range(1, nth_repetition):
            n_steps_played = (
                start_exploit_at_step_n + copy_weights_every_n_steps * i
            )
            evader = set_current_step_and_perform_one_optimization_step(
                evader, n_steps_played
            )

import random

import torch
from ray.rllib.agents import ppo, pg, dqn, a3c
from ray.rllib.agents.a3c.a3c_torch_policy import (
    A3CTorchPolicy,
    add_advantages,
)
from ray.rllib.agents.callbacks import DefaultCallbacks
from ray.rllib.agents.dqn.dqn_torch_policy import postprocess_nstep_and_prio
from ray.rllib.agents.pg.pg_torch_policy import post_process_advantages
from ray.rllib.agents.ppo.ppo_torch_policy import postprocess_ppo_gae
from ray.rllib.evaluation.sample_batch_builder import (
    MultiAgentSampleBatchBuilder,
)
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.utils import merge_dicts

from marltoolbox.algos.exploiters.influence_evader import (
    InfluenceEvaderTorchPolicy,
)
from marltoolbox.envs.matrix_sequential_social_dilemma import (
    IteratedPrisonersDilemma,
)
from marltoolbox.utils import postprocessing, miscellaneous

TEST_POLICIES = (
    (ppo.PPOTorchPolicy, postprocess_ppo_gae, ppo.DEFAULT_CONFIG),
    (dqn.DQNTorchPolicy, postprocess_nstep_and_prio, dqn.DEFAULT_CONFIG),
    (A3CTorchPolicy, add_advantages, a3c.DEFAULT_CONFIG),
    (pg.PGTorchPolicy, post_process_advantages, pg.DEFAULT_CONFIG),
)


def given_an_evader_policy():
    for policy_class, postprocessing_fn, default_config in TEST_POLICIES:
        print("policy_class", policy_class)
        coop_policy_class = policy_class.with_updates(
            postprocess_fn=miscellaneous.merge_policy_postprocessing_fn(
                postprocessing.welfares_postprocessing_fn(), postprocessing_fn
            )
        )

        if "target_network_update_freq" in default_config.keys():
            default_config["target_network_update_freq"] = 1

        config = merge_dicts(
            default_config,
            {
                "nested_policies": [
                    {
                        "Policy_class": coop_policy_class,
                        "config_update": {
                            postprocessing.ADD_UTILITARIAN_WELFARE: True
                        },
                    },
                    {"Policy_class": policy_class, "config_update": {}},
                ],
                "start_exploit_at_step_n": random.randint(1, 1000),
                "copy_weights_every_n_steps": random.randint(1, 1000),
                "welfare_key": postprocessing.WELFARE_UTILITARIAN,
            },
        )
        evader = InfluenceEvaderTorchPolicy(
            observation_space=IteratedPrisonersDilemma.OBSERVATION_SPACE,
            action_space=IteratedPrisonersDilemma.ACTION_SPACE,
            config=config,
        )
        yield evader, config["start_exploit_at_step_n"], config[
            "copy_weights_every_n_steps"
        ]


def when_reaching_the_start_of_exploitation(evader, start_exploit_at_step_n):
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    return evader


def get_fake_training_batch_for_ppo_in_ipd(policy):
    policy_id = "fake_player"
    players = {policy_id: policy}

    multi_agent_batch_builder = MultiAgentSampleBatchBuilder(
        policy_map={
            player_id: player for player_id, player in players.items()
        },
        clip_rewards=False,
        callbacks=DefaultCallbacks(),
    )

    n_steps_in_epi = 20
    for step_n in range(n_steps_in_epi):
        step_player_values = {
            SampleBatch.EPS_ID: 0,
            SampleBatch.OBS: 0,
            SampleBatch.NEXT_OBS: 0,
            SampleBatch.ACTIONS: 0,
            SampleBatch.REWARDS: random.randint(0, 10),
            SampleBatch.PREV_REWARDS: random.randint(0, 10),
            SampleBatch.VF_PREDS: random.randint(0, 10),
            SampleBatch.DONES: step_n == n_steps_in_epi - 1,
            SampleBatch.ACTION_DIST_INPUTS: [random.random(), random.random()],
            SampleBatch.ACTION_LOGP: random.random(),
        }
        multi_agent_batch_builder.add_values(
            agent_id=policy_id, policy_id=policy_id, **step_player_values
        )

    multiagent_batch = multi_agent_batch_builder.build_and_reset()
    return multiagent_batch.policy_batches[policy_id]


def assert_weights_of_both_policies_are_equal(evader):
    evading_weights, in_use_weights = get_models_weights(evader)
    for k, v in in_use_weights.items():
        assert torch.all(evading_weights[k] == v)


def get_models_weights(evader):
    evading_weights = {
        k: flatten_list_tensor(v)
        for k, v in evader.get_evading_weights().items()
    }
    in_use_weights = {
        k: flatten_list_tensor(v)
        for k, v in evader.get_in_use_weights().items()
    }
    return evading_weights, in_use_weights


def flatten_list_tensor(dict_tensor):
    list_tensor = [
        torch.reshape(w_tensor, shape=(-1,))
        for w_tensor in dict_tensor.values()
    ]
    return torch.cat(list_tensor, dim=0)


def when_before_the_start_of_exploitation(evader, start_exploit_at_step_n):
    n_steps_played = random.randint(0, start_exploit_at_step_n - 1)

    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def set_current_step_and_perform_one_optimization_step(
    evader, n_steps_played, samples=None
):
    evader.on_global_var_update({"timestep": n_steps_played})
    print("n_steps_played", n_steps_played)

    samples = (
        get_fake_training_batch_for_ppo_in_ipd(evader)
        if samples is None
        else samples
    )
    other_agent_batches = {
        "fake_opp": {1: get_fake_training_batch_for_ppo_in_ipd(evader)}
    }
    samples = evader.postprocess_trajectory(samples, other_agent_batches)
    evader.learn_on_batch(samples)
    if hasattr(evader, "update_target"):
        print("has update_target")
        evader.update_target()
    return evader


def assert_weights_of_both_policies_are_not_equal(evader):
    evading_weights, in_use_weights = get_models_weights(evader)
    for k, v in in_use_weights.items():
        print("k", k)
        assert torch.any(evading_weights[k] != v)


def when_every_n_step_after_start_of_exploitation(
    evader, start_exploit_at_step_n, copy_weights_every_n_steps, nth_repetition
):
    n_steps_played = (
        start_exploit_at_step_n + copy_weights_every_n_steps * nth_repetition
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def when_in_between_n_step_after_start_of_exploitation(
    evader, start_exploit_at_step_n, copy_weights_every_n_steps, nth_repetition
):
    play_intermediary_copies(
        evader,
        start_exploit_at_step_n,
        nth_repetition,
        copy_weights_every_n_steps,
    )

    additional_steps_be_be_in_between_two_copies = random.randint(
        1, copy_weights_every_n_steps - 1
    )
    n_steps_played = (
        start_exploit_at_step_n
        + copy_weights_every_n_steps * (nth_repetition - 1)
        + additional_steps_be_be_in_between_two_copies
    )
    evader = set_current_step_and_perform_one_optimization_step(
        evader, n_steps_played
    )
    return evader


def play_intermediary_copies(
    evader, start_exploit_at_step_n, nth_repetition, copy_weights_every_n_steps
):
    evader = set_current_step_and_perform_one_optimization_step(
        evader, start_exploit_at_step_n
    )
    if (nth_repetition - 1) > 0:
        for i in range(1, nth_repetition):
            n_steps_played = (
                start_exploit_at_step_n + copy_weights_every_n_steps * i
            )
            evader = set_current_step_and_perform_one_optimization_step(
                evader, n_steps_played
            )
