# Copyright 2019, Hudson and Thames Quantitative Research
# All rights reserved
# Read more: https://github.com/hudson-and-thames/mlfinlab/blob/master/LICENSE.txt

"""
Tests the Robust Bayesian Allocation (RBA) algorithm.
"""

import os
import unittest
import numpy as np
import pandas as pd
from mlfinlab.portfolio_optimization.bayesian import RobustBayesianAllocation
from mlfinlab.portfolio_optimization.modern_portfolio_theory import MeanVarianceOptimisation
np.random.seed(420)


class TestRobustBayesianAllocation(unittest.TestCase):
    # pylint: disable=too-many-public-methods, too-many-instance-attributes
    """
    Tests different functions of the RobustBayesianAllocation class.
    """

    def setUp(self):
        project_path = os.path.dirname(__file__)

        # Input parameters
        self.num_simulations = 15
        self.num_observations = 52
        self.num_assets = 5
        self.num_of_portfolios = 10
        self.asset_names = list(map(str, range(self.num_assets)))

        # True market parameters
        correlation_matrix = np.ones((self.num_assets, self.num_assets)) * 0.7
        np.fill_diagonal(correlation_matrix, 0.3)
        step_size = (0.4 - 0.1) / (self.num_assets - 1)
        steps = np.arange(0.1, 0.4 + step_size, step_size)
        s_diag = np.diag(np.diag(steps))
        self.true_covariance = s_diag * correlation_matrix * s_diag
        self.true_mean = (2.5 * self.true_covariance.dot(np.ones((self.num_assets, 1))) / self.num_assets).reshape((self.num_assets,))

        returns = pd.read_csv(project_path + '/test_data/bayesian_sample_returns.csv', index_col=0).values
        self.sample_mean = np.mean(returns, axis=0).reshape((self.num_assets, 1))
        self.sample_covariance = np.cov(returns, rowvar=False)
        self.prior_covariance = np.diag(np.diag(self.sample_covariance))
        self.prior_mean = 0.5 * self.sample_covariance.dot(np.ones((self.num_assets, 1))) / self.num_assets

    def test_bayesian_allocation_for_low_confidence(self):
        # pylint: disable=invalid-name, too-many-locals, broad-except
        """
        Test the efficient frontiers generated by RBA when investor has low confidence in prior.
        """

        sample_expected_returns = []
        sample_volatilities = []
        bayesian_expected_returns = []
        bayesian_volatilities = []
        robust_bayesian_expected_returns = []
        robust_bayesian_volatilities = []
        prior_expected_returns = []
        prior_volatilities = []

        for _ in range(self.num_simulations):
            try:
                returns = np.random.multivariate_normal(mean=self.true_mean, cov=self.true_covariance, size=self.num_observations)

                # Sample estimate
                sample_mean = np.mean(returns, axis=0).reshape((self.num_assets, 1))
                sample_covariance = np.cov(returns, rowvar=False)

                # Bayesian prior
                prior_covariance = np.diag(np.diag(sample_covariance))
                prior_mean = 0.5 * sample_covariance.dot(np.ones((self.num_assets, 1))) / self.num_assets

                # Prior efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(prior_mean,
                                                                                       prior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    prior_expected_returns.append(expected_returns[i])
                    prior_volatilities.append(volatilities[i])

                # Sample efficient frontier (True market frontier)
                _, volatilities, expected_returns = self._calculate_efficient_frontier(sample_mean,
                                                                                       sample_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    sample_expected_returns.append(expected_returns[i])
                    sample_volatilities.append(volatilities[i])

                # Do robust bayesian allocation
                bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
                bayes_allocator.allocate(
                    sample_mean=sample_mean,
                    sample_covariance=sample_covariance,
                    prior_mean=prior_mean,
                    prior_covariance=prior_covariance,
                    relative_confidence_in_prior_mean=1e-4,
                    relative_confidence_in_prior_covariance=1e-4,
                    max_volatility=0.8 * max(sample_mean),
                    sample_size=self.num_observations
                )


                # Bayesian efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(bayes_allocator.posterior_mean,
                                                                                       bayes_allocator.posterior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    bayesian_expected_returns.append(expected_returns[i])
                    bayesian_volatilities.append(volatilities[i])

                # Robust bayesian portfolio
                robust_bayesian_expected_returns.append(bayes_allocator.portfolio_return)
                robust_bayesian_volatilities.append(bayes_allocator.portfolio_risk)
            except Exception:
                continue

        assert np.mean(sample_expected_returns) >= np.mean(robust_bayesian_expected_returns)
        assert np.mean(robust_bayesian_expected_returns) >= np.mean(prior_expected_returns)
        assert np.abs(np.mean(sample_expected_returns) - np.mean(robust_bayesian_expected_returns)) < \
               np.abs(np.mean(robust_bayesian_expected_returns) - np.mean(prior_expected_returns))

    def test_bayesian_allocation_for_medium_confidence(self):
        # pylint: disable=invalid-name, too-many-locals, broad-except
        """
        Test the efficient frontiers generated by RBA when investor has medium confidence in prior.
        """

        sample_expected_returns = []
        sample_volatilities = []
        bayesian_expected_returns = []
        bayesian_volatilities = []
        robust_bayesian_expected_returns = []
        robust_bayesian_volatilities = []
        prior_expected_returns = []
        prior_volatilities = []

        for _ in range(self.num_simulations):
            try:
                returns = np.random.multivariate_normal(mean=self.true_mean, cov=self.true_covariance, size=self.num_observations)

                # Sample estimate
                sample_mean = np.mean(returns, axis=0).reshape((self.num_assets, 1))
                sample_covariance = np.cov(returns, rowvar=False)

                # Bayesian prior
                prior_covariance = np.diag(np.diag(sample_covariance))
                prior_mean = 0.5 * sample_covariance.dot(np.ones((self.num_assets, 1))) / self.num_assets

                # Prior efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(prior_mean,
                                                                                       prior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    prior_expected_returns.append(expected_returns[i])
                    prior_volatilities.append(volatilities[i])

                # Sample efficient frontier (True market frontier)
                _, volatilities, expected_returns = self._calculate_efficient_frontier(sample_mean,
                                                                                       sample_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    sample_expected_returns.append(expected_returns[i])
                    sample_volatilities.append(volatilities[i])

                # Do robust bayesian allocation
                bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
                bayes_allocator.allocate(
                    sample_mean=sample_mean,
                    sample_covariance=sample_covariance,
                    prior_mean=prior_mean,
                    prior_covariance=prior_covariance,
                    relative_confidence_in_prior_mean=1,
                    relative_confidence_in_prior_covariance=1,
                    max_volatility=0.8 * max(sample_mean),
                    sample_size=self.num_observations
                )

                # Bayesian efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(bayes_allocator.posterior_mean,
                                                                                       bayes_allocator.posterior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    bayesian_expected_returns.append(expected_returns[i])
                    bayesian_volatilities.append(volatilities[i])

                # Robust bayesian portfolio
                robust_bayesian_expected_returns.append(bayes_allocator.portfolio_return)
                robust_bayesian_volatilities.append(bayes_allocator.portfolio_risk)
            except Exception:
                continue

        assert np.mean(sample_expected_returns) >= np.mean(robust_bayesian_expected_returns)
        assert np.mean(robust_bayesian_expected_returns) >= np.mean(prior_expected_returns)
        assert np.mean(robust_bayesian_volatilities) <= np.mean(bayesian_volatilities)

    def test_bayesian_allocation_for_high_confidence(self):
        # pylint: disable=invalid-name, too-many-locals, broad-except
        """
        Test the efficient frontiers generated by RBA when investor has high confidence in prior.
        """

        sample_expected_returns = []
        sample_volatilities = []
        bayesian_expected_returns = []
        bayesian_volatilities = []
        robust_bayesian_expected_returns = []
        robust_bayesian_volatilities = []
        prior_expected_returns = []
        prior_volatilities = []

        for _ in range(self.num_simulations):
            try:
                returns = np.random.multivariate_normal(mean=self.true_mean, cov=self.true_covariance, size=self.num_observations)

                # Sample estimate
                sample_mean = np.mean(returns, axis=0).reshape((self.num_assets, 1))
                sample_covariance = np.cov(returns, rowvar=False)

                # Bayesian prior
                prior_covariance = np.diag(np.diag(sample_covariance))
                prior_mean = 0.5 * sample_covariance.dot(np.ones((self.num_assets, 1))) / self.num_assets

                # Prior efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(prior_mean,
                                                                                       prior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    prior_expected_returns.append(expected_returns[i])
                    prior_volatilities.append(volatilities[i])

                # Sample efficient frontier (True market frontier)
                _, volatilities, expected_returns = self._calculate_efficient_frontier(sample_mean,
                                                                                       sample_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    sample_expected_returns.append(expected_returns[i])
                    sample_volatilities.append(volatilities[i])

                # Do robust bayesian allocation
                bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
                bayes_allocator.allocate(
                    sample_mean=sample_mean,
                    sample_covariance=sample_covariance,
                    prior_mean=prior_mean,
                    prior_covariance=prior_covariance,
                    relative_confidence_in_prior_mean=3,
                    relative_confidence_in_prior_covariance=3,
                    max_volatility=0.8 * max(sample_mean),
                    sample_size=self.num_observations
                )

                # Bayesian efficient frontier
                _, volatilities, expected_returns = self._calculate_efficient_frontier(bayes_allocator.posterior_mean,
                                                                                       bayes_allocator.posterior_covariance,
                                                                                       self.asset_names,
                                                                                       self.num_of_portfolios)

                for i in range(self.num_of_portfolios):
                    bayesian_expected_returns.append(expected_returns[i])
                    bayesian_volatilities.append(volatilities[i])

                # Robust bayesian portfolio
                robust_bayesian_expected_returns.append(bayes_allocator.portfolio_return)
                robust_bayesian_volatilities.append(bayes_allocator.portfolio_risk)
            except Exception:
                continue

        assert np.mean(sample_expected_returns) >= np.mean(robust_bayesian_expected_returns)
        assert np.mean(robust_bayesian_expected_returns) >= np.mean(prior_expected_returns)
        assert np.mean(robust_bayesian_volatilities) <= np.mean(bayesian_volatilities)
        assert np.abs(np.mean(sample_expected_returns) - np.mean(robust_bayesian_expected_returns)) > \
               np.abs(np.mean(robust_bayesian_expected_returns) - np.mean(prior_expected_returns))

    def test_bayesian_allocation_dataframe_prior_covariance(self):
        # pylint: disable=invalid-name
        """
        Test bayesian allocation when prior covariance is a dataframe.
        """

        bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
        bayes_allocator.allocate(
            sample_mean=self.sample_mean,
            sample_covariance=self.sample_covariance,
            prior_mean=self.prior_mean,
            prior_covariance=pd.DataFrame(self.prior_covariance, columns=self.asset_names),
            relative_confidence_in_prior_mean=3,
            relative_confidence_in_prior_covariance=3,
            max_volatility=0.8 * max(self.sample_mean),
            sample_size=self.num_observations
        )
        weights = bayes_allocator.weights.values[0]
        assert len(weights) == self.num_assets
        assert round(weights[0], 3) == 0.733
        assert round(weights[1], 3) == 0.146
        assert round(weights[2], 3) == 0.094
        assert round(weights[3], 3) == 0.019
        assert round(weights[4], 3) == 0.008
        np.testing.assert_almost_equal(np.sum(weights), 1)

    def test_bayesian_allocation_dataframe_sample_covariance(self):
        # pylint: disable=invalid-name
        """
        Test bayesian allocation when sample covariance is a dataframe.
        """

        bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
        bayes_allocator.allocate(
            sample_mean=self.sample_mean,
            sample_covariance=pd.DataFrame(self.sample_covariance, columns=self.asset_names),
            prior_mean=self.prior_mean,
            prior_covariance=self.prior_covariance,
            relative_confidence_in_prior_mean=3,
            relative_confidence_in_prior_covariance=3,
            max_volatility=0.8 * max(self.sample_mean),
            sample_size=self.num_observations
        )
        weights = bayes_allocator.weights.values[0]
        assert len(weights) == self.num_assets
        assert round(weights[0], 3) == 0.733
        assert round(weights[1], 3) == 0.146
        assert round(weights[2], 3) == 0.094
        assert round(weights[3], 3) == 0.019
        assert round(weights[4], 3) == 0.008
        np.testing.assert_almost_equal(np.sum(weights), 1)

    def test_bayesian_allocation_for_asset_names(self):
        # pylint: disable=invalid-name
        """
        Test bayesian allocation when passing a list of custom asset names
        """

        bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
        bayes_allocator.allocate(
            sample_mean=self.sample_mean,
            sample_covariance=self.sample_covariance,
            prior_mean=self.prior_mean,
            prior_covariance=self.prior_covariance,
            relative_confidence_in_prior_mean=3,
            relative_confidence_in_prior_covariance=3,
            max_volatility=0.8 * max(self.sample_mean),
            asset_names=self.asset_names,
            sample_size=self.num_observations
        )
        weights = bayes_allocator.weights.values[0]
        assert len(weights) == self.num_assets
        assert round(weights[0], 3) == 0.733
        assert round(weights[1], 3) == 0.146
        assert round(weights[2], 3) == 0.094
        assert round(weights[3], 3) == 0.019
        assert round(weights[4], 3) == 0.008
        np.testing.assert_almost_equal(np.sum(weights), 1)

    def test_bayesian_allocation_for_no_asset_names(self):
        # pylint: disable=invalid-name
        """
        Test bayesian allocation when passing no asset names and the prior and sample covariances are not dataframes.
        """

        bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
        bayes_allocator.allocate(
            sample_mean=self.sample_mean,
            sample_covariance=self.sample_covariance,
            prior_mean=self.prior_mean,
            prior_covariance=self.prior_covariance,
            relative_confidence_in_prior_mean=3,
            relative_confidence_in_prior_covariance=3,
            max_volatility=0.8 * max(self.sample_mean),
            sample_size=self.num_observations
        )
        weights = bayes_allocator.weights.values[0]
        assert len(weights) == self.num_assets
        assert round(weights[0], 3) == 0.733
        assert round(weights[1], 3) == 0.146
        assert round(weights[2], 3) == 0.094
        assert round(weights[3], 3) == 0.019
        assert round(weights[4], 3) == 0.008
        np.testing.assert_almost_equal(np.sum(weights), 1)

    def test_value_error_for_no_weights_found(self):
        """
        Test ValueError when no robust portfolio is found.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                max_volatility=1e-5,
                sample_size=self.num_observations
            )

    def test_value_error_for_diff_prior_lengths(self):
        """
        Test ValueError when the lengths of prior mean and covariance do not match.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean[:2],
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=self.num_observations
            )

    def test_value_error_for_diff_sample_lengths(self):
        """
        Test ValueError when the lengths of sample mean and covariance do not match.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean[:2],
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=self.num_observations
            )

    def test_value_error_for_diff_means_lengths(self):
        """
        Test ValueError when the lengths of sample mean and prior mean do not match.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean[:2],
                sample_covariance=self.sample_covariance[:2, :2],
                prior_mean=self.prior_mean[:3],
                prior_covariance=self.prior_covariance[:3, :3],
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=self.num_observations
            )

    def test_value_error_for_negative_confidence(self):
        """
        Test ValueError when the investor specifies negative confidence in prior.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=-3,
                relative_confidence_in_prior_covariance=3,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=self.num_observations
            )

    def test_value_error_for_negative_estimation_risk_confidence(self):
        # pylint: disable=invalid-name
        """
        Test ValueError when the investor specifies negative posterior estimation risk confidence.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                posterior_covariance_estimation_risk_level=-1,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=self.num_observations
            )

    def test_value_error_for_negative_max_volatility(self):
        # pylint: disable=invalid-name
        """
        Test ValueError when the investor specifies negative maximum volatility.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                posterior_covariance_estimation_risk_level=1,
                max_volatility=-2,
                sample_size=self.num_observations
            )

    def test_value_error_for_single_sample_size(self):
        """
        Test ValueError when the investor specifies only sample size of 1.
        """

        with self.assertRaises(ValueError):
            bayes_allocator = RobustBayesianAllocation(discretisations=self.num_of_portfolios)
            bayes_allocator.allocate(
                sample_mean=self.sample_mean,
                sample_covariance=self.sample_covariance,
                prior_mean=self.prior_mean,
                prior_covariance=self.prior_covariance,
                relative_confidence_in_prior_mean=3,
                relative_confidence_in_prior_covariance=3,
                posterior_covariance_estimation_risk_level=1,
                max_volatility=0.8 * max(self.sample_mean),
                sample_size=1
            )

    @staticmethod
    def _calculate_efficient_frontier(mean, covariance, asset_names, num_portfolios):
        """
        Generate portfolios along the efficient frontier.

        :param mean: (Numpy array) Mean returns.
        :param covariance: (Numpy matrix) Covariance of returns.
        :param asset_names: (Python list) List of asset names in the portfolio.
        :param num_portfolios: (int) Number of portfolios to generate along the efficient frontier.
        :return: (list, list, list) Portfolios on the efficient frontier.
        """

        # Calculate minimum risk portfolio
        mvo = MeanVarianceOptimisation()
        mvo.allocate(expected_asset_returns=mean,
                     covariance_matrix=covariance,
                     asset_names=asset_names,
                     solution='min_volatility')
        min_volatility_weights = mvo.weights.values
        min_volatility_return = mvo.portfolio_return

        # Maximum return
        maximum_return = np.max(mean)

        # Get the target returns along the frontier
        steps = (maximum_return - min_volatility_return) / (num_portfolios - 1)
        target_returns = np.arange(min_volatility_return, maximum_return, steps)

        # Start calculating the portfolios along the frontier
        portfolio_weights, portfolio_volatility, portfolio_return = min_volatility_weights, mvo.portfolio_risk, mvo.portfolio_return

        bayesian_portfolios = [portfolio_weights]
        bayesian_portfolio_volatilities = [portfolio_volatility]
        bayesian_portfolio_returns = [portfolio_return]
        for target_return in target_returns:
            mvo.allocate(expected_asset_returns=mean,
                         covariance_matrix=covariance,
                         target_return=target_return,
                         asset_names=asset_names,
                         solution='efficient_risk')
            bayesian_portfolios.append(mvo.weights.values)
            bayesian_portfolio_volatilities.append(mvo.portfolio_risk)
            bayesian_portfolio_returns.append(mvo.portfolio_return)

        return bayesian_portfolios, bayesian_portfolio_volatilities, bayesian_portfolio_returns
