import unittest
import pytest
import numpy as np
import networkx as nx
from docplex.mp.model import Model
from openqaoa.problems import BPSP
from docplex.mp.constants import ObjectiveSense


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 fisher_yates_shuffle(arr, rng):
    """
    Perform an in-place shuffle of a list using the Fisher-Yates shuffle algorithm.

    This algorithm runs in O(n) time and ensures that every permutation of the array is equally likely.
    The shuffle is performed in-place, meaning that the input array `arr` is modified directly.

    Parameters:
    -----------
    arr : list
        A list of elements to be shuffled. This list will be modified in place.
    rng : numpy.random.Generator
        A random number generator instance from NumPy for generating random numbers.
        This allows for reproducibility and control over the random number generation.

    Returns:
    --------
    list
        The shuffled list. Note that the input list is shuffled in place, and the same list is returned.

    Example:
    --------
    >>> from numpy.random import default_rng
    >>> rng = default_rng(seed=42)
    >>> arr = [1, 2, 3, 4, 5]
    >>> fisher_yates_shuffle(arr, rng)
    [3, 4, 1, 5, 2]  # Output can vary based on the RNG seed and state.
    """

    # Iterate over the array from the last element down to the second element
    for i in range(len(arr) - 1, 0, -1):
        # Generate a random index to swap with the current element
        j = rng.integers(0, i + 1)

        # Swap the current element with the randomly chosen element
        arr[i], arr[j] = arr[j], arr[i]

    # Return the shuffled array
    return arr


class TestBPSP(unittest.TestCase):
    """
    Test suite for the BPSP problem.

    This test suite checks the functionality of methods related to the BPSP problem,
    including QUBO creation, car sequencing, solution methods, and more.
    """

    def test_bpsp_qubo(self):
        """
        Test the correct QUBO formation from a provided graph.

        This method validates that BPSP creates the expected QUBO by comparing
        its terms and weights with predefined expected values.
        """
        car_sequence = [3, 1, 0, 0, 2, 2, 4, 4, 3, 1]
        G = nx.Graph()
        G.add_weighted_edges_from(
            [[3, 1, -2], [3, 4, -1], [1, 0, -1], [0, 2, 1], [2, 4, 1]]
        )

        gr_edges = [[3, 1], [3, 4], [1, 0], [0, 2], [2, 4]]
        gr_weights = [-2, -1, -1, 1, 1]

        bpsp_prob_qubo = BPSP(car_sequence).qubo
        self.assertTrue(terms_list_equality(gr_edges, bpsp_prob_qubo.terms))
        self.assertEqual(gr_weights, bpsp_prob_qubo.weights)
        self.assertEqual(0.0, bpsp_prob_qubo.constant)

    def test_bpsp_car_sequence(self):
        """
        Test if ValueErrors are raised for invalid car sequences.

        Certain car sequences are expected to be invalid based on their content.
        This method checks if setting such sequences raises a ValueError.
        """
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        with self.assertRaises(ValueError):
            bpsp.car_sequence = np.array([0, 1, 0])
        with self.assertRaises(ValueError):
            bpsp.car_sequence = np.array([0, 0, 1, 2, 2, 3, 3, 3])

    def test_bpsp_random_instance(self):
        """
        Test the random instance generation of the Binary Paint Shop Problem (BPSP).

        This test verifies whether the `random_instance` method of the BPSP class
        correctly creates a randomized car sequence based on a given number of cars and seed.
        The car sequence should be a list containing two occurrences of each car ID,
        randomly shuffled using the Fisher-Yates shuffle algorithm.

        The function asserts that the sequence generated by the BPSP class's method
        is identical to the sequence produced by an explicitly defined Fisher-Yates shuffle process,
        ensuring both consistency and correctness in the shuffling method based on a fixed seed.
        """

        # Seed and number of cars for generating the instance
        seed = 1234
        num_cars = 10

        # Lambda function to create and shuffle car sequence using Fisher-Yates algorithm
        # The lambda function first duplicates each car ID using np.tile, then applies the shuffle.
        create_car_seq = lambda num_cars, seed: fisher_yates_shuffle(
            np.tile(np.arange(num_cars), 2), np.random.default_rng(seed)
        )

        # Generating a random BPSP instance using the class method
        bpsp = BPSP.random_instance(num_cars=num_cars, seed=seed)

        # Asserting if the generated car sequence in BPSP instance is identical
        # to the sequence produced by our lambda function.
        # This confirms the random instance generation works as expected.
        self.assertTrue(
            all(bpsp.car_sequence == create_car_seq(num_cars, seed)),
            "The generated car sequence does not match the expected shuffled sequence.",
        )

    def test_bpsp_car_pos(self):
        """Test the retrieval of car positions."""
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        self.assertEqual(bpsp.car_positions, {0: (0, 2), 1: (1, 3)})

    def test_bpsp_graph(self):
        """Test the generation of a graph representation of the BPSP instance."""
        car_sequence = [3, 1, 0, 0, 2, 2, 4, 4, 3, 1]
        G = nx.Graph()
        G.add_weighted_edges_from(
            [[3, 1, -2], [3, 4, -1], [1, 0, -1], [0, 2, 1], [2, 4, 1]]
        )

        bpsp_graph = BPSP(car_sequence).graph

        assert nx.is_isomorphic(
            G, bpsp_graph
        ), "The created graph and BPSP graph are not identical,"

    def test_bpsp_docplex_bpsp_model(self):
        """Test if the docplex model representation of BPSP is generated."""
        bpsp = BPSP.random_instance(num_cars=2, seed=1234)

        # Obtain the model
        mdl = bpsp.docplex_bpsp_model  # Access as attribute due to @property decorator

        # Verify that the function returns a valid model instance
        self.assertIsInstance(mdl, Model, "Model is not an instance of DOcplex Model")

        # Manually retrieve binary variables based on their expected names
        sequence = bpsp.car_sequence
        expected_var_names = [f"w_{w}_{i}" for i, w in enumerate(sequence)]
        binary_vars = [mdl.get_var_by_name(name) for name in expected_var_names]

        # Check that all variables are present and are binary
        self.assertEqual(len(binary_vars), 4, "Unexpected number of binary variables")
        for var in binary_vars:
            self.assertTrue(var is not None, f"Variable {var} not found in model")
            self.assertTrue(var.is_binary(), f"Variable {var} is not binary")

        # Check number of constraints
        self.assertEqual(
            mdl.number_of_constraints, 11, "Unexpected number of constraints"
        )

        # Check if objective function exists
        self.assertIsNotNone(mdl.objective_expr, "Objective function is missing")

        # Check if the objective is to minimize
        self.assertEqual(
            mdl.objective_sense,
            ObjectiveSense.Minimize,
            "Objective should be of type 'minimize'",
        )

    @pytest.mark.cplex
    def test_bpsp_solve_cplex(self):
        """
        Test the solution of the BPSP problem using CPLEX.

        This test assumes that CPLEX is installed and functional. It checks
        the length of the solution and the objective value.
        """
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        solution, objective_value = bpsp.solve_cplex()
        self.assertEqual(solution, [1, 1, 0, 0])
        self.assertEqual(objective_value, 1)

    def test_bpsp_paintseq_from_bits(self):
        """Test the solution of the BPSP problem using QAOA."""
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        sequence, color_swaps = bpsp.paintseq_from_bits("10")
        self.assertEqual(sequence, [1, 0, 0, 1])
        self.assertEqual(color_swaps, 2)

    def test_bpsp_solve_redfirst(self):
        """Test the solution of the BPSP problem using the Red-First method."""
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        sequence, color_swaps = bpsp.solve_redfirst()
        self.assertEqual(sequence, [1, 1, 0, 0])
        self.assertEqual(color_swaps, 1)

    def test_bpsp_solve_greedy(self):
        """Test the solution of the BPSP problem using the Greedy method."""
        bpsp = BPSP(np.array([0, 1, 0, 1]))
        sequence, color_swaps = bpsp.solve_greedy()
        self.assertEqual(sequence, [0, 0, 1, 1])
        self.assertEqual(color_swaps, 1)


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