from .utils import set_random_seed
from .ppo_nn import NnActorCritic
from .ppo_experience_buffer import PpoBuffer
from codebase.logger.log_epoch import EpochsLogger
from codebase.rudder.lstmcell_RUDDER_with_PPO import LstmCellRudder_with_PPO

import torch
import numpy as np


def run_ppo(env, spec, lstmcell_rudder: LstmCellRudder_with_PPO = None, hidden_dim=16, n_hidden_layers=1, lr=1e-3,
            weight_decay=0.0, gamma=0.99, lam=0.97, n_epoches=1000, steps_by_epoch=4000, target_kl=0.015,
            max_train_iters=80, reward_delayed=False, rew_factor=1.0, save_gap=5, seed=42, print_to_consol=True,
            device='cpu'):

    assert type(rew_factor) is float

    set_random_seed(env, seed)

    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n

    # Initialize agent network
    agent = NnActorCritic(state_size, action_size,
                          n_hidden_layers=n_hidden_layers,
                          hidden_dim=hidden_dim,
                          lr=lr,
                          weight_decay=weight_decay,
                          target_kl=target_kl,
                          max_train_pi_iters=max_train_iters,
                          device=device)

    # Load model
    # agent.load_model(env)

    if reward_delayed:
        print('Watch out ! Delayed rewards set in parameters ! Change to False if not wanted.')

    # Load LSTMCell model from LSTM (Rudder)
    if lstmcell_rudder is not None:
        if spec.selected_lstm_model_path is None:
            lstmcell_rudder.load_lstm_model(spec.experiment_path, spec.lstm_model_name)
        else:
            lstmcell_rudder.load_selected_lstm_model(spec.selected_lstm_model_path)


    # Initialize experience replay
    replay_buffer = PpoBuffer(steps_by_epoch, state_size, device, lstmcell_rudder=lstmcell_rudder)

    # Track trajectories info
    info_logger = EpochsLogger(n_epoches, save_gap=save_gap)

    # Track reward
    rewards_epoches_logger = []

    for epoch in range(n_epoches):

        s = torch.tensor(env.reset(), dtype=torch.float32, device=device)
        reward_logger = 0
        reward_tracker = []
        episode_tracker = 0
        reward_total_episode = 0

        # Reset hidden
        if lstmcell_rudder is not None:
            lstmcell_rudder.reset_cell_hidden_state()

        for t in range(steps_by_epoch):

            # Select action according to policy pi(a|s)
            a, v, log_prob_a = agent.step(s)

            # Obtain rewards r and observe next state s
            s_next, r, trajectory_done, _ = env.step(a)

            # log rewards
            reward_logger += r

            # If reward are delayed or not
            # if reward_delayed and trajectory_done:
            #     if env.unwrapped.spec.id == "CartPole-v1":
            #         r_modified = -torch.tanh(torch.tensor([reward_logger]))
            #     else:
            #         r_modified = reward_logger*rew_factor
            if reward_delayed and trajectory_done:
                r_modified = reward_logger
            elif reward_delayed:
                r_modified = 0
            else:
                r_modified = r

            s_next = torch.tensor(s_next, dtype=torch.float32, device=device)

            # Store information in buffer
            replay_buffer.store(s, a, r_modified, s_next, trajectory_done, v, log_prob_a)

            s = s_next

            if trajectory_done or t == steps_by_epoch - 1:

                # If trajectory not done, bootstrap
                if not trajectory_done:
                    _, last_v, _ = agent.step(s)
                else:
                    reward_tracker.append(reward_logger)
                    last_v = torch.tensor([0], dtype=torch.float32, device=device)
                    episode_tracker += 1
                    reward_logger = 0

                replay_buffer.epoch_ended(last_v, gamma, lam)
                s = torch.tensor(env.reset(), dtype=torch.float32, device=device)

                # Voir remarque de F-A ref task T13PRO-157
                if lstmcell_rudder is not None:
                    lstmcell_rudder.reset_cell_hidden_state()

        reward_mean = sum(reward_tracker)/episode_tracker

        trajectories_data = replay_buffer.get_trajectories()

        # Update critic v and actor policy
        loss_v = agent.update_critic(trajectories_data)
        loss_pi, approx_KL = agent.update_actor(trajectories_data, target_kl)

        # Save models
        # agent.save_model_data(info_logger)

        if print_to_consol:
            print(f'Epoch {epoch} :  e_avg_return: {reward_mean:.2f}, loss_pi = {loss_pi:.4f}, loss_v = {loss_v:.2f}, '
                  f'n_traject : {episode_tracker}')

        rewards_epoches_logger.append(reward_mean)

    return agent, rewards_epoches_logger
