# slightly modified copy of diophantine.py from https://github.com/katherinemwood/corvids
# see NOTES
#
# R. Perry, July 2018
#
"""
Diophantine is a python package for finding small solutions of systems of
diophantine equations (see https://en.wikipedia.org/wiki/Diophantine_equation).
It is based on  PHP code by Keith Matthews (see www.number-theory.org) that
implements the algorithm described in the included 'algorithm.pdf' (see
http://www.numbertheory.org/lll.html for a list of associated publications).
There are two branches of this code in the GitHub repository
(see https://github.com/tclose/Diophantine.git), 'master', which uses the
sympy library and therefore uses arbitrarily long integer representations, and
'numpy', which uses the numpy library, which is faster but can suffer from
integer overflow errors despite using int64 representations.
Diophantine is released under the MIT Licence (see Licence for details)
Author: Thomas G. Close (tom.g.close@gmail.com)
"""
# The MIT License (MIT)
#
# Copyright (c) 2015 Thomas G. Close
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
from copy import deepcopy
from math import copysign, sqrt, log10, floor
from fractions import gcd
from sympy import Matrix, zeros, ones, eye
from itertools import chain, product


class NoSolutionException(Exception):
    pass


class Diophantine:

    # Sign of a variable, which isn't included in math for some reason
    def sign(self, x):
        return copysign(1, x) if x else 0


    def nonzero(self, m):
        return [(i, j) for i, j in product(range(m.shape[0]), range(m.shape[1]))
                if m[i, j] != 0]


    def solve(self, A, b):
        """
        Finds small solutions to systems of diophantine equations, A x = b, where A
        is a M x N matrix of coefficents, b is a M x 1 vector and x is the
        N x 1 solution vector, e.g.
        >>> from sympy import Matrix
        >>> import diophantine
        >>> Dio = Diophantine()
        >>> A = Matrix([[1, 0, 0, 2], [0, 2, 3, 5], [2, 0, 3, 1], [-6, -1, 0, 2],
                        [0, 1, 1, 1], [-1, 2, 0,1], [-1, -2, 1, 0]]).T
        >>> b = Matrix([1, 1, 1, 1])
        >>> Dio.solve(A, b)
        [Matrix([
        [-1],
        [ 1],
        [ 0],
        [ 0],
        [-1],
        [-1],
        [-1]])]
        The returned solution vector will tend to be one with the smallest norms.
        If multiple solutions with the same norm are found they will all be
        returned. If there are no solutions the empty list will be returned.
        """
        A = Matrix(A)
        b = Matrix(b)
        if b.shape != (A.shape[0], 1):
            raise Exception("Length of b vector ({}) does not match number of rows"
                            " in A matrix ({})".format(b.shape[0], A.shape[0]))
        G = zeros(A.shape[1] + 1, A.shape[0] + 1)
        G[:-1, :-1] = A.T
        G[-1, :-1] = b.reshape(1, b.shape[0])
        G[-1, -1] = 1
        # A is m x n, b is m x 1, solving AX=b, X is n x 1+
        # Ab is the (n+1) x m transposed augmented matrix. G=[A^t|0] [b^t]1]
        hnf, P, rank = self.lllhermite(G)
        r = rank - 1  # For convenience
        if not any(chain(hnf[:r, -1], hnf[r, :-1])) and hnf[r, -1] == 1:
            nullity = hnf.shape[0] - rank
            if nullity:
                basis = P[rank:, :-1].col_join(-P[r, :-1])
                solutions = self.get_solutions(basis)
            else:
                solutions = []; solutions.append(-P[r,:-1].T) # RP, list of Matrix
                # raise NotImplementedError("Ax=B has unique solution in integers")
        else:
            solutions = []
        return solutions

    def getBasis(self, A, b):
        A = Matrix(A)
        b = Matrix(b)
        if b.shape != (A.shape[0], 1):
            raise Exception("Length of b vector ({}) does not match number of rows"
                            " in A matrix ({})".format(b.shape[0], A.shape[0]))
        G = zeros(A.shape[1] + 1, A.shape[0] + 1)
        G[:-1, :-1] = A.T
        G[-1, :-1] = b.reshape(1, b.shape[0])
        G[-1, -1] = 1
        # A is m x n, b is m x 1, solving AX=b, X is n x 1+
        # Ab is the (n+1) x m transposed augmented matrix. G=[A^t|0] [b^t]1]
        hnf, P, rank = self.lllhermite(G)
        r = rank - 1  # For convenience
        if not any(chain(hnf[:r, -1], hnf[r, :-1])) and hnf[r, -1] == 1:
            nullity = hnf.shape[0] - rank
            if nullity:
                basis = P[rank:, :-1].col_join(-P[r, :-1])
                solutions = basis
            else:
                return P[rank:, :-1].col_join(-P[r, :-1])
                raise NotImplementedError("Ax=B has unique solution in integers")
        else:
            solutions = []
        return solutions

    def lllhermite(self, G, m1=1, n1=1):
        """
        Input: integer mxn matrix A, nonzero, at least two rows+
        Output: small unimodular matrix B and HNF(A), such that BA=HNF(A)+
        The Havas, Majewski, Matthews LLL method is used+
        We usually take alpha=m1/n1, with (m1,n1)=(1,1) to get best results+
        """
        m = G.shape[0]
        n = G.shape[1]
        A, B, L, D = self.initialise_working_matrices(G)
        if self.first_nonzero_is_negative(A):
            B[m-1, m-1] = -1
            A[m-1, :] *= -1
        k = 1
        while k < m:
            col1, col2 = self.reduce_matrix(A, B, L, k, k - 1, D)
            u = n1 * (int(D[k - 1]) * int(D[k + 1]) +
                      int(L[k, k - 1]) * int(L[k, k - 1]))
            v = m1 * int(D[k]) * int(D[k])
            if col1 <= min(col2, n - 1) or (col1 == n and col2 == n and u < v):
                self.swap_rows(k, A, B, L, D)
                if k > 1:
                    k = k - 1
            else:
                for i in reversed(range(k - 1)):
                    self.reduce_matrix(A, B, L, k, i, D)
                k = k + 1
        try:
            rank = A.shape[0] - next(i for i in range(A.shape[0])
                                     if self.nonzero(A[i, :]))
        except StopIteration:
            assert False, "A matrix contains only zeros"
        hnf = A[::-1, :]
        unimodular_matrix = B[::-1, :]
        return hnf, unimodular_matrix, rank


    def initialise_working_matrices(self, G):
        """  G is a nonzero matrix with at least two rows.  """
        B = eye(G.shape[0])
        # Lower triang matrix
        L = zeros(G.shape[0], G.shape[0])
        D = ones(G.shape[0] + 1, 1)
        A = Matrix(G)
        return A, B, L, D


    def first_nonzero_is_negative(self, A):
        """
        returns 0 if the first nonzero column j of A contains more than one nonzero
        entry, or contains only one nonzero entry and which is positive+ returns 1
        if the first nonzero column j of A contains only one nonzero entry, which
        is negative+ This assumes A is a nonzero matrix with at least two rows+
        """
        nonzero_columns = list(zip(*self.nonzero(A)))[1]  # Should always be nonzero
        # Get the first nonzero column
        nonzero_col = A[:, min(nonzero_columns)]
        # Get the nonzero elements
        nonzero_elems = [e for e in nonzero_col if e != 0]
        # If there is only one and it is negative return 1 else 0
        return len(nonzero_elems) == 1 and nonzero_elems[0] < 0


    def reduce_matrix(self, A, B, L, k, i, D):
        nonzero_i_elems = list(zip(*self.nonzero(A[i, :])))
        if len(nonzero_i_elems):
            col1 = nonzero_i_elems[1][0]
            if A[i, col1] < 0:
                self.minus(i, L)
                A[i, :] *= -1
                B[i, :] *= -1
        else:
            col1 = A.shape[1]
        nonzero_k_elems = list(zip(*self.nonzero(A[k, :])))
        if len(nonzero_k_elems):
            col2 = nonzero_k_elems[1][0]
        else:
            col2 = A.shape[1]
        if col1 < A.shape[1]:
            q = A[k, col1] // A[i, col1]
        else:
            t = abs(L[k, i])
            t = 2 * t
            if t > D[i + 1]:
                q = self.lnearint(L[k, i], D[i + 1])
            else:
                q = 0
        if q != 0:
            A[k, :] -= q * A[i, :]
            B[k, :] -= q * B[i, :]
            L[k, i] -= q * D[i + 1]
            L[k, :i] -= q * L[i, :i]
        return col1, col2


    def minus(self, j, L):
        L[j, :] = -L[j, :]
        L[:, j] = -L[:, j]


    def swap_rows(self, k, A, B, L, D):
        # To avoid the interpretation of -1 as the last index of the matrix create
        # a reverse stop that ends past the negative of the length of the matrix
        reverse_stop = k - 2 if k > 1 else -(A.shape[0] + 1)
        # Swap rows of the matrices
        A[(k - 1):(k + 1), :] = A[k:reverse_stop:-1, :]
        B[(k - 1):(k + 1), :] = B[k:reverse_stop:-1, :]
        L[(k - 1):(k + 1), :(k - 1)] = L[k:reverse_stop:-1, :(k - 1)]
        t = (L[(k + 1):, k - 1] * D[k + 1] / D[k] -
             L[(k + 1):, k] * L[k, k - 1] / D[k])
        L[(k + 1):, k - 1] = (L[(k + 1):, k - 1] * L[k, k - 1] +
                              L[(k + 1):, k] * D[k - 1]) / D[k]
        L[(k + 1):, k] = t
        t = int(D[k - 1]) * int(D[k + 1]) + int(L[k, k - 1]) * int(L[k, k - 1])
        D[k] = t / D[k]


    def get_solutions(self, A):
        m = A.shape[0]
        n = A.shape[1]
        G = self.gram(A)
        N, D = self.cholesky(G)
        Qn, Qd = N, D
        m -= 1
        Nn = Qn[:m, m]
        Nd = Qd[:m, m]
        Cn = 0
        Cd = 1
        for i in range(m):
            num, den = self.multr(Nn[i], Nd[i], Nn[i], Nd[i])
            num, den = self.multr(num, den, Qn[i, i], Qd[i, i])
            Cn, Cd = self.addr(Cn, Cd, num, den)
        i = m - 1
        # List to hold working variables
        x = zeros(m, 1)
        UB = zeros(m, 1)
        Tn = zeros(m, 1)
        Td = zeros(m, 1)
        Un = zeros(m, 1)
        Ud = zeros(m, 1)
        Tn[i] = Cn
        Td[i] = Cd
        Un[i] = 0
        Ud[i] = 1
        solutions = []  # List to hold multipliers
        while True:
            # Calculate UB
            Zn, Zd = self.ratior(Tn[i], Td[i], Qn[i, i], Qd[i, i])
            num, den = self.subr(Nn[i], Nd[i], Un[i], Ud[i])
            UB[i] = self.introot(Zn, Zd, num, den)
            # Calculate x
            num, den = self.subr(Un[i], Ud[i], Nn[i], Nd[i])
            x[i] = -self.introot(Zn, Zd, num, den) - 1
            while True:
                x[i] += 1
                if x[i] <= UB[i]:
                    if i == 0:
                        lcv = self.lcasvector(A[:-1, :], x)
                        solution = A[m, :n] - lcv.reshape(1, lcv.shape[0])
                        solutions.append(solution.T)
                    else:
                        # now update U
                        Un[i - 1], Ud[i - 1] = 0, 1
                        for j in range(i, m):
                            # Loops from back of xs
                            num, den = self.multr(Qn[i - 1, j], Qd[i - 1, j], x[j], 1)
                            Un[i - 1], Ud[i - 1] = self.addr(Un[i - 1], Ud[i - 1], num,
                                                        den)
                        # now update T
                        num, den = self.addr(x[i], 1, Un[i], Ud[i])
                        num, den = self.subr(num, den, Nn[i], Nd[i])
                        num, den = self.multr(num, den, num, den)
                        num, den = self.multr(Qn[i, i], Qd[i, i], num, den)
                        Tn[i - 1], Td[i - 1] = self.subr(Tn[i], Td[i], num, den)
                        i = i - 1
                        break
                else:
                    i = i + 1
                    if i == m:
                        return solutions


    def cholesky(self, A):
        """
        # A is positive definite mxm
        """
        assert A.shape[0] == A.shape[1]
        # assert all(A.eigenvals() > 0)
        m = A.shape[0]
        N = deepcopy(A)
        D = ones(*A.shape)
        for i in range(m - 1):
            for j in range(i + 1, m):
                N[j, i] = N[i, j]
                D[j, i] = D[i, j]
                n, d = self.ratior(N[i, j], D[i, j], N[i, i], D[i, i])
                N[i, j], D[i, j] = n, d
            for k in range(i + 1, m):
                for l in range(k, m):
                    n, d = self.multr(N[k, i], D[k, i], N[i, l], D[i, l])
                    N[k, l], D[k, l] = self.subr(N[k, l], D[k, l], n, d)
        return N, D


    def gram(self, A):
        """
        Need to check for row and column operations
        """
        m = A.shape[0]
        B = zeros(m, m)
        for i in range(m):
            for j in range(m):
                B[i, j] = A[i, :].dot(A[j, :])  # dotproduct(A[i], A[j], n)
        return Matrix(B)


    def introot(self, a, b, c, d):
        """
        With Z=a/b, U=c/d, returns [sqrt(a/b)+c/d]. First ANSWER =
        [sqrt(Z)] + [U]. One then tests if Z < ([sqrt(Z)] + 1 -U)^2. If
        this does not hold, ANSWER += 1+ For use in fincke_pohst()+
        """
        y = c // d
        if a == 0:
            return y
        x = a // b
        assert x >= 0
        x_sqrt = int(floor(sqrt(x)))
        answer = x_sqrt + y
        num, den = self.subr(c, d, y, 1)
        num, den = self.subr(1, 1, num, den)
        num, den = self.addr(x_sqrt, 1, num, den)
        num, den = self.multr(num, den, num, den)
        t = self.comparer(num, den, a, b)
        if t <= 0:
            answer = answer + 1
        int_answer = int(answer)
        assert int_answer == answer
        return int_answer


    def lnearint(self, a, b):
        """
        left nearest integer
        returns y+1/2 if a/b=y+1/2, y integral+
        """
        y = a // b
        if b < 0:
            a = -a
            b = -b
        x = b * y
        z = a - x
        z = 2 * z
        if z > b:
            y = y + 1
        return y


    def ratior(self, a, b, c, d):
        """ returns (a/b)/(c/d)"""
        r = a * d
        s = b * c
        g = abs(gcd(r, s))
        if s < 0:
            g = -g
        return r / g, s / g


    def multr(self, a, b, c, d):
        # returns (a/b)(c/d)
        r = a * c
        s = b * d
        g = abs(gcd(r, s))
        return r / g, s / g


    def subr(self, a, b, c, d):
        t = a * d - b * c
        u = b * d
        g = abs(gcd(t, u))
        return t / g, u / g


    def addr(self, a, b, c, d):
        t = a * d + b * c
        u = b * d
        g = abs(gcd(t, u))
        return t / g, u / g


    def comparer(self, a, b, c, d):
        """Assumes b>0 and d>0.  Returns -1, 0 or 1 according as a/b <,=,> c/d+ """
        assert b > 0 and d > 0
        return self.sign(a * d - b * c)


    def lcasvector(self, A, x):
        """lcv[j]=X[1]A[1, j]+...+X[m]A[m, j], 1 <= j <= n+"""
        # global lcv
    #     printnp(x)
    #     printnp(A)
        n = A.shape[1]
        lcv = zeros(n, 1)
        for j in range(n):
            lcv[j] = x.dot(A[:, j])
        return lcv