#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import itertools
import logging

import numpy as np
import scipy as sp
import torch
from ml.rl.evaluation.cpe import CpeEstimate
from ml.rl.evaluation.evaluation_data_page import EvaluationDataPage


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


class WeightedSequentialDoublyRobustEstimator:
    NUM_SUBSETS_FOR_CB_ESTIMATES = 25
    CONFIDENCE_INTERVAL = 0.9
    NUM_BOOTSTRAP_SAMPLES = 50
    BOOTSTRAP_SAMPLE_PCT = 0.5

    def __init__(self, gamma):
        self.gamma = gamma

    def estimate(
        self,
        edp: EvaluationDataPage,
        num_j_steps,
        whether_self_normalize_importance_weights,
    ) -> CpeEstimate:
        # For details, visit https://arxiv.org/pdf/1604.00923.pdf Section 5, 7, 8
        (
            actions,
            rewards,
            logged_propensities,
            target_propensities,
            estimated_q_values,
        ) = WeightedSequentialDoublyRobustEstimator.transform_to_equal_length_trajectories(
            edp.mdp_id,
            edp.action_mask.cpu().numpy(),
            edp.logged_rewards.cpu().numpy().flatten(),
            edp.logged_propensities.cpu().numpy().flatten(),
            edp.model_propensities.cpu().numpy(),
            edp.model_values.cpu().numpy(),
        )

        num_trajectories = actions.shape[0]
        trajectory_length = actions.shape[1]

        j_steps = [float("inf")]

        if num_j_steps > 1:
            j_steps.append(-1)
        if num_j_steps > 2:
            interval = trajectory_length // (num_j_steps - 1)
            j_steps.extend([i * interval for i in range(1, num_j_steps - 1)])

        target_propensity_for_logged_action = np.sum(
            np.multiply(target_propensities, actions), axis=2
        )
        estimated_q_values_for_logged_action = np.sum(
            np.multiply(estimated_q_values, actions), axis=2
        )
        estimated_state_values = np.sum(
            np.multiply(target_propensities, estimated_q_values), axis=2
        )

        importance_weights = target_propensity_for_logged_action / logged_propensities
        importance_weights = np.cumprod(importance_weights, axis=1)
        importance_weights = WeightedSequentialDoublyRobustEstimator.normalize_importance_weights(
            importance_weights, whether_self_normalize_importance_weights
        )

        importance_weights_one_earlier = (
            np.ones([num_trajectories, 1]) * 1.0 / num_trajectories
        )
        importance_weights_one_earlier = np.hstack(
            [importance_weights_one_earlier, importance_weights[:, :-1]]
        )

        discounts = np.logspace(
            start=0, stop=trajectory_length - 1, num=trajectory_length, base=self.gamma
        )

        j_step_return_trajectories = []
        for j_step in j_steps:
            j_step_return_trajectories.append(
                WeightedSequentialDoublyRobustEstimator.calculate_step_return(
                    rewards,
                    discounts,
                    importance_weights,
                    importance_weights_one_earlier,
                    estimated_state_values,
                    estimated_q_values_for_logged_action,
                    j_step,
                )
            )
        j_step_return_trajectories = np.array(j_step_return_trajectories)

        j_step_returns = np.sum(j_step_return_trajectories, axis=1)

        if len(j_step_returns) == 1:
            weighted_doubly_robust = j_step_returns[0]
            weighted_doubly_robust_std_error = 0.0
        else:
            # break trajectories into several subsets to estimate confidence bounds
            infinite_step_returns = []
            num_subsets = int(
                min(
                    num_trajectories / 2,
                    WeightedSequentialDoublyRobustEstimator.NUM_SUBSETS_FOR_CB_ESTIMATES,
                )
            )
            interval = num_trajectories / num_subsets
            for i in range(num_subsets):
                trajectory_subset = np.arange(
                    int(i * interval), int((i + 1) * interval)
                )
                importance_weights = (
                    target_propensity_for_logged_action[trajectory_subset]
                    / logged_propensities[trajectory_subset]
                )
                importance_weights = np.cumprod(importance_weights, axis=1)
                importance_weights = WeightedSequentialDoublyRobustEstimator.normalize_importance_weights(
                    importance_weights, whether_self_normalize_importance_weights
                )
                importance_weights_one_earlier = (
                    np.ones([len(trajectory_subset), 1]) * 1.0 / len(trajectory_subset)
                )
                importance_weights_one_earlier = np.hstack(
                    [importance_weights_one_earlier, importance_weights[:, :-1]]
                )
                infinite_step_return = np.sum(
                    WeightedSequentialDoublyRobustEstimator.calculate_step_return(
                        rewards[trajectory_subset],
                        discounts,
                        importance_weights,
                        importance_weights_one_earlier,
                        estimated_state_values[trajectory_subset],
                        estimated_q_values_for_logged_action[trajectory_subset],
                        float("inf"),
                    )
                )
                infinite_step_returns.append(infinite_step_return)

            # Compute weighted_doubly_robust mean point estimate using all data
            weighted_doubly_robust = self.compute_weighted_doubly_robust_point_estimate(
                j_steps,
                num_j_steps,
                j_step_returns,
                infinite_step_returns,
                j_step_return_trajectories,
            )

            # Use bootstrapping to compute weighted_doubly_robust standard error
            bootstrapped_means = []
            sample_size = int(
                WeightedSequentialDoublyRobustEstimator.BOOTSTRAP_SAMPLE_PCT
                * num_subsets
            )
            for _ in range(
                WeightedSequentialDoublyRobustEstimator.NUM_BOOTSTRAP_SAMPLES
            ):
                random_idxs = np.random.choice(num_j_steps, sample_size, replace=False)
                random_idxs.sort()
                wdr_estimate = self.compute_weighted_doubly_robust_point_estimate(
                    j_steps=[j_steps[i] for i in random_idxs],
                    num_j_steps=sample_size,
                    j_step_returns=j_step_returns[random_idxs],
                    infinite_step_returns=infinite_step_returns,
                    j_step_return_trajectories=j_step_return_trajectories[random_idxs],
                )
                bootstrapped_means.append(wdr_estimate)
            weighted_doubly_robust_std_error = np.std(bootstrapped_means)

        episode_values = np.sum(np.multiply(rewards, discounts), axis=1)
        denominator = np.nanmean(episode_values)
        if abs(denominator) < 1e-6:
            return CpeEstimate(
                raw=0.0, normalized=0.0, raw_std_error=0.0, normalized_std_error=0.0
            )

        return CpeEstimate(
            raw=weighted_doubly_robust,
            normalized=weighted_doubly_robust / denominator,
            raw_std_error=weighted_doubly_robust_std_error,
            normalized_std_error=weighted_doubly_robust_std_error / denominator,
        )

    def compute_weighted_doubly_robust_point_estimate(
        self,
        j_steps,
        num_j_steps,
        j_step_returns,
        infinite_step_returns,
        j_step_return_trajectories,
    ):
        low_bound, high_bound = WeightedSequentialDoublyRobustEstimator.confidence_bounds(
            infinite_step_returns,
            WeightedSequentialDoublyRobustEstimator.CONFIDENCE_INTERVAL,
        )
        # decompose error into bias + variance
        j_step_bias = np.zeros([num_j_steps])
        where_lower = np.where(j_step_returns < low_bound)[0]
        j_step_bias[where_lower] = low_bound - j_step_returns[where_lower]
        where_higher = np.where(j_step_returns > high_bound)[0]
        j_step_bias[where_higher] = j_step_returns[where_higher] - high_bound

        covariance = np.cov(j_step_return_trajectories)
        error = covariance + j_step_bias.T * j_step_bias

        # minimize mse error
        constraint = {"type": "eq", "fun": lambda x: np.sum(x) - 1.0}

        x = np.zeros([len(j_steps)])
        res = sp.optimize.minimize(
            mse_loss,
            x,
            args=error,
            constraints=constraint,
            bounds=[(0, 1) for _ in range(x.shape[0])],
        )
        x = np.array(res.x)
        return float(np.dot(x, j_step_returns))

    @staticmethod
    def transform_to_equal_length_trajectories(
        mdp_ids,
        actions,
        rewards,
        logged_propensities,
        target_propensities,
        estimated_q_values,
    ):
        """
        Take in samples (action, rewards, propensities, etc.) and output lists
        of equal-length trajectories (episodes) accoriding to terminals.
        As the raw trajectories are of various lengths, the shorter ones are
        filled with zeros(ones) at the end.
        """
        num_actions = len(target_propensities[0])

        terminals = np.zeros(mdp_ids.shape[0])
        for x in range(0, mdp_ids.shape[0]):
            if x + 1 == mdp_ids.shape[0] or mdp_ids[x, 0] != mdp_ids[x + 1, 0]:
                terminals[x] = 1

        trajectories = []
        episode_start = 0
        episode_ends = np.nonzero(terminals)[0]
        if len(terminals) - 1 not in episode_ends:
            episode_ends = np.append(episode_ends, len(terminals) - 1)

        for episode_end in episode_ends:
            trajectories.append(np.arange(episode_start, episode_end + 1))
            episode_start = episode_end + 1

        action_trajectories = []
        reward_trajectories = []
        logged_propensity_trajectories = []
        target_propensity_trajectories = []
        Q_value_trajectories = []

        for trajectory in trajectories:
            action_trajectories.append(actions[trajectory])
            reward_trajectories.append(rewards[trajectory])
            logged_propensity_trajectories.append(logged_propensities[trajectory])
            target_propensity_trajectories.append(target_propensities[trajectory])
            Q_value_trajectories.append(estimated_q_values[trajectory])

        def to_equal_length(x, fill_value):
            x_equal_length = np.array(
                list(itertools.zip_longest(*x, fillvalue=fill_value))
            ).swapaxes(0, 1)
            return x_equal_length

        action_trajectories = to_equal_length(
            action_trajectories, np.zeros([num_actions])
        )
        reward_trajectories = to_equal_length(reward_trajectories, 0)
        logged_propensity_trajectories = to_equal_length(
            logged_propensity_trajectories, 1
        )
        target_propensity_trajectories = to_equal_length(
            target_propensity_trajectories, np.zeros([num_actions])
        )
        Q_value_trajectories = to_equal_length(
            Q_value_trajectories, np.zeros([num_actions])
        )

        return (
            action_trajectories,
            reward_trajectories,
            logged_propensity_trajectories,
            target_propensity_trajectories,
            Q_value_trajectories,
        )

    @staticmethod
    def normalize_importance_weights(
        importance_weights, whether_self_normalize_importance_weights
    ):
        if whether_self_normalize_importance_weights:
            sum_importance_weights = np.sum(importance_weights, axis=0)
            where_zeros = np.where(sum_importance_weights == 0.0)[0]
            sum_importance_weights[where_zeros] = len(importance_weights)
            importance_weights[:, where_zeros] = 1.0
            importance_weights /= sum_importance_weights
            return importance_weights
        else:
            importance_weights /= importance_weights.shape[0]
            return importance_weights

    @staticmethod
    def calculate_step_return(
        rewards,
        discounts,
        importance_weights,
        importance_weights_one_earlier,
        estimated_state_values,
        estimated_q_values,
        j_step,
    ):
        trajectory_length = len(rewards[0])
        num_trajectories = len(rewards)
        j_step = int(min(j_step, trajectory_length - 1))

        weighted_discounts = np.multiply(discounts, importance_weights)
        weighted_discounts_one_earlier = np.multiply(
            discounts, importance_weights_one_earlier
        )

        importance_sampled_cumulative_reward = np.sum(
            np.multiply(weighted_discounts[:, : j_step + 1], rewards[:, : j_step + 1]),
            axis=1,
        )

        if j_step < trajectory_length - 1:
            direct_method_value = (
                weighted_discounts_one_earlier[:, j_step + 1]
                * estimated_state_values[:, j_step + 1]
            )
        else:
            direct_method_value = np.zeros([num_trajectories])

        control_variate = np.sum(
            np.multiply(
                weighted_discounts[:, : j_step + 1], estimated_q_values[:, : j_step + 1]
            )
            - np.multiply(
                weighted_discounts_one_earlier[:, : j_step + 1],
                estimated_state_values[:, : j_step + 1],
            ),
            axis=1,
        )

        j_step_return = (
            importance_sampled_cumulative_reward + direct_method_value - control_variate
        )

        return j_step_return

    @staticmethod
    def confidence_bounds(x, confidence):
        n = len(x)
        m, se = np.mean(x), sp.stats.sem(x)
        h = se * sp.stats.t._ppf((1 + confidence) / 2.0, n - 1)
        return m - h, m + h


def mse_loss(x, error):
    return np.dot(np.dot(x, error), x.T)

[((14, 9, 14, 36), 'logging.getLogger', 'logging.getLogger', ({(14, 27, 14, 35): '__name__'}, {}), '(__name__)', False, 'import logging\n'), ((71, 29, 71, 67), 'numpy.cumprod', 'np.cumprod', (), '', True, 'import numpy as np\n'), ((79, 41, 81, 9), 'numpy.hstack', 'np.hstack', ({(80, 12, 80, 72): '[importance_weights_one_earlier, importance_weights[:, :-1]]'}, {}), '([importance_weights_one_earlier, importance_weights[:, :-1]])', True, 'import numpy as np\n'), ((83, 20, 85, 9), 'numpy.logspace', 'np.logspace', (), '', True, 'import numpy as np\n'), ((100, 37, 100, 73), 'numpy.array', 'np.array', ({(100, 46, 100, 72): 'j_step_return_trajectories'}, {}), '(j_step_return_trajectories)', True, 'import numpy as np\n'), ((102, 25, 102, 67), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((179, 22, 179, 48), 'numpy.nanmean', 'np.nanmean', ({(179, 33, 179, 47): 'episode_values'}, {}), '(episode_values)', True, 'import numpy as np\n'), ((185, 15, 190, 9), 'ml.rl.evaluation.cpe.CpeEstimate', 'CpeEstimate', (), '', False, 'from ml.rl.evaluation.cpe import CpeEstimate\n'), ((205, 22, 205, 45), 'numpy.zeros', 'np.zeros', ({(205, 31, 205, 44): '[num_j_steps]'}, {}), '([num_j_steps])', True, 'import numpy as np\n'), ((211, 21, 211, 55), 'numpy.cov', 'np.cov', ({(211, 28, 211, 54): 'j_step_return_trajectories'}, {}), '(j_step_return_trajectories)', True, 'import numpy as np\n'), ((225, 12, 225, 27), 'numpy.array', 'np.array', ({(225, 21, 225, 26): 'res.x'}, {}), '(res.x)', True, 'import numpy as np\n'), ((245, 20, 245, 46), 'numpy.zeros', 'np.zeros', ({(245, 29, 245, 45): 'mdp_ids.shape[0]'}, {}), '(mdp_ids.shape[0])', True, 'import numpy as np\n'), ((330, 29, 330, 71), 'numpy.multiply', 'np.multiply', ({(330, 41, 330, 50): 'discounts', (330, 52, 330, 70): 'importance_weights'}, {}), '(discounts, importance_weights)', True, 'import numpy as np\n'), ((331, 41, 333, 9), 'numpy.multiply', 'np.multiply', ({(332, 12, 332, 21): 'discounts', (332, 23, 332, 53): 'importance_weights_one_earlier'}, {}), '(discounts, importance_weights_one_earlier)', True, 'import numpy as np\n'), ((374, 18, 374, 34), 'numpy.dot', 'np.dot', ({(374, 25, 374, 26): 'x', (374, 28, 374, 33): 'error'}, {}), '(x, error)', True, 'import numpy as np\n'), ((61, 12, 61, 53), 'numpy.multiply', 'np.multiply', ({(61, 24, 61, 43): 'target_propensities', (61, 45, 61, 52): 'actions'}, {}), '(target_propensities, actions)', True, 'import numpy as np\n'), ((64, 12, 64, 52), 'numpy.multiply', 'np.multiply', ({(64, 24, 64, 42): 'estimated_q_values', (64, 44, 64, 51): 'actions'}, {}), '(estimated_q_values, actions)', True, 'import numpy as np\n'), ((67, 12, 67, 64), 'numpy.multiply', 'np.multiply', ({(67, 24, 67, 43): 'target_propensities', (67, 45, 67, 63): 'estimated_q_values'}, {}), '(target_propensities, estimated_q_values)', True, 'import numpy as np\n'), ((176, 47, 176, 73), 'numpy.std', 'np.std', ({(176, 54, 176, 72): 'bootstrapped_means'}, {}), '(bootstrapped_means)', True, 'import numpy as np\n'), ((178, 32, 178, 63), 'numpy.multiply', 'np.multiply', ({(178, 44, 178, 51): 'rewards', (178, 53, 178, 62): 'discounts'}, {}), '(rewards, discounts)', True, 'import numpy as np\n'), ((181, 19, 183, 13), 'ml.rl.evaluation.cpe.CpeEstimate', 'CpeEstimate', (), '', False, 'from ml.rl.evaluation.cpe import CpeEstimate\n'), ((206, 22, 206, 58), 'numpy.where', 'np.where', ({(206, 31, 206, 57): '(j_step_returns < low_bound)'}, {}), '(j_step_returns < low_bound)', True, 'import numpy as np\n'), ((208, 23, 208, 60), 'numpy.where', 'np.where', ({(208, 32, 208, 59): '(j_step_returns > high_bound)'}, {}), '(j_step_returns > high_bound)', True, 'import numpy as np\n'), ((226, 21, 226, 46), 'numpy.dot', 'np.dot', ({(226, 28, 226, 29): 'x', (226, 31, 226, 45): 'j_step_returns'}, {}), '(x, j_step_returns)', True, 'import numpy as np\n'), ((252, 23, 252, 44), 'numpy.nonzero', 'np.nonzero', ({(252, 34, 252, 43): 'terminals'}, {}), '(terminals)', True, 'import numpy as np\n'), ((280, 33, 280, 56), 'numpy.zeros', 'np.zeros', ({(280, 42, 280, 55): '[num_actions]'}, {}), '([num_actions])', True, 'import numpy as np\n'), ((287, 44, 287, 67), 'numpy.zeros', 'np.zeros', ({(287, 53, 287, 66): '[num_actions]'}, {}), '([num_actions])', True, 'import numpy as np\n'), ((290, 34, 290, 57), 'numpy.zeros', 'np.zeros', ({(290, 43, 290, 56): '[num_actions]'}, {}), '([num_actions])', True, 'import numpy as np\n'), ((306, 37, 306, 71), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((336, 12, 336, 86), 'numpy.multiply', 'np.multiply', ({(336, 24, 336, 59): 'weighted_discounts[:, :j_step + 1]', (336, 61, 336, 85): 'rewards[:, :j_step + 1]'}, {}), '(weighted_discounts[:, :j_step + 1], rewards[:, :j_step + 1])', True, 'import numpy as np\n'), ((346, 34, 346, 62), 'numpy.zeros', 'np.zeros', ({(346, 43, 346, 61): '[num_trajectories]'}, {}), '([num_trajectories])', True, 'import numpy as np\n'), ((368, 16, 368, 26), 'numpy.mean', 'np.mean', ({(368, 24, 368, 25): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((368, 28, 368, 43), 'scipy.stats.sem', 'sp.stats.sem', ({(368, 41, 368, 42): 'x'}, {}), '(x)', True, 'import scipy as sp\n'), ((369, 17, 369, 63), 'scipy.stats.t._ppf', 'sp.stats.t._ppf', ({(369, 33, 369, 55): '((1 + confidence) / 2.0)', (369, 57, 369, 62): '(n - 1)'}, {}), '((1 + confidence) / 2.0, n - 1)', True, 'import scipy as sp\n'), ((77, 12, 77, 42), 'numpy.ones', 'np.ones', ({(77, 20, 77, 41): '[num_trajectories, 1]'}, {}), '([num_trajectories, 1])', True, 'import numpy as np\n'), ((125, 37, 125, 75), 'numpy.cumprod', 'np.cumprod', (), '', True, 'import numpy as np\n'), ((132, 49, 134, 17), 'numpy.hstack', 'np.hstack', ({(133, 20, 133, 80): '[importance_weights_one_earlier, importance_weights[:, :-1]]'}, {}), '([importance_weights_one_earlier, importance_weights[:, :-1]])', True, 'import numpy as np\n'), ((166, 30, 166, 87), 'numpy.random.choice', 'np.random.choice', (), '', True, 'import numpy as np\n'), ((257, 32, 257, 73), 'numpy.arange', 'np.arange', ({(257, 42, 257, 55): 'episode_start', (257, 57, 257, 72): '(episode_end + 1)'}, {}), '(episode_start, episode_end + 1)', True, 'import numpy as np\n'), ((307, 26, 307, 65), 'numpy.where', 'np.where', ({(307, 35, 307, 64): '(sum_importance_weights == 0.0)'}, {}), '(sum_importance_weights == 0.0)', True, 'import numpy as np\n'), ((349, 12, 351, 13), 'numpy.multiply', 'np.multiply', ({(350, 16, 350, 51): 'weighted_discounts[:, :j_step + 1]', (350, 53, 350, 88): 'estimated_q_values[:, :j_step + 1]'}, {}), '(weighted_discounts[:, :j_step + 1], estimated_q_values[:, :\n    j_step + 1])', True, 'import numpy as np\n'), ((352, 14, 355, 13), 'numpy.multiply', 'np.multiply', ({(353, 16, 353, 63): 'weighted_discounts_one_earlier[:, :j_step + 1]', (354, 16, 354, 55): 'estimated_state_values[:, :j_step + 1]'}, {}), '(weighted_discounts_one_earlier[:, :j_step + 1],\n    estimated_state_values[:, :j_step + 1])', True, 'import numpy as np\n'), ((215, 53, 215, 62), 'numpy.sum', 'np.sum', ({(215, 60, 215, 61): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((275, 21, 275, 68), 'itertools.zip_longest', 'itertools.zip_longest', (), '', False, 'import itertools\n')]