#!/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)

[((305, 332), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (322, 332), False, 'import logging\n'), ((2219, 2257), 'numpy.cumprod', 'np.cumprod', (['importance_weights'], {'axis': '(1)'}), '(importance_weights, axis=1)\n', (2229, 2257), True, 'import numpy as np\n'), ((2604, 2675), 'numpy.hstack', 'np.hstack', (['[importance_weights_one_earlier, importance_weights[:, :-1]]'], {}), '([importance_weights_one_earlier, importance_weights[:, :-1]])\n', (2613, 2675), True, 'import numpy as np\n'), ((2719, 2811), 'numpy.logspace', 'np.logspace', ([], {'start': '(0)', 'stop': '(trajectory_length - 1)', 'num': 'trajectory_length', 'base': 'self.gamma'}), '(start=0, stop=trajectory_length - 1, num=trajectory_length,\n    base=self.gamma)\n', (2730, 2811), True, 'import numpy as np\n'), ((3379, 3415), 'numpy.array', 'np.array', (['j_step_return_trajectories'], {}), '(j_step_return_trajectories)\n', (3387, 3415), True, 'import numpy as np\n'), ((3442, 3484), 'numpy.sum', 'np.sum', (['j_step_return_trajectories'], {'axis': '(1)'}), '(j_step_return_trajectories, axis=1)\n', (3448, 3484), True, 'import numpy as np\n'), ((7126, 7152), 'numpy.nanmean', 'np.nanmean', (['episode_values'], {}), '(episode_values)\n', (7136, 7152), True, 'import numpy as np\n'), ((7336, 7549), 'ml.rl.evaluation.cpe.CpeEstimate', '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)'}), '(raw=weighted_doubly_robust, normalized=weighted_doubly_robust /\n    denominator, raw_std_error=weighted_doubly_robust_std_error,\n    normalized_std_error=weighted_doubly_robust_std_error / denominator)\n', (7347, 7549), False, 'from ml.rl.evaluation.cpe import CpeEstimate\n'), ((8085, 8108), 'numpy.zeros', 'np.zeros', (['[num_j_steps]'], {}), '([num_j_steps])\n', (8093, 8108), True, 'import numpy as np\n'), ((8410, 8444), 'numpy.cov', 'np.cov', (['j_step_return_trajectories'], {}), '(j_step_return_trajectories)\n', (8416, 8444), True, 'import numpy as np\n'), ((8851, 8866), 'numpy.array', 'np.array', (['res.x'], {}), '(res.x)\n', (8859, 8866), True, 'import numpy as np\n'), ((9494, 9520), 'numpy.zeros', 'np.zeros', (['mdp_ids.shape[0]'], {}), '(mdp_ids.shape[0])\n', (9502, 9520), True, 'import numpy as np\n'), ((12708, 12750), 'numpy.multiply', 'np.multiply', (['discounts', 'importance_weights'], {}), '(discounts, importance_weights)\n', (12719, 12750), True, 'import numpy as np\n'), ((12792, 12846), 'numpy.multiply', 'np.multiply', (['discounts', 'importance_weights_one_earlier'], {}), '(discounts, importance_weights_one_earlier)\n', (12803, 12846), True, 'import numpy as np\n'), ((14105, 14121), 'numpy.dot', 'np.dot', (['x', 'error'], {}), '(x, error)\n', (14111, 14121), True, 'import numpy as np\n'), ((1792, 1833), 'numpy.multiply', 'np.multiply', (['target_propensities', 'actions'], {}), '(target_propensities, actions)\n', (1803, 1833), True, 'import numpy as np\n'), ((1919, 1959), 'numpy.multiply', 'np.multiply', (['estimated_q_values', 'actions'], {}), '(estimated_q_values, actions)\n', (1930, 1959), True, 'import numpy as np\n'), ((2031, 2083), 'numpy.multiply', 'np.multiply', (['target_propensities', 'estimated_q_values'], {}), '(target_propensities, estimated_q_values)\n', (2042, 2083), True, 'import numpy as np\n'), ((7003, 7029), 'numpy.std', 'np.std', (['bootstrapped_means'], {}), '(bootstrapped_means)\n', (7009, 7029), True, 'import numpy as np\n'), ((7063, 7094), 'numpy.multiply', 'np.multiply', (['rewards', 'discounts'], {}), '(rewards, discounts)\n', (7074, 7094), True, 'import numpy as np\n'), ((7208, 7293), 'ml.rl.evaluation.cpe.CpeEstimate', 'CpeEstimate', ([], {'raw': '(0.0)', 'normalized': '(0.0)', 'raw_std_error': '(0.0)', 'normalized_std_error': '(0.0)'}), '(raw=0.0, normalized=0.0, raw_std_error=0.0,\n    normalized_std_error=0.0)\n', (7219, 7293), False, 'from ml.rl.evaluation.cpe import CpeEstimate\n'), ((8131, 8167), 'numpy.where', 'np.where', (['(j_step_returns < low_bound)'], {}), '(j_step_returns < low_bound)\n', (8139, 8167), True, 'import numpy as np\n'), ((8269, 8306), 'numpy.where', 'np.where', (['(j_step_returns > high_bound)'], {}), '(j_step_returns > high_bound)\n', (8277, 8306), True, 'import numpy as np\n'), ((8888, 8913), 'numpy.dot', 'np.dot', (['x', 'j_step_returns'], {}), '(x, j_step_returns)\n', (8894, 8913), True, 'import numpy as np\n'), ((9755, 9776), 'numpy.nonzero', 'np.nonzero', (['terminals'], {}), '(terminals)\n', (9765, 9776), True, 'import numpy as np\n'), ((10949, 10972), 'numpy.zeros', 'np.zeros', (['[num_actions]'], {}), '([num_actions])\n', (10957, 10972), True, 'import numpy as np\n'), ((11269, 11292), 'numpy.zeros', 'np.zeros', (['[num_actions]'], {}), '([num_actions])\n', (11277, 11292), True, 'import numpy as np\n'), ((11385, 11408), 'numpy.zeros', 'np.zeros', (['[num_actions]'], {}), '([num_actions])\n', (11393, 11408), True, 'import numpy as np\n'), ((11860, 11894), 'numpy.sum', 'np.sum', (['importance_weights'], {'axis': '(0)'}), '(importance_weights, axis=0)\n', (11866, 11894), True, 'import numpy as np\n'), ((12937, 13009), 'numpy.multiply', 'np.multiply', (['weighted_discounts[:, :j_step + 1]', 'rewards[:, :j_step + 1]'], {}), '(weighted_discounts[:, :j_step + 1], rewards[:, :j_step + 1])\n', (12948, 13009), True, 'import numpy as np\n'), ((13303, 13331), 'numpy.zeros', 'np.zeros', (['[num_trajectories]'], {}), '([num_trajectories])\n', (13311, 13331), True, 'import numpy as np\n'), ((13941, 13951), 'numpy.mean', 'np.mean', (['x'], {}), '(x)\n', (13948, 13951), True, 'import numpy as np\n'), ((13953, 13968), 'scipy.stats.sem', 'sp.stats.sem', (['x'], {}), '(x)\n', (13965, 13968), True, 'import scipy as sp\n'), ((13986, 14032), 'scipy.stats.t._ppf', 'sp.stats.t._ppf', (['((1 + confidence) / 2.0)', '(n - 1)'], {}), '((1 + confidence) / 2.0, n - 1)\n', (14001, 14032), True, 'import scipy as sp\n'), ((2497, 2527), 'numpy.ones', 'np.ones', (['[num_trajectories, 1]'], {}), '([num_trajectories, 1])\n', (2504, 2527), True, 'import numpy as np\n'), ((4435, 4473), 'numpy.cumprod', 'np.cumprod', (['importance_weights'], {'axis': '(1)'}), '(importance_weights, axis=1)\n', (4445, 4473), True, 'import numpy as np\n'), ((4887, 4958), 'numpy.hstack', 'np.hstack', (['[importance_weights_one_earlier, importance_weights[:, :-1]]'], {}), '([importance_weights_one_earlier, importance_weights[:, :-1]])\n', (4896, 4958), True, 'import numpy as np\n'), ((6381, 6438), 'numpy.random.choice', 'np.random.choice', (['num_j_steps', 'sample_size'], {'replace': '(False)'}), '(num_j_steps, sample_size, replace=False)\n', (6397, 6438), True, 'import numpy as np\n'), ((9976, 10017), 'numpy.arange', 'np.arange', (['episode_start', '(episode_end + 1)'], {}), '(episode_start, episode_end + 1)\n', (9985, 10017), True, 'import numpy as np\n'), ((11921, 11960), 'numpy.where', 'np.where', (['(sum_importance_weights == 0.0)'], {}), '(sum_importance_weights == 0.0)\n', (11929, 11960), True, 'import numpy as np\n'), ((13379, 13467), 'numpy.multiply', 'np.multiply', (['weighted_discounts[:, :j_step + 1]', 'estimated_q_values[:, :j_step + 1]'], {}), '(weighted_discounts[:, :j_step + 1], estimated_q_values[:, :\n    j_step + 1])\n', (13390, 13467), True, 'import numpy as np\n'), ((13509, 13612), 'numpy.multiply', 'np.multiply', (['weighted_discounts_one_earlier[:, :j_step + 1]', 'estimated_state_values[:, :j_step + 1]'], {}), '(weighted_discounts_one_earlier[:, :j_step + 1],\n    estimated_state_values[:, :j_step + 1])\n', (13520, 13612), True, 'import numpy as np\n'), ((8585, 8594), 'numpy.sum', 'np.sum', (['x'], {}), '(x)\n', (8591, 8594), True, 'import numpy as np\n'), ((10756, 10803), 'itertools.zip_longest', 'itertools.zip_longest', (['*x'], {'fillvalue': 'fill_value'}), '(*x, fillvalue=fill_value)\n', (10777, 10803), False, 'import itertools\n')]