import pytest
param = pytest.mark.parametrize

from shutil import rmtree
from pathlib import Path
from functools import partial

import torch
from torch import cat
from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score, extract_grpo_data
from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper

from memmap_replay_buffer import ReplayBuffer

from einops import rearrange

# functions

def exists(v):
    return v is not None

# test

@param('use_binary_mapper_variant, switch_per_latent_dim', [
    (False, False),
    (False, True),
    (True, False)
])
@param('action_discrete', (False, True))
@param('variable_length', (False, True))
def test_metacontroller(
    use_binary_mapper_variant,
    switch_per_latent_dim,
    action_discrete,
    variable_length
):

    state = torch.randn(2, 128, 384)
    episode_lens = torch.tensor([64, 64]) if variable_length else None

    if action_discrete:
        actions = torch.randint(0, 4, (2, 128))
        action_embed_readout = dict(num_discrete = 4)
        assert_shape = (4,)
    else:
        actions = torch.randn(2, 128, 8)
        action_embed_readout = dict(num_continuous = 8)
        assert_shape = (8, 2)

    # behavioral cloning phase

    model = Transformer(
        dim = 512,
        action_embed_readout = action_embed_readout,
        state_embed_readout = dict(num_continuous = 384),
        lower_body = dict(depth = 2,),
        upper_body = dict(depth = 2,),
    )

    state_clone_loss, action_clone_loss = model(state, actions, episode_lens = episode_lens)
    (state_clone_loss + 0.5 * action_clone_loss).backward()

    # discovery and internal rl phase with meta controller

    if not use_binary_mapper_variant:
        meta_controller = MetaController(
            dim_model = 512,
            dim_meta_controller = 256,
            dim_latent = 128,
            switch_per_latent_dim = switch_per_latent_dim
        )
    else:
        meta_controller = MetaControllerWithBinaryMapper(
            dim_model = 512,
            dim_meta_controller = 256,
            switch_per_code = switch_per_latent_dim,
            dim_code_bits = 8, # 2 ** 8 = 256 codes
        )

    # discovery phase

    (action_recon_loss, kl_loss, switch_loss) = model(state, actions, meta_controller = meta_controller, discovery_phase = True, episode_lens = episode_lens)
    (action_recon_loss + kl_loss * 0.1 + switch_loss * 0.2).backward()

    # internal rl - done iteratively

    # replay buffer

    test_folder = './test-buffer-for-grpo'

    replay_buffer = ReplayBuffer(
        test_folder,
        max_episodes = 3,
        max_timesteps = 256,
        fields = meta_controller.replay_buffer_field_dict,
        meta_fields = dict(
            advantages = 'float'
        )
    )

    # simulate grpo

    all_episodes = []
    all_rewards = []

    for _ in range(3): # group of 3
        subset_state = state[:1]

        cache = None
        past_action_id = None

        grpo_data_list = []

        for one_state in subset_state.unbind(dim = 1):
            one_state = rearrange(one_state, 'b d -> b 1 d')

            logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, cache = cache, return_cache = True)

            past_action_id = model.action_readout.sample(logits)

            # extract grpo data and store

            grpo_data = extract_grpo_data(meta_controller, cache)
            grpo_data_list.append(grpo_data)

        # accumulate across time for the episode data

        states, actions, log_probs, switch_betas = zip(*grpo_data_list)

        all_episodes.append((
            cat(states, dim = 1),
            cat(log_probs, dim = 1),
            cat(switch_betas, dim = 1),
            cat(actions, dim = 1)
        ))

        all_rewards.append(torch.randn(1))

    # calculate advantages using z-score

    rewards = cat(all_rewards)
    group_advantages = z_score(rewards)

    assert group_advantages.shape == (3,)

    # simulate a policy loss update over the entire group

    group_states, group_log_probs, group_switch_betas, group_latent_actions = map(partial(cat, dim = 0), zip(*all_episodes))
    
    # parallel verification

    parallel_action_dist = meta_controller.get_action_dist_for_internal_rl(group_states)
    parallel_log_probs = meta_controller.log_prob(parallel_action_dist, group_latent_actions)

    assert torch.allclose(parallel_log_probs, group_log_probs, atol = 1e-5), 'parallel log probs do not match stored log probs'

    for states, log_probs, switch_betas, latent_actions, advantages in zip(group_states, group_log_probs, group_switch_betas, group_latent_actions, group_advantages):
        replay_buffer.store_episode(
            states = states,
            log_probs = log_probs,
            switch_betas = switch_betas,
            latent_actions = latent_actions,
            advantages = advantages
        )

    dl = replay_buffer.dataloader(batch_size = 3)

    batch = next(iter(dl))

    loss = meta_controller.policy_loss(
        batch['states'],
        batch['log_probs'],
        batch['latent_actions'],
        batch['advantages'],
        batch['switch_betas'] == 1.,
        episode_lens = batch['_lens']
    )

    loss.backward()

    # evolutionary strategies over grpo

    model.meta_controller = meta_controller
    model.evolve(1, lambda _: 1., noise_population_size = 2)

    # saving and loading

    meta_controller.save('./meta_controller.pt')

    meta_controller_klass = meta_controller.__class__
    rehydrated_meta_controller = meta_controller_klass.init_and_load('./meta_controller.pt')

    model.save('./trained.pt')

    rehydrated_model = Transformer.init_and_load('./trained.pt', strict = False)

    Path('./meta_controller.pt').unlink()
    Path('./trained.pt').unlink()

    rmtree(test_folder, ignore_errors = True)
