"""
Implement the deploy model
Add evaluation code in here

This works with new dataset generated by William
"""

import os

import torch
import torch.nn as nn
import torch.nn.functional as F
from .model import MLPPolicy

from torch.distributions.categorical import Categorical

import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader

norm_names = ['grade_norm', 'pre-score_norm', 'stage_norm', 'failed_attempts_norm',
         'pos_norm', 'neg_norm', 'help_norm', 'anxiety_norm']

raw_names = ['grade', 'pre', 'stage', 'failed_attempts', 'pos', 'neg', 'help', 'anxiety']

def generate_features(norm_ob_row, unnormed_ob_row):
    assert len(unnormed_ob_row.shape) == 1

    norm_obs_dict = dict(zip(norm_names, norm_ob_row))
    raw_obs_dict = dict(zip(raw_names, unnormed_ob_row))

    features = [norm_obs_dict['stage_norm'],
                norm_obs_dict['failed_attempts_norm'],
                norm_obs_dict['pos_norm'],
                norm_obs_dict['neg_norm'],
                norm_obs_dict['help_norm'],
                norm_obs_dict['anxiety_norm'],
                norm_obs_dict['grade_norm'],
                int(raw_obs_dict['pre']),
                int(raw_obs_dict['anxiety']),
                int(raw_obs_dict['stage'])]

    return features

def normalize(ob):
    return (ob - [3, 4, 3, 10, 0.5, 0.5, 0.5, 27])/ [1, 4, 3, 10, 0.5, 0.5, 0.5, 18]

def unnormalize(norm_ob):
    if type(norm_ob) == list:
        norm_ob = np.array(norm_ob)

    return norm_ob * np.array([1, 4, 3, 10, 0.5, 0.5, 0.5, 18]) + np.array([3, 4, 3, 10, 0.5, 0.5, 0.5, 27])

def load_pytorch_agent():
    bc_policy = MLPPolicy([10, 4, 4])
    policy = DeployPolicy(bc_policy)

    policy.load_model("./bc_policy_4_4_full_data.ckpt")

    # if not os.path.exists('./store'):
    #     print("\n******************* Create an empty store folder *******************\n")
    #     os.mkdir('./store')
    #
    # if os.path.exists('./store/checkpoint'):
    #     policy.load_model("./store/bc_policy_4_4_full_data.ckpt")
    #     print("\n******************* Load the model from the last checkpoint *******************\n")
    #     # check for nan
    #     w = policy.get_weights()
    #     print(w)
    # else:
    #     print("\n!!!!!!!!!!!!!!!!!!!!! Store Folder is Empty !!!!!!!!!!!!!!!!!!!!!\n")

    return policy

class DeployPolicy(nn.Module):
    def __init__(self, mlp_model: nn.Module):
        super().__init__()
        self.policy = mlp_model

    def reset(self):
        pass

    def get_weights(self):
        weights = {}
        for name, param in self.policy.named_parameters():
            weights[name] = param.data
        return weights

    def store_model(self, ckpt_path):
        torch.save(self.policy.state_dict(), ckpt_path)

    def load_model(self, ckpt_path):
        self.policy.load_state_dict(torch.load(ckpt_path))

    def get_action(self, obs, temperature=1):
        # Crucial note: this takes in "normalized" data
        # Not-batched

        # input obs: numpy array
        # ['grade', 'pre', 'stage', 'failed_attempts', 'pos', 'neg', 'help', 'anxiety']

        features = np.zeros(10)
        features[0] = obs[2]
        features[1] = obs[3]
        features[2] = obs[4]
        features[3] = obs[5]
        features[4] = obs[6]
        features[5] = obs[7]
        features[6] = obs[0]
        features[7] = obs[1] * 4 + 4
        features[8] = obs[7] * 18 + 27
        features[9] = obs[2] * 3 + 3

        # features = [obs[2], obs[3], obs[4], obs[5], obs[6], obs[7], obs[0], obs[1] * 4 + 4,
        #             obs[7] * 18 + 27,
        #             obs[2] * 3 + 3]
        # features = np.concatenate(features)

        features = torch.from_numpy(np.array(features)).float()

        probs = self.policy.get_action_probability(features, no_grad=True, temperature=temperature)
        print(probs)

        # Step 2, let's sample it
        dist = Categorical(probs)
        action = dist.sample()

        # integer
        return action.item()

    def get_action_probability(self, obs, no_grad=True, action_mask=None):
        # handle batch

        features = [obs[:, 2], obs[:, 3], obs[:, 4], obs[:, 5], obs[:, 6], obs[:, 7], obs[:, 0], obs[:, 1] * 4 + 4, obs[:, 7] * 18 + 27,
                        obs[:, 2] * 3 + 3]
        features = np.vstack(features).transpose(1, 0)

        features = torch.from_numpy(features).float()

        probs = self.policy.get_action_probability(features, no_grad=True)

        return probs

    def observe(self, preprocessed_states, actions, internals, rewards, next_states, terminals, env_id=None):
        pass

    def update(self):
        pass


MAX_TIME = 28

# raw_names = ['grade', 'pre', 'stage', 'failed_attempts', 'pos', 'neg', 'help', 'anxiety']

# All of these are normalized in "children_probs.csv"
feature_names = ['grade', 'pre_score', 'stage', 'failed_attempts', 'pos', 'neg', 'help', 'anxiety']

target_names = ["p_hint", "p_nothing", "p_encourage", "p_question"]

action_names = ['hint', 'nothing', 'encourage', 'question']

knn_target_names = []

# Let's unify it here...
def compute_is_weights_for_deploy_policy(behavior_df, eval_policy, eps=0.05, temp=0.1,
                                     reward_column='adjusted_score', no_grad=True,
                                     gr_safety_thresh=0.0, is_train=True, normalize_reward=False, use_knn=False,
                                     model_round_as_feature=False):
    # is_weights with batch processing
    df = behavior_df
    user_ids = df['user_id'].unique()
    n = len(user_ids)

    # now max_time is dynamic, finally!!
    MAX_TIME = max(behavior_df.groupby('user_id').size())

    assert reward_column in ['adjusted_improvement', 'reward']

    pies = torch.zeros((n, MAX_TIME))  # originally 20
    pibs = torch.zeros((n, MAX_TIME))
    rewards = torch.zeros((n, MAX_TIME))
    lengths = np.zeros((n))  # we are not doing anything to this

    # compute train split reward mean and std
    # (n,): and not average
    user_rewards = df.groupby("user_id")[reward_column].mean()
    train_reward_mu = user_rewards.mean()
    train_reward_std = user_rewards.std()

    action_names = ['hint', 'nothing', 'encourage', 'question']
    df['action_names'] = df['action']
    df['action'] = df['action_names'].map(action_names.index)

    for idx, user_id in enumerate(user_ids):
        data = df[df['user_id'] == user_id]
        # get features, targets
        if not model_round_as_feature:
            features = np.asarray(data[feature_names]).astype(float)
        else:
            features = np.asarray(data[feature_names + ['model_round']]).astype(float)
        targets = np.asarray(data[target_names]).astype(float)
        actions = np.asarray(data['action']).astype(int)

        length = features.shape[0]
        lengths[idx] = length

        T = targets.shape[0]
        # shape: (1, T)
        beh_probs = torch.from_numpy(np.array([targets[i, a] for i, a in enumerate(actions)])).float()
        pibs[idx, :T] = beh_probs

        gr_mask = None
        if gr_safety_thresh > 0:
            if not is_train:
                # we only use KNN during validation
                if not use_knn:
                    gr_mask = None
                else:
                    knn_targets = np.asarray(data[knn_target_names]).astype(float)
                    assert knn_targets.shape[0] == targets.shape[0]
                    beh_action_probs = torch.from_numpy(knn_targets)

                    gr_mask = beh_action_probs >= gr_safety_thresh
            else:
                beh_action_probs = torch.from_numpy(targets)

                # gr_mask = beh_probs >= gr_safety_thresh
                gr_mask = beh_action_probs >= gr_safety_thresh

            # need to renormalize behavior policy as well?

        # assign rewards (note adjusted_score we only assign to last step)
        # reward, we assign to all
        if reward_column == 'adjusted_improvement':
            reward = np.asarray(data[reward_column])[-1]
            # only normalize reward during training
            if normalize_reward and is_train:
                # might not be the best -- we could just do a plain shift instead
                # like -1 shift
                reward = (reward - train_reward_mu) / train_reward_std
            rewards[idx, T - 1] = reward
        else:
            # normal reward
            # rewards[idx, :T] = torch.from_numpy(np.asarray(data[reward_column])).float()
            raise Exception("We currrently do not offer training in this mode")

        # last thing: model prediction
        eval_action_probs = eval_policy.get_action_probability(features, no_grad,
                                                               action_mask=gr_mask)

        pies[idx, :T] = torch.hstack([eval_action_probs[i, a] for i, a in enumerate(actions)])

    return pibs, pies, rewards, lengths.astype(int), MAX_TIME

def ope_evaluate_deploy(policy, df, ope_method, gr_safety_thresh, is_train=True, reward_column='adjusted_improvement', return_weights=False, use_knn=False):
    with torch.no_grad():
        pibs, pies, rewards, lengths, max_time = compute_is_weights_for_deploy_policy(df, policy, is_train=is_train,
                                                                                  gr_safety_thresh=gr_safety_thresh,
                                                                                  reward_column=reward_column,
                                                                                  use_knn=use_knn)

        ope_score, weights = ope_method(pibs, pies, rewards, lengths, max_time=max_time)
        ess = 1 / (torch.sum(weights ** 2, dim=0))

    if return_weights:
        return ope_score, ess, weights
    else:
        return ope_score, ess