"""
Tests for functions in utv.py and qr.py.
to run all tests from the command line: $ python utv_tests.py

To run the tests instantiated in 'def suite()' near the end of this file
from a Colab or Jupyter Notebook:
    *in the notebook*
    import utv_tests
    import unittest

    def run_tests():
        suite = utv_tests.suite()
        runner = unittest.TextTestRunner()
        runner.run(suite)
    run_tests()

To add the tests in a subclass to 'def suite()', so that they will also
be run from the Colab or Jupyter notebook:
    *in this file*
    class NewTestSubClass(unittest.TestCase):

    def suite():
        suite = unittest.TestSuite()
        *one of the below for each subclass*
        suite.addTests(unittest.makeSuite(TestSubclass, 'test'))

A base class GaussianMatrixTest(unittest.TestCase) is provided to
loop over Gaussian matrices, treating each of several permutations of shape and
dtype as a different subtest. See that class' docstring for details.

Tests are currently defined that mean to do each of the following:
    TODO: FILL THESE IN
    test_replace_diagonal: Checks that matutils.replace_diagonal(A, D)
                           correctly returns a matrix identical to A, with
                           first D.size entries replaced by those of D, and
                           the rest zeroed out.
"""

import jax.numpy as jnp
import jax
import numpy as np
import unittest
import itertools
import math
import dfact.utv as utv
import dfact.qr as qr
import dfact.matutils as matutils
from dfact.matutils import dag

###############################################################################
# BASE CLASSES AND UTILITIES
###############################################################################
def errstring(arr1, name1, arr2, name2):
    """
    Utility function to generate failure messages in a few tests. Returns
    the Frobenius norm of the difference between arr1 and arr2, and a
    string errstring to be output if desired.
    """
    error = matutils.frob(arr1, arr2)
    errormsg = "Error = " + str(error) + "\n"
    errormsg += name1 + ":\n " + str(arr1) + "\n"
    errormsg += name2 + ":\n " + str(arr2)
    return (error, errormsg)

def errorstack(errtups, passed=True, msg="\n", thresh=1E-6):
    """
    Combines the output from multiples calls to errstring into a single
    pass-or-fail condition, based on comparison of each errtups[i][0] to
    thresh. Returns this flag along with a single error message, concatenated
    from
    those of each individual call.
    """
    for error, errormsg in errtups:
        if error < thresh:
            passed = False
            msg += "**********************************************************"
            msg += errormsg
            msg += "**********************************************************"
    return (passed, msg)

def manyclose(pairs):
    """
    Loops over pairs of numpy arrays and runs allclose on each. Returns True
    if allclose passes for each pair and False otherwise.
    """
    passed = True
    for pair in pairs:
        passed = passed and jnp.allclose(pair[0], pair[1])
    return passed

class GaussianMatrixTest(unittest.TestCase):
    """
    Concrete base class providing a set of random m x n matrices of differing
    data types to operate on.

    The parameter lists (ms, ns, dtypes) defined in the body of __init__
    will
    be used to generate a shape (m, n), dtype=dtype matrix, as if from
    nested for loops. Each will define a new subtest. The parameters to be
    looped over are fixed for any particular subclass of GaussianMatrixTest.
    Each subclass defines its choice of parameters by specializing
    __init__.

    Subclasses of GaussianMatrixTest should thus employ the following maneuvers
    to define tests:

    class MySubClass(GaussianMatrixTest):
        def __init__(self, *args, **kwargs):
            *specialize this class only if you want to change the default*
            *parameters*
            self.ns = (*start, stop, step*)
            self.ms = (*start, stop, step*)
            self.dtypes = [jnp.float32, jnp.complex64...]
            super().__init__(*args, **kwargs)

        def test_something(self, **kwargs):
            def impl(A, paramtup):
                m, n, dtype = paramtup
                ***body of test acting on the random input matrix A***
            self.iterloop(impl)

    *in def suite(): defined near the end of this file)*
    def suite():
        suite.addTests(unittest.makeSuite(MySubClass, 'test')

    *in a colab notebook from which you wish to run tests*

    """
    def __init__(self, *args, ns=(1, 6, 2),
                 ms=(6, 18, 6),
                 dtypes=[jnp.float32, jnp.complex64],
                 **kwargs):
        self.ns = range(*ns)
        self.ms = range(*ms)
        self.dtypes = dtypes
        super().__init__(*args, **kwargs)

    def setUp(self):
        self.matrices = [matutils.gaussian_random(shape=(m, n), dtype=dtype)
                         for m, n, dtype
                         in itertools.product(self.ms, self.ns, self.dtypes)]

    def iterloop(self, func):
        """
        Iterates over the parameters defined in setUpImpl, stores them in
        paramtup, generates a random
        matrix A for each, and calls func(A, paramtup) as a new subtest.
        func(A, paramtup) should thus store the body of the test, and is
        usually defined within class methods as 'impl'.
        iterloop would most naturally be a decorator, but my Python
        isn't up to getting the interpreter to treat a class method as such.
        """
        params = itertools.product(self.ms, self.ns, self.dtypes)
        for A, paramtup in zip(self.matrices, params):
            m, n, dtype = paramtup
            with self.subTest(m=m, n=n, dtype=dtype):
                func(A, paramtup)

###############################################################################
# QR DECOMPOSITION TESTS
###############################################################################
class ExplicitQRTests(unittest.TestCase):
    """
    These tests check whether our QR code functions correctly, by creating
    a matrix with explicitly known input (self.testA) and ensuring we
    retrieve explicitly known results.
    """
    def setUp(self):
        self.testA = jnp.array([[1., -4.],
                                [2., 3.],
                                [2., 2.]])

    def test_householder_generation_on_explicit_input(self):
        """
        Checks that qr.house gives the correct output for known input.
        With x = [1, 2, 2] we should have v = [1, -1, -1]^T and
        beta = 2/3.

        ***THIS TEST SEEMS TO BE WRONG***
        """
        test_me = self.testA[:, 0]
        v, beta = qr.house(test_me)
        success = jnp.allclose(jnp.array([beta]), jnp.array([2./3.]))
        self.assertTrue(success, msg="beta[0] = "+str(beta)+" was wrong.")
        success = jnp.allclose(v, jnp.array([1., -1., -1.]))
        self.assertTrue(success, msg="v[0] = "+str(v)+" was wrong.")

    def test_factored_qr_on_explicit_input(self):
        """
        Checks that qr.house_qr(mode="factored") gives the correct output.

        ***THIS TEST SEEMS TO BE WRONG***
        """
        test_me = self.testA
        # print("\n")
        # print("*****************************")
        H, beta = qr.house_qr(test_me, mode="factored")
        print("H:\n ", H)
        print("beta: ", beta)
        Hnp, betanp = np.linalg.qr(test_me, mode="raw")
        print("Hnp:\n ", Hnp)
        print("betanp: ", betanp)
        correct_beta = jnp.array([2./3., 8./5.])
        # print("beta:", beta)
        # print("1/beta:", 1/beta)
        # print("correct_beta:", correct_beta)
        # print("*****************************")
        self.assertTrue(jnp.allclose(beta, correct_beta),
                        msg="beta="+str(beta)+" was wrong.")
        correct_H = jnp.array([[3., 2.],
                               [-1., 5.],
                               [-1., 0.5]])
        self.assertTrue(jnp.allclose(H, correct_H),
                        msg="\nH=\n"+str(H)+" was wrong.")

        correct_Q = jnp.array([[1./3, -14./15, -2./15],
                               [2./3, 1./3, -2./3],
                               [2./3, 2./15, 11./15]
                               ])
        Q, R = qr.factored_to_QR(H, beta)
        self.assertTrue(jnp.allclose(Q, correct_Q),
                        msg="\nQ=\n"+str(Q)+" was wrong.")

        QR = jnp.dot(Q, R)
        self.assertTrue(jnp.allclose(QR, test_me),
                        msg="\nQR=\n"+str(QR)+" was wrong.")

class TestHouseholderVectorProperties(GaussianMatrixTest):
    """
    Tests the code to compute and apply Householder reflections.
    """
    def __init__(self, *args, **kwargs):
        ns = (1, 2, 1)
        ms = (1, 5, 1)
        super().__init__(*args, ns=ns, ms=ms, **kwargs)

    def test_householder_unitarity(self, thresh=1E-6):
        """
        Random (m,) vectors are generated, and Householder reflections
        (v, beta) computed from them. The dense matrix
        P = I_m - beta v otimes dag(v) is formed, and its unitarity
        (orthogonality) is confirmed.
        """
        def impl(A, paramtup):
            v, beta = qr.house(A)
            P = qr.form_dense_P([v, beta])
            Pd = dag(P)
            unitary1 = jnp.dot(P, Pd)
            unitary2 = jnp.dot(Pd, P)
            Id = jnp.eye(v.size, dtype=P.dtype)
            err1, errormsg1 = errstring(unitary1, "P Pd", Id, "I")
            err2, errormsg2 = errstring(unitary2, "Pd P", Id, "I")
            errormsg = ""
            passed = True
            if err1 > thresh:
                passed = False
                errormsg += "\n" + errormsg1
            if err2 > thresh:
                passed = False
                errormsg += "\n" + errormsg2
            self.assertTrue(passed, msg=errormsg)
        self.iterloop(impl)

    # def test_householder(self, thresh=1E-6):
        # """
        # Random (m,) vectors are generated, and Householder reflections
        # (v, beta) computed from them. The dense matrix
        # P = I_m - beta v otimes dag(v) is formed, and its unitarity
        # (orthogonality) is confirmed.
        # """
        # def impl(A, paramtup):
            # v, beta = qr.house(A)

            # x0 = A.ravel()[0]
            # r = jnp.abs(x0)
            # theta = jnp.angle(x0)
            # vp = x +

            # P = qr.form_dense_P([v, beta])
            # Pd = dag(P)
            # unitary1 = jnp.dot(P, Pd)
            # unitary2 = jnp.dot(Pd, P)
            # Id = jnp.eye(v.size, dtype=P.dtype)
            # err1, errormsg1 = errstring(unitary1, "P Pd", Id, "I")
            # err2, errormsg2 = errstring(unitary2, "Pd P", Id, "I")
            # errormsg = ""
            # passed = True
            # if err1 > thresh:
                # passed = False
                # errormsg += "\n" + errormsg1
            # if err2 > thresh:
                # passed = False
                # errormsg += "\n" + errormsg2
            # self.assertTrue(passed, msg=errormsg)
        # self.iterloop(impl)

class TestComputeAndApplyHouseholderReflectors(GaussianMatrixTest):
    """
    Tests the code to compute and apply Householder reflections.
    """
    def __init__(self, *args, **kwargs):
        ns = (1, 6, 1)
        ms = (1, 6, 1)
        super().__init__(*args, ns=ns, ms=ms, **kwargs)

    def test_house_leftmult(self, thresh=1E-6):
        """
        Random (m,n) matrices A are generated, along with length-m vectors x.
        Householder reflections
        (v, beta) are computed from each x. The dense matrix
        P = I_m - beta v otimes dag(v) is formed. It is confirmed
        that P A and house_leftmult(A, v, beta) yield the same result.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            x = matutils.gaussian_random(shape=(m,), dtype=dtype)
            v, beta = qr.house(x)
            PA_h = qr.house_leftmult(A, v, beta)
            P = qr.form_dense_P([v, beta])
            PA = jnp.dot(P, A)
            err, errmsg = errstring(PA, "PA", PA_h, "PA_h")
            self.assertTrue(err < thresh, msg=errmsg)
        self.iterloop(impl)

    def test_house_rightmult(self, thresh=1E-6):
        """
        Random (m,n) matrices A are generated, along with length-m vectors x.
        Householder reflections
        (v, beta) are computed from each x. The dense matrix
        P = I_n - beta v otimes dag(v) is formed. It is confirmed
        that A P and house_rightmult(A, v, beta) yield the same result.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            x = matutils.gaussian_random(shape=(n,), dtype=dtype)
            v, beta = qr.house(x)
            AP_h = qr.house_rightmult(A, v, beta)
            P = qr.form_dense_P([v, beta])
            AP = jnp.dot(A, P)
            err, errmsg = errstring(AP, "AP", AP_h, "AP_h")
            self.assertTrue(err < thresh, msg=errmsg)
        self.iterloop(impl)

class GaussianSVDTests(GaussianMatrixTest):
    """
    Tests of the randSVD decomposition that loop over Gaussian random matrices.
    """
    def __init__(self, *args, **kwargs):
        ns = (1, 6, 1)
        ms = (1, 6, 1)
        dtypes = [jnp.float32]  # , jnp.complex64]
        super().__init__(*args, ns=ns, ms=ms, dtypes=dtypes, **kwargs)

    def test_blockpowerSVD_svs(self, thresh=1E-5):
        """
        Checks that block power SVD gets the right singular values.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            U, S, Vh = jnp.linalg.svd(A)
            for sigma in range(1, n-1):
                with self.subTest(sigma=sigma):
                    Uc, Sc, Vc  = qr.block_power_svd(A, sigma, tol=thresh)
                    error, errormsg = errstring(S[:sigma],
                                                "NP SVD", jnp.abs(Sc),
                                                "Chase SVD")
                    self.assertTrue(error < thresh, msg=errormsg)
        self.iterloop(impl)

class GaussianQRTests(GaussianMatrixTest):
    """
    These tests check whether the QR decomposition routines function correctly,
    by generating Gaussian random input and ensuring results meet various
    conditions.

    """
    def __init__(self, *args, **kwargs):
        ns = (1, 6, 1)
        ms = (1, 6, 1)
        dtypes = [jnp.float32]  # , jnp.complex64]
        super().__init__(*args, ns=ns, ms=ms, dtypes=dtypes, **kwargs)

    def test_recursive_QR(self, thresh=1E-6):
        """
        Checks that recursive_QR correctly reconstructs A for various matrix
        and block sizes.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            # Qtest, Rtest = jnp.linalg.QR(A, mode="reduced")
            for block_size in range(1, n):
                with self.subTest(block_size=block_size):
                    Qblock, Rblock = qr.recursiveQR(A, block_size)
                    #  print("***********DONE!********")
                    #  print("A:", A.shape)
                    #  print("Q:", Qblock.shape)
                    #  print("R:", Rblock.shape)
                    A_recon = Qblock @ Rblock
                    #print("QR:", A_recon.shape)
                    error, errormsg = errstring(A, "A", A_recon, "QR A")
                    self.assertTrue(error < thresh, msg=errormsg)
        self.iterloop(impl)

    #  def test_forward_vs_backward_accumulation(self, thresh=1E-6):
    #      """
    #      Checks that Q computed from the factored representation gives the
    #      same result when using either forward or backward accumulation.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  H, betas = qr.house_qr(A, mode="factored")
    #              return
    #          H, betas = qr.house_qr(A, mode="factored")
    #          Im = jnp.eye(m, dtype=A.dtype)
    #          Qforward = qr.factored_rightmult(Im, H, betas)
    #          Qbackward, R = qr.factored_to_QR(H, betas)
    #          err, errmsg = errstring(Qforward, "Qforward", Qbackward,
    #                                  "Qbackward")
    #          self.assertTrue(err < thresh, msg=errmsg)
    #      self.iterloop(impl)

    #  def test_factored_mult(self, thresh=1E-5):
    #      """
    #      A = QR -> [H, tau] is computed. R is extracted. We compare
    #      C * A with C * Q * R without forming Q explicitly.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup

    #          C = matutils.gaussian_random(shape=(n, m), dtype=dtype)
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  H, betas = qr.house_qr(A, mode="factored")
    #              return
    #          H, betas = qr.house_qr(A, mode="factored")
    #          R = jnp.triu(H)

    #          CA = jnp.dot(C, A)
    #          CQ = qr.factored_rightmult(C, H, betas)
    #          CQR = CQ@R
    #          err, errmsg = errstring(CA, "CA", CQR, "CQR")
    #          self.assertTrue(err < thresh, msg=errmsg)
    #      self.iterloop(impl)

    #  def test_factored_to_dense_Q(self, thresh=1E-6):
    #      """
    #      Runs the qr decomposition in 'factored' mode. Factored mode returns
    #      matrices H and tau that record the Householder transformations
    #      from which Q and R are formed.

    #      Specifically, R is the upper triangle of H, the Householder vectors
    #      mapping A to R are the lower triangle, and the normalizations of those
    #      vectors in a certain sense are stored in tau.

    #      This routine explicitly forms
    #      Q from these outputs using qr.factored_to_Q, checks that Q is
    #      unitary, and that QR = A
    #      to within Frobenius norm 'thresh'.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  H, betas = qr.house_qr(A, mode="factored")
    #              return
    #          H, betas = qr.house_qr(A, mode="factored")
    #          jaxQ, jaxR = qr.factored_to_QR(H, betas)
    #          Id = jnp.eye(jaxQ.shape[0], dtype=A.dtype)
    #          errormsg = ""
    #          success = True

    #          unitary_check1 = jnp.dot(jaxQ, dag(jaxQ))
    #          error1, errormsg1 = errstring(unitary_check1, "Q Qdag", Id, "I")
    #          if error1 > thresh:
    #              success = False
    #              errormsg += "Q wasn't unitary. \n" + errormsg1 + "\n"

    #          unitary_check2 = jnp.dot(dag(jaxQ), jaxQ)
    #          error2, errormsg2 = errstring(unitary_check2, "Qdag Q", Id, "I")
    #          if error2 > thresh:
    #              success = False
    #              errormsg += "Q wasn't unitary. \n" + errormsg2 + "\n"

    #          nullopcheck = jnp.dot(jaxQ, jaxR)
    #          error3, errormsg3 = errstring(nullopcheck, "QR", A, "A")
    #          if error3 > thresh:
    #              errormsg += "QR != A. \n" + errormsg3 + "\n"
    #              success = False

    #          self.assertTrue(success, msg=errormsg)
    #      self.iterloop(impl)

    #  def test_WY_Q_properties(self, thresh=1E-6):
    #      """
    #      Runs the qr decomposition in 'WY' mode. WY mode returns
    #      matrices W and Y, storing the same Householder transformations as
    #      'factored' mode in a 'blocked' representation permitting their
    #      application using Level 3 BLAS operations.

    #      This routine explicitly forms
    #      Q from these outputs using qr.factored_to_Q. It checks that Q is
    #      unitary, and that QR = A
    #      to within Frobenius norm 'thresh'.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  H, betas = qr.house_qr(A, mode="WY")
    #              return
    #          W, YH, _ = qr.house_qr(A, mode="WY")
    #          jaxQ = qr.WY_to_Q(W, YH)
    #          jaxQdag = dag(jaxQ)
    #          QQdag = jaxQ @ jaxQdag
    #          Id = jnp.eye(QQdag.shape[0], dtype=QQdag.dtype)
    #          err, errmsg = errstring(QQdag, "Qdag", Id, "I")
    #          self.assertLessEqual(err, thresh, msg=errmsg)

    #      self.iterloop(impl)

    #  def test_WY_reconstruction(self, thresh=1E-6):
    #      """
    #      Runs the qr decomposition in 'WY' mode. WY mode returns
    #      matrices W and Y, storing the same Householder transformations as
    #      'factored' mode in a 'blocked' representation permitting their
    #      application using Level 3 BLAS operations.

    #      This routine explicitly forms
    #      Q from these outputs using qr.factored_to_Q. It checks that Q is
    #      unitary, and that QR = A
    #      to within Frobenius norm 'thresh'.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  H, betas = qr.house_qr(A, mode="WY")
    #              return
    #          W, YH, R = qr.house_qr(A, mode="WY")
    #          Q = qr.WY_to_Q(W, YH)
    #          A_recon = Q @ R
    #          err, errmsg = errstring(A, "A", A_recon, "QR")
    #          self.assertLessEqual(err, thresh, msg=errmsg)

    #      self.iterloop(impl)

    #  def test_WY_to_Q(self, thresh=1E-6):
    #      """
    #      Makes sure that retrieval of Q from WY^H, Q = I - WY^H, works
    #      correctly.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          W = A
    #          YH = matutils.gaussian_random(shape=(n, m), dtype=dtype)
    #          Id = jnp.eye(m, dtype=dtype)
    #          Q = Id - W @ YH
    #          Q2 = qr.WY_to_Q(W, YH)
    #          err, errmsg = errstring(Q, "Q", Q2, "I-WY^H")
    #          self.assertLessEqual(err, thresh, msg=errmsg)

    #  def test_B_times_Q_WY(self, thresh=1E-6):
    #      """
    #      Makes sure that B * Q = B * (I - W Y^H) for Q = I - WY^H, where
    #      the RHS is computed implicitly from W and YH.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          W = A
    #          YH = matutils.gaussian_random(shape=(n, m), dtype=dtype)
    #          B = matutils.gaussian_random(shape=(m, n), dtype=dtype)
    #          Id = jnp.eye(m, dtype=dtype)
    #          Q = Id - W @ YH
    #          BQ = B@Q
    #          BQ_WY = qr.B_times_Q_WY(B, W, YH)
    #          err, errmsg = errstring(BQ, "BQ", BQ_WY, "B(I-WY^T)")
    #          self.assertLessEqual(err, thresh, msg=errmsg)

    #  def test_Qdag_WY_times_B(self, thresh=1E-6):
    #      """
    #      Makes sure that Q^H@B = (I - W Y^H)^H @ B  for Q = I - WY^H, where
    #      the RHS is computed implicitly from W and YH.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          W = A
    #          YH = matutils.gaussian_random(shape=(n, m), dtype=dtype)
    #          B = matutils.gaussian_random(shape=(m, n), dtype=dtype)
    #          Id = jnp.eye(m, dtype=dtype)
    #          Q = Id - W @ YH

    #          QHB = dag(Q)@B
    #          QHB_WY = qr.Qdag_WY_times_B(B, W, YH)
    #          err, errmsg = errstring(QHB, "QHB", QHB_WY, "(I-WY^T)^H @ B")
    #          self.assertLessEqual(err, thresh, msg=errmsg)

class TestRandSVD(GaussianMatrixTest):
    """
    Tests of the randSVD decomposition that loop over Gaussian random matrices.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def test_randSVD(self, thresh=1E-2, k=5):
        """
        Tests the singular values from utv.randSVD against the numpy
        implementation.
        """
        def impl(A, paramtup):
            out_jnp = jnp.linalg.svd(A)
            out_rand = utv.randSVD(A, k=k)
            svd_jnp = out_jnp[1][:k]
            svd_rand = out_rand[1]
            error, errormsg = errstring(svd_jnp, "Numpy SVs", svd_rand,
                                        "randSVs")
            self.assertTrue(error < thresh, msg=errormsg)
        self.iterloop(impl)

    def test_randSVD_reconstruction(self, thresh=1E-3):
        """
        Checks that randSVD correctly reconstructs its input.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            U, S, Vh = utv.randSVD(A)
            A_recon = matutils.trimultdag(U, S, Vh)
            error, errormsg = errstring(A, "Input A", A_recon, "randSVD A")
            self.assertTrue(error < thresh, msg=errormsg)
        self.iterloop(impl)

class TestUTV(GaussianMatrixTest):
    """
    Tests of the UTV decomposition that loop over Gaussian random matrices.
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def setUp(self):
        ns = (3, 4, 1)
        ms = (3, 4, 1)
        dtypes = [jnp.float32]#, jnp.complex64]
        self.ns = range(*ns)
        self.ms = range(*ms)
        self.dtypes = dtypes
        self.matrices = [matutils.gaussian_random(shape=(m, n), dtype=dtype)
                         for m, n, dtype
                         in itertools.product(self.ms, self.ns, self.dtypes)]
        #  self.matrices = [jnp.ones((m,n), dtype=dtype) for m, n, dtype
        #                   in itertools.product(self.ms, self.ns, self.dtypes)]

    ###########################################################################
    # randUTV
    ###########################################################################
    #  def test_randUTV_svs(self, thresh=0.1):
    #      """
    #      Tests the singular values from randUTV against those from
    #      stepUTV_slow, using blocksize = number of columns.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          U, slow_sv, Vh = jnp.linalg.svd(A)
    #          for b in range(1, n+1, 1):
    #              with self.subTest(b=b):
    #                  out_fast = utv.randUTV(A, b=b, q=1)
    #                  fast_sv = jnp.diag(out_fast[1])
    #                  error, errormsg = errstring(slow_sv, "slow UTV SVs",
    #                                              fast_sv, "rand_UTV SVs")
    #                  self.assertTrue(error < thresh, msg=errormsg)
    #      self.iterloop(impl)

    def test_randUTV_reconstruction(self, thresh=1E-5):
        """
        Tests that A can be recovered from randUTV_slow, using various
        blocksizes.
        """
        def impl(A, paramtup):
            m, n, dtype = paramtup
            #for b in range(1, n+1, 1):
            for b in range(1, n+1, 1):
                with self.subTest(b=b):
                    U, T, V = utv.randUTV(A, b=b, q=2, p=5)
                    A_UTV = matutils.trimultmat(U, T, dag(V))
                    error, errormsg = errstring(A, "A", A_UTV,
                                                "UTV A")
                    self.assertTrue(error < thresh, msg=errormsg)
        self.iterloop(impl)
    ###########################################################################
    # stepUTV
    ###########################################################################
    #  def test_stepUTV_slow_svs(self, thresh=1E-5):
    #      """
    #      Tests the singular values from utv.stepUTV_slow against those from a
    #      numpy SVD.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  out_rand = utv.stepUTV_slow(A)
    #              return
    #          out_rand = utv.stepUTV_slow(A)
    #          out_jnp = jnp.linalg.svd(A)

    #          svd_sv = out_jnp[1]
    #          utv_sv = jnp.diag(out_rand[1])
    #          error, errormsg = errstring(svd_sv, "Numpy SVs", utv_sv,
    #                                      "rand_UTV SVs")
    #          self.assertTrue(error < thresh, msg=errormsg)
    #      self.iterloop(impl)

    #  def test_stepUTV_reconstruction(self, thresh=1E-5):
    #      """
    #      Checks that stepUTV correctly reconstructs its input.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  out = utv.stepUTV_slow(A)
    #              return
    #          U, T, V = utv.stepUTV_slow(A)
    #          A_UTV = matutils.trimultmat(U, T, dag(V))
    #          error, errormsg = errstring(A, "Input A", A_UTV, "stepUTV A")
    #          self.assertTrue(error < thresh, msg=errormsg)
    #      self.iterloop(impl)

    #  ###########################################################################
    #  # randUTV_slow
    #  ###########################################################################
    #  def test_randUTVslow_svs(self, thresh=1E-5):
    #      """
    #      Tests the singular values from randUTV against those from
    #      stepUTV_slow, using blocksize = number of columns.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          if n > m:
    #              with self.assertRaises(NotImplementedError):
    #                  _ = utv.stepUTV_slow(A)
    #              return
    #          out_slow = utv.stepUTV_slow(A)
    #          out_fast = utv.randUTV_slow(A, n, 1)
    #          slow_sv = out_slow[1]
    #          fast_sv = out_fast[1]
    #          error, errormsg = errstring(slow_sv, "slow UTV SVs", fast_sv,
    #                                      "rand_UTV SVs")
    #          self.assertTrue(error < thresh, msg=errormsg)
    #      self.iterloop(impl)

    #  def test_randUTVslow_reconstruction(self, thresh=1E-5):
    #      """
    #      Tests that A can be recovered from randUTV_slow, using various
    #      blocksizes.
    #      """
    #      def impl(A, paramtup):
    #          m, n, dtype = paramtup
    #          for b in range(1, n+1, 1):
    #              with self.subTest(b=b):
    #                  if n > m:
    #                      with self.assertRaises(NotImplementedError):
    #                          U, T, V = utv.randUTV_slow(A, b, 1)
    #                      return
    #                  U, T, V = utv.randUTV_slow(A, b, 1)
    #                  A_UTV = matutils.trimultmat(U, T, dag(V))
    #                  error, errormsg = errstring(A, "A", A_UTV,
    #                                              "UTV A")
    #                  Us, Ds, Vhs = jnp.linalg.svd(A)
    #                  # print("U: \n", U)
    #                  # print("U svd: \n", Us)
    #                  # print("V: \n", dag(V))
    #                  # print("V svd: \n", Vhs)
    #                  # print("SVDS:", Ds)
    #                  # print("Error: ", error)
    #                  # print("***")
    #                  self.assertTrue(error < thresh, msg=errormsg)
    #      self.iterloop(impl)

class TestUtils(GaussianMatrixTest):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def test_replace_diagonal(self, thresh=1E-6):
        def impl(A, paramtup):
            m, n, dtype = paramtup
            for i in range(1, n):
                D = matutils.gaussian_random(shape=(i,), dtype=dtype)
                with self.subTest(i=i):
                    result = matutils.replace_diagonal(A, D)
                    err = matutils.frob(jnp.triu(result, k=1),
                                        jnp.triu(A, k=1))
                    err += matutils.frob(jnp.tril(result, k=-1),
                                         jnp.tril(A, k=-1))
                    err += matutils.frob(D, jnp.diag(result)[:i])
                    err += matutils.frob(jnp.diag(result)[i:],
                                         jnp.zeros(jnp.diag(result)[i:].shape))
                    errstr = "\nA: \n " + str(A) + "\nRes: \n " + str(result)
                    errstr += "\nD: \n " + str(D)
                    errstr += "\nErr: " + str(err)
                    self.assertTrue(err < thresh, msg=errstr)
        self.iterloop(impl)

###############################################################################
# Functions to call tests
###############################################################################

def suite():
    suite = unittest.TestSuite()
    #suite.addTests(unittest.makeSuite(TestUtils, 'test'))
    #suite.addTests(unittest.makeSuite(TestUTV, 'test'))
    # suite.addTests(unittest.makeSuite(TestRandSVD, 'test'))
    # suite.addTests(unittest.makeSuite(TestHouseholderVectorProperties, 'test'))
    # suite.addTests(unittest.makeSuite(TestComputeAndApplyHouseholderReflectors,
                                      # 'test'))
    # 'ExplicitQRTests' is commented out because it is wrong.
    # suite.addTests(unittest.makeSuite(ExplicitQRTests, 'test'))
    suite.addTests(unittest.makeSuite(GaussianQRTests, 'test'))
    suite.addTests(unittest.makeSuite(GaussianSVDTests, 'test'))
    return suite

if __name__ == '__main__':
    unittest.main()


import os
import pickle
import numpy as np
from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
import torch
from torch import Tensor

class SimpleDataset(Dataset):
    """SimpleDataset [summary]

    [extended_summary]

    :param path_to_pkl: Path to PKL file with Images
    :type path_to_pkl: str
    :param path_to_labels: path to file with labels
    :type path_to_labels: str
    """
    def __init__(self, path_to_pkl, path_to_labels):
        ## TODO: Add code to read csv and load data.
        ## You should store the data in a field.
        # Eg (on how to read .csv files):
        # with open('path/to/.csv', 'r') as f:
        #   lines = ...
        ## Look up how to read .csv files using Python. This is common for datasets in projects.

        # Load images
        with open(path_to_pkl, 'rb') as file:
            self.images = pickle.load(file)
        # Load labels
        with open(path_to_labels, 'rb') as file:
            self.labels = pickle.load(file)

    def __len__(self):
        """__len__ [summary]

        [extended_summary]
        """
        ## TODO: Returns the length of the dataset.
        return len(self.images)

    def __getitem__(self, index):
        """__getitem__ [summary]

        [extended_summary]

        :param index: [description]
        :type index: [type]
        """
        ## TODO: This returns only ONE sample from the dataset, for a given index.
        ## The returned sample should be a tuple (x, y) where x is your input
        ## vector and y is your label
        ## Before returning your sample, you should check if there is a transform
        ## sepcified, and pply that transform to your sample
        # Eg:
        # if self.transform:
        #   sample = self.transform(sample)
        ## Remember to convert the x and y into torch tensors.

        x = np.array(self.images[index])

        y = np.array(self.labels[index])
        y = torch.tensor(y, dtype=torch.float)
        x = Tensor(x).view(1, 28, 28).float()
        return x , y

def get_data_loaders(path_to_pkl,
                     path_to_labels,
                     train_val_test=[0.8, 0.2, 0.2],
                     batch_size=32):
    """get_data_loaders [summary]

    [extended_summary]

    :param path_to_csv: [description]
    :type path_to_csv: [type]
    :param train_val_test: [description], defaults to [0.8, 0.2, 0.2]
    :type train_val_test: list, optional
    :param batch_size: [description], defaults to 32
    :type batch_size: int, optional
    :return: [description]
    :rtype: [type]
    """
    # First we create the dataset given the path to the .csv file
    dataset = SimpleDataset(path_to_pkl, path_to_labels)

    # Then, we create a list of indices for all samples in the dataset.
    dataset_size = len(dataset)
    indices = list(range(dataset_size))

    ## TODO: Rewrite this section so that the indices for each dataset split
    ## are formed. You can take your code from last time

    ## BEGIN: YOUR CODE
    train_split, validation_split, test_split = train_val_test[0], train_val_test[1], train_val_test[2]
    # Generate split between train and test data
    train_test_split = int(np.floor(train_split * dataset_size))
    # Shuffle indices
    np.random.shuffle(indices)
    # Generate split within train data for train and val data (train_val_split has 80% of train_val_data for training, 20% of train_val_data for val data
    train_val_split = int(np.floor(train_split * train_test_split))
    # Generate list of indices for train and test data
    train_val_indices = indices[:train_test_split]
    train_indices = train_val_indices[:train_val_split]
    val_indices = train_val_indices[train_val_split:]
    test_indices = indices[train_test_split:]
    ## END: YOUR CODE

    # Now, we define samplers for each of the train, val and test data
    train_sampler = SubsetRandomSampler(train_indices)
    train_loader = DataLoader(dataset, batch_size=batch_size, sampler=train_sampler)

    val_sampler = SubsetRandomSampler(val_indices)
    val_loader = DataLoader(dataset, batch_size=batch_size, sampler=val_sampler)

    test_sampler = SubsetRandomSampler(test_indices)
    test_loader = DataLoader(dataset, batch_size=batch_size, sampler=test_sampler)

    return train_loader, val_loader, test_loader

if __name__ == '__main__':
    # Testing purposes
    train_loader, val_loader, test_data = get_data_loaders('data/processedImages.pkl', 'data/processedLabels.pkl')

    for batch_index, (x, y) in enumerate(train_loader):
        print(f"Batch {batch_index}")
        print(f"X: {x.shape}")
        print(f"Y: {y.shape}")

        if batch_index == 1000:
            break


"""Module for compiling codegen output, and wrap the binary for use in
python.

.. note:: To use the autowrap module it must first be imported

   >>> from sympy.utilities.autowrap import autowrap

This module provides a common interface for different external backends, such
as f2py, fwrap, Cython, SWIG(?) etc. (Currently only f2py and Cython are
implemented) The goal is to provide access to compiled binaries of acceptable
performance with a one-button user interface, i.e.

    >>> from sympy.abc import x,y
    >>> expr = ((x - y)**(25)).expand()
    >>> binary_callable = autowrap(expr)
    >>> binary_callable(1, 2)
    -1.0

The callable returned from autowrap() is a binary python function, not a
SymPy object.  If it is desired to use the compiled function in symbolic
expressions, it is better to use binary_function() which returns a SymPy
Function object.  The binary callable is attached as the _imp_ attribute and
invoked when a numerical evaluation is requested with evalf(), or with
lambdify().

    >>> from sympy.utilities.autowrap import binary_function
    >>> f = binary_function('f', expr)
    >>> 2*f(x, y) + y
    y + 2*f(x, y)
    >>> (2*f(x, y) + y).evalf(2, subs={x: 1, y:2})
    0.e-110

The idea is that a SymPy user will primarily be interested in working with
mathematical expressions, and should not have to learn details about wrapping
tools in order to evaluate expressions numerically, even if they are
computationally expensive.

When is this useful?

    1) For computations on large arrays, Python iterations may be too slow,
       and depending on the mathematical expression, it may be difficult to
       exploit the advanced index operations provided by NumPy.

    2) For *really* long expressions that will be called repeatedly, the
       compiled binary should be significantly faster than SymPy's .evalf()

    3) If you are generating code with the codegen utility in order to use
       it in another project, the automatic python wrappers let you test the
       binaries immediately from within SymPy.

    4) To create customized ufuncs for use with numpy arrays.
       See *ufuncify*.

When is this module NOT the best approach?

    1) If you are really concerned about speed or memory optimizations,
       you will probably get better results by working directly with the
       wrapper tools and the low level code.  However, the files generated
       by this utility may provide a useful starting point and reference
       code. Temporary files will be left intact if you supply the keyword
       tempdir="path/to/files/".

    2) If the array computation can be handled easily by numpy, and you
       don't need the binaries for another project.

"""

from __future__ import print_function, division

import sys
import os
import shutil
import tempfile
from subprocess import STDOUT, CalledProcessError, check_output
from string import Template
from warnings import warn

from sympy.core.cache import cacheit
from sympy.core.compatibility import range, iterable
from sympy.core.function import Lambda
from sympy.core.relational import Eq
from sympy.core.symbol import Dummy, Symbol
from sympy.tensor.indexed import Idx, IndexedBase
from sympy.utilities.codegen import (make_routine, get_code_generator,
                                     OutputArgument, InOutArgument,
                                     InputArgument, CodeGenArgumentListError,
                                     Result, ResultBase, C99CodeGen)
from sympy.utilities.lambdify import implemented_function
from sympy.utilities.decorator import doctest_depends_on

_doctest_depends_on = {'exe': ('f2py', 'gfortran', 'gcc'),
                       'modules': ('numpy',)}

class CodeWrapError(Exception):
    pass

class CodeWrapper(object):
    """Base Class for code wrappers"""
    _filename = "wrapped_code"
    _module_basename = "wrapper_module"
    _module_counter = 0

    @property
    def filename(self):
        return "%s_%s" % (self._filename, CodeWrapper._module_counter)

    @property
    def module_name(self):
        return "%s_%s" % (self._module_basename, CodeWrapper._module_counter)

    def __init__(self, generator, filepath=None, flags=[], verbose=False):
        """
        generator -- the code generator to use
        """
        self.generator = generator
        self.filepath = filepath
        self.flags = flags
        self.quiet = not verbose

    @property
    def include_header(self):
        return bool(self.filepath)

    @property
    def include_empty(self):
        return bool(self.filepath)

    def _generate_code(self, main_routine, routines):
        routines.append(main_routine)
        self.generator.write(
            routines, self.filename, True, self.include_header,
            self.include_empty)

    def wrap_code(self, routine, helpers=[]):

        workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
        if not os.access(workdir, os.F_OK):
            os.mkdir(workdir)
        oldwork = os.getcwd()
        os.chdir(workdir)
        try:
            sys.path.append(workdir)
            self._generate_code(routine, helpers)
            self._prepare_files(routine)
            self._process_files(routine)
            mod = __import__(self.module_name)
        finally:
            sys.path.remove(workdir)
            CodeWrapper._module_counter += 1
            os.chdir(oldwork)
            if not self.filepath:
                try:
                    shutil.rmtree(workdir)
                except OSError:
                    # Could be some issues on Windows
                    pass

        return self._get_wrapped_function(mod, routine.name)

    def _process_files(self, routine):
        command = self.command
        command.extend(self.flags)
        try:
            retoutput = check_output(command, stderr=STDOUT)
        except CalledProcessError as e:
            raise CodeWrapError(
                "Error while executing command: %s. Command output is:\n%s" % (
                    " ".join(command), e.output.decode()))
        if not self.quiet:
            print(retoutput)

class DummyWrapper(CodeWrapper):
    """Class used for testing independent of backends """

    template = """# dummy module for testing of SymPy
def %(name)s():
    return "%(expr)s"
%(name)s.args = "%(args)s"
%(name)s.returns = "%(retvals)s"
"""

    def _prepare_files(self, routine):
        return

    def _generate_code(self, routine, helpers):
        with open('%s.py' % self.module_name, 'w') as f:
            printed = ", ".join(
                [str(res.expr) for res in routine.result_variables])
            # convert OutputArguments to return value like f2py
            args = filter(lambda x: not isinstance(
                x, OutputArgument), routine.arguments)
            retvals = []
            for val in routine.result_variables:
                if isinstance(val, Result):
                    retvals.append('nameless')
                else:
                    retvals.append(val.result_var)

            print(DummyWrapper.template % {
                'name': routine.name,
                'expr': printed,
                'args': ", ".join([str(a.name) for a in args]),
                'retvals': ", ".join([str(val) for val in retvals])
            }, end="", file=f)

    def _process_files(self, routine):
        return

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

class CythonCodeWrapper(CodeWrapper):
    """Wrapper that uses Cython"""

    setup_template = """\
try:
    from setuptools import setup
    from setuptools import Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension
from Cython.Build import cythonize
cy_opts = {cythonize_options}
{np_import}
ext_mods = [Extension(
    {ext_args},
    include_dirs={include_dirs},
    library_dirs={library_dirs},
    libraries={libraries},
    extra_compile_args={extra_compile_args},
    extra_link_args={extra_link_args}
)]
setup(ext_modules=cythonize(ext_mods, **cy_opts))
"""

    pyx_imports = (
        "import numpy as np\n"
        "cimport numpy as np\n\n")

    pyx_header = (
        "cdef extern from '{header_file}.h':\n"
        "    {prototype}\n\n")

    pyx_func = (
        "def {name}_c({arg_string}):\n"
        "\n"
        "{declarations}"
        "{body}")

    std_compile_flag = '-std=c99'

    def __init__(self, *args, **kwargs):
        """Instantiates a Cython code wrapper.

        The following optional parameters get passed to ``distutils.Extension``
        for building the Python extension module. Read its documentation to
        learn more.

        Parameters
        ==========
        include_dirs : [list of strings]
            A list of directories to search for C/C++ header files (in Unix
            form for portability).
        library_dirs : [list of strings]
            A list of directories to search for C/C++ libraries at link time.
        libraries : [list of strings]
            A list of library names (not filenames or paths) to link against.
        extra_compile_args : [list of strings]
            Any extra platform- and compiler-specific information to use when
            compiling the source files in 'sources'.  For platforms and
            compilers where "command line" makes sense, this is typically a
            list of command-line arguments, but for other platforms it could be
            anything. Note that the attribute ``std_compile_flag`` will be
            appended to this list.
        extra_link_args : [list of strings]
            Any extra platform- and compiler-specific information to use when
            linking object files together to create the extension (or to create
            a new static Python interpreter). Similar interpretation as for
            'extra_compile_args'.
        cythonize_options : [dictionary]
            Keyword arguments passed on to cythonize.

        """

        self._include_dirs = kwargs.pop('include_dirs', [])
        self._library_dirs = kwargs.pop('library_dirs', [])
        self._libraries = kwargs.pop('libraries', [])
        self._extra_compile_args = kwargs.pop('extra_compile_args', [])
        self._extra_compile_args.append(self.std_compile_flag)
        self._extra_link_args = kwargs.pop('extra_link_args', [])
        self._cythonize_options = kwargs.pop('cythonize_options', {})

        self._need_numpy = False

        super(CythonCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        command = [sys.executable, "setup.py", "build_ext", "--inplace"]
        return command

    def _prepare_files(self, routine, build_dir=os.curdir):
        # NOTE : build_dir is used for testing purposes.
        pyxfilename = self.module_name + '.pyx'
        codefilename = "%s.%s" % (self.filename, self.generator.code_extension)

        # pyx
        with open(pyxfilename, 'w') as f:
            self.dump_pyx([routine], f, self.filename)

        # setup.py
        ext_args = [repr(self.module_name), repr([pyxfilename, codefilename])]
        if self._need_numpy:
            np_import = 'import numpy as np\n'
            self._include_dirs.append('np.get_include()')
        else:
            np_import = ''

        with open(os.path.join(build_dir, 'setup.py'), 'w') as f:
            includes = str(self._include_dirs).replace("'np.get_include()'",
                                                       'np.get_include()')
            f.write(self.setup_template.format(
                ext_args=", ".join(ext_args),
                np_import=np_import,
                include_dirs=includes,
                library_dirs=self._library_dirs,
                libraries=self._libraries,
                extra_compile_args=self._extra_compile_args,
                extra_link_args=self._extra_link_args,
                cythonize_options=self._cythonize_options
            ))

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name + '_c')

    def dump_pyx(self, routines, f, prefix):
        """Write a Cython file with python wrappers

        This file contains all the definitions of the routines in c code and
        refers to the header file.

        Arguments
        ---------
        routines
            List of Routine instances
        f
            File-like object to write the file to
        prefix
            The filename prefix, used to refer to the proper header file.
            Only the basename of the prefix is used.
        """
        headers = []
        functions = []
        for routine in routines:
            prototype = self.generator.get_prototype(routine)

            # C Function Header Import
            headers.append(self.pyx_header.format(header_file=prefix,
                                                  prototype=prototype))

            # Partition the C function arguments into categories
            py_rets, py_args, py_loc, py_inf = self._partition_args(routine.arguments)

            # Function prototype
            name = routine.name
            arg_string = ", ".join(self._prototype_arg(arg) for arg in py_args)

            # Local Declarations
            local_decs = []
            for arg, val in py_inf.items():
                proto = self._prototype_arg(arg)
                mat, ind = val
                local_decs.append("    cdef {0} = {1}.shape[{2}]".format(proto, mat, ind))
            local_decs.extend(["    cdef {0}".format(self._declare_arg(a)) for a in py_loc])
            declarations = "\n".join(local_decs)
            if declarations:
                declarations = declarations + "\n"

            # Function Body
            args_c = ", ".join([self._call_arg(a) for a in routine.arguments])
            rets = ", ".join([str(r.name) for r in py_rets])
            if routine.results:
                body = '    return %s(%s)' % (routine.name, args_c)
                if rets:
                    body = body + ', ' + rets
            else:
                body = '    %s(%s)\n' % (routine.name, args_c)
                body = body + '    return ' + rets

            functions.append(self.pyx_func.format(name=name, arg_string=arg_string,
                    declarations=declarations, body=body))

        # Write text to file
        if self._need_numpy:
            # Only import numpy if required
            f.write(self.pyx_imports)
        f.write('\n'.join(headers))
        f.write('\n'.join(functions))

    def _partition_args(self, args):
        """Group function arguments into categories."""
        py_args = []
        py_returns = []
        py_locals = []
        py_inferred = {}
        for arg in args:
            if isinstance(arg, OutputArgument):
                py_returns.append(arg)
                py_locals.append(arg)
            elif isinstance(arg, InOutArgument):
                py_returns.append(arg)
                py_args.append(arg)
            else:
                py_args.append(arg)
        # Find arguments that are array dimensions. These can be inferred
        # locally in the Cython code.
            if isinstance(arg, (InputArgument, InOutArgument)) and arg.dimensions:
                dims = [d[1] + 1 for d in arg.dimensions]
                sym_dims = [(i, d) for (i, d) in enumerate(dims) if
                            isinstance(d, Symbol)]
                for (i, d) in sym_dims:
                    py_inferred[d] = (arg.name, i)
        for arg in args:
            if arg.name in py_inferred:
                py_inferred[arg] = py_inferred.pop(arg.name)
        # Filter inferred arguments from py_args
        py_args = [a for a in py_args if a not in py_inferred]
        return py_returns, py_args, py_locals, py_inferred

    def _prototype_arg(self, arg):
        mat_dec = "np.ndarray[{mtype}, ndim={ndim}] {name}"
        np_types = {'double': 'np.double_t',
                    'int': 'np.int_t'}
        t = arg.get_datatype('c')
        if arg.dimensions:
            self._need_numpy = True
            ndim = len(arg.dimensions)
            mtype = np_types[t]
            return mat_dec.format(mtype=mtype, ndim=ndim, name=arg.name)
        else:
            return "%s %s" % (t, str(arg.name))

    def _declare_arg(self, arg):
        proto = self._prototype_arg(arg)
        if arg.dimensions:
            shape = '(' + ','.join(str(i[1] + 1) for i in arg.dimensions) + ')'
            return proto + " = np.empty({shape})".format(shape=shape)
        else:
            return proto + " = 0"

    def _call_arg(self, arg):
        if arg.dimensions:
            t = arg.get_datatype('c')
            return "<{0}*> {1}.data".format(t, arg.name)
        elif isinstance(arg, ResultBase):
            return "&{0}".format(arg.name)
        else:
            return str(arg.name)

class F2PyCodeWrapper(CodeWrapper):
    """Wrapper that uses f2py"""

    def __init__(self, *args, **kwargs):

        ext_keys = ['include_dirs', 'library_dirs', 'libraries',
                    'extra_compile_args', 'extra_link_args']
        msg = ('The compilation option kwarg {} is not supported with the f2py '
               'backend.')

        for k in ext_keys:
            if k in kwargs.keys():
                warn(msg.format(k))
            kwargs.pop(k, None)

        super(F2PyCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        filename = self.filename + '.' + self.generator.code_extension
        args = ['-c', '-m', self.module_name, filename]
        command = [sys.executable, "-c", "import numpy.f2py as f2py2e;f2py2e.main()"]+args
        return command

    def _prepare_files(self, routine):
        pass

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

# Here we define a lookup of backends -> tuples of languages. For now, each
# tuple is of length 1, but if a backend supports more than one language,
# the most preferable language is listed first.
_lang_lookup = {'CYTHON': ('C99', 'C89', 'C'),
                'F2PY': ('F95',),
                'NUMPY': ('C99', 'C89', 'C'),
                'DUMMY': ('F95',)}     # Dummy here just for testing

def _infer_language(backend):
    """For a given backend, return the top choice of language"""
    langs = _lang_lookup.get(backend.upper(), False)
    if not langs:
        raise ValueError("Unrecognized backend: " + backend)
    return langs[0]

def _validate_backend_language(backend, language):
    """Throws error if backend and language are incompatible"""
    langs = _lang_lookup.get(backend.upper(), False)
    if not langs:
        raise ValueError("Unrecognized backend: " + backend)
    if language.upper() not in langs:
        raise ValueError(("Backend {0} and language {1} are "
                          "incompatible").format(backend, language))

@cacheit
@doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
def autowrap(expr, language=None, backend='f2py', tempdir=None, args=None,
             flags=None, verbose=False, helpers=None, code_gen=None, **kwargs):
    """Generates python callable binaries based on the math expression.

    Parameters
    ----------
    expr
        The SymPy expression that should be wrapped as a binary routine.
    language : string, optional
        If supplied, (options: 'C' or 'F95'), specifies the language of the
        generated code. If ``None`` [default], the language is inferred based
        upon the specified backend.
    backend : string, optional
        Backend used to wrap the generated code. Either 'f2py' [default],
        or 'cython'.
    tempdir : string, optional
        Path to directory for temporary files. If this argument is supplied,
        the generated code and the wrapper input files are left intact in the
        specified path.
    args : iterable, optional
        An ordered iterable of symbols. Specifies the argument sequence for the
        function.
    flags : iterable, optional
        Additional option flags that will be passed to the backend.
    verbose : bool, optional
        If True, autowrap will not mute the command line backends. This can be
        helpful for debugging.
    helpers : iterable, optional
        Used to define auxillary expressions needed for the main expr. If the
        main expression needs to call a specialized function it should be put
        in the ``helpers`` iterable. Autowrap will then make sure that the
        compiled main expression can link to the helper routine. Items should
        be tuples with (<funtion_name>, <sympy_expression>, <arguments>). It
        is mandatory to supply an argument sequence to helper routines.
    code_gen : CodeGen instance
        An instance of a CodeGen subclass. Overrides ``language``.
    include_dirs : [string]
        A list of directories to search for C/C++ header files (in Unix form
        for portability).
    library_dirs : [string]
        A list of directories to search for C/C++ libraries at link time.
    libraries : [string]
        A list of library names (not filenames or paths) to link against.
    extra_compile_args : [string]
        Any extra platform- and compiler-specific information to use when
        compiling the source files in 'sources'.  For platforms and compilers
        where "command line" makes sense, this is typically a list of
        command-line arguments, but for other platforms it could be anything.
    extra_link_args : [string]
        Any extra platform- and compiler-specific information to use when
        linking object files together to create the extension (or to create a
        new static Python interpreter).  Similar interpretation as for
        'extra_compile_args'.

    Examples
    --------
    >>> from sympy.abc import x, y, z
    >>> from sympy.utilities.autowrap import autowrap
    >>> expr = ((x - y + z)**(13)).expand()
    >>> binary_func = autowrap(expr)
    >>> binary_func(1, 4, 2)
    -1.0
    """
    if language:
        if not isinstance(language, type):
            _validate_backend_language(backend, language)
    else:
        language = _infer_language(backend)

    helpers = [helpers] if helpers else ()
    args = list(args) if iterable(args, exclude=set) else args

    if code_gen is None:
        code_gen = get_code_generator(language, "autowrap")

    CodeWrapperClass = {
        'F2PY': F2PyCodeWrapper,
        'CYTHON': CythonCodeWrapper,
        'DUMMY': DummyWrapper
    }[backend.upper()]
    code_wrapper = CodeWrapperClass(code_gen, tempdir, flags if flags else (),
                                    verbose, **kwargs)

    helps = []
    for name_h, expr_h, args_h in helpers:
        helps.append(make_routine(name_h, expr_h, args_h))

    for name_h, expr_h, args_h in helpers:
        if expr.has(expr_h):
            name_h = binary_function(name_h, expr_h, backend='dummy')
            expr = expr.subs(expr_h, name_h(*args_h))
    try:
        routine = make_routine('autofunc', expr, args)
    except CodeGenArgumentListError as e:
        # if all missing arguments are for pure output, we simply attach them
        # at the end and try again, because the wrappers will silently convert
        # them to return values anyway.
        new_args = []
        for missing in e.missing_args:
            if not isinstance(missing, OutputArgument):
                raise
            new_args.append(missing.name)
        routine = make_routine('autofunc', expr, args + new_args)

    return code_wrapper.wrap_code(routine, helpers=helps)

@doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
def binary_function(symfunc, expr, **kwargs):
    """Returns a sympy function with expr as binary implementation

    This is a convenience function that automates the steps needed to
    autowrap the SymPy expression and attaching it to a Function object
    with implemented_function().

    Parameters
    ----------
    symfunc : sympy Function
        The function to bind the callable to.
    expr : sympy Expression
        The expression used to generate the function.
    kwargs : dict
        Any kwargs accepted by autowrap.

    Examples
    --------
    >>> from sympy.abc import x, y
    >>> from sympy.utilities.autowrap import binary_function
    >>> expr = ((x - y)**(25)).expand()
    >>> f = binary_function('f', expr)
    >>> type(f)
    <class 'sympy.core.function.UndefinedFunction'>
    >>> 2*f(x, y)
    2*f(x, y)
    >>> f(x, y).evalf(2, subs={x: 1, y: 2})
    -1.0
    """
    binary = autowrap(expr, **kwargs)
    return implemented_function(symfunc, binary)

#################################################################
#                           UFUNCIFY                            #
#################################################################

_ufunc_top = Template("""\
#include "Python.h"
#include "math.h"
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/halffloat.h"
#include ${include_file}

static PyMethodDef ${module}Methods[] = {
        {NULL, NULL, 0, NULL}
};""")

_ufunc_outcalls = Template("*((double *)out${outnum}) = ${funcname}(${call_args});")

_ufunc_body = Template("""\
static void ${funcname}_ufunc(char **args, npy_intp *dimensions, npy_intp* steps, void* data)
{
    npy_intp i;
    npy_intp n = dimensions[0];
    ${declare_args}
    ${declare_steps}
    for (i = 0; i < n; i++) {
        ${outcalls}
        ${step_increments}
    }
}
PyUFuncGenericFunction ${funcname}_funcs[1] = {&${funcname}_ufunc};
static char ${funcname}_types[${n_types}] = ${types}
static void *${funcname}_data[1] = {NULL};""")

_ufunc_bottom = Template("""\
#if PY_VERSION_HEX >= 0x03000000
static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "${module}",
    NULL,
    -1,
    ${module}Methods,
    NULL,
    NULL,
    NULL,
    NULL
};

PyMODINIT_FUNC PyInit_${module}(void)
{
    PyObject *m, *d;
    ${function_creation}
    m = PyModule_Create(&moduledef);
    if (!m) {
        return NULL;
    }
    import_array();
    import_umath();
    d = PyModule_GetDict(m);
    ${ufunc_init}
    return m;
}
#else
PyMODINIT_FUNC init${module}(void)
{
    PyObject *m, *d;
    ${function_creation}
    m = Py_InitModule("${module}", ${module}Methods);
    if (m == NULL) {
        return;
    }
    import_array();
    import_umath();
    d = PyModule_GetDict(m);
    ${ufunc_init}
}
#endif\
""")

_ufunc_init_form = Template("""\
ufunc${ind} = PyUFunc_FromFuncAndData(${funcname}_funcs, ${funcname}_data, ${funcname}_types, 1, ${n_in}, ${n_out},
            PyUFunc_None, "${module}", ${docstring}, 0);
    PyDict_SetItemString(d, "${funcname}", ufunc${ind});
    Py_DECREF(ufunc${ind});""")

_ufunc_setup = Template("""\
def configuration(parent_package='', top_path=None):
    import numpy
    from numpy.distutils.misc_util import Configuration

    config = Configuration('',
                           parent_package,
                           top_path)
    config.add_extension('${module}', sources=['${module}.c', '${filename}.c'])

    return config

if __name__ == "__main__":
    from numpy.distutils.core import setup
    setup(configuration=configuration)""")

class UfuncifyCodeWrapper(CodeWrapper):
    """Wrapper for Ufuncify"""

    def __init__(self, *args, **kwargs):

        ext_keys = ['include_dirs', 'library_dirs', 'libraries',
                    'extra_compile_args', 'extra_link_args']
        msg = ('The compilation option kwarg {} is not supported with the numpy'
               ' backend.')

        for k in ext_keys:
            if k in kwargs.keys():
                warn(msg.format(k))
            kwargs.pop(k, None)

        super(UfuncifyCodeWrapper, self).__init__(*args, **kwargs)

    @property
    def command(self):
        command = [sys.executable, "setup.py", "build_ext", "--inplace"]
        return command

    def wrap_code(self, routines, helpers=None):
        # This routine overrides CodeWrapper because we can't assume funcname == routines[0].name
        # Therefore we have to break the CodeWrapper private API.
        # There isn't an obvious way to extend multi-expr support to
        # the other autowrap backends, so we limit this change to ufuncify.
        helpers = helpers if helpers is not None else []
        # We just need a consistent name
        funcname = 'wrapped_' + str(id(routines) + id(helpers))

        workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
        if not os.access(workdir, os.F_OK):
            os.mkdir(workdir)
        oldwork = os.getcwd()
        os.chdir(workdir)
        try:
            sys.path.append(workdir)
            self._generate_code(routines, helpers)
            self._prepare_files(routines, funcname)
            self._process_files(routines)
            mod = __import__(self.module_name)
        finally:
            sys.path.remove(workdir)
            CodeWrapper._module_counter += 1
            os.chdir(oldwork)
            if not self.filepath:
                try:
                    shutil.rmtree(workdir)
                except OSError:
                    # Could be some issues on Windows
                    pass

        return self._get_wrapped_function(mod, funcname)

    def _generate_code(self, main_routines, helper_routines):
        all_routines = main_routines + helper_routines
        self.generator.write(
            all_routines, self.filename, True, self.include_header,
            self.include_empty)

    def _prepare_files(self, routines, funcname):

        # C
        codefilename = self.module_name + '.c'
        with open(codefilename, 'w') as f:
            self.dump_c(routines, f, self.filename, funcname=funcname)

        # setup.py
        with open('setup.py', 'w') as f:
            self.dump_setup(f)

    @classmethod
    def _get_wrapped_function(cls, mod, name):
        return getattr(mod, name)

    def dump_setup(self, f):
        setup = _ufunc_setup.substitute(module=self.module_name,
                                        filename=self.filename)
        f.write(setup)

    def dump_c(self, routines, f, prefix, funcname=None):
        """Write a C file with python wrappers

        This file contains all the definitions of the routines in c code.

        Arguments
        ---------
        routines
            List of Routine instances
        f
            File-like object to write the file to
        prefix
            The filename prefix, used to name the imported module.
        funcname
            Name of the main function to be returned.
        """
        if (funcname is None) and (len(routines) == 1):
            funcname = routines[0].name
        elif funcname is None:
            msg = 'funcname must be specified for multiple output routines'
            raise ValueError(msg)
        functions = []
        function_creation = []
        ufunc_init = []
        module = self.module_name
        include_file = "\"{0}.h\"".format(prefix)
        top = _ufunc_top.substitute(include_file=include_file, module=module)

        name = funcname

        # Partition the C function arguments into categories
        # Here we assume all routines accept the same arguments
        r_index = 0
        py_in, _ = self._partition_args(routines[0].arguments)
        n_in = len(py_in)
        n_out = len(routines)

        # Declare Args
        form = "char *{0}{1} = args[{2}];"
        arg_decs = [form.format('in', i, i) for i in range(n_in)]
        arg_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
        declare_args = '\n    '.join(arg_decs)

        # Declare Steps
        form = "npy_intp {0}{1}_step = steps[{2}];"
        step_decs = [form.format('in', i, i) for i in range(n_in)]
        step_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
        declare_steps = '\n    '.join(step_decs)

        # Call Args
        form = "*(double *)in{0}"
        call_args = ', '.join([form.format(a) for a in range(n_in)])

        # Step Increments
        form = "{0}{1} += {0}{1}_step;"
        step_incs = [form.format('in', i) for i in range(n_in)]
        step_incs.extend([form.format('out', i, i) for i in range(n_out)])
        step_increments = '\n        '.join(step_incs)

        # Types
        n_types = n_in + n_out
        types = "{" + ', '.join(["NPY_DOUBLE"]*n_types) + "};"

        # Docstring
        docstring = '"Created in SymPy with Ufuncify"'

        # Function Creation
        function_creation.append("PyObject *ufunc{0};".format(r_index))

        # Ufunc initialization
        init_form = _ufunc_init_form.substitute(module=module,
                                                funcname=name,
                                                docstring=docstring,
                                                n_in=n_in, n_out=n_out,
                                                ind=r_index)
        ufunc_init.append(init_form)

        outcalls = [_ufunc_outcalls.substitute(
            outnum=i, call_args=call_args, funcname=routines[i].name) for i in
            range(n_out)]

        body = _ufunc_body.substitute(module=module, funcname=name,
                                      declare_args=declare_args,
                                      declare_steps=declare_steps,
                                      call_args=call_args,
                                      step_increments=step_increments,
                                      n_types=n_types, types=types,
                                      outcalls='\n        '.join(outcalls))
        functions.append(body)

        body = '\n\n'.join(functions)
        ufunc_init = '\n    '.join(ufunc_init)
        function_creation = '\n    '.join(function_creation)
        bottom = _ufunc_bottom.substitute(module=module,
                                          ufunc_init=ufunc_init,
                                          function_creation=function_creation)
        text = [top, body, bottom]
        f.write('\n\n'.join(text))

    def _partition_args(self, args):
        """Group function arguments into categories."""
        py_in = []
        py_out = []
        for arg in args:
            if isinstance(arg, OutputArgument):
                py_out.append(arg)
            elif isinstance(arg, InOutArgument):
                raise ValueError("Ufuncify doesn't support InOutArguments")
            else:
                py_in.append(arg)
        return py_in, py_out

@cacheit
@doctest_depends_on(exe=('f2py', 'gfortran', 'gcc'), modules=('numpy',))
def ufuncify(args, expr, language=None, backend='numpy', tempdir=None,
             flags=None, verbose=False, helpers=None, **kwargs):
    """Generates a binary function that supports broadcasting on numpy arrays.

    Parameters
    ----------
    args : iterable
        Either a Symbol or an iterable of symbols. Specifies the argument
        sequence for the function.
    expr
        A SymPy expression that defines the element wise operation.
    language : string, optional
        If supplied, (options: 'C' or 'F95'), specifies the language of the
        generated code. If ``None`` [default], the language is inferred based
        upon the specified backend.
    backend : string, optional
        Backend used to wrap the generated code. Either 'numpy' [default],
        'cython', or 'f2py'.
    tempdir : string, optional
        Path to directory for temporary files. If this argument is supplied,
        the generated code and the wrapper input files are left intact in
        the specified path.
    flags : iterable, optional
        Additional option flags that will be passed to the backend.
    verbose : bool, optional
        If True, autowrap will not mute the command line backends. This can
        be helpful for debugging.
    helpers : iterable, optional
        Used to define auxillary expressions needed for the main expr. If
        the main expression needs to call a specialized function it should
        be put in the ``helpers`` iterable. Autowrap will then make sure
        that the compiled main expression can link to the helper routine.
        Items should be tuples with (<funtion_name>, <sympy_expression>,
        <arguments>). It is mandatory to supply an argument sequence to
        helper routines.
    kwargs : dict
        These kwargs will be passed to autowrap if the `f2py` or `cython`
        backend is used and ignored if the `numpy` backend is used.

    Note
    ----
    The default backend ('numpy') will create actual instances of
    ``numpy.ufunc``. These support ndimensional broadcasting, and implicit type
    conversion. Use of the other backends will result in a "ufunc-like"
    function, which requires equal length 1-dimensional arrays for all
    arguments, and will not perform any type conversions.

    References
    ----------
    [1] http://docs.scipy.org/doc/numpy/reference/ufuncs.html

    Examples
    ========

    >>> from sympy.utilities.autowrap import ufuncify
    >>> from sympy.abc import x, y
    >>> import numpy as np
    >>> f = ufuncify((x, y), y + x**2)
    >>> type(f)
    <class 'numpy.ufunc'>
    >>> f([1, 2, 3], 2)
    array([  3.,   6.,  11.])
    >>> f(np.arange(5), 3)
    array([  3.,   4.,   7.,  12.,  19.])

    For the 'f2py' and 'cython' backends, inputs are required to be equal length
    1-dimensional arrays. The 'f2py' backend will perform type conversion, but
    the Cython backend will error if the inputs are not of the expected type.

    >>> f_fortran = ufuncify((x, y), y + x**2, backend='f2py')
    >>> f_fortran(1, 2)
    array([ 3.])
    >>> f_fortran(np.array([1, 2, 3]), np.array([1.0, 2.0, 3.0]))
    array([  2.,   6.,  12.])
    >>> f_cython = ufuncify((x, y), y + x**2, backend='Cython')
    >>> f_cython(1, 2)  # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    TypeError: Argument '_x' has incorrect type (expected numpy.ndarray, got int)
    >>> f_cython(np.array([1.0]), np.array([2.0]))
    array([ 3.])
    """

    if isinstance(args, Symbol):
        args = (args,)
    else:
        args = tuple(args)

    if language:
        _validate_backend_language(backend, language)
    else:
        language = _infer_language(backend)

    helpers = helpers if helpers else ()
    flags = flags if flags else ()

    if backend.upper() == 'NUMPY':
        # maxargs is set by numpy compile-time constant NPY_MAXARGS
        # If a future version of numpy modifies or removes this restriction
        # this variable should be changed or removed
        maxargs = 32
        helps = []
        for name, expr, args in helpers:
            helps.append(make_routine(name, expr, args))
        code_wrapper = UfuncifyCodeWrapper(C99CodeGen("ufuncify"), tempdir,
                                           flags, verbose)
        if not isinstance(expr, (list, tuple)):
            expr = [expr]
        if len(expr) == 0:
            raise ValueError('Expression iterable has zero length')
        if (len(expr) + len(args)) > maxargs:
            msg = ('Cannot create ufunc with more than {0} total arguments: '
                   'got {1} in, {2} out')
            raise ValueError(msg.format(maxargs, len(args), len(expr)))
        routines = [make_routine('autofunc{}'.format(idx), exprx, args) for
                    idx, exprx in enumerate(expr)]
        return code_wrapper.wrap_code(routines, helpers=helps)
    else:
        # Dummies are used for all added expressions to prevent name clashes
        # within the original expression.
        y = IndexedBase(Dummy())
        m = Dummy(integer=True)
        i = Idx(Dummy(integer=True), m)
        f = implemented_function(Dummy().name, Lambda(args, expr))
        # For each of the args create an indexed version.
        indexed_args = [IndexedBase(Dummy(str(a))) for a in args]
        # Order the arguments (out, args, dim)
        args = [y] + indexed_args + [m]
        args_with_indices = [a[i] for a in indexed_args]
        return autowrap(Eq(y[i], f(*args_with_indices)), language, backend,
                        tempdir, args, flags, verbose, helpers, **kwargs)


"""
This file contains code that will kick off training and testing processes
"""
import os
import json
import numpy as np
from sklearn.model_selection import train_test_split
from experiments.UNetExperiment import UNetExperiment
from data_prep.HippocampusDatasetLoader import LoadHippocampusData

class Config:
    """
    Holds configuration parameters
    """
    def __init__(self):
        self.name = "Basic_unet"
        self.root_dir = r"/home/workspace/src/data"
        self.n_epochs = 10
        self.learning_rate = 0.0002
        self.batch_size = 8
        self.patch_size = 64
        self.test_results_dir = r"..\out"

if __name__ == "__main__":
    # Get configuration

    # TASK: Fill in parameters of the Config class and specify directory where the data is stored and
    # directory where results will go
    c = Config()

    # Load data
    print("Loading data...")

    # TASK: LoadHippocampusData is not complete. Go to the implementation and complete it.
    data = LoadHippocampusData(c.root_dir, y_shape = c.patch_size, z_shape = c.patch_size)

    # Create test-train-val split
    # In a real world scenario you would probably do multiple splits for
    # multi-fold training to improve your model quality

    keys = range(len(data))

    # Here, random permutation of keys array would be useful in case if we do something like
    # a k-fold training and combining the results.

    split = dict()

    # TASK: create three keys in the dictionary: "train", "val" and "test". In each key, store
    # the array with indices of training volumes to be used for training, validation
    # and testing respectively.
    # <YOUR CODE GOES HERE>

    # 60:20:20 split using train_test_split()
    train, test = train_test_split(keys, test_size=0.2, shuffle=True, random_state=1)
    train, val = train_test_split(train, test_size=0.25, shuffle=True, random_state=1)

    print("Train Size:", len(train), "; Validation Size:", len(val), "; Test Size:", len(test))
    split={'train': np.array(train),'val': np.array(val),'test': np.array(test)}

    # Set up and run experiment

    # TASK: Class UNetExperiment has missing pieces. Go to the file and fill them in
    exp = UNetExperiment(c, split, data)

    # You could free up memory by deleting the dataset
    # as it has been copied into loaders
    # del dataset
    del test; del train; del val

    # run training
    exp.run()

    # prep and run testing

    # TASK: Test method is not complete. Go to the method and complete it
    results_json = exp.run_test()

    results_json["config"] = vars(c)

    with open(os.path.join(exp.out_dir, "results.json"), 'w') as out_file:
        json.dump(results_json, out_file, indent=2, separators=(',', ': '))


from matplotlib.pyplot import figure, plot, savefig, xlabel, ylabel
from numpy import genfromtxt
from sys import argv

input_file = str(argv[1])
output_file = str(argv[2])

data = genfromtxt(input_file, delimiter=',')
x = data[0]
ufinal = data[1]

figure()
plot(x, ufinal, '-')
xlabel('x')
ylabel('u(x, T)')
savefig(output_file)

print('Generated ' + output_file)