"""
Tools for use in Sage.
"""
import string
from ..sage_helper import _within_sage, sage_method

if _within_sage:
    from ..sage_helper import ZZ, vector, matrix, block_matrix, identity_matrix, gcd, prod, det
    from ..sage_helper import MatrixSpace, PolynomialRing, LaurentPolynomialRing
    from ..sage_helper import SageObject, AbelianGroup, GroupAlgebra
    from .polished_reps import polished_holonomy, MatrixRepresentation
    Id2 = MatrixSpace(ZZ, 2)(1)
else:
    SageObject = object
    ZZ, Id2 = None, None


def search_for_low_rank_triangulation(M, tries=100, target_lower_bound=0):
    rank_lower_bound = max(M.homology().rank(), target_lower_bound)

    rank_upper_bound = M.fundamental_group().num_generators()
    N = M.copy()
    curr_best_tri = N.copy()
    for i in range(tries):
        if rank_upper_bound == rank_lower_bound:
            break
        N.randomize()
        new_rank = N.fundamental_group().num_generators()
        if new_rank < rank_upper_bound:
            rank_upper_bound = new_rank
            curr_best_tri = N.copy()

    return rank_upper_bound, rank_lower_bound, curr_best_tri


# ----------------------------------------------------------------
#
#  Abelianization of the fundamental group
#
# ----------------------------------------------------------------

def abelianize_word(word, gens):
    return vector(ZZ, [ word.count(g) - word.count(g.swapcase()) for g in gens])


class MapToAbelianization(SageObject):
    """
    sage: M = Manifold('v2037')
    sage: ab = MapToAbelianization(M.fundamental_group())
    sage: ab.range()
    Multiplicative Abelian group isomorphic to C2 x C4 x Z
    sage: ab('ab').order()
    4
    sage: ab('abc').order()
    +Infinity

    sage: U = Manifold('dLQbcbchecv')
    sage: ab = MapToAbelianization(U.fundamental_group())
    sage: ab.range()
    Multiplicative Abelian group isomorphic to Z
    sage: ab('aaa')
    t^3
    """
    def __init__(self, fund_group):
        self.domain_gens = fund_group.generators()
        ab_words = [abelianize_word(R, self.domain_gens) for R in fund_group.relators()]
        if not ab_words:
            n = fund_group.num_generators()
            self.elementary_divisors = n*[0,]
            self.U = identity_matrix(n)
        else:
            R = matrix(ZZ, ab_words).transpose()
            D, U, V = R.smith_form()
            m = U.nrows()
            assert m == D.nrows()
            d = min(D.nrows(), D.ncols())
            diag = D.diagonal()
            num_ones = diag.count(1)
            self.elementary_divisors = diag[num_ones:] + [0,]*(m - d)
            self.U = U[num_ones:]

        tor = [d for d in self.elementary_divisors if d != 0]
        free = [d for d in self.elementary_divisors if d == 0]
        names = []
        if len(tor) == 1:
            names.append('u')
        else:
            names += ['u%d' % i for i in range(len(tor))]
        if len(free) == 1:
            names.append('t')
        else:
            names += ['t%d' % i for i in range(len(free))]
        self._range = AbelianGroup(self.elementary_divisors, names=names)

    def range(self):
        return self._range

    def _normalize_exponents(self, exponents):
        D = self.elementary_divisors
        return [v % d if d > 0 else v for (v, d) in zip(exponents, D)]

    def _exponents_of_word(self, word):
        exponents = self.U*abelianize_word(word, self.domain_gens)
        return self._normalize_exponents(exponents)

    def __call__(self, word):
        return self._range(self._exponents_of_word(word))


class MapToFreeAbelianization(MapToAbelianization):
    def range(self):
        return ZZ**self.elementary_divisors.count(0)

    def __call__(self, word):
        D = self.elementary_divisors
        v = self.U*abelianize_word(word, self.domain_gens)
        return vector(ZZ, [v[i] for i in range(len(D)) if D[i] == 0])

# Finding the longitude


@sage_method
def homological_longitude(manifold, cusp=None):
    """
    Returns the peripheral curve in the given cusp, if any, which is
    homologically trivial (with rational coefficients) in the manifold::

        sage: M = Manifold('m015')
        sage: M.homological_longitude()
        (2, -1)

    If no cusp is specified, the default is the first unfilled cusp;
    if all cusps are filled, the default is the first cusp::

        sage: M = Manifold('L5a1(3,4)(0,0)')
        sage: M.homological_longitude()
        (0, 1)

    The components of the next link have nontrivial linking number
    so there is no such curve::

        sage: W = Manifold('L7a2')
        sage: W.homological_longitude(cusp=1) is None
        True

    If every curve in the given cusp is trivial in the rational homology of
    the manifold, an exception is raised::

        sage: M = Manifold('4_1(1,0)')
        sage: M.homological_longitude()
        Traceback (most recent call last):
        ...
        ValueError: Every curve on cusp is homologically trivial
    """
    if cusp is None:
        unfilled = [i for i, status in
                    enumerate(manifold.cusp_info('complete?')) if status]
        if len(unfilled):
            cusp = unfilled[0]
        else:
            cusp = 0
    G = manifold.fundamental_group()
    f = MapToFreeAbelianization(G)
    m, l = G.peripheral_curves()[cusp]
    kernel_basis = matrix(ZZ, [f(m), f(l)]).left_kernel().basis()
    if len(kernel_basis) >= 2:
        raise ValueError('Every curve on cusp is homologically trivial')
    if len(kernel_basis) == 0:
        return None
    return kernel_basis[0]

# --------------------------------------------------------------
#
#  Computing the Alexander polynomial
#
# --------------------------------------------------------------


class MapToGroupRingOfAbelianization(MapToAbelianization):
    """
    sage: M = Manifold('m003')
    sage: G = M.fundamental_group()
    sage: psi = MapToGroupRingOfAbelianization(G)
    sage: psi('ab') + psi('AAAAB')
    u*t + u^4*t^-4
    """
    def __init__(self, fund_group, base_ring=ZZ):
        MapToAbelianization.__init__(self, fund_group)
        self.H = H = self._range  # The abelian group
        self.R = GroupAlgebra(H, base_ring)

    def range(self):
        return self.R

    def __call__(self, word):
        v = MapToAbelianization.__call__(self, word)
        return self.R.monomial(v)


class MapToGroupRingOfFreeAbelianization(MapToFreeAbelianization):
    def __init__(self, fund_group, base_ring=ZZ):
        MapToFreeAbelianization.__init__(self, fund_group)
        n = self.elementary_divisors.count(0)
        self.R = LaurentPolynomialRing(base_ring, list(string.ascii_lowercase[:n]))

    def range(self):
        return self.R

    def __call__(self, word):
        v = MapToFreeAbelianization.__call__(self, word)
        return prod([ g**v[i] for i, g in enumerate(self.R.gens())])


def setup_fox_derivative(word, phi, var, involute=False):
    R = phi.range()
    if len(word) == 0:
        return R.zero()

    # Cache things for speed

    gens = list(set( (var + word).lower() ))
    gens += [g.upper() for g in gens]

    phi_ims = {}
    fox_ders = {}
    for g in gens:
        phi_ims[g] = phi(g) if not involute else phi(g.swapcase())
        if g == g.lower():
            fox_ders[g] = R.zero() if g != var else R.one()
        else:
            fox_ders[g] = R.zero() if g.lower() != var else -phi_ims[var.upper()]

    return R, phi_ims, fox_ders


def fox_derivative(word, phi, var):
    """
    Given a homorphism phi of a group pi, computes
    phi( d word / d var), i.e. the image of the fox derivative
    of the word with respect to var.
    """

    R, phi_ims, fox_ders = setup_fox_derivative(word, phi, var)
    D = 0
    for w in reverse_word(word):
        D = fox_ders[w] + phi_ims[w] * D
    return D


def fox_derivative_with_involution(word, phi, var):
    """
    The group ring Z[pi] has a natural involution iota sends
    g in pi to g^-1 and respects addition.  This function
    computes

        phi( iota( d word / d var) )
    """
    R, phi_ims, fox_ders = setup_fox_derivative(word, phi, var, involute=True)
    D = 0
    for w in reverse_word(word):
        D = fox_ders[w] + D * phi_ims[w]
    return D


# It's best to deal with matrices of polynomials rather than Laurent
# polynomials, so we need to be able to clear denominators.  This add
# to the complexity of the code.


def join_lists(list_of_lists):
    for L in list_of_lists:
        yield from L


def uniform_poly_exponents(poly):
    if poly.parent().ngens() == 1:
        return [(e,) for e in poly.exponents()]
    return poly.exponents()


def minimum_exponents(elts):
    exps = iter(join_lists(uniform_poly_exponents(p) for p in elts))
    mins = list(next(exps))
    n = len(mins)
    for e in exps:
        for i in range(n):
            if e[i] < mins[i]:
                mins[i] = e[i]
    return mins


def convert_laurent_to_poly(elt, minexp, P):
    if P.ngens() == 1:
        f = minexp[0]
        return P({e - f: c for e, c in elt.dict().items()})
    return P({tuple(e - f for e, f in zip(exps, minexp)): c
              for exps, c in elt.dict().items()})


def alexander_polynomial_basic(G, phi, d=0, pos_leading_coeff=False):
    """
    Compute the d-th Alexander polynomial, i.e. the gcd of the
    (d + 1)-st elementrary ideal of A(M).
    """
    R = phi.range()
    P = R.polynomial_ring()
    if G.num_relators() == 0:
        return P(1) if G.num_generators() <= d + 1 else P(0)
    M = [[fox_derivative(rel, phi, var) for rel in G.relators()]
         for var in G.generators()]
    entries = sum(M, [])
    if all(e == 0 for e in entries):
        M = matrix(P, M)
    else:
        minexp = minimum_exponents(entries)
        M = matrix(P, [[ convert_laurent_to_poly(p, minexp, P) for p in row]
                       for row in M])
    alex_poly = gcd(M.minors(G.num_generators() - d - 1))
    if alex_poly == 0:
        return alex_poly

    # Normalize it
    alex_poly = convert_laurent_to_poly(alex_poly, minimum_exponents([alex_poly]), P)
    if pos_leading_coeff:
        e = max(alex_poly.exponents())
        if alex_poly[e] < 0:
            alex_poly = -alex_poly

    return alex_poly


def alexander_polynomial_group(G):
    phi = MapToGroupRingOfFreeAbelianization(G)
    return alexander_polynomial_basic(G, phi, d=0, pos_leading_coeff=True)


@sage_method
def alexander_polynomial(manifold, **kwargs):
    """
    Computes the multivariable Alexander polynomial of the manifold::

        sage: M = Manifold('K12n123')
        sage: M.alexander_polynomial()
        2*a^6 - 14*a^5 + 34*a^4 - 45*a^3 + 34*a^2 - 14*a + 2

        sage: N = Triangulation('v1539(5,1)')
        sage: N.alexander_polynomial()
        a^2*b + a*b^2 + a*b + a + b

    Any provided keyword arguments are passed to
    :meth:`fundamental_group <snappy.Triangulation.fundamental_group>` and
    so affect the group presentation used in the computation.
    """
    if manifold.homology().order() != 'infinite':
        raise ValueError(
            "Alexander polynomial only defined for manifolds with "
            "infinite homology.")

    return alexander_polynomial_group(manifold.fundamental_group(**kwargs))


# --------------------------------------------------------------
#
#  Computing the twisted torsion polynomials
#     for deficiency one presentations.
#
# --------------------------------------------------------------


class PhiAlpha():
    def __init__(self, phi, alpha):
        self.base_ring = phi.range()
        self.image_ring = MatrixSpace(self.base_ring, 2)
        self.phi, self.alpha = phi, alpha

    def range(self):
        return self.image_ring

    def __call__(self, word):
        a = self.phi(word)
        A = self.alpha(word)
        M = self.image_ring(0)
        M[0,0], M[0,1], M[1,0], M[1,1] = a*A[0,0], a*A[0,1], a*A[1,0], a*A[1,1]
        return M


def free_monoid_elt_to_string(elt):
    vars = elt.parent().variable_names()
    return "".join([e*vars[v] for v, e in elt._element_list])


def inverse_word(word):
    L = list(word.swapcase())
    L.reverse()
    return "".join(L)


def reverse_word(word):
    L = list(word)
    L.reverse()
    return "".join(L)


def clean_RR(r, error):
    return 0 if abs(r) < error else r


def clean_CC(z, error):
    CC = z.parent()
    return CC( clean_RR(z.real(),error), clean_RR(z.imag(), error) ) if not CC.is_exact() else z


def univ_exponents(p):
    try:
        return [a[0] for a in p.exponents()]
    except TypeError:
        return p.exponents()


def clean_laurent(p, error):
    R = p.parent()
    t = R.gen()
    new_coeffs = [clean_CC(z, error) for z in p.coefficients()]
    return sum( [a*t**n for a , n in zip(new_coeffs, univ_exponents(p))], R.zero())


def clean_laurent_to_poly(p, error=10**-8):
    R = p.parent()
    P = PolynomialRing(R.base_ring(), R.variable_names())
    t = P.gen()
    cp = clean_laurent(p,error)
    exponents = univ_exponents(cp)
    if len(exponents) == 0:
        return P(0)
    m = min(exponents)
    return sum( [a*t**(n-m) for a,n in zip(cp.coefficients(), exponents)], P(0))


def univ_abs(z):
    """
    Compute a reasonable choice for the absolute value of z.
    """
    try:
        return z.abs()
    except (TypeError, AttributeError):
        if hasattr(z, 'coefficients'):
            return max([0,] + [univ_abs(c) for c in z.coefficients()])
        else:
            return 0 if z == 0 else Infinity


def univ_matrix_norm(A):
    return max([0] + [univ_abs(a) for a in A.list()])


def matrix_has_small_entries(A, bound):
    if A.base_ring().is_exact():
        return A == 0
    else:
        return univ_matrix_norm(A) < bound


def last_square_submatrix(A):
    r, c = A.nrows(), A.ncols()
    if r <= c:
        return A.matrix_from_columns( range(c - r, c) )
    else:
        return A.matrix_from_rows( range( r - c, r) )


def first_square_submatrix(A):
    r, c = A.nrows(), A.ncols()
    if r <= c:
        return A.matrix_from_columns( range(r) )
    else:
        return A.matrix_from_rows( range(c) )


class TorsionComputationError(Exception):
    pass


def fast_determinant_of_laurent_poly_matrix(A):
    """
    Return the determinant of the given matrix up to
    a power of t^n, using the faster algorithm for
    polynomial entries.
    """
    R = A.base_ring()
    P = R.polynomial_ring()
    if A == 0:
        return P(0)
    minexp = minimum_exponents(A.list())
    MS = A.parent().change_ring(P)
    Ap = MS([convert_laurent_to_poly(p, minexp, P) for p in A.list()])
    return Ap.det()


def compute_torsion(G, bits_prec, alpha=None, phi=None, phialpha=None,
                    return_parts=False, return_as_poly=True,
                    wada_conventions=False, symmetry_test=True):
    if alpha:
        F = alpha('a').base_ring()
    elif phialpha:
        F = phialpha('a').base_ring().base_ring()
    epsilon = ZZ(2)**(-bits_prec//3) if not F.is_exact() else None
    big_epsilon = ZZ(2)**(-bits_prec//5) if not F.is_exact() else None
    gens, rels = G.generators(), G.relators()

    if len(rels) != len(gens) - 1:
        raise ValueError(
            "Algorithm to compute torsion requires a group presentation "
            "with deficiency one")

    k = len(gens)
    if phi is None:
        phi = MapToGroupRingOfFreeAbelianization(G, F)

    if len(phi.range().gens()) != 1:
        raise ValueError(
            "Algorithm to compute torsion requires betti number 1")

    # Want the first variable to be homologically nontrivial
    i0 = [i for i, g in enumerate(gens) if phi(g) != 1][0]
    gens = gens[i0:] + gens[:i0]
    if phialpha is None:
        phialpha = PhiAlpha(phi, alpha)

    # Boundary maps for chain complex

    if not wada_conventions:   # Using the conventions of our paper.
        d2 = [ [fox_derivative_with_involution(R, phialpha, g) for R in rels] for g in gens]
        d2 = block_matrix(d2, nrows=k, ncols=k-1)
        d1 = [phialpha(g.swapcase()) - 1 for g in gens]
        d1 = block_matrix(d1, nrows=1, ncols=k)
        dsquared = d1 * d2

    else: # Using those implicit in Wada's paper.
        d2 = [ [fox_derivative(R, phialpha, g) for g in gens] for R in rels]
        d2 = block_matrix(sum(d2, []), nrows=k-1, ncols=k)
        d1 = [phialpha(g) - 1 for g in gens]
        d1 = block_matrix(d1, nrows=k, ncols=1)
        dsquared = d2 * d1

    if not matrix_has_small_entries( dsquared, epsilon ):
        raise TorsionComputationError("(boundary)^2 != 0")

    T = last_square_submatrix(d2)
    if return_as_poly:
        T = fast_determinant_of_laurent_poly_matrix(T)
    else:
        T = det(T)
    B = first_square_submatrix(d1)
    B = det(B)
    if return_as_poly:
        T = clean_laurent_to_poly(T, epsilon)
        B = clean_laurent_to_poly(B, epsilon)
    else:
        T = clean_laurent(T, epsilon)
        B = clean_laurent(B, epsilon)

    if return_parts:
        return (T, B)

    q, r = T.quo_rem(B)
    ans = clean_laurent_to_poly(q, epsilon)
    if (not F.is_exact() and univ_abs(r) > epsilon) or (F.is_exact() and r != 0):
        raise TorsionComputationError("Division failed")

    # Now do a quick sanity check

    if symmetry_test:
        coeffs = ans.coefficients()
        error = max([univ_abs(a-b) for a,b in zip(coeffs, reversed(coeffs))], default=0)
        if (not F.is_exact() and error > epsilon) or (F.is_exact() and error != 0):
            raise TorsionComputationError("Torsion polynomial doesn't seem symmetric")

    return ans

# --------------------------------------------------------------
#
#  Computing the torsion polynomials
#     for the *adjoint* representation
#     i.e. the poly of Dubois-Yamaguichi
# --------------------------------------------------------------


def adjoint_action(A):
    a, b, c, d = A.list()
    return matrix( [[a**2, 2*a*b, b**2],[a*c,b*c+a*d,b*d],[c**2,2*c*d,d**2]] )


def SL2_to_SLN(A, N):
    F = A.base_ring()
    R = PolynomialRing(F, ['x', 'y'])
    x, y = R.gens()
    X, Y = A * vector(R, (x, y))
    monomials = [x**(N - 1 - i) * y**i for i in range(N)]
    image_vectors = [m(X, Y) for m in monomials]
    return matrix(F, [[v.monomial_coefficient(m) for m in monomials] for v in image_vectors])


class PhiAlpha3():
    def __init__(self, phi, alpha):
        self.base_ring = phi.range()
        self.image_ring = MatrixSpace(self.base_ring, 3)
        self.phi, self.alpha = phi, alpha

    def range(self):
        return self.image_ring

    def __call__(self, word):
        a = self.phi(word)
        A = adjoint_action(self.alpha(word))
        M = self.image_ring(0)
        for i in range(3):
            for j in range(3):
                M[i,j] = a*A[i,j]

        return M


class PhiAlphaN():
    def __init__(self, phi, alpha, N):
        self.base_ring = phi.range()
        self.image_ring = MatrixSpace(self.base_ring, N)
        self.phi, self.alpha, self.N = phi, alpha, N

    def range(self):
        return self.image_ring

    def __call__(self, word):
        a = self.phi(word)
        A = SL2_to_SLN(self.alpha(word), self.N)
        M = self.image_ring(0)
        for i in range(self.N):
            for j in range(self.N):
                M[i,j] = a*A[i,j]

        return M


def test_rep(G, phialpha):
    def manually_apply_word(w):
        return prod(phialpha(x) for x in w)
    return max([univ_matrix_norm(manually_apply_word(R) - 1) for R in G.relators()])


@sage_method
def hyperbolic_torsion(manifold, bits_prec=100, all_lifts=False, wada_conventions=False, phi=None):
    """
    Computes the hyperbolic torsion polynomial as defined in
    `[DFJ] <http://arxiv.org/abs/1108.3045>`_::

        sage: M = Manifold('K11n42')
        sage: M.alexander_polynomial()
        1
        sage: tau = M.hyperbolic_torsion(bits_prec=200)
        sage: tau.degree()
        6
    """
    if manifold.homology().betti_number() != 1:
        raise ValueError('Algorithm needs H^1(M; Z) = Z to be able to compute torsion')
    H = manifold.fundamental_group()
    if H.num_generators() != H.num_relators() + 1:
        raise ValueError('Algorithm to compute torsion requires a group presentation '
                         'with deficiency one')
    G = alpha = polished_holonomy(manifold, bits_prec=bits_prec, lift_to_SL2=True)
    if not all_lifts:
        return compute_torsion(G, bits_prec, alpha, phi, wada_conventions=wada_conventions)
    else:
        return [compute_torsion(G, bits_prec, beta, phi, wada_conventions=wada_conventions)
                for beta in alpha.all_lifts_to_SL2C()]


@sage_method
def hyperbolic_SLN_torsion(manifold, N, bits_prec=100):
    """
    Compute the torsion polynomial of the holonomy representation lifted
    to SL(2, C) and then followed by the irreducible representation
    from SL(2, C) -> SL(N, C)::

        sage: M = Manifold('m016')
        sage: [M.hyperbolic_SLN_torsion(N).degree() for N in [2, 3, 4]]
        [18, 27, 36]
    """
    if manifold.homology().betti_number() != 1:
        raise ValueError('Algorithm needs H^1(M; Z) = Z to be able to compute torsion')
    H = manifold.fundamental_group()
    if H.num_generators() != H.num_relators() + 1:
        raise ValueError('Algorithm to compute torsion requires a group presentation '
                         'with deficiency one')
    G = alpha = polished_holonomy(manifold, bits_prec)
    phi = MapToGroupRingOfFreeAbelianization(G, alpha('a').base_ring())
    phialpha = PhiAlphaN(phi, alpha, N)
    if not test_rep(G, phialpha) < ZZ(2)**(bits_prec // 2):
        raise RuntimeError("Invalid representation")
    return compute_torsion(G, bits_prec, phialpha=phialpha, symmetry_test=False)


@sage_method
def hyperbolic_adjoint_torsion(manifold, bits_prec=100):
    """
    Computes the torsion polynomial of the adjoint representation
    a la Dubois-Yamaguichi.   This is not a sign-refined computation
    so the result is only defined up to sign, not to mention a power
    of the variable 'a'::

        sage: M = Manifold('K11n42')
        sage: tau = M.hyperbolic_adjoint_torsion()
        sage: tau.parent()
        Univariate Polynomial Ring in a over Complex Field with 100 bits of precision
        sage: tau.degree()
        7
    """
    return hyperbolic_SLN_torsion(manifold, 3, bits_prec)
