"""Module responsible for sloving an MILP to find best tour around the given landmarks."""
import logging
from collections import defaultdict, deque
import yaml
import numpy as np
import pulp as pl

from ..structs.landmark import Landmark
from ..utils.get_time_distance import get_time
from ..constants import OPTIMIZER_PARAMETERS_PATH


# Silence the pupl logger
logging.getLogger('pulp').setLevel(level=logging.CRITICAL)


class Optimizer:
    """
    Optimizes the balance between the efficiency of a tour and the inclusion of landmarks. 

    The `Optimizer` class is responsible for calculating the best possible detour adjustments 
    to a tour based on specific parameters such as detour time, walking speed, and the maximum 
    number of landmarks to visit. It helps refine a tour by determining whether adding additional 
    landmarks would significantly reduce the overall efficiency.

    Responsibilities:
    - Calculates the maximum detour time allowed for a given tour.
    - Considers the detour factor, which accounts for real-world walking paths versus straight-line distance.
    - Takes into account the average walking speed to estimate walking times.
    - Limits the number of landmarks that can be added to the tour to prevent excessive detouring.
    - Allows some overflow (overshoot) in the maximum detour time to accommodate for slight inefficiencies.

    Attributes:
        logger (logging.Logger): Logger for capturing relevant events and errors.
        detour (int): The accepted maximum detour time in minutes.
        detour_factor (float): The ratio between straight-line distance and actual walking distance in cities.
        average_walking_speed (float): The average walking speed of an adult (in meters per second or kilometers per hour).
        max_landmarks (int): The maximum number of landmarks to include in the tour.
        overshoot (float): The overshoot allowance for exceeding the maximum detour time in a restrictive manner.
    """
    logger = logging.getLogger(__name__)

    detour: int = None              # accepted max detour time (in minutes)
    detour_factor: float            # detour factor of straight line vs real distance in cities
    average_walking_speed: float    # average walking speed of adult
    max_landmarks: int              # max number of landmarks to visit
    overshoot: float                # overshoot to allow maxtime to overflow. Optimizer is a bit restrictive

    def __init__(self) :

        # load parameters from file
        with OPTIMIZER_PARAMETERS_PATH.open('r') as f:
            parameters = yaml.safe_load(f)
            self.detour_factor = parameters['detour_factor']
            self.average_walking_speed = parameters['average_walking_speed']
            self.max_landmarks = parameters['max_landmarks']
            self.overshoot = parameters['overshoot']
            self.time_limit = parameters['time_limit']
            self.gap_rel = parameters['gap_rel']
            self.max_iter = parameters['max_iter']


    def init_ub_time(self, prob: pl.LpProblem, x: pl.LpVariable, L: int, landmarks: list[Landmark], max_time: int):
        """
        Initialize the objective function and inequality constraints for the linear program.

        This function sets up the objective to maximize the attractiveness of visiting landmarks,
        while ensuring that the total time (including travel and visit duration) does not exceed
        the maximum allowed time. It calculates the pairwise travel times between landmarks and
        incorporates visit duration to form the inequality constraints.

        The objective is to maximize sightseeing by selecting the most attractive landmarks within
        the time limit.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints and the objective will be added.
            x (pl.LpVariable): A decision variable representing whether a landmark is visited.
            L (int): The number of landmarks.
            landmarks (list[Landmark]): List of landmarks to visit.
            max_time (int): Maximum allowable time for sightseeing, including travel and visit duration.

        Returns:
            None: Adds the objective function and constraints to the LP problem directly.
            constraint coefficients, and the right-hand side of the inequality constraint.
        """
        L = len(landmarks)

        # Objective function coefficients. a*x1 + b*x2 + c*x3 + ...
        c = np.zeros(L, dtype=np.int16)

        # inequality matrix and vector
        A_ub = np.zeros(L*L, dtype=np.int16)
        b_ub = round(max_time*(1.1+max_time*self.overshoot))

        for i, spot1 in enumerate(landmarks) :
            c[i] = spot1.attractiveness
            for j in range(i+1, L) :
                if i !=j :
                    t = get_time(spot1.location, landmarks[j].location)
                    A_ub[i*L + j] = t + spot1.duration
                    A_ub[j*L + i] = t + landmarks[j].duration

        # Expand 'c' to L*L for every decision variable and ad
        c = np.tile(c, L)

        # Now sort and modify A_ub for each row
        if L > 22 :
            for i in range(L):
                # Get indices of the 4 smallest values in row i
                row_values = A_ub[i*L:i*L+L]
                closest_indices = np.argpartition(row_values, 22)[:22]

                # Create a mask for non-closest landmarks
                mask = np.ones(L, dtype=bool)
                mask[closest_indices] = False

                # Set non-closest landmarks to 32765
                row_values[mask] = 32765
                A_ub[i*L:i*L+L] = row_values

        # Add the objective and the 1 distance constraint
        prob += pl.lpSum([c[j] * x[j] for j in range(L*L)])
        prob += (pl.lpSum([A_ub[j] * x[j] for j in range(L*L)]) <= b_ub)


    def respect_number(self, prob: pl.LpProblem, x: pl.LpVariable, L: int, max_landmarks: int):
        """
        Generate constraints to ensure each landmark is visited at most once and cap the total number of visited landmarks.

        This function adds the following constraints to the linear program:
        1. Each landmark is visited at most once by creating L-2 constraints (one for each landmark).
        2. The total number of visited landmarks is capped by the specified maximum number (`max_landmarks`) plus 2.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable indicating whether a landmark is visited.
            L (int): The total number of landmarks.
            max_landmarks (int): The maximum number of landmarks that can be visited.

        Returns:
            None: This function directly modifies the `prob` object by adding constraints.
        """
        # L-2 constraints: each landmark is visited exactly once
        for i in range(1, L-1):
            prob += (pl.lpSum([x[L*i + j] for j in range(L)]) <= 1)

        # 1 constraint: cap the total number of visits
        prob += (pl.lpSum([1 * x[j] for j in range(L*L)]) <= max_landmarks+2)


    def break_sym(self, prob: pl.LpProblem, x: pl.LpVariable, L: int):
        """
        Generate constraints to prevent simultaneous travel between two landmarks
        in both directions. This constraint ensures that, for any pair of landmarks,
        travel from landmark i to landmark j (dij) and travel from landmark j to landmark i (dji)
        cannot happen simultaneously.

        This method adds constraints to break symmetry, specifically to prevent
        cyclic paths with only two elements. It does not prevent cyclic paths involving more than two elements.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable representing travel between landmarks.
            L (int): The total number of landmarks.

        Returns:
            None: This function modifies the `prob` object by adding constraints in-place.
        """
        upper_ind = np.triu_indices(L, 0, L)  # Get the upper triangular indices
        up_ind_x = upper_ind[0]
        up_ind_y = upper_ind[1]

        # Loop over the upper triangular indices, excluding diagonal elements
        for i, up_ind in enumerate(up_ind_x):
            if up_ind != up_ind_y[i]:
                # Add (L*L-L)/2 constraints to break symmetry
                prob += (x[up_ind*L + up_ind_y[i]] + x[up_ind_y[i]*L + up_ind] <= 1)


    def init_eq_not_stay(self, prob: pl.LpProblem, x: pl.LpVariable, L: int):
        """
        Generate constraints to prevent staying at the same position during travel. 
        Specifically, it removes travel from a landmark to itself (e.g., d11, d22, d33, etc.).

        This function adds one equality constraint to the optimization problem that ensures 
        no decision variable corresponding to staying at the same landmark is included 
        in the solution. This helps in ensuring that the path does not include self-loops.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable representing travel between landmarks.
            L (int): The total number of landmarks.

        Returns:
            None: This function modifies the `prob` object by adding an equality constraint in-place.
        """
        A_eq = np.zeros((L, L), dtype=np.int8)

        # Set diagonal elements to 1 (to prevent staying in the same position)
        np.fill_diagonal(A_eq, 1)
        A_eq = A_eq.flatten()

        # First equality constraint
        prob += (pl.lpSum([A_eq[j] * x[j] for j in range(L*L)]) == 0)


    def respect_start_finish(self, prob: pl.LpProblem, x: pl.LpVariable, L: int):
        """
        Generate constraints to ensure that the optimization starts at the designated 
        start landmark and finishes at the goal landmark.

        Specifically, this function adds three equality constraints:
        1. Ensures that the path starts at the designated start landmark (row 0).
        2. Ensures that the path finishes at the designated goal landmark (row 1).
        3. Prevents any arrivals at the start landmark or departures from the goal landmark (row 2).

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable representing travel between landmarks.
            L (int): The total number of landmarks.

        Returns:
            None: This function modifies the `prob` object by adding three equality constraints in-place.
        """
        # Fill-in row 0.
        A_eq = np.zeros((3,L*L), dtype=np.int8)
        A_eq[0, :L] = np.ones(L, dtype=np.int8)        # sets departures only for start (horizontal ones)
        for k in range(L-1) :
            if k != 0 :
                # Fill-in row 1
                A_eq[1, k*L+L-1] = 1                   # sets arrivals only for finish (vertical ones)
            # Fill-in row 1
            A_eq[2, k*L] = 1

        A_eq[2, L*(L-1):] = np.ones(L, dtype=np.int8)  # prevents arrivals at start and departures from goal
        b_eq= [1, 1, 0]

        # Add the constraints to pulp
        for i in range(3) :
            prob += (pl.lpSum([A_eq[i][j] * x[j] for j in range(L*L)]) == b_eq[i])


    def respect_order(self, prob: pl.LpProblem, x: pl.LpVariable, L: int):
        """
        Generate constraints to tie the optimization problem together and prevent 
        stacked ones, although this does not fully prevent circles.

        This function adds constraints to the optimization problem that prevent 
        simultaneous travel between landmarks in a way that would result in stacked ones.
        However, it does not fully prevent circular paths.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable representing travel between landmarks.
            L (int): The total number of landmarks.

        Returns:
            None: This function modifies the `prob` object by adding L-2 equality constraints in-place.
        """
        # Loop through rows 1 to L-2 to prevent stacked ones
        for i in range(1, L-1):
            # Add the constraint that sums across each "row" or "block" in the decision variables
            row_sum = -pl.lpSum(x[i + j*L] for j in range(L)) + pl.lpSum(x[i*L:(i+1)*L])
            prob += (row_sum == 0)


    def respect_user_must(self, prob: pl.LpProblem, x: pl.LpVariable, L: int, landmarks: list[Landmark]) :
        """
        Generate constraints to ensure that landmarks marked as 'must_do' are included in the optimization.

        This function adds constraints to the optimization problem to ensure that landmarks marked as 
        'must_do' are included in the solution. It precomputes the constraints and adds them to the 
        problem accordingly.

        Args:
            prob (pl.LpProblem): The linear programming problem where constraints will be added.
            x (pl.LpVariable): Decision variable representing travel between landmarks.
            L (int): The total number of landmarks.
            landmarks (list[Landmark]): List of landmarks, where some are marked as 'must_do'.

        Returns:
            None: This function modifies the `prob` object by adding equality constraints in-place.
        """
        ones = np.ones(L, dtype=np.int8)
        A_eq = np.zeros(L*L, dtype=np.int8)

        for i, elem in enumerate(landmarks) :
            if elem.must_do is True and i not in [0, L-1]:
                A_eq[i*L:i*L+L] = ones
                prob += (pl.lpSum([A_eq[j] * x[j] for j in range(L*L)]) == 1)
            if elem.must_avoid is True and i not in [0, L-1]:
                A_eq[i*L:i*L+L] = ones
                prob += (pl.lpSum([A_eq[j] * x[j] for j in range(L*L)]) == 2)


    def prevent_circle(self, prob: pl.LpProblem, x: pl.LpVariable, circle_vertices: list, L: int) :
        """
        Prevent circular paths by adding constraints to the optimization.

        This function ensures that circular paths in both directions (i.e., forward and reverse) 
        between landmarks are avoided in the optimization problem by adding the corresponding constraints.

        Args:
            prob (pl.LpProblem): The linear programming problem instance to which the constraints will be added.
            x (pl.LpVariable): Decision variable representing the travel between landmarks in the problem.
            circle_vertices (list): List of indices representing the landmarks that form a circular path.
            L (int): The total number of landmarks.

        Returns:
            None: This function modifies the `prob` object by adding two equality constraints that 
                prevent circular paths in both directions for the specified circle vertices.
        """
        l = np.zeros((2, L*L), dtype=np.int8)

        for i, node in enumerate(circle_vertices[:-1]) :
            next = circle_vertices[i+1]

            l[0, node*L + next] = 1
            l[1, next*L + node] = 1

        s = circle_vertices[0]
        g = circle_vertices[-1]

        l[0, g*L + s] = 1
        l[1, s*L + g] = 1

        # Add the constraints
        prob += (pl.lpSum([l[0][j] * x[j] for j in range(L*L)]) == 0)
        prob += (pl.lpSum([l[1][j] * x[j] for j in range(L*L)]) == 0)


    def is_connected(self, resx) :
        """
        Determine the order of visits and detect any circular paths in the given configuration.

        Args:
            resx (list): List of edge weights.

        Returns:
            tuple[list[int], Optional[list[list[int]]]]: A tuple containing the visit order and a list of any detected circles.
        """
        resx = np.round(resx).astype(np.int8)  # round all elements and cast them to int

        N = len(resx)               # length of res
        L = int(np.sqrt(N))         # number of landmarks. CAST INTO INT but should not be a problem because N = L**2 by def.

        nonzeroind = np.nonzero(resx)[0] # the return is a little funny so I use the [0]
        nonzero_tup = np.unravel_index(nonzeroind, (L,L))

        ind_a = nonzero_tup[0]
        ind_b = nonzero_tup[1]

        # Extract all journeys
        all_journeys_nodes = []
        visited_nodes = set()

        for node in ind_a:
            if node not in visited_nodes:
                journey_nodes = self.get_journey(node, ind_a, ind_b)
                all_journeys_nodes.append(journey_nodes)
                visited_nodes.update(journey_nodes)

        for l in all_journeys_nodes :
            if 0 in l :
                all_journeys_nodes.remove(l)
                break

        if not all_journeys_nodes :
            return None

        return all_journeys_nodes


    def get_journey(self, start, ind_a, ind_b):
        """
        Trace the journey starting from a given node and follow the connections between landmarks.
        This method constructs a graph from two lists of landmark connections, `ind_a` and `ind_b`,
        where each element in `ind_a` is connected to the corresponding element in `ind_b`.
        It then performs a depth-first search (DFS) starting from the `start` node to determine
        the path (journey) by following the connections.

        Args:
            start (int): The starting node of the journey.
            ind_a (list[int]): List of "from" nodes, representing the starting points of each connection.
            ind_b (list[int]): List of "to" nodes, representing the endpoints of each connection.
        
        Returns:
            list[int]: A list of nodes representing the order of the journey, starting from the `start` node.

        Example:
            If `ind_a = [0, 1, 2]` and `ind_b = [1, 2, 3]`, starting from node 0, the journey would be `[0, 1, 2, 3]`.
        """
        graph = defaultdict(list)
        for a, b in zip(ind_a, ind_b):
            graph[a].append(b)

        journey_nodes = []
        visited = set()
        stack = deque([start])

        while stack:
            node = stack.pop()
            if node not in visited:
                visited.add(node)
                journey_nodes.append(node)
                for neighbor in graph[node]:
                    if neighbor not in visited:
                        stack.append(neighbor)

        return journey_nodes


    def get_order(self, resx):
        """
        Determine the order of visits given the result of the optimization.

        Args:
            resx (list): List of edge weights.

        Returns:
            list[int]: A list containing the visit order.
        """
        resx = np.round(resx).astype(np.uint8)  # must contain only 0 and 1

        N = len(resx)               # length of res
        L = int(np.sqrt(N))         # number of landmarks. CAST INTO INT but should not be a problem because N = L**2 by def.

        nonzeroind = np.nonzero(resx)[0] # the return is a little funny so I use the [0]
        nonzero_tup = np.unravel_index(nonzeroind, (L,L))

        ind_a = nonzero_tup[0].tolist()
        ind_b = nonzero_tup[1].tolist()

        order = [0]
        current = 0
        used_indices = set()  # Track visited index pairs

        while True:
            # Find index of the current node in ind_a
            try:
                i = ind_a.index(current)
            except ValueError:
                break  # No more links, stop the search

            if i in used_indices:
                break  # Prevent infinite loops

            used_indices.add(i)  # Mark this index as visited
            next_node = ind_b[i]  # Get the corresponding node in ind_b
            order.append(next_node)  # Add it to the path

            # Switch roles, now look for next_node in ind_a
            try:
                current = next_node
            except ValueError:
                break  # No further connections, end the path

        return order


    def link_list(self, order: list[int], landmarks: list[Landmark])->list[Landmark] :
        """
        Compute the time to reach from each landmark to the next and create a list of landmarks with updated travel times.

        Args:
            order (list[int]): List of indices representing the order of landmarks to visit.
            landmarks (list[Landmark]): List of all landmarks.

        Returns:
            list[Landmark]]: The updated linked list of landmarks with travel times
        """
        L =  []
        j = 0
        while j < len(order)-1 :
            # get landmarks involved
            elem = landmarks[order[j]]
            next = landmarks[order[j+1]]

            # get attributes
            elem.time_to_reach_next = get_time(elem.location, next.location)
            elem.must_do = True
            elem.location = (round(elem.location[0], 5), round(elem.location[1], 5))
            elem.next_uuid = next.uuid
            L.append(elem)
            j += 1

        next.location = (round(next.location[0], 5), round(next.location[1], 5))
        next.must_do = True
        L.append(next)

        return L


    def warm_start(self, x: list[pl.LpVariable], L: int) :
        """
        This function sets the initial values of the decision variables to a feasible solution.
        This can help the solver start with a feasible or heuristic solution,
        potentially speeding up convergence.

        Args:
        x (list[pl.LpVariable]): A list of PuLP decision variables (binary variables).
        L (int): The size parameter, representing a dimension (likely related to a grid or matrix).

        Returns:
        list[pl.LpVariable]: The modified list of PuLP decision variables with initial values set.
        """
        for i in range(L*L) :
            x[i].setInitialValue(0)

        x[1].setInitialValue(1)
        x[2*L-1].setInitialValue(1)

        return x


    def pre_processing(self, L: int, landmarks: list[Landmark], max_time: int, max_landmarks: int | None) :
        """
        Preprocesses the optimization problem by setting up constraints and variables for the tour optimization.

        This method initializes and prepares the linear programming problem to optimize a tour that includes landmarks, 
        while respecting various constraints such as time limits, the number of landmarks to visit, and user preferences. 
        The pre-processing step sets up the problem before solving it using a linear programming solver.

        Responsibilities:
        - Defines the optimization problem using linear programming (LP) with the objective to maximize the tour value.
        - Creates binary decision variables for each potential transition between landmarks.
        - Sets up inequality constraints to respect the maximum time available for the tour and the maximum number of landmarks.
        - Implements equality constraints to ensure the tour respects the start and finish positions, avoids staying in the same place, 
        and adheres to a visit order.
        - Forces inclusion or exclusion of specific landmarks based on user preferences.

        Attributes:
            prob (pl.LpProblem): The linear programming problem to be solved.
            x (list): A list of binary variables representing transitions between landmarks.
            L (int): The total number of landmarks considered in the optimization.
            landmarks (list[Landmark]): The list of landmarks to be visited in the tour.
            max_time (int): The maximum allowable time for the entire tour.
            max_landmarks (int | None): The maximum number of landmarks to visit in the tour, or None if no limit is set.

        Returns:
            prob (pl.LpProblem): The linear programming problem setup for optimization.
            x (list): The list of binary variables for transitions between landmarks in the tour.
        """

        if max_landmarks is None :
            max_landmarks = self.max_landmarks

        # Initalize the optimization problem
        prob = pl.LpProblem("OptimizationProblem", pl.LpMaximize)

        # Define the problem
        x_bounds = [(0, 1)]*L*L
        x = [pl.LpVariable(f"x_{i}", lowBound=x_bounds[i][0], upBound=x_bounds[i][1], cat='Binary') for i in range(L*L)]

        # Setup the inequality constraints
        self.init_ub_time(prob, x, L, landmarks, max_time)   # Adds the distances from each landmark to the other.
        self.respect_number(prob, x, L, max_landmarks)                      # Respects max number of visits (no more possible stops than landmarks).
        self.break_sym(prob, x, L)                           # Breaks the 'zig-zag' symmetry. Avoids d12 and d21 but not larger cirlces.

        # Setup the equality constraints
        self.init_eq_not_stay(prob, x, L)                # Force solution not to stay in same place
        self.respect_start_finish(prob, x, L)            # Force start and finish positions
        self.respect_order(prob, x, L)                   # Respect order of visit (only works when max_time is limiting factor)
        self.respect_user_must(prob, x, L, landmarks)    # Force to do/avoid landmarks set by user.

        # return prob, self.warm_start(x, L)
        return prob, x


    def solve_optimization(self, max_time: int, landmarks: list[Landmark], max_landmarks: int = None) -> list[Landmark]:
        """
        Main optimization pipeline to solve the landmark visiting problem.

        This method sets up and solves a linear programming problem with constraints to find an optimal tour of landmarks,
        considering user-defined must-visit landmarks, start and finish points, and ensuring no cycles are present.

        Args:
            max_time (int): Maximum time allowed for the tour in minutes.
            landmarks (list[Landmark]): List of landmarks to visit.
            max_landmarks (int): Maximum number of landmarks visited
        Returns:
            list[Landmark]: The optimized tour of landmarks with updated travel times, or None if no valid solution is found.
        """
        # Setup the optimization proplem.
        L = len(landmarks)
        prob, x = self.pre_processing(L, landmarks, max_time, max_landmarks)

        # Solve the problem and extract results.
        try :
            prob.solve(pl.PULP_CBC_CMD(msg=False, timeLimit=self.time_limit+1, gapRel=self.gap_rel))
        except Exception as exc :
            raise Exception(f"No solution found: {str(exc)}") from exc
        status = pl.LpStatus[prob.status]
        solution = [pl.value(var) for var in x]  # The values of the decision variables (will be 0 or 1)

        self.logger.debug("First results are out. Looking out for circles and correcting...")

        # Raise error if no solution is found. FIXME: for now this throws the internal server error
        if status != 'Optimal' :
            self.logger.warning("The problem is overconstrained, no solution on first try.")
            raise ArithmeticError("No solution could be found. Please try again with more time or different preferences.")

        # If there is a solution, we're good to go, just check for connectiveness
        circles = self.is_connected(solution)

        i = 0
        while circles is not None :
            i += 1
            if i == self.max_iter :
                self.logger.warning(f'Timeout: No solution found after {self.max_iter} iterations.')
                raise TimeoutError(f"Optimization took too long. No solution found after {self.max_iter} iterations.")

            for circle in circles :
                self.prevent_circle(prob, x, circle, L)

            # Solve the problem again
            try :
                prob.solve(pl.PULP_CBC_CMD(msg=False, timeLimit=self.time_limit, gapRel=self.gap_rel))
            except Exception as exc :
                self.logger.warning("No solution found: {str(exc)")
                raise Exception(f"No solution found: {str(exc)}") from exc

            solution = [pl.value(var) for var in x]

            if pl.LpStatus[prob.status] != 'Optimal' :
                self.logger.warning("The problem is overconstrained, no solution after {i} cycles.")
                raise ArithmeticError("No solution could be found. Please try again with more time or different preferences.")

            circles = self.is_connected(solution)
            if circles is None :
                break

        # Sort the landmarks in the order of the solution
        order = self.get_order(solution)
        tour =  [landmarks[i] for i in order]

        self.logger.info(f"Re-optimized {i} times, objective value : {int(pl.value(prob.objective))}")
        return tour
