import sympy as sp
import time
import datetime
import random
import readline

errs = (ValueError, TypeError)

def genmatrix(rowcol, intmax, dif) :
    # generate a random matrices
    a = sp.randMatrix(rowcol, rowcol, -intmax, intmax)
    # if determinant non-zero, radnom value less than difficulty, set a to its inverse
    if a.det() != 0 and random.random() <= dif :
        a = a.inv()
    return a

# in function to take arbitrary matrix input from user
def inmat(rowcol) :
    # initialise list of lists
    mat=[[]]*rowcol
    # take input from user into list
    print("Enter line by line with entries separated by space:")
    for i in range(rowcol):
        mat[i]=input().split(" ")
        for j in range(len(mat[i])) :
            # convert list using nsimplify in order to take rational number symbolically
            try :
                mat[i][j] = sp.nsimplify(mat[i][j])
            except errs : return False
    try : mat = sp.Matrix(mat)
    except errs : return False
    return mat

# multcheck function to check if two matrices are multiplied together
def multcheck(a, b, rowcol, intmax) :
    sp.pprint(b)
    print("Multiply these two matrices together")
    # return bool based on if input equals the two matrices multiplied together
    return (inmat(rowcol) == a*b)

# detcheck function to check if the determinant of the matrix is correct
def detcheck(a, b, rowcol, intmax):
    det = sp.det(a)
    # return bool based on if input equals determinant
    try : d = sp.nsimplify(input("What is the determinant of this matrix?: "))
    except errs : return False
    return (det == d)

# invcheck function to check if the inverse of the matrix is correct
def invcheck(a, b, rowcol, intmax):
    # return bool based on if input equals inverse of matrix
    print("What is the inverse of this matrix?")
    return (inmat(rowcol) == a.inv())

# eigcheck function to check if the eigenvalues are correct
def eigcheck(a, b, rowcol, intmax):
    eigs = a.eigenvals()
    for i in range(len(eigs)) :
        try :
            val = sp.nsimplify(input("Input eigenvalue: "))
            algm = sp.nsimplify(input("Input its algebraic multiplicity: "))
        except errs : return False
        if not (val in eigs and eigs[val] == algm) :
            # return false if the eigenvalue not in dictionary and wrong alg multiplicity
            return False
    return True

def diagcheck(a, b, rowcol, intmax):
    return (a.diagonalize()[1] == inmat(rowcol))

# pascalmat function to generate pascal matrices (used for generating unimodular matrices)
def pascalmat(n) :
    mat = []
    # generate rows of binomial coefficients
    for i in range(n) :
        mat.append(sp.binomial_coefficients_list(i) + [0]*(n-1-i))
    mat = sp.Matrix(mat)
    # make a lower triangular pascal matrix
    lmat = mat.transpose()
    rand = random.random()
    # randomly pick either lower or upper triangular pascal matrix
    if rand <= 1/2 :
        mat = lmat
    return mat

def triangcheck(a, b, rowcol, intmax) :
    print("Input triangularised matrix,")
    T = inmat(rowcol)
    print("Input basis change matrix (from standard basis)")
    P_inv = inmat(rowcol)
    # check if P_inv is invertible
    if P_inv.det() == 0 : return False
    # return false if T is not upper triangular and not similar to the original matrix
    return (T.is_upper and a == P_inv.inv()*T*P_inv)

# practice function
def practice(t) :
    count = 0

    # choose the function of the program that you want
    if t == "mult" :
        f = multcheck
    elif t == "det" :
        f = detcheck
    elif t == "inv" :
        f = invcheck
    elif t == "eig" :
        f = eigcheck
    elif t == "diag" :
        f = diagcheck
    elif t == "triang" :
        f = triangcheck
    else :
        exit()

    while True :
        try :
            rowcol = abs(int(input("What size of matrix do you want to practice with? ")))
            intmax = abs(int(input("What maximum size of integer do you want the matrix to be made out of? ")))
            dif = float(input("What difficulty (probability for a matrix of rational values between 0, 1) do you want? "))
            break
        except errs:
            continue

    # initialise time measurement
    tic = time.perf_counter()

    # infinite loop of practice
    while True :
        a = genmatrix(rowcol, intmax, dif)
        b = genmatrix(rowcol, intmax, dif)
        if t == "triang" :
            while True :
                # generate matrices with smaller integers
                a = genmatrix(rowcol, intmax//rowcol, dif)
                # ensure a is invertible
                if a.det() == 0 : continue
                # find a non-fractional upper trianagular matrix (ensures triangularisability)
                c = a.LUdecompositionFF()[3]
                croots = {k: k for k in sp.roots(c.charpoly()) if k not in sp.roots(c.charpoly(), filter='Z')}
                # make sure only integer roots and c is not diagonalisable
                if not (len(croots) != 0 or c.is_diagonalizable()) :
                    a = c
                    pmat = pascalmat(rowcol)
                    # multiply by a pascal matrix to get a similar matrix with nice numbers
                    a = pmat.inv()*a*pmat
                    if a.is_upper: continue
                    break

        # infinite loop until user succeeds 
        while True :
            # if diagcheck, make sure the matrix is diagonalizable
            if t == "diag" :
                while not a.is_diagonalizable :
                    a = genmatrix(rowcol, intmax, dif)
            if t == "inv" :
                while a.det() == 0 :
                    a = genmatrix(rowcol, intmax, dif)
            sp.pprint(a)
            # if return of function is True
            if f(a, b, rowcol, intmax) :
                # take time interval
                toc = time.perf_counter()
                tictoc=toc-tic
                # increment count
                count += 1
                # print success message
                print("Correct! You have solved", count, "problems in", str(datetime.timedelta(seconds=round(tictoc))))
                break
            else :
                # print failure message
                print("Try again")
                continue

# take input from user on the type of practice they want to do
t = input("What do you want to practice? (mult, det, inv, eig, diag, triang) ")
practice(t)