import numpy as np
import random

from src.ac_agent import MAgentDDPG, GaussianProcess
from src.utils import transpose_to_tensor

def instanciate():
    #


    RND_SEED = 123
    # Problem
    N_EPISODES = 1
    MAX_STEPS = 20
    SOLVED_AT = -.5
    GAMMA = 1  # Discount factor

    # Noise
    NOISE_MU = 0
    NOISE_SIGMA = 0.1
    NOISE_DECAY = 1.
    NOISE_MIN_WEIGHT = 1.

    # Agent
    ACT_HID_LAYERS = (8, 4)
    CRIT_HID_LAYERS = (8, 4)
    ACT_ADD_BN = (False)
    CRIT_ADD_BN = (False)
    GRAD_CLIP = (False, False)  # (actor, critic)
    BATCH_SIZE = 1
    LEARNING_RATES = (1e-3, 1e-3)  # (actor, critic)
    WEIGHT_DECAY = (0, 0)  # (actor, critic)
    SOFT_UPD_PARAM = 1e-3
    UPDATE_EVERY = 1
    BUFFER_SIZE = int(1e6)
    LEARN_EVERY = 1
    LEARN_NUM = 1

    g_noise = GaussianProcess(action_size, RND_SEED, mu=NOISE_MU, sigma=NOISE_SIGMA)

    maddpg = MAgentDDPG(num_agents=num_agents, state_size=state_size, action_size=action_size, gamma=GAMMA,
                         actor_hidden_layers=num_agents*[ACT_HID_LAYERS],
                         critic_hidden_layers=num_agents*[CRIT_HID_LAYERS],
                         actor_add_bn=num_agents*[ACT_ADD_BN],
                         critic_add_bn=num_agents*[CRIT_ADD_BN],
                         grad_clipping=num_agents*[GRAD_CLIP],
                         learning_rates=num_agents*[LEARNING_RATES],
                         weight_decay=num_agents*[WEIGHT_DECAY],
                         batch_size=BATCH_SIZE,
                         soft_upd_param=SOFT_UPD_PARAM,
                         update_every=UPDATE_EVERY,
                         buffer_size=BUFFER_SIZE,
                         noise=g_noise,
                         learn_every=LEARN_EVERY,
                         learn_num=LEARN_NUM,
                         seed=RND_SEED)

    return maddpg

def act(obs):
    actions = maddpg.act(obs, add_noise=True, noise_weight=1.)
    return actions

def step(states, actions, rewards, next_states, dones, t_step=1):
    maddpg.step(states, actions, rewards, next_states, dones, t_step)


def tensor_to_lst(obs_full, num_agents, space_size):
    obs_full_lst = []
    for agent_idx in range(num_agents):
        col_idx_left = agent_idx * space_size
        col_idx_right = (agent_idx + 1) * space_size
        obs_full_lst.append(obs_full[:, col_idx_left:col_idx_right])

    return obs_full_lst


class Env(object):
    def __init__(self, num_agents, state_size, action_size, seed=123):
        self.num_agents = num_agents
        self.state_size = state_size
        self.action_size = action_size
        self.seed = seed
        random.seed(seed)

        self.t_steps = 0
        self.n_episodes = 0

    def step(self, actions):
        next_states = np.random.randn(self.num_agents, self.state_size)
        rewards = self.num_agents*[0.]
        dones = self.num_agents*[False]
        env_info = (next_states, rewards, dones)

        return env_info

    def reset(self):
        return np.random.randn(self.num_agents, self.state_size)


if __name__ == '__main__':
    num_agents = 2
    action_size = 2
    state_size = 24

    env = Env(num_agents, state_size, action_size)
    single_obs = env.reset()
    maddpg = instanciate()

    ag1_local_params, ag2_local_params = [], []
    trans_hist = []

    t_step = 0
    actions = act(single_obs)
    single_next_obs, rewards, dones = env.step(actions)
    trans_hist.append((single_obs, actions))
    step(single_obs, actions, rewards, single_next_obs, dones, t_step)
    single_obs = single_next_obs

    t_step = 1
    actions = act(single_obs)
    single_next_obs, rewards, dones = env.step(actions)
    trans_hist.append((single_obs, actions))

    ###
    buffer = maddpg.shared_memory
    experiences = buffer.sample_shared()


    obs_full, actions_full, rewards, next_obs_full, dones = experiences

    agent_idx = 0
    obs_full_lst = tensor_to_lst(obs_full, num_agents, state_size)
    np.testing.assert_almost_equal(
        obs_full.numpy()[:, :state_size].reshape(-1),
        trans_hist[0][0][agent_idx, :]
    )
    np.testing.assert_almost_equal(
        obs_full_lst[0].view(-1),
        trans_hist[0][0][agent_idx, :]
    )
    np.testing.assert_almost_equal(
        actions_full.numpy()[:, :action_size].reshape(-1),
        trans_hist[0][1][agent_idx, :]
    )

    agent_idx = 1
    np.testing.assert_almost_equal(
        obs_full_lst[1].view(-1),
        trans_hist[0][0][agent_idx, :]
    )
    np.testing.assert_almost_equal(
        obs_full.numpy()[:, state_size:].reshape(-1),
        trans_hist[0][0][agent_idx, :]
    )
    np.testing.assert_almost_equal(
        actions_full.numpy()[:, action_size:].reshape(-1),
        trans_hist[0][1][agent_idx, :]
    )

    ###

    step(single_obs, actions, rewards, single_next_obs, dones, t_step)
    single_obs = single_next_obs

    t_step = 2
    actions = act(single_obs)
    single_next_obs, rewards, dones = env.step(actions)
    step(single_obs, actions, rewards, single_next_obs, dones, t_step)
    single_obs = single_next_obs



    t_step = 3
    actions = act(single_obs)
    single_next_obs, rewards, dones = env.step(actions)
    step(single_obs, actions, rewards, single_next_obs, dones, t_step)
    single_obs = single_next_obs
    ag1_local_params.append(maddpg.maddpg_agent[0].actor_local.fc1.weight)
    ag2_local_params.append(maddpg.maddpg_agent[1].actor_local.fc1.weight)