"""
Numerical computation of newforms
"""

#*****************************************************************************
#       Copyright (C) 2004-2006 William Stein <wstein@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from six import integer_types

from sage.structure.sage_object  import SageObject
from sage.structure.sequence     import Sequence
from sage.structure.richcmp import richcmp_method, richcmp
from sage.modular.modsym.all     import ModularSymbols
from sage.modular.arithgroup.all import Gamma0
from sage.modules.all            import vector
from sage.misc.misc              import verbose
from sage.rings.all import CDF, Integer, QQ
from sage.arith.all import next_prime, prime_range
from sage.misc.prandom           import randint
from sage.matrix.constructor     import matrix

# This variable controls importing the SciPy library sparingly
scipy=None

@richcmp_method
class NumericalEigenforms(SageObject):
    """
    numerical_eigenforms(group, weight=2, eps=1e-20, delta=1e-2, tp=[2,3,5])

    INPUT:

    - ``group`` - a congruence subgroup of a Dirichlet character of
      order 1 or 2

    - ``weight`` - an integer >= 2

    - ``eps`` - a small float; abs( ) < eps is what "equal to zero" is
      interpreted as for floating point numbers.

    - ``delta`` - a small-ish float; eigenvalues are considered distinct
      if their difference has absolute value at least delta

    - ``tp`` - use the Hecke operators T_p for p in tp when searching
      for a random Hecke operator with distinct Hecke eigenvalues.

    OUTPUT:

    A numerical eigenforms object, with the following useful methods:

    - :meth:`ap` - return all eigenvalues of $T_p$

    - :meth:`eigenvalues` - list of eigenvalues corresponding
      to the given list of primes, e.g.,::

          [[eigenvalues of T_2],
           [eigenvalues of T_3],
           [eigenvalues of T_5], ...]

    - :meth:`systems_of_eigenvalues` - a list of the systems of
      eigenvalues of eigenforms such that the chosen random linear
      combination of Hecke operators has multiplicity 1 eigenvalues.

    EXAMPLES::

        sage: n = numerical_eigenforms(23)
        sage: n == loads(dumps(n))
        True
        sage: n.ap(2)  # rel tol 2e-15
        [3.0, 0.6180339887498941, -1.618033988749895]
        sage: n.systems_of_eigenvalues(7)  # rel tol 2e-15
        [
        [-1.618033988749895, 2.23606797749979, -3.23606797749979],
        [0.6180339887498941, -2.2360679774997902, 1.2360679774997883],
        [3.0, 4.0, 6.0]
        ]
        sage: n.systems_of_abs(7)
        [
        [0.6180339887..., 2.236067977..., 1.236067977...],
        [1.6180339887..., 2.236067977..., 3.236067977...],
        [3.0, 4.0, 6.0]
        ]
        sage: n.eigenvalues([2,3,5])  # rel tol 2e-15
        [[3.0, 0.6180339887498941, -1.618033988749895],
         [4.0, -2.2360679774997902, 2.23606797749979],
         [6.0, 1.2360679774997883, -3.23606797749979]]
    """
    def __init__(self, group, weight=2, eps=1e-20,
                 delta=1e-2, tp=[2,3,5]):
        """
        Create a new space of numerical eigenforms.

        EXAMPLES::

            sage: numerical_eigenforms(61) # indirect doctest
            Numerical Hecke eigenvalues for Congruence Subgroup Gamma0(61) of weight 2
        """
        if isinstance(group, integer_types + (Integer,)):
            group = Gamma0(Integer(group))
        self._group  = group
        self._weight = Integer(weight)
        self._tp = tp
        if self._weight < 2:
            raise ValueError("weight must be at least 2")
        self._eps = eps
        self._delta = delta

    def __richcmp__(self, other, op):
        """
        Compare two spaces of numerical eigenforms.

        They are considered equal if and only if they come from the
        same space of modular symbols.

        EXAMPLES::

            sage: n = numerical_eigenforms(23)
            sage: n == loads(dumps(n))
            True
        """
        if not isinstance(other, NumericalEigenforms):
            return NotImplemented
        return richcmp(self.modular_symbols(), other.modular_symbols(), op)

    def level(self):
        """
        Return the level of this set of modular eigenforms.

        EXAMPLES::

            sage: n = numerical_eigenforms(61) ; n.level()
            61
        """
        return self._group.level()

    def weight(self):
        """
        Return the weight of this set of modular eigenforms.

        EXAMPLES::

            sage: n = numerical_eigenforms(61) ; n.weight()
            2
        """
        return self._weight

    def _repr_(self):
        """
        Print string representation of self.

        EXAMPLES::

            sage: n = numerical_eigenforms(61) ; n
            Numerical Hecke eigenvalues for Congruence Subgroup Gamma0(61) of weight 2

            sage: n._repr_()
            'Numerical Hecke eigenvalues for Congruence Subgroup Gamma0(61) of weight 2'
        """
        return "Numerical Hecke eigenvalues for %s of weight %s"%(
            self._group, self._weight)

    def modular_symbols(self):
        """
        Return the space of modular symbols used for computing this
        set of modular eigenforms.

        EXAMPLES::

            sage: n = numerical_eigenforms(61) ; n.modular_symbols()
            Modular Symbols space of dimension 5 for Gamma_0(61) of weight 2 with sign 1 over Rational Field
        """
        try:
            return self.__modular_symbols
        except AttributeError:
            M = ModularSymbols(self._group,
                    self._weight, sign=1)
            if M.base_ring() != QQ:
                raise ValueError("modular forms space must be defined over QQ")
            self.__modular_symbols = M
            return M

    def _eigenvectors(self):
        r"""
        Find numerical approximations to simultaneous eigenvectors in
        self.modular_symbols() for all T_p in self._tp.

        EXAMPLES::

            sage: n = numerical_eigenforms(61)
            sage: n._eigenvectors() # random order
            [              1.0    0.289473640239    0.176788851952    0.336707726757  2.4182243084e-16]
            [                0  -0.0702748344418    0.491416161212    0.155925712173    0.707106781187]
            [                0    0.413171180356    0.141163094698   0.0923242547901    0.707106781187]
            [                0    0.826342360711    0.282326189397     0.18464850958 6.79812569682e-16]
            [                0      0.2402380858    0.792225196393    0.905370774276 4.70805946682e-16]

        TESTS:

        This tests if this routine selects only eigenvectors with
        multiplicity one.  Two of the eigenvalues are
        (roughly) -92.21 and -90.30 so if we set ``eps = 2.0``
        then they should compare as equal, causing both eigenvectors
        to be absent from the matrix returned.  The remaining eigenvalues
        (ostensibly unique) are visible in the test, which should be
        independent of which eigenvectors are returned, but it does presume
        an ordering of these eigenvectors for the test to succeed.
        This exercises a correction in :trac:`8018`. ::

            sage: n = numerical_eigenforms(61, eps=2.0)
            sage: evectors = n._eigenvectors()
            sage: evalues = diagonal_matrix(CDF, [-283.0, 108.522012456, 142.0])
            sage: diff = n._hecke_matrix*evectors - evectors*evalues
            sage: sum([abs(diff[i,j]) for i in range(5) for j in range(3)]) < 1.0e-9
            True
        """
        try:
            return self.__eigenvectors
        except AttributeError:
            pass
        verbose('Finding eigenvector basis')
        M = self.modular_symbols()
        N = self.level()

        tp = self._tp
        p = tp[0]
        t = M.T(p).matrix()
        for p in tp[1:]:
            t += randint(-50,50)*M.T(p).matrix()

        self._hecke_matrix = t

        global scipy
        if scipy is None:
            import scipy
        import scipy.linalg
        evals,eig = scipy.linalg.eig(self._hecke_matrix.numpy(), right=True, left=False)
        B = matrix(eig)
        v = [CDF(evals[i]) for i in range(len(evals))]

        # Determine the eigenvectors with eigenvalues of multiplicity
        # one, with equality controlled by the value of eps
        # Keep just these eigenvectors
        eps = self._eps
        w = []
        for i in range(len(v)):
            e = v[i]
            uniq = True
            for j in range(len(v)):
                if uniq and i != j and abs(e-v[j]) < eps:
                    uniq = False
            if uniq:
                w.append(i)
        self.__eigenvectors = B.matrix_from_columns(w)
        return self.__eigenvectors

    def _easy_vector(self):
        """
        Return a very sparse vector v such that v times the eigenvector matrix
        has all entries nonzero.

        ALGORITHM:

        1. Choose row with the most nonzero entries.   (put 1 there)

        2. Consider submatrix of columns corresponding to zero entries
           in row chosen in 1.

        3. Find row of submatrix with most nonzero entries, and add
           appropriate multiple.  Repeat.

        EXAMPLES::

            sage: n = numerical_eigenforms(37)
            sage: n._easy_vector()                 # slightly random output
            (1.0, 1.0, 0)
            sage: n = numerical_eigenforms(43)
            sage: n._easy_vector()                 # slightly random output
            (1.0, 0, 1.0, 0)
            sage: n = numerical_eigenforms(125)
            sage: n._easy_vector()                 # slightly random output
            (0, 0, 0, 1.0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
        """
        try:
            return self.__easy_vector
        except AttributeError:
            pass
        E = self._eigenvectors()
        delta = self._delta
        x = (CDF**E.nrows()).zero_vector()
        if E.nrows() == 0:
            return x



        def best_row(M):
            """
            Find the best row among rows of M, i.e. the row
            with the most entries supported outside [-delta, delta].

            EXAMPLES::

                sage: numerical_eigenforms(61)._easy_vector() # indirect doctest
                (1.0, 1.0, 0.0, 0.0, 0.0)
            """
            R = M.rows()
            v = [len(support(r, delta)) for r in R]
            m = max(v)
            i = v.index(m)
            return i, R[i]

        i, e = best_row(E)

        x[i] = 1

        while True:
            s = set(support(e, delta))
            zp = [i for i in range(e.degree()) if not i in s]
            if len(zp) == 0:
                break
            C = E.matrix_from_columns(zp)
            # best row
            i, f = best_row(C)
            x[i] += 1   # simplistic
            e = x*E

        self.__easy_vector = x
        return x

    def _eigendata(self):
        """
        Return all eigendata for self._easy_vector().

        EXAMPLES::

            sage: numerical_eigenforms(61)._eigendata() # random order
            ((1.0, 0.668205013164, 0.219198805797, 0.49263343893, 0.707106781187), (1.0, 1.49654668896, 4.5620686498, 2.02990686579, 1.41421356237), [0, 1], (1.0, 1.0))
        """
        try:
            return self.__eigendata
        except AttributeError:
            pass
        x = self._easy_vector()

        B = self._eigenvectors()
        def phi(y):
            """
            Take coefficients and a basis, and return that
            linear combination of basis vectors.

            EXAMPLES::

                sage: n = numerical_eigenforms(61) # indirect doctest
                sage: n._eigendata() # random order
                ((1.0, 0.668205013164, 0.219198805797, 0.49263343893, 0.707106781187), (1.0, 1.49654668896, 4.5620686498, 2.02990686579, 1.41421356237), [0, 1], (1.0, 1.0))
            """
            return y.element() * B

        phi_x = phi(x)
        V = phi_x.parent()
        phi_x_inv = V([a**(-1) for a in phi_x])
        eps = self._eps
        nzp = support(x, eps)
        x_nzp = vector(CDF, x.list_from_positions(nzp))
        self.__eigendata = (phi_x, phi_x_inv, nzp, x_nzp)
        return self.__eigendata

    def ap(self, p):
        """
        Return a list of the eigenvalues of the Hecke operator `T_p`
        on all the computed eigenforms.  The eigenvalues match up
        between one prime and the next.

        INPUT:

        - ``p`` - integer, a prime number

        OUTPUT:

        - ``list`` - a list of double precision complex numbers

        EXAMPLES::

            sage: n = numerical_eigenforms(11,4)
            sage: n.ap(2) # random order
            [9.0, 9.0, 2.73205080757, -0.732050807569]
            sage: n.ap(3) # random order
            [28.0, 28.0, -7.92820323028, 5.92820323028]
            sage: m = n.modular_symbols()
            sage: x = polygen(QQ, 'x')
            sage: m.T(2).charpoly('x').factor()
            (x - 9)^2 * (x^2 - 2*x - 2)
            sage: m.T(3).charpoly('x').factor()
            (x - 28)^2 * (x^2 + 2*x - 47)
        """
        p = Integer(p)
        if not p.is_prime():
            raise ValueError("p must be a prime")
        try:
            return self._ap[p]
        except AttributeError:
            self._ap = {}
        except KeyError:
            pass
        a = Sequence(self.eigenvalues([p])[0], immutable=True)
        self._ap[p] = a
        return a

    def eigenvalues(self, primes):
        """
        Return the eigenvalues of the Hecke operators corresponding
        to the primes in the input list of primes.   The eigenvalues
        match up between one prime and the next.

        INPUT:

        - ``primes`` - a list of primes

        OUTPUT:

        list of lists of eigenvalues.

        EXAMPLES::

            sage: n = numerical_eigenforms(1,12)
            sage: n.eigenvalues([3,5,13])  # rel tol 2.4e-10
            [[177148.0, 252.00000000001896], [48828126.0, 4830.000000001376], [1792160394038.0, -577737.9999898539]]
        """
        primes = [Integer(p) for p in primes]
        for p in primes:
            if not p.is_prime():
                raise ValueError('each element of primes must be prime.')
        phi_x, phi_x_inv, nzp, x_nzp = self._eigendata()
        B = self._eigenvectors()
        def phi(y):
            """
            Take coefficients and a basis, and return that
            linear combination of basis vectors.

            EXAMPLES::

                sage: n = numerical_eigenforms(1,12)  # indirect doctest
                sage: n.eigenvalues([3,5,13])  # rel tol 2.4e-10
                [[177148.0, 252.00000000001896], [48828126.0, 4830.000000001376], [1792160394038.0, -577737.9999898539]]
            """
            return y.element() * B

        ans = []
        m = self.modular_symbols().ambient_module()
        for p in primes:
            t = m._compute_hecke_matrix_prime(p, nzp)
            w = phi(x_nzp*t)
            ans.append([w[i]*phi_x_inv[i] for i in range(w.degree())])
        return ans

    def systems_of_eigenvalues(self, bound):
        """
        Return all systems of eigenvalues for self for primes
        up to bound.

        EXAMPLES::

            sage: numerical_eigenforms(61).systems_of_eigenvalues(10)  # rel tol 6e-14
            [
            [-1.4811943040920152, 0.8060634335253695, 3.1563251746586642, 0.6751308705666477],
            [-1.0, -2.0000000000000027, -3.000000000000003, 1.0000000000000044],
            [0.3111078174659775, 2.903211925911551, -2.525427560843529, -3.214319743377552],
            [2.170086486626034, -1.7092753594369208, -1.63089761381512, -0.46081112718908984],
            [3.0, 4.0, 6.0, 8.0]
            ]
        """
        P = prime_range(bound)
        e = self.eigenvalues(P)
        v = Sequence([], cr=True)
        if len(e) == 0:
            return v
        for i in range(len(e[0])):
            v.append([e[j][i] for j in range(len(e))])
        v.sort()
        v.set_immutable()
        return v

    def systems_of_abs(self, bound):
        """
        Return the absolute values of all systems of eigenvalues for
        self for primes up to bound.

        EXAMPLES::

            sage: numerical_eigenforms(61).systems_of_abs(10)  # rel tol 6e-14
            [
            [0.3111078174659775, 2.903211925911551, 2.525427560843529, 3.214319743377552],
            [1.0, 2.0000000000000027, 3.000000000000003, 1.0000000000000044],
            [1.4811943040920152, 0.8060634335253695, 3.1563251746586642, 0.6751308705666477],
            [2.170086486626034, 1.7092753594369208, 1.63089761381512, 0.46081112718908984],
            [3.0, 4.0, 6.0, 8.0]
            ]
        """
        P = prime_range(bound)
        e = self.eigenvalues(P)
        v = Sequence([], cr=True)
        if len(e) == 0:
            return v
        for i in range(len(e[0])):
            v.append([abs(e[j][i]) for j in range(len(e))])
        v.sort()
        v.set_immutable()
        return v

def support(v, eps):
    """
    Given a vector `v` and a threshold eps, return all
    indices where `|v|` is larger than eps.

    EXAMPLES::

        sage: sage.modular.modform.numerical.support( numerical_eigenforms(61)._easy_vector(), 1.0 )
        []

        sage: sage.modular.modform.numerical.support( numerical_eigenforms(61)._easy_vector(), 0.5 )
        [0, 1]

    """
    return [i for i in range(v.degree()) if abs(v[i]) > eps]
