# (C) Copyright 2018- ECMWF.
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

from typing import List
import numpy as np
from loki.ir import Loop
from loki.expression import (
    symbols as sym,
    FindVariables,
    simplify,
    is_constant,
    accumulate_polynomial_terms,
)
from loki.tools import as_tuple

__all__ = ["Polyhedron"]


class Polyhedron:
    """
    Halfspace representation of a (convex) polyhedron.

    A polyhedron `P c R^d` is described by a set of inequalities, in matrix form
    ```
    P = { x=[x1,...,xd]^T c R^d | Ax <= b }
    ```
    with n-by-d matrix `A` and d-dimensional right hand side `b`.

    In loop transformations, polyhedrons are used to represent iteration spaces of
    d-deep loop nests.


    Parameters
    ----------
    A : numpy.array
        The representation matrix A.
    b : numpy.array
        The right hand-side vector b.
    variables : list, optional
        List of variables representing the dimensions in the polyhedron.

    Attributes
    ----------
    A : numpy.array
        The representation matrix A.
    b : numpy.array
        The right hand-side vector b.
    variables : list, optional, default = None
    """

    def __init__(self, A, b, variables=None):
        A = np.array(A, dtype=np.dtype(int))
        b = np.array(b, dtype=np.dtype(int))
        assert A.ndim == 2 and b.ndim == 1
        assert A.shape[0] == b.shape[0]
        self.A = A
        self.b = b

        self.variables = None
        self.variable_names = None
        if variables is not None:
            assert len(variables) == A.shape[1]
            self.variables = variables
            self.variable_names = [v.name.lower() for v in self.variables]

    def __str__(self):
        str_A = "[" + ", ".join([str(row) for row in self.A]) + "]"
        str_b = f"[{', '.join(map(str, self.b))}]"
        str_variable_names = (
            f"[{', '.join(map(str, self.variable_names))}]"
            if self.variable_names
            else "[]"
        )

        return f"Polyhedron(\n\tA={str_A}, \n\tb={str_b}, \n\tvariables={str_variable_names}\n)"

    def __repr__(self):
        return str(self)

    def _has_satisfiable_constant_restrictions(self):
        """
        Check whether the constant restrictions of the polyhedron are satisfiable.

        This method checks if 0 <= b, assuming that A x = 0.

        Returns:
        bool: True if all constant restrictions are satisfiable, False otherwise.

        """

        return (0 <= self.b).all()

    def is_empty(self):
        """
        Determine whether a polyhedron is empty.

        A polyhedron is considered empty under the following conditions:
        1. It contains no inequalities.
        2. It spans no space, which is a nontrivial problem. The simplest case is when it has an empty
        matrix A and does not satisfy the constant restrictions 0 <= b.

        Notes
        -----
        An empty polyhedron implies that it has no valid solutions or feasible points within its boundaries.
        This function is expected to be called only for polyhedrons with an empty matrix.

        Returns
        -------
        bool
            True if the polyhedron is empty; False if it is not.
        """
        if self.A.size == 0:
            return self.b.size == 0 or not self._has_satisfiable_constant_restrictions()

        raise RuntimeError(
            """
            Checking if a polyhedron with a non-empty matrix spans no space is a nontrivial problem.
            This function is expected to be only called upon polyhedrons with an empty matrix!
            """
        )

    def variable_to_index(self, variable):
        if self.variable_names is None:
            raise RuntimeError("No variables list associated with polyhedron.")
        if isinstance(variable, sym.TypedSymbol):
            variable = variable.name.lower()
        assert isinstance(variable, str)
        return self.variable_names.index(variable)

    @staticmethod
    def _to_literal(value):
        if value < 0:
            return sym.Product((-1, sym.IntLiteral(abs(value))))
        return sym.IntLiteral(value)

    def lower_bounds(self, index_or_variable, ignore_variables=None):
        """
        Return all lower bounds imposed on a variable.

        The lower bounds for the variable `j` are given by the index set:
        ```
        L = {i | A_ij < 0, i in {0, ..., d-1}}
        ```
        Parameters
        ----------
        index_or_variable : int or str or sym.Array or sym.Scalar
            The index, name, or expression symbol for which the lower bounds are produced.
        ignore_variables : list or None, optional
            A list of variable names, indices, or symbols for which constraints should be ignored
            if they depend on one of them.

        Returns
        -------
        list
            The bounds for the specified variable.
        """
        if isinstance(index_or_variable, int):
            j = index_or_variable
        else:
            j = self.variable_to_index(index_or_variable)

        if ignore_variables:
            ignore_variables = [
                i if isinstance(i, int) else self.variable_to_index(i)
                for i in ignore_variables
            ]

        bounds = []
        for i in range(self.A.shape[0]):
            if self.A[i, j] < 0:
                if ignore_variables and any(
                    self.A[i, k] != 0 for k in ignore_variables
                ):
                    # Skip constraint that depends on any of the ignored variables
                    continue

                components = [
                    self._to_literal(self.A[i, k]) * self.variables[k]
                    for k in range(self.A.shape[1])
                    if k != j and self.A[i, k] != 0
                ]
                if not components:
                    lhs = sym.IntLiteral(0)
                elif len(components) == 1:
                    lhs = components[0]
                else:
                    lhs = sym.Sum(as_tuple(components))
                bounds += [
                    simplify(
                        sym.Quotient(
                            self._to_literal(self.b[i]) - lhs,
                            self._to_literal(self.A[i, j]),
                        )
                    )
                ]
        return bounds

    def upper_bounds(self, index_or_variable, ignore_variables=None):
        """
        Return all upper bounds imposed on a variable.

        The upper bounds for the variable `j` are given by the index set:
        ```
        U = {i | A_ij > 0, i in {0, ..., d-1}}
        ```
        Parameters
        ----------
        index_or_variable : int or str or sym.Array or sym.Scalar
            The index, name, or expression symbol for which the upper bounds are produced.
        ignore_variables : list or None, optional
            A list of variable names, indices, or symbols for which constraints should be ignored
            if they depend on one of them.

        Returns
        -------
        list
            The bounds for the specified variable.
        """
        if isinstance(index_or_variable, int):
            j = index_or_variable
        else:
            j = self.variable_to_index(index_or_variable)

        if ignore_variables:
            ignore_variables = [
                i if isinstance(i, int) else self.variable_to_index(i)
                for i in ignore_variables
            ]

        bounds = []
        for i in range(self.A.shape[0]):
            if self.A[i, j] > 0:
                if ignore_variables and any(
                    self.A[i, k] != 0 for k in ignore_variables
                ):
                    # Skip constraint that depends on any of the ignored variables
                    continue

                components = [
                    self._to_literal(self.A[i, k]) * self.variables[k]
                    for k in range(self.A.shape[1])
                    if k != j and self.A[i, k] != 0
                ]
                if not components:
                    lhs = sym.IntLiteral(0)
                elif len(components) == 1:
                    lhs = components[0]
                else:
                    lhs = sym.Sum(as_tuple(components))
                bounds += [
                    simplify(
                        sym.Quotient(
                            self._to_literal(self.b[i]) - lhs,
                            self._to_literal(self.A[i, j]),
                        )
                    )
                ]
        return bounds

    @staticmethod
    def generate_entries_for_lower_bound(bound, variables, index):
        """
        Helper routine to generate matrix and right-hand side entries for a given lower bound.

        Note that this routine can only handle affine bounds, which means expressions that are
        constant or can be reduced to a linear polynomial.

        Upper bounds can be derived from this by multiplying the left-hand side and right-hand side
        with -1.

        Parameters
        ----------
        bound : int or str or sym.Array or sym.Scalar
            The expression representing the lower bound.
        variables : list of str
            The list of variable names.
        index : int
            The index of the variable constrained by this bound.

        Returns
        -------
        tuple(np.array, np.array)
            The pair ``(lhs, rhs)`` of the row in the matrix inequality, where `lhs` is the left-hand side
            and `rhs` is the right-hand side.
        """
        supported_types = (sym.TypedSymbol, sym.MetaSymbol, sym.Sum, sym.Product)
        if not (is_constant(bound) or isinstance(bound, supported_types)):
            raise ValueError(f"Cannot derive inequality from bound {str(bound)}")
        summands = accumulate_polynomial_terms(bound)
        b = -summands.pop(1, 0)  # Constant term or 0
        A = np.zeros([1, len(variables)], dtype=np.dtype(int))
        A[0, index] = -1
        for base, coef in summands.items():
            if not len(base) == 1:
                raise ValueError(f"Non-affine bound {str(bound)}")
            A[0, variables.index(base[0].name.lower())] = coef
        return A, b

    @classmethod
    def from_loop_ranges(cls, loop_variables, loop_ranges):
        """
        Create polyhedron from a list of loop ranges and associated variables.
        """
        assert len(loop_ranges) == len(loop_variables)

        # Add any variables that are not loop variables to the vector of variables
        variables = list(loop_variables)
        variable_names = [v.name.lower() for v in variables]
        for v in sorted(
            FindVariables().visit(loop_ranges), key=lambda v: v.name.lower()
        ):
            if v.name.lower() not in variable_names:
                variables += [v]
                variable_names += [v.name.lower()]

        n = 2 * len(loop_ranges)
        d = len(variables)
        A = np.zeros([n, d], dtype=np.dtype(int))
        b = np.zeros([n], dtype=np.dtype(int))

        for i, (loop_variable, loop_range) in enumerate(
            zip(loop_variables, loop_ranges)
        ):
            assert loop_range.step is None or loop_range.step == "1"
            j = variables.index(loop_variable.name.lower())

            # Create inequality from lower bound
            lhs, rhs = cls.generate_entries_for_lower_bound(
                loop_range.start, variable_names, j
            )
            A[2 * i, :] = lhs
            b[2 * i] = rhs

            # Create inequality from upper bound
            lhs, rhs = cls.generate_entries_for_lower_bound(
                loop_range.stop, variable_names, j
            )
            A[2 * i + 1, :] = -lhs
            b[2 * i + 1] = -rhs

        return cls(A, b, variables)

    @classmethod
    def from_nested_loops(cls, nested_loops: List[Loop]):
        """
        Helper function, for creating a polyhedron from a list of loops.
        """
        return cls.from_loop_ranges(
            [l.variable for l in nested_loops], [l.bounds for l in nested_loops]
        )
