import unittest
import numpy as np
from openqaoa.problems import PortfolioOptimization


def terms_list_equality(terms_list1, terms_list2):
    """
    Check the terms equality between two terms list
    where the order of edges do not matter.
    """
    if len(terms_list1) != len(terms_list2):
        bool = False
    else:
        for term1, term2 in zip(terms_list1, terms_list2):
            bool = True if (term1 == term2 or term1 == term2[::-1]) else False

    return bool


def terms_list_isclose(terms_list1, terms_list2):
    """
    Check if the distance between two terms list
    where the order of edges do not matter.
    """
    if len(terms_list1) != len(terms_list2):
        bool = False
    else:
        for term1, term2 in zip(terms_list1, terms_list2):
            bool = (
                True
                if np.isclose(term1, term2) or np.isclose(term1, term2[::-1])
                else False
            )

    return bool

class TestPortfolioOptimization(unittest.TestCase):
    """Tests for PortfoliaOptimization class"""
    def test_portfoliooptimization_terms_weights_constant(self):
        """Test terms,weights,constant of QUBO generated by PortfolioOptimization class"""

        po_terms = [[0, 1], [0, 2], [1, 2], [0], [1], [2]]
        po_weights = [0.505, 0.505, 0.505, 0.535, 0.535, 0.535]
        po_constant = 0.8799999999999999

        mu = [0.1, 0.1, 0.1]
        sigma = [[0.01, 0.01, 0.01], [0.01, 0.01, 0.01], [0.01, 0.01, 0.01]]
        risk_factor = 0.1
        budget = 2
        penalty = 1

        qubo = PortfolioOptimization(mu, sigma, risk_factor, budget, penalty).qubo
        terms, weights = qubo.terms, qubo.weights
        constant = qubo.constant

        self.assertEqual(weights, po_weights)
        self.assertEqual(terms, po_terms)
        self.assertEqual(po_constant, constant)

    def test_portfoliooptimization_random_instance(self):
        """Test random instance method of PortfolioOptimization problem class"""
        seed = 1234
        np.random.seed(seed)
        num_assets = 3
        risk_factor = 0.1
        budget = 2
        penalty = 1
        mu_bounds = [-0.1, 0.1]
        sigma_bounds = [-0.01, 0.01]
        mu = [
            (mu_bounds[1] - mu_bounds[0]) * np.random.rand() + mu_bounds[0]
            for _ in range(num_assets)
        ]
        sigma = [[0 for i in range(num_assets)] for j in range(num_assets)]
        for i in range(num_assets):
            for j in range(num_assets):
                sigma[i][j] = (
                    sigma_bounds[1] - sigma_bounds[0]
                ) * np.random.rand() + sigma_bounds[0]

        qubo = PortfolioOptimization(mu, sigma, risk_factor, budget, penalty).qubo

        qubo_random = PortfolioOptimization.random_instance(
            mu_bounds=mu_bounds,
            sigma_bounds=sigma_bounds,
            risk_factor=risk_factor,
            budget=budget,
            seed=seed,
            penalty=penalty,
            num_assets=num_assets,
        ).qubo
        self.assertEqual(qubo.weights, qubo_random.weights)
        self.assertEqual(qubo.terms, qubo_random.terms)
        self.assertEqual(qubo.constant, qubo_random.constant)

    def test_portfoliooptimization_classical_sol(self):
        """Test the portfolio optimization set random instance method classical solution"""

        seed = 1234
        np.random.seed(seed)
        po_sol = PortfolioOptimization.random_instance(
            num_assets=10, seed=seed
        ).classical_solution()

        sol = {
            "asset_0": 0,
            "asset_1": 0,
            "asset_2": 0,
            "asset_3": 1,
            "asset_4": 1,
            "asset_5": 0,
            "asset_6": 0,
            "asset_7": 1,
            "asset_8": 1,
            "asset_9": 1,
        }

        self.assertEqual(po_sol, sol)

    def test_portfoliooptimization_plot(self):
        """Test portfolio optimization random instance method"""
        import matplotlib.pyplot as plt

        seed = 1234
        portfoliooptimization_random_prob = PortfolioOptimization.random_instance(
            num_assets=12, budget=8, seed=seed
        )
        sol = portfoliooptimization_random_prob.classical_solution()
        figure = portfoliooptimization_random_prob.plot_solution(sol)
        self.assertTrue(isinstance(figure, plt.Figure))
        fig, ax = plt.subplots(figsize=(5, 5))
        self.assertTrue(
            portfoliooptimization_random_prob.plot_solution(sol, ax=ax) == None
        )


if __name__ == "__main__":
    unittest.main()