"""
Utility functions for working with CVXOPT matrices (instances of the
class:`cvxopt.base.matrix` class).
"""

from math import sqrt
from cvxopt import matrix
from cvxopt.lapack import gees, syevr

from . import options


def append_col(left, right):
    """
    Append two matrices side-by-side.

    Parameters
    ----------

    left, right : matrix
        The two matrices to append to one another.

    Returns
    -------

    matrix
        A new matrix consisting of ``right`` appended to the right
        of ``left``.

    Examples
    --------

        >>> A = matrix([1,2,3,4], (2,2))
        >>> B = matrix([5,6,7,8,9,10], (2,3))
        >>> print(A)
        [ 1  3]
        [ 2  4]
        <BLANKLINE>
        >>> print(B)
        [  5   7   9]
        [  6   8  10]
        <BLANKLINE>
        >>> print(append_col(A,B))
        [  1   3   5   7   9]
        [  2   4   6   8  10]
        <BLANKLINE>

    """
    return matrix([left.trans(), right.trans()]).trans()


def append_row(top, bottom):
    """
    Append two matrices top-to-bottom.

    Parameters
    ----------

    top, bottom : matrix
        The two matrices to append to one another.

    Returns
    -------

    matrix
        A new matrix consisting of ``bottom`` appended below ``top``.

    Examples
    --------

        >>> A = matrix([1,2,3,4], (2,2))
        >>> B = matrix([5,6,7,8,9,10], (3,2))
        >>> print(A)
        [ 1  3]
        [ 2  4]
        <BLANKLINE>
        >>> print(B)
        [  5   8]
        [  6   9]
        [  7  10]
        <BLANKLINE>
        >>> print(append_row(A,B))
        [  1   3]
        [  2   4]
        [  5   8]
        [  6   9]
        [  7  10]
        <BLANKLINE>

    """
    return matrix([top, bottom])


def eigenvalues(symmat):
    """
    Return the eigenvalues of the given symmetric real matrix.

    On the surface, this appears redundant to the :func:`eigenvalues_re`
    function. However, if we know in advance that our input is
    symmetric, a better algorithm can be used.

    Parameters
    ----------

    symmat : matrix
        The real symmetric matrix whose eigenvalues you want.

    Returns
    -------

    list of float
       A list of the eigenvalues (in no particular order) of ``symmat``.

    Raises
    ------

    TypeError
        If the input matrix is not symmetric.

    Examples
    --------

        >>> A = matrix([[2,1],[1,2]], tc='d')
        >>> eigenvalues(A)
        [1.0, 3.0]

    If the input matrix is not symmetric, it may not have real
    eigenvalues, and we don't know what to do::

        >>> A = matrix([[1,2],[3,4]])
        >>> eigenvalues(A)
        Traceback (most recent call last):
        ...
        TypeError: input must be a symmetric real matrix

    """
    if not norm(symmat.trans() - symmat) < options.ABS_TOL:
        # Ensure that ``symmat`` is symmetric (and thus square).
        raise TypeError('input must be a symmetric real matrix')

    domain_dim = symmat.size[0]
    eigs = matrix(0, (domain_dim, 1), tc='d')
    syevr(symmat, eigs)
    return list(eigs)


def eigenvalues_re(anymat):
    """
    Return the real parts of the eigenvalues of the given square matrix.

    Parameters
    ----------

    anymat : matrix
        The square matrix whose eigenvalues you want.

    Returns
    -------

    list of float
       A list of the real parts (in no particular order) of the
       eigenvalues of ``anymat``.

    Raises
    ------

    TypeError
        If the input matrix is not square.

    Examples
    --------

    This is symmetric and has two real eigenvalues:

        >>> A = matrix([[2,1],[1,2]], tc='d')
        >>> sorted(eigenvalues_re(A))
        [1.0, 3.0]

    But this rotation matrix has eigenvalues `i` and `-i`, both of whose
    real parts are zero:

        >>> A = matrix([[0,-1],[1,0]])
        >>> eigenvalues_re(A)
        [0.0, 0.0]

    If the input matrix is not square, it doesn't have eigenvalues::

        >>> A = matrix([[1,2],[3,4],[5,6]])
        >>> eigenvalues_re(A)
        Traceback (most recent call last):
        ...
        TypeError: input matrix must be square

    """
    if not anymat.size[0] == anymat.size[1]:
        raise TypeError('input matrix must be square')

    domain_dim = anymat.size[0]
    eigs = matrix(0, (domain_dim, 1), tc='z')

    # Create a copy of ``anymat`` here for two reasons:
    #
    #   1. ``gees`` clobbers its input.
    #   2. We need to ensure that the type code of ``dummy`` is 'd' or 'z'.
    #
    dummy = matrix(anymat, anymat.size, tc='d')

    gees(dummy, eigs)
    return [eig.real for eig in eigs]


def identity(domain_dim, typecode='i'):
    """
    Create an identity matrix of the given dimensions.

    Parameters
    ----------

    domain_dim : int
        The dimension of the vector space on which the identity will act.

    typecode : {'i', 'd', 'z'}, optional
        The type code for the returned matrix, defaults to 'i' for integers.
        Can also be 'd' for real double, or 'z' for complex double.

    Returns
    -------

    matrix
        A ``domain_dim``-by-``domain_dim`` dense integer identity matrix.

    Raises
    ------

    ValueError
        If you ask for the identity on zero or fewer dimensions.

    Examples
    --------

       >>> print(identity(3))
       [ 1  0  0]
       [ 0  1  0]
       [ 0  0  1]
       <BLANKLINE>

    """
    if domain_dim <= 0:
        raise ValueError('domain dimension must be positive')

    entries = [int(i == j)
               for i in range(domain_dim)
               for j in range(domain_dim)]
    return matrix(entries, (domain_dim, domain_dim), tc=typecode)


def inner_product(vec1, vec2):
    """
    Compute the Euclidean inner product of two vectors.

    Parameters
    ----------

    vec1, vec2 : matrix
        The two vectors whose inner product you want.

    Returns
    -------

    float
        The inner product of ``vec1`` and ``vec2``.

    Raises
    ------

    TypeError
        If the lengths of ``vec1`` and ``vec2`` differ.

    Examples
    --------

        >>> x = [1,2,3]
        >>> y = [3,4,1]
        >>> inner_product(x,y)
        14

        >>> x = matrix([1,1,1])
        >>> y = matrix([2,3,4], (1,3))
        >>> inner_product(x,y)
        9

        >>> x = [1,2,3]
        >>> y = [1,1]
        >>> inner_product(x,y)
        Traceback (most recent call last):
        ...
        TypeError: the lengths of vec1 and vec2 must match

    """
    if not len(vec1) == len(vec2):
        raise TypeError('the lengths of vec1 and vec2 must match')

    return sum([x*y for (x, y) in zip(vec1, vec2)])


def norm(matrix_or_vector):
    """
    Return the Frobenius norm of a matrix or vector.

    When the input is a vector, its matrix-Frobenius norm is the same
    thing as its vector-Euclidean norm.

    Parameters
    ----------

    matrix_or_vector : matrix
        The matrix or vector whose norm you want.

    Returns
    -------

    float
        The norm of ``matrix_or_vector``.

    Examples
    --------

        >>> v = matrix([1,1])
        >>> print('{:.5f}'.format(norm(v)))
        1.41421

        >>> A = matrix([1,1,1,1], (2,2))
        >>> norm(A)
        2.0

    """
    return sqrt(inner_product(matrix_or_vector, matrix_or_vector))


def vec(mat):
    """
    Create a long vector in column-major order from ``mat``.

    Parameters
    ----------

    mat : matrix
        Any sort of real matrix that you want written as a long vector.

    Returns
    -------

    matrix
        An ``len(mat)``-by-``1`` long column vector containign the
        entries of ``mat`` in column major order.

    Examples
    --------

        >>> A = matrix([[1,2],[3,4]])
        >>> print(A)
        [ 1  3]
        [ 2  4]
        <BLANKLINE>

        >>> print(vec(A))
        [ 1]
        [ 2]
        [ 3]
        [ 4]
        <BLANKLINE>

    Note that if ``mat`` is a vector, this function is a no-op:

        >>> v = matrix([1,2,3,4], (4,1))
        >>> print(v)
        [ 1]
        [ 2]
        [ 3]
        [ 4]
        <BLANKLINE>
        >>> print(vec(v))
        [ 1]
        [ 2]
        [ 3]
        [ 4]
        <BLANKLINE>

    """
    return matrix(mat, (len(mat), 1))