"""
Copyright 2013 Steven Diamond, 2017 Robin Verschueren

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from typing import Tuple

import numpy as np
import scipy.sparse as sp
from scipy import linalg as LA

from cvxpy.atoms.affine.wraps import psd_wrap
from cvxpy.atoms.atom import Atom
from cvxpy.expressions.expression import Expression
from cvxpy.interface.matrix_utilities import is_sparse
from cvxpy.utilities.linalg import sparse_cholesky
from cvxpy.utilities.warn import warn


class CvxPyDomainError(Exception):
    pass


class QuadForm(Atom):
    _allow_complex = True
    block_indices = None  # For compatibility with SymbolicQuadForm

    def __init__(self, x, P) -> None:
        """Atom representing :math:`x^T P x`."""
        super(QuadForm, self).__init__(x, P)

    def numeric(self, values):
        prod = values[1].dot(values[0])
        if self.args[0].is_complex():
            quad = np.dot(np.conj(values[0]).T, prod)
        else:
            quad = np.dot(np.transpose(values[0]), prod)
        return np.real(quad)

    def validate_arguments(self) -> None:
        super(QuadForm, self).validate_arguments()
        n = self.args[1].shape[0]
        if self.args[1].shape[1] != n or self.args[0].shape not in [(n, 1), (n,)]:
            raise ValueError("Invalid dimensions for arguments to quad_form.")
        if not self.args[1].is_hermitian():
            raise ValueError("Quadratic form matrices must be symmetric/Hermitian.")

    def sign_from_args(self) -> Tuple[bool, bool]:
        """Returns sign (is positive, is negative) of the expression.
        """
        return (self.is_atom_convex(), self.is_atom_concave())

    def is_atom_convex(self) -> bool:
        """Is the atom convex?
        """
        P = self.args[1]
        return P.is_constant() and P.is_psd()

    def is_atom_concave(self) -> bool:
        """Is the atom concave?
        """
        P = self.args[1]
        return P.is_constant() and P.is_nsd()

    def is_atom_log_log_convex(self) -> bool:
        """Is the atom log-log convex?
        """
        return True

    def is_atom_log_log_concave(self) -> bool:
        """Is the atom log-log concave?
        """
        return False

    def is_incr(self, idx) -> bool:
        """Is the composition non-decreasing in argument idx?
        """
        if idx == 0:
            # ∇_x f = 2Px: nonneg when (x≥0, P≥0) or (x≤0, P≤0)
            return (self.args[0].is_nonneg() and self.args[1].is_nonneg()) or \
                   (self.args[0].is_nonpos() and self.args[1].is_nonpos())
        elif idx == 1:
            # ∂f/∂P_{ij} = x_i x_j: nonneg when x is all nonneg or all nonpos
            return self.args[0].is_nonneg() or self.args[0].is_nonpos()
        return False

    def is_decr(self, idx) -> bool:
        """Is the composition non-increasing in argument idx?
        """
        if idx == 0:
            # ∇_x f = 2Px: nonpos when (x≥0, P≤0) or (x≤0, P≥0)
            return (self.args[0].is_nonneg() and self.args[1].is_nonpos()) or \
                   (self.args[0].is_nonpos() and self.args[1].is_nonneg())
        return False

    def is_quadratic(self) -> bool:
        """Is the atom quadratic?
        """
        return True

    def has_quadratic_term(self) -> bool:
        """Always a quadratic term.
        """
        return True

    def is_pwl(self) -> bool:
        """Is the atom piecewise linear?
        """
        return False

    def name(self) -> str:
        return f"{type(self).__name__}({self.args[0]}, {self.args[1]})"

    def format_labeled(self) -> str:
        if self._label is not None:
            return self._label
        return (
            f"{type(self).__name__}({self.args[0].format_labeled()}, "
            f"{self.args[1].format_labeled()})"
        )

    def _grad(self, values):
        x = np.array(values[0])
        P = np.array(values[1])
        D = (P + np.conj(P.T)) @ x
        return [sp.csc_array([D.ravel(order="F")]).T]

    def shape_from_args(self) -> Tuple[int, ...]:
        return tuple()


class SymbolicQuadForm(Atom):
    """
    Symbolic form of QuadForm when quadratic matrix is not known (yet).

    Parameters
    ----------
    x : Variable or Expression
        The input expression.
    P : ndarray or sparse matrix
        The quadratic matrix.
    expr : Expression
        The original expression that this represents.
    block_indices : list of np.ndarray, optional
        For non-scalar outputs, maps each output element j to input indices.
        block_indices[j] is an array of indices that output[j] depends on.
        Supports both contiguous and non-contiguous blocks.
        If None, uses existing scalar/diagonal behavior.
    """
    def __init__(self, x, P, expr, block_indices=None) -> None:
        self.original_expression = expr
        self.block_indices = block_indices
        super(SymbolicQuadForm, self).__init__(x, P)
        self.P = self.args[1]

    def get_data(self):
        return [self.original_expression, self.block_indices]

    def _grad(self, values):
        raise NotImplementedError()

    def is_atom_concave(self) -> bool:
        return self.original_expression.is_atom_concave()

    def is_atom_convex(self) -> bool:
        return self.original_expression.is_atom_convex()

    def is_decr(self, idx) -> bool:
        return self.original_expression.is_decr(idx)

    def is_incr(self, idx) -> bool:
        return self.original_expression.is_incr(idx)

    def shape_from_args(self) -> Tuple[int, ...]:
        return self.original_expression.shape_from_args()

    def sign_from_args(self) -> Tuple[bool, bool]:
        return self.original_expression.sign_from_args()

    def is_quadratic(self) -> bool:
        return True


def decomp_quad(P, cond=None, rcond=None, lower=True, check_finite: bool = True):
    """
    Compute a matrix decomposition.

    Compute scale, M1, M2 such that P = scale * (dot(M1, M1.T) - dot(M2, M2.T)).
    Uses LDL decomposition for dense matrices, which is faster than
    eigendecomposition while providing the same factorization structure.

    Parameters
    ----------
    P : matrix or ndarray
        A real symmetric positive or negative (semi)definite input matrix
    cond, rcond : float, optional
        Cutoff for small pivot values.
        Pivot values smaller than rcond * largest_pivot
        are considered negligible.
        If None or -1, suitable machine precision is used (default).
    lower : bool, optional
        Whether the array data is taken from the lower or upper triangle of P.
        The default is to take it from the lower triangle.
    check_finite : bool, optional
        Whether to check that the input matrix contains only finite numbers.
        The default is True; disabling may give a performance gain
        but may result in problems (crashes, non-termination) if the inputs
        contain infinities or NaNs.

    Returns
    -------
    scale : float
        largest absolute pivot value from the LDL decomposition of P
    M1, M2 : 2d ndarray
        A rectangular ndarray such that P = scale * (dot(M1, M1.T) - dot(M2, M2.T))

    """
    if is_sparse(P):
        # TODO: consider using QDLDL instead, if available.
        try:
            sign, L, p = sparse_cholesky(P)
            if sign > 0:
                return 1.0, L[p, :], np.empty((0, 0))
            else:
                return 1.0, np.empty((0, 0)), L[p, :]
        except (ValueError, ModuleNotFoundError):
            P = np.array(P.todense())  # make dense (needs to happen for ldl).
    lu, d, _perm = LA.ldl(P, lower=lower, check_finite=check_finite)

    # Extract effective diagonal values from D, handling any 2x2 blocks.
    # For PSD/NSD matrices D is diagonal (no 2x2 blocks). For indefinite
    # matrices, Bunch-Kaufman pivoting may introduce 2x2 blocks which we
    # resolve by batching all 2x2 blocks into a single eigh call.
    sub_diag = np.diag(d, -1)
    block_starts = np.nonzero(sub_diag)[0]
    diag_vals = np.real(np.diag(d)).copy()
    if len(block_starts) > 0:
        bs = block_starts
        idx = bs[:, None] + np.arange(2)[None, :]  # (k, 2)
        blocks = d[idx[:, :, None], idx[:, None, :]]  # (k, 2, 2)
        eigvals, eigvecs = np.linalg.eigh(blocks)  # (k, 2), (k, 2, 2)
        diag_vals[bs] = eigvals[:, 0]
        diag_vals[bs + 1] = eigvals[:, 1]
        # Apply eigenvector rotations to lu columns
        lu_i = lu[:, bs].copy()
        lu_ip1 = lu[:, bs + 1].copy()
        lu[:, bs] = lu_i * eigvecs[:, 0, 0] + lu_ip1 * eigvecs[:, 1, 0]
        lu[:, bs + 1] = lu_i * eigvecs[:, 0, 1] + lu_ip1 * eigvecs[:, 1, 1]

    if rcond is not None:
        cond = rcond
    if cond in (None, -1):
        t = lu.dtype.char.lower()
        factor = {'f': 1e3, 'd': 1e6}
        cond = factor[t] * np.finfo(t).eps

    scale = max(np.absolute(diag_vals))
    if scale == 0:
        d_scaled = diag_vals
    else:
        d_scaled = diag_vals / scale
    maskp = d_scaled > cond
    maskn = d_scaled < -cond
    # TODO: allow indefinite quad_form
    if np.any(maskp) and np.any(maskn):
        warn("Forming a nonconvex expression quad_form(x, indefinite).")
    M1 = lu[:, maskp] * np.sqrt(d_scaled[maskp])
    M2 = lu[:, maskn] * np.sqrt(-d_scaled[maskn])
    return scale, M1, M2


def quad_form(x, P, assume_PSD: bool = False):
    """ Alias for :math:`x^T P x`.

    Parameters
    ----------
    x : vector argument.
    P : matrix argument.
    assume_PSD : P is assumed to be PSD without checking.
    """
    x, P = map(Expression.cast_to_const, (x, P))
    # Check dimensions.
    if not P.ndim == 2 or P.shape[0] != P.shape[1] or max(x.shape, (1,))[0] != P.shape[0]:
        raise Exception("Invalid dimensions for arguments to quad_form.")
    if x.is_constant():
        return x.T.conjugate() @ P @ x
    elif P.is_constant():
        if assume_PSD:
            P = psd_wrap(P)
        return QuadForm(x, P)
    else:
        raise Exception(
            "At least one argument to quad_form must be non-variable."
        )
