############################################
# This script generates a random Toeplitz matrix and a random vector, and multiplies them together.
# The matrix and vector can be binary or not.
# The script is used to generate test cases for the hardware implementation of the Toeplitz matrix-vector multiplication.
# The generated test cases are saved in a file.
# The script is used as follows:
#
#
# usage: testcase-generator.py [-h] [-n CASES] [-j MINDIM] [-k MAXDIM] [-o OUTPUT] [-b BINARY] [-v]
# Generate test cases for the Toeplitz matrix-vector multiplication 
# options:
#  -h, --help            show this help message and exit
#  -n CASES, --cases CASES
#                        Number of test cases to generate
#  -j MINDIM, --mindim MINDIM
#                        Minimum dimension of the Toeplitz matrix
#  -k MAXDIM, --maxdim MAXDIM
#                        Maximum dimension of the Toeplitz matrix
#  -o OUTPUT, --output OUTPUT
#                        Name of the file to save the test cases
#  -b BINARY, --binary BINARY
#                        Boolean indicating if the matrix and vector should be binary
#  -v, --verbose         Boolean indicating if the matrices should be printed
############################################


import numpy as np
import sys
import argparse
 
"""The following function generates a random n x m binary Toeplitz matrix. Takes the parameter 'binary' to indicate if the vector should be binary."""
def generate_toeplitz(n, m, binary=True):
    # Generate a random vector of length n
    vector = np.random.randint(2, size=n) if binary else np.random.randint(100, size=n)
    
    # Generate the Toeplitz matrix
    matrix = np.zeros((n, m))
    for i in range(n):
        for j in range(m):
            if i > j:
                matrix[i][j] = vector[i-j]
            else:
                matrix[i][j] = vector[j-i]
    return matrix


"""Generate a random n x 1 vector. Takes the parameter binary to indicate if the vector should be binary"""
def generate_vector(n, binary=True):
    if binary:
        return np.random.randint(0, 2, n)
    else:
        return np.random.randint(0, 255, n)




"""Helper function to pretty print matrices"""
def print_matrix(matrix):
    for row in matrix:
        print(row)
    print("\n")

"""Multiply a matrix by a vector. Takes the parameter binary to indicate if the matrix and vector are binary"""
def multiply_matrix_vector(matrix, vector, binary=True):
    if not binary:
        return np.matmul(matrix, vector)
    else:
        return np.matmul(matrix, vector) % 2
    

def verbose_print(verbose, data):
    if verbose:
        print(data)


if __name__ == "__main__":
    # Parse the command line arguments
    parser = argparse.ArgumentParser(prog='testcase-generator.py', description='Generate test cases for the Toeplitz matrix-vector multiplication', epilog='Enjoy the program! :)')
    parser.add_argument('-n', '--cases', type=int, default=1, help='Number of test cases to generate')
    parser.add_argument('-j', '--mindim', type=int, default=2, help='Minimum dimension of the Toeplitz matrix')
    parser.add_argument('-k', '--maxdim', type=int, default=3, help='Maximum dimension of the Toeplitz matrix')
    parser.add_argument('-o', '--output', type=str, default="testcases.txt", help='Name of the file to save the test cases')
    parser.add_argument('-b', '--binary', type=int, default=False, help='Boolean indicating if the matrix and vector should be binary')
    parser.add_argument('-v', '--verbose', action='store_true', default=False, help='Boolean indicating if the matrices should be printed')
    args = parser.parse_args()

    with open(args.output, "w") as f:
        for i in range(args.cases):
            # Generate a random Toeplitz matrix
            n = np.random.randint(args.mindim, args.maxdim)
            m = np.random.randint(args.mindim, args.maxdim)
            matrix = generate_toeplitz(n, m, args.binary)
            # Generate a random vector
            vector = generate_vector(n, args.binary)
            # Multiply the matrix by the vector
            result = multiply_matrix_vector(matrix, vector, args.binary)
            # Write the test case to the file
            np.savetxt(f, vector, fmt="%d")
            f.write("\n")
            np.savetxt(f, matrix, fmt="%d")
            f.write("\n")
            np.savetxt(f, result, fmt="%d")
            f.write("\n")
            f.write("---\n")

            # Print the test case if verbose is set
            if args.verbose:
                print_matrix(vector)
                print_matrix(matrix)
                print_matrix(result)
                print("---\n")