import gym
import time
import pickle
import argparse
import numpy as np
import tensorflow as tf
from typing import Callable, Union, Tuple, List
from models.models import actor_fc_discrete_network, actor_critic_fc_discrete_network
from algorithms.imitation.utils import plot_training_results
from util.replay_buffer import ReplayBuffer


# Set up
GAMMA = 0.99
LEARNING_RATE = 0.0001


class ImitationAgent:
    def __init__(self,
                 environment: gym.Env,
                 model_fn: Callable[..., tf.keras.Model],
                 optimizer: tf.keras.optimizers,
                 run_dagger: bool,
                 expert_policy,
                 expert_data_path,
                 replay_buffer: ReplayBuffer,
                 model_kwargs: dict = None,
                 train_kwargs: dict = None,
                 save_dir: str = None) -> None:
        # Env vars
        self.env = environment
        self.state_dims = model_kwargs.get('state_dims')
        self.num_actions = model_kwargs.get('num_actions')

        num_hidden_layers = model_kwargs.get("num_hidden_layers")
        hidden_size = model_kwargs.get("hidden_size")

        # Algorithm
        self.run_dagger = run_dagger

        # Expert
        self.expert_policy = expert_policy
        self.expert_data = ImitationAgent.load_expert_data(expert_data_path)

        # Actor model
        self.model = model_fn(state_dims=self.state_dims,
                              num_actions=self.num_actions,
                              num_hidden_layers=num_hidden_layers,
                              hidden_size=hidden_size)

        self.optimizer = optimizer
        self.loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)  # Discrete action space only

        # Replay buffer
        self.replay_buffer = replay_buffer

        # Training vars
        self.cur_episode = 0
        self.total_steps = 0
        self.max_ep_len = train_kwargs.get("max_ep_len")
        self.batch_size = train_kwargs.get("batch_size")  # Batch size of data collection from buffer
        self.train_batch_size = train_kwargs.get('train_batch_size')  # Batch size for training models
        self.eval_batch_size = train_kwargs.get('eval_batch_size')  # Batch size for eval
        self.num_agent_train_steps_per_iter = train_kwargs.get('num_agent_train_steps_per_iter')  # Grad updates per run

        # Save directories
        self.save_dir = save_dir

    def save_models(self) -> None:
        self.model.save(self.save_dir)

    def load_models(self) -> tf.keras.Model:
        self.model = tf.keras.models.load_model(self.save_dir)
        return self.model

    @staticmethod
    def load_expert_data(path):
        with open(path, 'rb') as f:
            expert_data = pickle.load(f)
        return expert_data

    def sample_random_trajectory(self) -> Tuple[List[Tuple], Union[int, float]]:
        """
        Sample 1 trajectory.

        :param max_path_length: the maximum number of steps to take in the trajectory
        :param random: whether or not to sample actions randomly or using MPC
        :return:
        """
        state = tf.expand_dims(tf.convert_to_tensor(self.env.reset()), 0)
        num_steps = 0
        total_rewards = 0
        transitions = []  # transition tuples (s,a,r,s',d)
        while True:
            num_steps += 1
            action_prob = self.model(state)
            action = np.random.choice(self.num_actions, p=np.squeeze(action_prob))
            next_state, reward, done, _ = self.env.step(action)
            next_state = tf.reshape(next_state, [1, self.state_dims])

            total_rewards += reward

            if done or num_steps > self.max_ep_len:
                transitions.append((state, action, reward, next_state, 1))
                break

            transitions.append((state, action, reward, next_state, 0))
            state = next_state

        return transitions, total_rewards

    def sample_n_trajectories(self) -> Tuple[List, List, int]:
        """
        Sample `self.batch_size` trajectories. Each trajectory should be no longer than
        `max_path_length` steps/transitions. Note that transitions are different than trajectories!
         A transition is a tuple (s,a,r,s',d) and a trajectory is made up of 1 to `max_path_length` transitions.

        :param batch_size: The number of transitions to sample.
        :param max_path_length: The maximum steps/transitions per trajectory
        :param random: Boolean to indicate whether or not to sample actions randomly or via MPC
        :return:
        """
        num_steps_this_batch = 0
        trajectory_rewards = []
        transitions = []
        while num_steps_this_batch < self.batch_size:
            traj, rews = self.sample_random_trajectory()
            num_steps_this_batch += len(traj)
            trajectory_rewards.append(rews)
            # Note that we're extending, not appending, because we don't care about trajectories, we care about
            #   the transitions. If we appended, it would be ([[(tran 1), (tran 2)], ..., [(tran n), (tran n+1)]],
            #   where each sublist is a trajectory. But by extending, it's instead ([(tran 1), ..., (tran n)]
            transitions.extend(traj)
        return transitions, trajectory_rewards, num_steps_this_batch

    def relabel_actions_with_expert(self, transitions: List[Tuple]) -> List[Tuple]:
        """
        Given a batch of transition tuples, query the Expert Policy and update the action based on
        the Expert. This is the key difference between vanilla behavioral cloning and DAgger. This
        step is equivalent to asking a human expert to label our dataset with actions the correct actions.
        """
        updated_transitions = []
        for transition in transitions:
            state, action, reward, next_state, done = transition
            action_prob, _ = self.expert_policy(state)
            expert_action = np.argmax(np.squeeze(action_prob))
            updated_transitions.append((state, expert_action, reward, next_state, done))
        return updated_transitions

    def train_episode(self) -> List:
        # Step 1: Sample trajectories
        if self.cur_episode == 0:
            # Load expert_data
            transitions = self.expert_data
        else:
            # Or sample trajectories using current policy
            transitions, _, _ = self.sample_n_trajectories()

        # Step 2: For DAgger only, ask expert policy to label data with actions
        if self.run_dagger and self.cur_episode > 0:
            transitions = self.relabel_actions_with_expert(transitions)

        # Step 3: Store the sampled transitions in the replay buffer
        self.replay_buffer.store_transitions_batch(transitions)

        # Step 4: Train model!
        losses = []
        for train_step in range(self.num_agent_train_steps_per_iter):
            # Sample a random batch of data from the replay buffer
            states, actions, _, _, _ = self.replay_buffer.sample(batch_size=self.train_batch_size)

            with tf.GradientTape() as tape:
                action_prob = self.model(states)
                loss = self.loss_fn(actions, action_prob)
            grads = tape.gradient(loss, self.model.trainable_variables)
            self.optimizer.apply_gradients(zip(grads, self.model.trainable_variables))

            losses.append(loss)
        self.cur_episode += 1
        return losses

    def run_agent(self, render=False) -> Tuple[float, int]:
        total_reward, total_steps = 0, 0
        state = self.env.reset()
        done = False

        while not done:
            if render:
                self.env.render()

            # Select action
            action_prob = self.model(tf.expand_dims(state, axis=0))
            action = np.argmax(np.squeeze(action_prob))

            # Interact with environment
            state, reward, done, _ = self.env.step(action)

            # Bookkeeping
            total_reward += reward
            total_steps += 1
        return total_reward, total_steps


def main() -> None:
    # Check input params
    if args.run_dagger:
        assert args.epochs > 1, "DAgger needs more than 1 iteration of training, where each iter" \
                                "we query the expert and train"
    else:
        assert args.epochs == 1, "Vanilla behavior cloning collects expert data only once and does traditional" \
                                 "supervised learning on that dataset."

    # Create environment
    env = gym.make(args.env)

    # Set seeds
    if args.seed:
        np.random.seed(args.seed)
        tf.random.set_seed(args.seed)
        env.seed(args.seed)

    # Create helper vars for model creation
    _state_dims = len(env.observation_space.high)
    _action_dims = 1
    _num_actions = env.action_space.n

    # Create Replay Buffer
    buffer = ReplayBuffer(state_dims=_state_dims, action_dims=_action_dims)

    # Instantiate optimizer
    opt = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)

    # Instantiate expert policy from file
    # TODO >> I think it's a bit cleaner to load the entire model instead of just the weights
    #   but I'm getting a TF error that I think was fixed in a later version. I should probably
    #   try updating the version and seeing if it fixes itself.
    expert = actor_critic_fc_discrete_network(state_dims=_state_dims,
                                              num_actions=_num_actions,
                                              num_hidden_layers=2,
                                              hidden_size=128)
    expert.load_weights(args.expert_policy_file)

    # Create agent
    agent = ImitationAgent(environment=env,
                           model_fn=actor_fc_discrete_network,
                           optimizer=opt,
                           replay_buffer=buffer,
                           run_dagger=args.run_dagger,
                           expert_policy=expert,
                           expert_data_path=args.expert_data,
                           model_kwargs=dict(state_dims=_state_dims,
                                             num_actions=_num_actions,
                                             num_hidden_layers=2,
                                             hidden_size=256),
                           train_kwargs=dict(max_ep_len=args.max_ep_len,
                                             batch_size=args.batch_size,
                                             train_batch_size=args.train_batch_size,
                                             eval_batch_size=args.eval_batch_size,
                                             num_agent_train_steps_per_iter=args.num_agent_train_steps_per_iter)
                           )

    # Run training
    ep_mean_rewards_history, ep_max_rewards_history, ep_min_rewards_history = [], [], []
    ep_mean_loss_history, ep_max_loss_history, ep_min_loss_history = [], [], []
    ep_steps_history = []
    ep_wallclock_history = []
    start = time.time()
    for e in range(args.epochs):
        # Run one episode
        ep_loss = agent.train_episode()
        ep_rew, ep_steps = agent.run_agent()

        # Prepare for logging
        mean_ep_rew, max_ep_rew, min_ep_rew, std_ep_rew = np.mean(ep_rew), np.max(ep_rew), np.min(ep_rew), np.std(ep_rew)
        mean_ep_loss, max_ep_loss, min_ep_loss = np.mean(ep_loss), np.max(ep_loss), np.min(ep_loss)
        ep_wallclock_history.append(time.time() - start)

        ep_mean_rewards_history.append(mean_ep_rew)
        ep_max_rewards_history.append(max_ep_rew)
        ep_min_rewards_history.append(min_ep_rew)

        ep_mean_loss_history.append(mean_ep_loss)
        ep_max_loss_history.append(max_ep_loss)
        ep_min_loss_history.append(min_ep_loss)

        ep_steps_history.append(ep_steps)

        template = "EPISODE {} | mean ep reward: {:.2f} - max ep reward: {:.2f}" \
                   " - min ep reward: {:.2f} - std ep reward: {:.2f} - mean ep loss {:.2f}"
        print(template.format(e, mean_ep_rew, max_ep_rew, min_ep_rew, std_ep_rew, mean_ep_loss))

    # Now that we've completed training, let's plot the results
    print(f"Training time elapsed (sec): {round(time.time() - start, 2)}")

    # Let's evaluate the performance of the trained agent
    print("Beginning evaluation of trained agent!")
    eval_rew = []
    for i in range(50):
        ep_rew, ep_steps = agent.run_agent()
        eval_rew.append(ep_rew)
    print(f"Evaluation rewards: mean - {np.mean(eval_rew)} | min - {np.min(eval_rew)} | max - {np.max(eval_rew)}")

    # Plot summary of results
    plot_training_results(mean_rewards_history=ep_mean_rewards_history,
                          max_rew_history=ep_max_rewards_history,
                          min_rew_history=ep_min_rewards_history,
                          mean_loss_history=ep_mean_loss_history,
                          max_loss_history=ep_max_loss_history,
                          min_loss_history=ep_min_loss_history,
                          steps_history=ep_steps_history,
                          wallclock_history=ep_wallclock_history,
                          save_dir="./results.png")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--env", type=str, default="CartPole-v0")
    parser.add_argument('--expert_policy_file', type=str,  default='./checkpoints/expert_model_weights')
    parser.add_argument('--expert_data', type=str, default='expert_data.pkl')
    # parser.add_argument("--run_dagger", action="store_false")
    parser.add_argument("--run_dagger", type=bool, default=False)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument('--max_ep_len', type=int, default=100)  # max trajectory length

    parser.add_argument('--num_agent_train_steps_per_iter', type=int, default=20)  # number of grad updates per iter
    parser.add_argument('--batch_size', type=int, default=1000)  # num steps/transitions to sample for itr 1+
    parser.add_argument('--train_batch_size', type=int, default=512)  # training batch size per model
    parser.add_argument('--eval_batch_size', type=int, default=400)  # steps collected per eval iteration
    args = parser.parse_args()

    main()

tensorflow.GradientTape tensorflow.random.set_seed tensorflow.expand_dims tensorflow.reshape tensorflow.keras.losses.SparseCategoricalCrossentropy tensorflow.keras.models.load_model tensorflow.keras.optimizers.Adam