<reponame>EkremBayar/bayar
import pytest
import numpy as np
from scipy.optimize import quadratic_assignment, OptimizeWarning
from scipy.optimize._qap import _calc_score as _score
from numpy.testing import assert_equal, assert_, assert_warns

################
# Common Tests #
################

def chr12c():
    A = [
        [0, 90, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [90, 0, 0, 23, 0, 0, 0, 0, 0, 0, 0, 0],
        [10, 0, 0, 0, 43, 0, 0, 0, 0, 0, 0, 0],
        [0, 23, 0, 0, 0, 88, 0, 0, 0, 0, 0, 0],
        [0, 0, 43, 0, 0, 0, 26, 0, 0, 0, 0, 0],
        [0, 0, 0, 88, 0, 0, 0, 16, 0, 0, 0, 0],
        [0, 0, 0, 0, 26, 0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 0, 0, 16, 0, 0, 0, 96, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 29, 0],
        [0, 0, 0, 0, 0, 0, 0, 96, 0, 0, 0, 37],
        [0, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 37, 0, 0],
    ]
    B = [
        [0, 36, 54, 26, 59, 72, 9, 34, 79, 17, 46, 95],
        [36, 0, 73, 35, 90, 58, 30, 78, 35, 44, 79, 36],
        [54, 73, 0, 21, 10, 97, 58, 66, 69, 61, 54, 63],
        [26, 35, 21, 0, 93, 12, 46, 40, 37, 48, 68, 85],
        [59, 90, 10, 93, 0, 64, 5, 29, 76, 16, 5, 76],
        [72, 58, 97, 12, 64, 0, 96, 55, 38, 54, 0, 34],
        [9, 30, 58, 46, 5, 96, 0, 83, 35, 11, 56, 37],
        [34, 78, 66, 40, 29, 55, 83, 0, 44, 12, 15, 80],
        [79, 35, 69, 37, 76, 38, 35, 44, 0, 64, 39, 33],
        [17, 44, 61, 48, 16, 54, 11, 12, 64, 0, 70, 86],
        [46, 79, 54, 68, 5, 0, 56, 15, 39, 70, 0, 18],
        [95, 36, 63, 85, 76, 34, 37, 80, 33, 86, 18, 0],
    ]
    A, B = np.array(A), np.array(B)
    n = A.shape[0]

    opt_perm = np.array([7, 5, 1, 3, 10, 4, 8, 6, 9, 11, 2, 12]) - [1] * n

    return A, B, opt_perm

class QAPCommonTests(object):
    """
    Base class for `quadratic_assignment` tests.
    """
    def setup_method(self):
        np.random.seed(0)

    # Test global optima of problem from Umeyama IVB
    # https://pcl.sitehost.iu.edu/rgoldsto/papers/weighted%20graph%20match2.pdf
    # Graph matching maximum is in the paper
    # QAP minimum determined by brute force
    def test_accuracy_1(self):
        # besides testing accuracy, check that A and B can be lists
        A = [[0, 3, 4, 2],
             [0, 0, 1, 2],
             [1, 0, 0, 1],
             [0, 0, 1, 0]]

        B = [[0, 4, 2, 4],
             [0, 0, 1, 0],
             [0, 2, 0, 2],
             [0, 1, 2, 0]]

        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, "maximize": False})
        assert_equal(res.fun, 10)
        assert_equal(res.col_ind, np.array([1, 2, 3, 0]))

        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, "maximize": True})

        if self.method == 'faq':
            # Global optimum is 40, but FAQ gets 37
            assert_equal(res.fun, 37)
            assert_equal(res.col_ind, np.array([0, 2, 3, 1]))
        else:
            assert_equal(res.fun, 40)
            assert_equal(res.col_ind, np.array([0, 3, 1, 2]))

        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, "maximize": True})

    # Test global optima of problem from Umeyama IIIB
    # https://pcl.sitehost.iu.edu/rgoldsto/papers/weighted%20graph%20match2.pdf
    # Graph matching maximum is in the paper
    # QAP minimum determined by brute force
    def test_accuracy_2(self):

        A = np.array([[0, 5, 8, 6],
                      [5, 0, 5, 1],
                      [8, 5, 0, 2],
                      [6, 1, 2, 0]])

        B = np.array([[0, 1, 8, 4],
                      [1, 0, 5, 2],
                      [8, 5, 0, 5],
                      [4, 2, 5, 0]])

        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, "maximize": False})
        if self.method == 'faq':
            # Global optimum is 176, but FAQ gets 178
            assert_equal(res.fun, 178)
            assert_equal(res.col_ind, np.array([1, 0, 3, 2]))
        else:
            assert_equal(res.fun, 176)
            assert_equal(res.col_ind, np.array([1, 2, 3, 0]))

        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, "maximize": True})
        assert_equal(res.fun, 286)
        assert_equal(res.col_ind, np.array([2, 3, 0, 1]))

    def test_accuracy_3(self):

        A, B, opt_perm = chr12c()

        # basic minimization
        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0})
        assert_(11156 <= res.fun < 21000)
        assert_equal(res.fun, _score(A, B, res.col_ind))

        # basic maximization
        res = quadratic_assignment(A, B, method=self.method,
                                   options={"rng": 0, 'maximize': True})
        assert_(74000 <= res.fun < 85000)
        assert_equal(res.fun, _score(A, B, res.col_ind))

        # check ofv with strictly partial match
        seed_cost = np.array([4, 8, 10])
        seed = np.asarray([seed_cost, opt_perm[seed_cost]]).T
        res = quadratic_assignment(A, B, method=self.method,
                                   options={'partial_match': seed})
        assert_(11156 <= res.fun < 21000)
        assert_equal(res.col_ind[seed_cost], opt_perm[seed_cost])

        # check performance when partial match is the global optimum
        seed = np.asarray([np.arange(len(A)), opt_perm]).T
        res = quadratic_assignment(A, B, method=self.method,
                                   options={'partial_match': seed})
        assert_equal(res.col_ind, seed[:, 1].T)
        assert_equal(res.fun, 11156)
        assert_equal(res.nit, 0)

        # check performance with zero sized matrix inputs
        empty = np.empty((0, 0))
        res = quadratic_assignment(empty, empty, method=self.method,
                                   options={"rng": 0})
        assert_equal(res.nit, 0)
        assert_equal(res.fun, 0)

    def test_unknown_options(self):
        A, B, opt_perm = chr12c()

        def f():
            quadratic_assignment(A, B, method=self.method,
                                 options={"ekki-ekki": True})
        assert_warns(OptimizeWarning, f)

class TestFAQ(QAPCommonTests):
    method = "faq"

    def test_options(self):
        # cost and distance matrices of QAPLIB instance chr12c
        A, B, opt_perm = chr12c()
        n = len(A)

        # check that max_iter is obeying with low input value
        res = quadratic_assignment(A, B,
                                   options={'maxiter': 5})
        assert_equal(res.nit, 5)

        # test with shuffle
        res = quadratic_assignment(A, B,
                                   options={'shuffle_input': True})
        assert_(11156 <= res.fun < 21000)

        # test with randomized init
        res = quadratic_assignment(A, B,
                                   options={'rng': 1, 'P0': "randomized"})
        assert_(11156 <= res.fun < 21000)

        # check with specified P0
        K = np.ones((n, n)) / float(n)
        K = _doubly_stochastic(K)
        res = quadratic_assignment(A, B,
                                   options={'P0': K})
        assert_(11156 <= res.fun < 21000)

    def test_specific_input_validation(self):

        A = np.identity(2)
        B = A

        # method is implicitly faq

        # ValueError Checks: making sure single value parameters are of
        # correct value
        with pytest.raises(ValueError, match="Invalid 'P0' parameter"):
            quadratic_assignment(A, B, options={'P0': "random"})
        with pytest.raises(
                ValueError, match="'maxiter' must be a positive integer"):
            quadratic_assignment(A, B, options={'maxiter': -1})
        with pytest.raises(ValueError, match="'tol' must be a positive float"):
            quadratic_assignment(A, B, options={'tol': -1})

        # TypeError Checks: making sure single value parameters are of
        # correct type
        with pytest.raises(TypeError):
            quadratic_assignment(A, B, options={'maxiter': 1.5})

        # test P0 matrix input
        with pytest.raises(
                ValueError,
                match="`P0` matrix must have shape m' x m', where m'=n-m"):
            quadratic_assignment(
                np.identity(4), np.identity(4),
                options={'P0': np.ones((3, 3))}
            )

        K = [[0.4, 0.2, 0.3],
             [0.3, 0.6, 0.2],
             [0.2, 0.2, 0.7]]
        # matrix that isn't quite doubly stochastic
        with pytest.raises(
                ValueError, match="`P0` matrix must be doubly stochastic"):
            quadratic_assignment(
                np.identity(3), np.identity(3), options={'P0': K}
            )

class Test2opt(QAPCommonTests):
    method = "2opt"

    def test_deterministic(self):
        # np.random.seed(0) executes before every method
        n = 20

        A = np.random.rand(n, n)
        B = np.random.rand(n, n)
        res1 = quadratic_assignment(A, B, method=self.method)

        np.random.seed(0)

        A = np.random.rand(n, n)
        B = np.random.rand(n, n)
        res2 = quadratic_assignment(A, B, method=self.method)

        assert_equal(res1.nit, res2.nit)

    def test_partial_guess(self):
        n = 5
        A = np.random.rand(n, n)
        B = np.random.rand(n, n)

        res1 = quadratic_assignment(A, B, method=self.method,
                                    options={'rng': 0})
        guess = np.array([np.arange(5), res1.col_ind]).T
        res2 = quadratic_assignment(A, B, method=self.method,
                                    options={'rng': 0, 'partial_guess': guess})
        fix = [2, 4]
        match = np.array([np.arange(5)[fix], res1.col_ind[fix]]).T
        res3 = quadratic_assignment(A, B, method=self.method,
                                    options={'rng': 0, 'partial_guess': guess,
                                             'partial_match': match})
        assert_(res1.nit != n*(n+1)/2)
        assert_equal(res2.nit, n*(n+1)/2)      # tests each swap exactly once
        assert_equal(res3.nit, (n-2)*(n-1)/2)  # tests free swaps exactly once

    def test_specific_input_validation(self):
        # can't have more seed nodes than cost/dist nodes
        _rm = _range_matrix
        with pytest.raises(
                ValueError,
                match="`partial_guess` can have only as many entries as"):
            quadratic_assignment(np.identity(3), np.identity(3),
                                 method=self.method,
                                 options={'partial_guess': _rm(5, 2)})
        # test for only two seed columns
        with pytest.raises(
                ValueError, match="`partial_guess` must have two columns"):
            quadratic_assignment(
                np.identity(3), np.identity(3), method=self.method,
                options={'partial_guess': _range_matrix(2, 3)}
            )
        # test that seed has no more than two dimensions
        with pytest.raises(
                ValueError, match="`partial_guess` must have exactly two"):
            quadratic_assignment(
                np.identity(3), np.identity(3), method=self.method,
                options={'partial_guess': np.random.rand(3, 2, 2)}
            )
        # seeds cannot be negative valued
        with pytest.raises(
                ValueError, match="`partial_guess` must contain only pos"):
            quadratic_assignment(
                np.identity(3), np.identity(3), method=self.method,
                options={'partial_guess': -1 * _range_matrix(2, 2)}
            )
        # seeds can't have values greater than number of nodes
        with pytest.raises(
                ValueError,
                match="`partial_guess` entries must be less than number"):
            quadratic_assignment(
                np.identity(5), np.identity(5), method=self.method,
                options={'partial_guess': 2 * _range_matrix(4, 2)}
            )
        # columns of seed matrix must be unique
        with pytest.raises(
                ValueError,
                match="`partial_guess` column entries must be unique"):
            quadratic_assignment(
                np.identity(3), np.identity(3), method=self.method,
                options={'partial_guess': np.ones((2, 2))}
            )

class TestQAPOnce():
    def setup_method(self):
        np.random.seed(0)

    # these don't need to be repeated for each method
    def test_common_input_validation(self):
        # test that non square matrices return error
        with pytest.raises(ValueError, match="`A` must be square"):
            quadratic_assignment(
                np.random.random((3, 4)),
                np.random.random((3, 3)),
            )
        with pytest.raises(ValueError, match="`B` must be square"):
            quadratic_assignment(
                np.random.random((3, 3)),
                np.random.random((3, 4)),
            )
        # test that cost and dist matrices have no more than two dimensions
        with pytest.raises(
                ValueError, match="`A` and `B` must have exactly two"):
            quadratic_assignment(
                np.random.random((3, 3, 3)),
                np.random.random((3, 3, 3)),
            )
        # test that cost and dist matrices of different sizes return error
        with pytest.raises(
                ValueError,
                match="`A` and `B` matrices must be of equal size"):
            quadratic_assignment(
                np.random.random((3, 3)),
                np.random.random((4, 4)),
            )
        # can't have more seed nodes than cost/dist nodes
        _rm = _range_matrix
        with pytest.raises(
                ValueError,
                match="`partial_match` can have only as many seeds as"):
            quadratic_assignment(np.identity(3), np.identity(3),
                                 options={'partial_match': _rm(5, 2)})
        # test for only two seed columns
        with pytest.raises(
                ValueError, match="`partial_match` must have two columns"):
            quadratic_assignment(
                np.identity(3), np.identity(3),
                options={'partial_match': _range_matrix(2, 3)}
            )
        # test that seed has no more than two dimensions
        with pytest.raises(
                ValueError, match="`partial_match` must have exactly two"):
            quadratic_assignment(
                np.identity(3), np.identity(3),
                options={'partial_match': np.random.rand(3, 2, 2)}
            )
        # seeds cannot be negative valued
        with pytest.raises(
                ValueError, match="`partial_match` must contain only pos"):
            quadratic_assignment(
                np.identity(3), np.identity(3),
                options={'partial_match': -1 * _range_matrix(2, 2)}
            )
        # seeds can't have values greater than number of nodes
        with pytest.raises(
                ValueError,
                match="`partial_match` entries must be less than number"):
            quadratic_assignment(
                np.identity(5), np.identity(5),
                options={'partial_match': 2 * _range_matrix(4, 2)}
            )
        # columns of seed matrix must be unique
        with pytest.raises(
                ValueError,
                match="`partial_match` column entries must be unique"):
            quadratic_assignment(
                np.identity(3), np.identity(3),
                options={'partial_match': np.ones((2, 2))}
            )

def _range_matrix(a, b):
    mat = np.zeros((a, b))
    for i in range(b):
        mat[:, i] = np.arange(a)
    return mat

def _doubly_stochastic(P, tol=1e-3):
    # cleaner implementation of btaba/sinkhorn_knopp

    max_iter = 1000
    c = 1 / P.sum(axis=0)
    r = 1 / (P @ c)
    P_eps = P

    for it in range(max_iter):
        if ((np.abs(P_eps.sum(axis=1) - 1) < tol).all() and
                (np.abs(P_eps.sum(axis=0) - 1) < tol).all()):
            # All column/row sums ~= 1 within threshold
            break

        c = 1 / (r @ P)
        r = 1 / (P @ c)
        P_eps = r[:, None] * P * c

    return P_eps


'''
    Compute average distances for each metric on covers80
'''

def segment(y, s, rs_size, kmin, kmax, filter):
    """structurally segments the selected audio

        ds_size: side length to which combined matrix is going to be resampled to
        [kmin, kmax]: min and maximum approximation ranks
        filtering: True or False, whether memory stacking, timelag and path enhance are going to be used

    returns set of low rank approximations"""

    #compute cqt
    C = librosa.amplitude_to_db(np.abs(librosa.cqt(y=y, sr=sr,
                                        hop_length=512,
                                        bins_per_octave=12*3,
                                        n_bins=7*12*3)),
                                        ref=np.max)

    #beat synch cqt
    Csync = cv2.resize(C, (int(C.shape[1]/10), C.shape[0]))

    #stack memory
    if filter:
        Csync = librosa.feature.stack_memory(Csync, 4)

    #Affinity matrix
    R = librosa.segment.recurrence_matrix(Csync, width=3, mode='affinity', sym=True)

    #Filtering
    if filter:
        df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)
        R = df(R, size=(1, 7))
        R = librosa.segment.path_enhance(R, 15)

    #mfccs
    mfcc = librosa.feature.mfcc(y=y, sr=sr)

    #downsample like CQT, compress time by 10
    Msync = cv2.resize(C, (int(mfcc.shape[1]/10), mfcc.shape[0]))

    #weighted sequence
    path_distance = np.sum(np.diff(Msync, axis=1)**2, axis=0)
    sigma = np.median(path_distance)
    path_sim = np.exp(-path_distance / sigma)
    R_path = np.diag(path_sim, k=1) + np.diag(path_sim, k=-1)

    #weighted combination of affinity matrix and mfcc diagonal
    deg_path = np.sum(R_path, axis=1)
    deg_rec = np.sum(R, axis=1)

    mu = deg_path.dot(deg_path + deg_rec) / np.sum((deg_path + deg_rec)**2)

    A = mu * R + (1 - mu) * R_path

    #resampling
    A_d = cv2.resize(A, (rs_size, rs_size))

    #laplacian
    L = scipy.sparse.csgraph.laplacian(A_d, normed=True)

    #eigendecomposition
    evals, evecs = scipy.linalg.eigh(L)
    #eigenvector filtering
    evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))

    #normalization
    Cnorm = np.cumsum(evecs**2, axis=1)**0.5

    #temporary replacement for bug
    a_min_value = 3.6934424e-08
    Cnorm[Cnorm == 0.0] = a_min_value
    if (np.isnan(np.sum(Cnorm))):
        print("WOOOOOAH")

    # print("Cnorm shape:",Cnorm.shape)
    # plt.matshow(Cnorm)
    # plt.savefig(filedir[-10:-4])

    #approximations
    dist_set = []
    for k in range(kmin, kmax):

        # #debug
        # print(np.all(Cnorm[:, k-1:k]))
        # divisor = Cnorm[:, k-1:k]
        # if not np.all(divisor):
        #     print("0 divisor")

        Xs = evecs[:, :k] / Cnorm[:, k-1:k]

        #debug
        if np.isnan(np.sum(Xs)):
            print('woops')
            # fig, axs = plt.subplots(1, approx[1]-approx[0], figsize=(20, 20))
            # for i in range(approx[1]-approx[0]):
            #     axs[i].matshow(struct[i])
            # plt.savefig(filedir[-10:-1])

        distance = scipy.spatial.distance.squareform(scipy.spatial.distance.pdist(Xs, metric='euclidean'))
        dist_set.append(distance)
    dist_set = np.asarray(dist_set)

    #return
    return(dist_set)

#Importing
import librosa
import numpy as np
import scipy
from scipy.spatial.distance import pdist, squareform
from scipy.interpolate import interp2d
from scipy.sparse.csgraph import laplacian
from scipy.spatial.distance import directed_hausdorff
from scipy.cluster import hierarchy
from scipy.linalg import eigh
from scipy.ndimage import median_filter
import cv2
from sklearn import metrics
import matplotlib.pyplot as plt
import dill
import sys
import glob
import os
import random
import csv

#--supress warnings--#
import warnings
warnings.filterwarnings("ignore")

#--reading--#

all_dirs = []
all_names = []
all_roots = []
all_audio = []
max_files = 40000
for root, dirs, files in os.walk('/home/ismir/Documents/ISMIR/Datasets/covers80/'):
        for name in files:
            if (('.wav' in name) or ('.aif' in name) or ('.mp3' in name)):
                filepath = os.path.join(root, name)
                all_dirs.append(filepath)
                all_names.append(name[:-4])
                all_roots.append(root)

                if len(all_dirs)>=max_files:
                    break
        if len(all_dirs)>=max_files:
            break
file_no = len(all_dirs)

#load audio
for f in range(file_no):
    y, sr = librosa.load(all_dirs[f], sr=16000, mono=True)
    #bug: empty mel bins
    all_audio.append((y,sr))

    #progress
    sys.stdout.write("\rLoading %i/%s pieces." % ((f+1), str(file_no)))
    sys.stdout.flush()
print('')

#--cover (True) vs non-cover (False)--#
covers = np.zeros((file_no, file_no), dtype=np.bool_)
for i in range(file_no):
    for j in range(file_no):
        if (all_roots[i] == all_roots[j]):
            covers[i][j] = True
        else:
            covers[i][j] = False

#--Distance dictionary--#
"""Terminology
distances: L1, fro, dtw, hau, pair, sh2, sh3
format: rs_size-approx[0]-approx[1]-distance e.g. 128-2-8-L1
"""
distances = {}

#--Score dictionary--#
"""Terminology
distances: L1, fro, dtw, hau, pair, sh2, sh3
format: (filt-)rs_size-approx[0]-approx[1]-distance e.g. filt-128-2-8-L1
"""
scores = {}

#--traverse parameters, compute segmentations, save evaluation--#

#resampling parameters
#for rs_size in [32]:
for rs_size in [128]:
    #approximations
    #for approx in [[2,6]]:
    for approx in [[2,11]]:
        for filtering in [True]:

            #string for keys to indicate filtering
            if filtering:
                filt = 'filt-'
            else:
                filt = ''

            #hold all structures and their formats
            all_struct = [] #kmax-kmin sets each with a square matrix
            all_flat = [] #kmax-kmin sets each with a flattened matrix
            all_merged = [] #single concatenated vector with all flattened matrices
            all_shingled2 = [] #shingled pairs of flat approximations
            all_shingled3 = [] #shingled triples of flat approximations

            print("--------------------")
            print("Resampling size:", str(rs_size))
            print("Approximation range: [" + str(approx[0]) + ',' + str(approx[1]) + ']')
            print("Filtering:", str(filtering))

            #songs
            for f in range(file_no):
                #structure segmentation
                struct = segment(all_audio[f][0], all_audio[f][1],
                                rs_size, approx[0], approx[1], filtering)
                all_struct.append(struct)

                # #debug
                # fig, axs = plt.subplots(1, approx[1]-approx[0], figsize=(20, 20))
                # for i in range(approx[1]-approx[0]):
                #     axs[i].matshow(struct[i])
                # plt.savefig(all_names[f])

                #formatting
                flat_approximations = []
                merged_approximations = np.empty((0))
                for j in range(approx[1]-approx[0]):
                    flat_approximations.append(struct[j].flatten())
                    merged_approximations = np.concatenate((merged_approximations, flat_approximations[j]))
                all_flat.append(np.asarray(flat_approximations))
                all_merged.append(merged_approximations)

                #shingling per 2
                shingled = []
                for j in range(approx[1]-approx[0]-1):
                    shingled.append(np.concatenate((all_flat[f][j],all_flat[f][j+1]),axis=None))
                    #shingled.append(np.concatenate((struct[all_names[f]]['OG'][1][j],struct[all_names[f]]['OG'][1][j+1]),axis=None))
                all_shingled2.append(np.asarray(shingled))

                #shingling per 3
                shingled = []
                for j in range(approx[1]-approx[0]-2):
                    shingled.append(np.concatenate((all_flat[f][j],all_flat[f][j+1],all_flat[f][j+2]),axis=None))
                    #shingled.append(np.concatenate((struct[all_names[f]]['OG'][1][j],struct[all_names[f]]['OG'][1][j+1],struct[all_names[f]]['OG'][1][j+2]), axis=None))
                all_shingled3.append(np.asarray(shingled))

                #progress
                sys.stdout.write("\rSegmented %i/%s pieces." % ((f+1), str(file_no)))
                sys.stdout.flush()
            print('')

            # #plot approximations
            # fig, axs = plt.subplots(1, approx[1]-approx[0], figsize=(20, 20))
            # for i in range(approx[1]-approx[0]):
            #     axs[i].matshow(all_struct[0][i])
            # plt.savefig('approximations'+str(rs_size))

            #list to numpy array
            all_struct = np.asarray(all_struct)
            all_flat = np.asarray(all_flat)
            all_merged = np.asarray(all_merged)

            rows = [['', 'mean', 'max']]

            #L1 norm
            L1_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    L1_distances[i][j] = np.linalg.norm(all_merged[i]-all_merged[j], ord=1)

            rows.append(['L1', np.mean(L1_distances), np.amax(L1_distances)])
            print("Computed L1 distances.")

            #Frobenius norm
            fro_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    fro_distances[i][j] = np.linalg.norm(all_merged[i]-all_merged[j])

            rows.append(['Frobenius', np.mean(fro_distances), np.amax(fro_distances)])
            print("Computed Frobenius distsances.")

            #Sub-sequence Dynamic Time Warping cost
            dtw_cost = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    costs = []
                    for k in range(approx[1]-approx[0]):
                        costs.append(librosa.sequence.dtw(all_struct[i][k], all_struct[j][k], subseq=False, metric='euclidean')[0][rs_size-1,rs_size-1])
                    dtw_cost[i][j] = sum(costs)/len(costs)

            rows.append(['DTW', np.mean(dtw_cost), np.amax(dtw_cost)])
            print("Computed DTW costs.")

            #Directed Hausdorff distance
            hausdorff_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    hausdorff_distances[i][j] = (directed_hausdorff(all_flat[i], all_flat[j]))[0]

            rows.append(['Hausdorff', np.mean(hausdorff_distances), np.amax(hausdorff_distances)])
            print("Computed directed Hausdorff distances.")

            #Minimum distance across all pairs
            min_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    dists = []
                    for n in range(approx[1]-approx[0]):
                        for m in range(approx[1]-approx[0]):
                            dists.append(np.linalg.norm(all_struct[i][n]-all_struct[j][m]))
                    min_distances[i][j] = min(dists)

            rows.append(['Pair', np.mean(min_distances), np.amax(min_distances)])
            print("Computed minimum paiwise distances.")

            #Directed Hausdorff distance of shingled pairs
            shingled2_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    shingled2_distances[i][j] = (directed_hausdorff(all_shingled2[i], all_shingled2[j]))[0]

            rows.append(['Shingled 2', np.mean(shingled2_distances), np.amax(shingled2_distances)])
            print("Copmuter directed Hausforff distances for bi-grams")

            #Directed Hausdorff distance of shingled triples
            shingled3_distances = np.zeros((file_no, file_no))
            for i in range(file_no):
                for j in range(file_no):
                    shingled3_distances[i][j] = (directed_hausdorff(all_shingled3[i], all_shingled3[j]))[0]

            rows.append(['Shingled 3', np.mean(shingled3_distances), np.amax(shingled3_distances)])
            print("Computed directed Hausforff distances for tri-grams.")

with open('/home/ismir/Documents/ISMIR/figures/deformations_run2/mean_max.csv', mode='w') as f:
    writer = csv.writer(f)
    writer.writerows(rows)
print('Stats computed.')


import os
import glob
import scipy.misc as misc
import numpy as np
import imageio

from io import BytesIO

def normalize_image(img):
    """
    Make image zero centered and in between (0, 1)
    """
    normalized = img / 255.
    return normalized

def read_split_image(img):
    mat = misc.imread(img).astype(np.float)
    side = int(mat.shape[1] / 2)
    assert side * 2 == mat.shape[1]
    img_A = mat[:, :side]  # target
    img_B = mat[:, side:]  # source

    return img_A, img_B

def bytes_to_file(bytes_img):
    return BytesIO(bytes_img)

def read_split_image(img):
    mat = misc.imread(img).astype(np.float)
    side = int(mat.shape[1] / 2)
    assert side * 2 == mat.shape[1]
    img_A = mat[:, :side]  # target
    img_B = mat[:, side:]  # source

    return img_A, img_B

def shift_and_resize_image(img, shift_x, shift_y, nw, nh):
    w, h, _ = img.shape
    enlarged = misc.imresize(img, [nw, nh])
    return enlarged[shift_x:shift_x + w, shift_y:shift_y + h]

def scale_back(images):
    return (images + 1.) / 2.

def merge(images, size):
    h, w = images.shape[1], images.shape[2]
    img = np.zeros((h * size[0], w * size[1], 3))
    for idx, image in enumerate(images):
        i = idx % size[1]
        j = idx // size[1]
        img[j * h:j * h + h, i * w:i * w + w, :] = image

    return img

def save_concat_images(imgs, img_path):
    concated = np.concatenate(imgs, axis=1)
    misc.imsave(img_path, concated)

def compile_frames_to_gif(frame_dir, gif_file):
    frames = sorted(glob.glob(os.path.join(frame_dir, "*.png")))
    print(frames)
    images = [misc.imresize(imageio.imread(f), interp='nearest', size=0.33) for f in frames]
    imageio.mimsave(gif_file, images, duration=0.1)
    return gif_file


<gh_stars>0
'''
FFT_testing.py

Benchmark of the FFT method with CST files
'''
import numpy as np
import matplotlib.pyplot as plt
import time
import sys
import os
from scipy.constants import c, mu_0, pi
from scipy.special import iv
from copy import copy
import pickle as pk

# UNIT=1e-3 #mm to m

# # Gaussian bunch
# t0=0.53e-10         #injection time [s]
# WL=50*UNIT          #wakelength [m]
# sigmaz=1.0*UNIT     # [m]
# sigmat=sigmaz/c     # [s]
# q=1e-9              # [C]

# #--- define time
# N=10000
# tau=np.linspace(1e-12, WL/c, N)
# dt=tau[2]-tau[1]

# #--- define s
# s=np.linspace(-t0*c, WL, N)
# ds=s[2]-s[1]

# # bunch=np.exp(-((tau-2*t0)**2)/(2*(sigmat**2))) #*1/(sigmat*np.sqrt(2*pi))
# bunch=q*np.exp(-((s)**2)/(2*(sigmaz**2)))*1/(sigmaz*np.sqrt(2*pi))

# # Wakefunction

# F=1.0           #form factor
# L=300*UNIT      #length [m]
# b=1.0*UNIT      #radius [m]
# Zo=376.73       #vacuum impedance [Ohm]
# sigma_el=1.0e6   #electric conductivity [S/m]
# #s_wf=np.linspace(1.0e-5, WL, 1000)
# WF=np.zeros_like(s)

# #WF=-F*L/(4*pi*b)*np.sqrt(Zo/(pi*c*sigma_el))*(1/np.power(tau,3/2))
# mask = s > 0
# WF[mask]=F*L*c/(4*pi*b)*np.sqrt(Zo/(pi*sigma_el))*(1/np.power(abs(s[mask]),3/2))
# #WF[np.logical_not(mask)] = 0.0

# # Wakepotential

# #--- with convolution

# '''
# WFf, f = Wsol.DFT(WF, ds/c, fmax=fmax, Nf=1000)
# bunchf, f=Wsol.DFT(bunch, ds/c, fmax=fmax, Nf=1000)

# convf = WFf*bunchf
# WP = np.fft.ifft(Wf)
# '''

# #WP_conv=(1/(q*1e12))*np.convolve( bunch[mask] , WF[mask] )   #convolution of Wakefunction and bunch charge distribution [V/pC]
# WP_conv=(1/(q*1e12))*np.convolve( bunch , WF )
# s_conv=np.linspace(0, WL, len(WP_conv))
# WP_conv = np.interp(s, s_conv, WP_conv)

# #--- from theory
# x=s/(2*sigmaz)
# WPth=-c*L/(4*pi*b*np.power(sigmaz, 3/2))*np.sqrt(Zo/sigma_el)*np.power(abs(x), 3/2)*np.exp(-x**2)*(iv(-3/4, x*x)-iv(1/4, x*x)+np.sign(s)*(iv(3/4, x*x)-iv(-1/4, x*x)))
# WPth=WPth*1e-12 #[V/pC]

# # Impedance

# fmax=1/(3*sigmat)

# #--- with FFT
# bunchf, f2=Wsol.FFT(bunch/q, ds/c, fmax=2*fmax, flag_zeropadding=False)
# WPf, f=Wsol.FFT(WPth*1e12, ds/c, flag_zeropadding=False)

# bunchf = np.interp(f,f2,bunchf)

# #--- with DFT
# #WPf, f=Wsol.DFT(WPth, ds/c, fmax=fmax, Nf=1000)
# #bunchf, f=Wsol.DFT(bunch/q, ds/c, fmax=fmax, Nf=1000)
# #WPf_conv, f=Wsol.DFT(WP_conv, ds/c, fmax=fmax, Nf=1000)

# Z = - WPf / bunchf
# Z_abs= abs(Z)
# Z_re=np.real(Z)
# Z_im=np.imag(Z)

# #Z_conv = - WPf_conv / bunchf

# #--- from theory
# fth=np.linspace(0, fmax, 1000)
# Zth=(1-1j*np.sign(fth))*L/(2*pi*b)*np.sqrt(Zo*2*pi*fth/(2*c*sigma_el))

# Zth_abs=abs(Zth)
# Zth_re=np.real(Zth)
# Zth_im=np.imag(Zth)

# '''
# # Plot WP
# fig = plt.figure(2, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(s, WP_conv, color='red', label='Wake potential from convolution [norm]')
# ax.plot(s, WPth, color='red', ls='--', label='Wake potential from theory')
# ax.plot(s, bunch/max(bunch)*max(abs(WPth)), color='orange', label='lambda(s)')
# ax.set(title='Resistive wall Wake potential W//(s)',
#         xlabel='s [m]',
#         ylabel='WP [V/pC]',
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()
# '''

# '''
# # Plot WF
# fig = plt.figure(2, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(s, WF, color='blue', ls='--', label='Wake function from theory')
# ax.plot(s, bunch, color='orange', label='lambda(s)')
# ax.set(title='Resistive wall Wake function',
#         xlabel='s [m]',
#         ylabel='WF [V]',
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()
# '''

# # Plot Z vs Zth
# factor=sum(Zth_abs)/len(Zth_abs)/(sum(Z_abs)/len(Z_abs))
# fig = plt.figure(2, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# #ax.plot(fth, Zth_abs, color='green', ls='--', label='|Z(f)| from theory')
# ax.plot(f, Z_abs, color='green', label='|Z(f)| from FFT')
# #ax.plot(f, abs(Z_conv), color='m', label='|Z(f)| from convolution')
# ax.plot(f, Z_re, color='red', label='Zre(f) from FFT')
# ax.plot(f, Z_im, color='blue', label='Zim(f) imaginary from FFT')
# ax.set(title='Resistive wall impedance Z//(f)', #' \n Zth / Z = ' + str(round(factor,3)),
#         xlabel='f [Hz]',
#         ylabel='Z [Ohm]',
#         xlim=(0, 1/(3*sigmat)),
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()

# '''
# # Plot Z vs Zth Re and Im
# fig = plt.figure(2, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(f, Z_re, color='red', label='Zre(f) from FFT')
# ax.plot(f, Z_im, color='blue', label='Zim(f) imaginary from FFT')
# ax.plot(f, Zth_re, color='green', ls='--', label='Zre(f) from theory')
# ax.plot(f, Zth_im, color='m', ls='--', label='Zim(f) from theory')
# ax.set(title='Resistive wall impedance',
#         xlabel='f [Hz]',
#         ylabel='Z [Real / Imag]',
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()
# '''

# #--- read the cst dictionary
# with open('cst_out.txt', 'rb') as handle:
#   cst_data = pk.loads(handle.read())

# print(cst_data.keys())
# charge_dist_time=cst_data.get('charge_dist_time')
# t=cst_data.get('t_charge_dist')
# t0=cst_data.get('init_time')
# dt=t[2]-t[1]

# charge_dist=cst_data.get('charge_dist')
# s_charge_dist=cst_data.get('s_charge_dist')
# ds_charge_dist=s_charge_dist[2]-s_charge_dist[1]
# spectrum=cst_data.get('charge_dist_spectrum')
# df=0.0094650788*1e9
# f_spectrum=np.arange(0, 1001*df, df)
# q=1e-9 #[C]

# #'''
# Wake_potential=cst_data.get('WP_cst')
# s_cst=cst_data.get('s_cst')
# Z_cst=cst_data.get('Z_cst')
# freq_cst=cst_data.get('freq_cst')
# sigmaz=cst_data.get('sigmaz')
# #'''
# '''
# Wake_potential_cst=cst_data.get('WPx_dipolar_cst')
# s_cst=cst_data.get('s_cst_dipolar')
# Z_cst=cst_data.get('Zx_dipolar_cst')
# freq_cst=cst_data.get('freq_cst_dipolar')
# '''
# '''
# Wake_potential_cst=cst_data.get('WPy_quadrupolar_cst')
# s_cst=cst_data.get('s_cst_quadrupolar')
# Z_cst=cst_data.get('Zy_quadrupolar_cst')
# freq_cst=cst_data.get('freq_cst_quadrupolar')
# '''

# #--- Auxiliary variables
# ds=s_cst[2]-s_cst[1]
# s=np.arange(np.min(s_cst),np.max(s_cst),ds) #constant ds vector

# #--- Obtain impedance Z with Fourier transform numpy.fft.fft
# # MAKE A SYMMETRIC SIGNAL

# # Interpolate charge distribution
# # INTERPOLATE TO HAVE A CONSTANT ds. PLOT CST DS DISTRIBUTION
# charge_dist_interp=np.interp(s, s_charge_dist, charge_dist)
# Wake_potential_interp=np.interp(s, s_cst, Wake_potential)

# #lambdaf, f=Wsol.FFT(charge_dist_interp, ds/c, fmax=np.max(freq_cst), r=10.0)
# #WPf, f=Wsol.FFT(Wake_potential_interp, ds/c, fmax=np.max(freq_cst), r=10.0)

# #lambdaf, f2=Wsol.DFT(charge_dist/q, ds_charge_dist/c, fmax=max(freq_cst), Nf=2001)
# WPf, f=Wsol.DFT(Wake_potential*1e12, ds/c, fmax=max(freq_cst), Nf=2001)
# #WPf=WPf*sum(Wake_potential*1e12)*ds/c/np.sqrt(pi)

# #lambdaf=np.interp(f,f2,lambdaf)
# lambdaf=np.interp(f, f_spectrum, spectrum/q)*c

# # Compute the impedance
# Z = abs(- WPf / lambdaf) # * 2/(t_sample*ds/np.sqrt(pi)) #normalized according to CST wakesolver manual

# # Plot Impedance and maximum frequency
# fig = plt.figure(1, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()

# # add CST fft result
# ax.plot(freq_cst*1.0e-9, Z_cst, lw=1.2, color='black', label='Z// from CST')

# # add numpy.fft result (normalized)
# factor=np.max(Z_cst)/np.max(Z)
# ax.plot(f*1.0e-9, Z, lw=1.2, color='red', label='Z// from numpy FFT')

# ax.set(title='Longitudinal impedance Z from CST magnitude',
#         xlabel='frequency [GHz]',
#         ylabel='Z//(s) [Ohm]',         #ylim=(-8.0e4,8.0e4)
#         xlim=(0.,np.max(freq_cst)*1e-9)
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()

# # Plot charge dist and wake potential
# '''
# fig = plt.figure(2, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(s*1e3, charge_dist_interp, color='red', label='$\lambda$(s)')
# ax.plot(s*1e3, Wake_potential_cst, color='orange', label='W||(s)')
# ax.set(title='Wake potential and charge distribution',
#         xlabel='s [mm]',
#         ylabel='W||(s) [V/pC]',         #ylim=(-8.0e4,8.0e4)
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()
# '''

# fig = plt.figure(1, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(f/1e9, abs(lambdaf), color='red', label='$\lambda$(w)')
# ax.plot(f/1e9, abs(WPf), color='orange', label='W||(w)')
# ax.set(title='Wake potential and charge distribution',
#         xlabel='f [GHz]',
#         ylabel='W||(s) [V/pC]',         #ylim=(-8.0e4,8.0e4)
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()

# #VALUE IN 0 SHOULD BE THE SAME AS THE INTEGRAL FOR THE CHARGE DIST

# spectrum_fft,f=Wsol.DFT(charge_dist_time, dt, fmax=max(f_spectrum), Nf=2001)
# spectrum_fft=spectrum_fft*sum(charge_dist_time)*dt/np.sqrt(pi)

# fig = plt.figure(1, figsize=(6,4), dpi=200, tight_layout=True)
# ax=fig.gca()
# ax.plot(f/1e9, abs(spectrum_fft), color='red', label='DFT')
# ax.plot(f_spectrum/1e9, spectrum, color='blue', label='CST')
# ax.set(title='Charge distribution spectrum',
#         xlabel='f [GHz]',
#         ylabel='Spectrum',         #ylim=(-8.0e4,8.0e4)
#         )
# ax.legend(loc='best')
# ax.grid(True, color='gray', linewidth=0.2)
# plt.show()

#--- read the cst dictionary
with open('cst_out.txt', 'rb') as handle:
  cst_data = pk.loads(handle.read())

bunch=cst_data.get('charge_dist')
bunch_time=cst_data.get('charge_dist_time')
spectrum=cst_data.get('charge_dist_spectrum')
current=cst_data.get('current')
s_bunch=cst_data.get('s_charge_dist')
ds_bunch=s_bunch[2]-s_bunch[1]
df=0.0094650788*1e9
f_spectrum=np.arange(0, 1001*df, df)
dt=0.0013634439*1e-9
t_current=np.arange(0,7.5371175*1e-9, dt)
q=1e-9 #[C]

WP=cst_data.get('WP_cst')
Z_cst=cst_data.get('Z_cst')
WP_dip=cst_data.get('WPx_dipolarX_cst')
Z_dip_cst=cst_data.get('Zx_dipolarX_cst')
WP_quad=cst_data.get('WPx_quadrupolarX_cst')
Z_quad_cst=cst_data.get('Zx_quadrupolarX_cst')

f_cst=cst_data.get('freq_cst_dipolar')
s_cst=cst_data.get('s_cst')
ds = s_cst[2]-s_cst[1]
df = f_cst[2]-f_cst[1]

bunch_i=np.interp(s_cst, s_bunch, bunch)

'''
lambdaf, f2=Wsol.DFT(bunch_i/q, ds/c, fmax=max(f_spectrum), Nf=2001)
WPf, f=Wsol.DFT(WP*1e12, ds/c, fmax=max(f_cst), Nf=2001)
WPf_dip, f=Wsol.DFT(WP_dip*1e12, ds/c, fmax=max(f_cst), Nf=2001)
WPf_quad, f=Wsol.DFT(WP_quad*1e12, ds/c, fmax=max(f_cst), Nf=2001)
'''
lambdafft = np.fft.fft(bunch_i/q*c, n=200000)
WPfft = np.fft.fft(WP*1e12, n=200000)
WPfft_dip = np.fft.fft(WP_dip*1e12, n=200000)
WPfft_quad = np.fft.fft(WP_quad*1e12, n=200000)
ffft=np.fft.fftfreq(len(WPfft), ds/c)

mask  = np.logical_and(ffft >= 0 , ffft < 5.5*1e9)
WPf = WPfft[mask]*ds
WPf_dip = WPfft_dip[mask]*ds
WPf_quad = WPfft_quad[mask]*ds
lambdaf = lambdafft[mask]*ds
f = ffft[mask]            # Positive frequencies

# Compute the impedance
Z = abs(- WPf / lambdaf)
Z_dip = abs(1j* WPf_dip / lambdaf)
Z_quad = abs(1j* WPf_quad / lambdaf)

# Plot Impedance and maximum frequency
fig = plt.figure(1, figsize=(6,4), dpi=200, tight_layout=True)
ax=fig.gca()

# add CST fft result
'''
ax.plot(f*1.0e-9, Z, lw=1.2, color='red', label='Z// from numpy FFT')
ax.plot(f_cst*1.0e-9, Z_cst, lw=1.2, color='black', ls='--', label='Z// from CST')
'''
'''

ax.plot(f*1.0e-9, Z_dip, lw=1.2, color='red', label='Z dipolar from numpy FFT')
ax.plot(f_cst*1.0e-9, Z_dip_cst, lw=1.2, color='black', ls='--', label='Z dipolar from CST')
'''

ax.plot(f*1.0e-9, Z_quad, lw=1.2, color='red', label='Z quadrupolar from numpy FFT')
ax.plot(f_cst*1.0e-9, Z_quad_cst, lw=1.2, color='black', ls='--', label='Z quadrupolar from CST')

#ax.plot(f_cst*1.0e-9, Z_cst/Z_dip_cst, lw=1.2, color='black', label='Z// from CST')
#ax.plot(f*1.0e-9, Z/Z_dip, lw=1.2, color='red', label='Z// from numpy FFT')
#ax.plot(s_cst*1e3, WP_quad, lw=1.2, color='blue', label='Z// from numpy FFT')

#ax.plot(f_cst*1.0e-9, abs(WPf_dip)/Z_dip_cst, lw=1.2, color='blue', label='Z// from numpy FFT')
#ax.plot(f*1.0e-9, abs(lambdaf_cst), lw=1.2, color='blue', label='Z// from numpy FFT')

ax.set(title='Longitudinal impedance Z from CST magnitude',
        xlabel='frequency [GHz]',
        ylabel='Z//(s) [Ohm]',         #ylim=(-8.0e4,8.0e4)
        #xlim=(0.,np.max(f_cst)*1e-9)
        )
ax.legend(loc='best')
ax.grid(True, color='gray', linewidth=0.2)
plt.show()


# -*- coding: utf-8 -*-
"""
My implementation of a linear algebra library in Python. Consists of
a Vector and Matrix class that can be compatible with one another, as
well as useful methods for interacting with both concepts/classes.

github repository: https://github.com/ulloaluis/linear-algebra
"""

from math import gcd, pow, sqrt, isclose
from linear_lib.linear_tests import *
from fractions import Fraction

__author__ = "<NAME>"
__license__ = "MIT"
__version__ = "0.0.2"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"

class Vector:
    """
    The Vector class imitates the m x 1 vector from linear algebra and
    contains many useful functions for dealing and interacting with Vectors.

    Getting values directly from the vector should be done using the get(index)
    function since the comp vector location in memory may change with functions
    like mag() or zero().

    class Vector
    __init__(comp)          - takes in a list of components or a valid mx1 Matrix
    resize(length)          - while preserving current elements or filling with 0's, changes current vector length
    set(comp, index=-1)     - sets entire list at once or one specific index/value
    get(index)              - returns item at specified index of vector
    zero()                  - turns the current vector into a zero vector and returns it
    mag()                   - returns the magnitude of current vector
    normalize(change=False) - returns normalized current vector, if change=True, internal vector is updated
    same_normalized(other)  - returns True/False depending on equality of the two vectors once normalized
    dot(other)              - returns the dot product of th two vectors
    cross(other)            - returns the cross product of u x v (u is current vector, v is other)
    perp(other)             - returns True/False if current and other are/aren't perpendicular
    parallel(other)         - returns True/False if current and other are/aren't parallel
    indep(other)          - returns True/False if curr vector and other vector(s) are linearly independent
    operator +              - returns sum of two vectors, component wise addition
    operator -              - returns difference of two vectors, component wise subtraction
    operator *              - alternate for dot product, or can use for vector scaling
    operator **             - returns original vector with its components raised to power
    operator ==             - checks to see if lists are equal
    to string method        - format: "<elem1, elem2, ...>"
    len() method            - can use len() to get vector length
    get and set []          - user can get and set values in vector with index based operations []

    comp = vector composition, list of components
    length = number of components in vector
    rows = same as length, used with cols for backwards compatibility with Matrix
    cols = 1 (num of columns)
    """

    def __init__(self, comp=[]):
        """
        Initializes the vector with either a list containing its components
        or an appropriate Matrix with mx1 dimensions. Defaults to an empty
        vector if not specified.

        :param comp: a list of the elements to be included in the vector;
                        the initial components, defaulted to an empty list.
                        could also be an appropriately sized Matrix
        :type comp: list, Matrix
        :return: none
        :raises: ValueError when Matrix is invalid size for Vector conversion
        """
        if isinstance(comp, Matrix):
            if comp.cols == 1:
                self.comp = [row[0] for row in comp.comp]
            else:
                raise ValueError("Cannot convert Matrix with greater than 1 column to Vector.")
        else:
            self.comp = comp             # user should never change comp directly; use set()

        self.length = len(self.comp)     # user should never change length directly; use resize()
        self.rows = self.length          # rows and cols included for backwards compatibility as
        self.cols = 1                    # a Matrix and for use in matrix-vector product

    def resize(self, length):
        """
        Re-sizes the vector to the specified length. If length is greater than
        current size, the new components are initialized with 0, otherwise if length
        is less than current size, the last few components are lost.

        :param length: new length of vector
        :type length: int
        :return: current vector, now resized
        :rtype: Vector
        """
        assert(length >= 0)                # no negative lengths
        dist = length - self.length

        if dist < 0:
            self.comp = self.comp[:dist]  # notice how this fails when dist = 0, but else correctly handles it
        else:
            self.comp = self.comp + [0]*dist

        self.length = length
        return self

    def set(self, comp, index=-1):
        """
        Set/change the values of the current vector. Can either pass in a new
        list to replace the internal list, or can specify an index in vector
        to change just that value, in which case comp can be a single value.
        No errors are thrown if user re-sizes the list.

        :param comp: list to replace whole vector or value to replace single component
        :param index: optional parameter that specifies the index of the value to be replaced
        :type comp: list, int, float
        :type index: int, float (that is whole ex. 1.0)
        :return: current vector, now updated
        :rtype: Vector
        :raises: index error if out of bounds index
        """
        if index < 0:                   # default None and index=0 calls would conflict
            self.comp = comp
            self.length = self.rows = len(comp)
        else:
            if index >= self.length:
                raise IndexError("Index out of bounds in vector.")
            self.comp[index] = comp
        return self

    def get(self, index):
        """
        :param index: index of value
        :type index: int
        :return: element at specified index
        :rtype: int, float
        :raises: IndexError if index not in vector
        """
        if 0 <= index < self.length:
            return self.comp[index]
        else:
            raise IndexError("Specified index is not in vector.")

    def zero(self):
        """
        Zeroes out the current vector by replacing each component with a 0.

        :return: returns current vector, which is now a zero vector
        :rtype: Vector
        """
        self.comp = [0]*self.length
        return self

    def mag(self):
        """
        Will get the magnitude of a vector.

        :return: the magnitude of a vector (sqrt(sum of components squared))
        :rtype: int, float
        """
        return sqrt(sum([pow(x, 2) for x in self.comp]))

    def normalize(self, change=False):
        """
        Normalizes a vector (acts on internal vector, does not take in a vector)

        :param change: if True, internal vector components are changed in addition
                          to returning vector
                 if False, vector says the same but normalized vector is returned;
                           default is false
        :type change: bool
        :return: another Vector but with the normalized components (False)
                 current Vector but with normalized components (True)
        :rtype: Vector
        """
        magnitude = self.mag()
        if magnitude == 0:      # already zero vector
            return self

        if change:
            self.comp = [elem / magnitude for elem in self.comp]
            return self
        else:
            return Vector([x / magnitude for x in self.comp])

    def same_normalized(self, other):
        """
        This function states whether the current vector is the same as other
        vector when normalized.

        :param other: other vector to be compared
        :type other: Vector
        :return: True if they have same normalized version, False otherwise
        :rtype: bool
        """
        return self.normalize() == other.normalize()

    def dot(self, other):
        """
        This function returns a scalar (number) value representing the dot
        product of the current vector and the other vector.

        :param other: the b vector in a dot b
        :type other: Vector
        :return: The dot product of the current vector and other vector.
        :rtype: int, float
        :raises: ValueError when vectors are not the same length
        """
        if len(self.comp) == len(other.comp):
            return sum([x * y for x, y in zip(self.comp, other.comp)])
        else:
            raise ValueError("Invalid vectors - must be of same length.")

    def cross(self, other):
        """
        For 3-dimensional vectors (3x1), this function allows you to take
        the cross product, which produces a vector i.e. orthogonal to both.

        :param other: 3D Vector (b in a X b)
        :return: Vector representing cross product of current and other
        :rtype: Vector
        :raises: Value error if vectors are not 3 dimensional
        """

        # Simplified version, after determinants: u is current vector v is other
        # u x v = (u2v3 - u3v2)i - (u1v3-u3v1)j + (u1v2-u2v1)k
        if self.length == 3 and other.length == 3:
            i_hat = self.comp[1]*other.comp[2] - self.comp[2]*other.comp[1]
            j_hat = -1 * (self.comp[0]*other.comp[2] - self.comp[2]*other.comp[0])
            k_hat = self.comp[0]*other.comp[1] - self.comp[1]*other.comp[0]
            return Vector([i_hat, j_hat, k_hat])
        else:
            raise ValueError("Invalid vectors - Can only take the cross product of 3D vectors.")

    def perp(self, other):
        """
        Boolean function for whether two vectors are perpendicular/orthogonal to each other.

        :param other: the other vector
        :type other: Vector
        :return: Will return True if current vector and other vector are perpendicular, false otherwise.
        :rtype: bool
        """

        return self.dot(other) == 0

    def parallel(self, other):
        """
        Boolean function for whether two vectors are parallel to each other.

        :param other: the other vector
        :type other: Vector
        :return: Will return True if current vector and other vector are parallel, false otherwise.
        :rtype: bool
        """

        return self.cross(other).mag() == 0   # could also check dot prod = |A*B|

    def indep(self, other):
        """
        Determines whether current vector and one or more vectors are linearly independent.

        Note: User should make sure to pass in vectors of correct dimension.

        :param other: list of vectors or a vector to be compared to current
        :type other: List, Vector
        :return: boolean true/false if given vectors are linearly independent
        :rtype: bool
        :raises: ValueError if other is not a valid type
        """

        if isinstance(other, Vector):   # make 'other' a list if it's a vector
            other = [other]

        if isinstance(other, list) and len(other) > 0:
            other.append(self)
            m, n = len(other), len(other[0])    # m is num vectors, n is vector dimension

            if m == n:              # Place list into matrix and check if determinant is 0
                return Matrix([vec.comp for vec in other]).det() != 0
            elif m < n:
                row_reduced = Matrix([vec.comp for vec in other]).row_reduce()
                return Vector(row_reduced[-1]).mag() != 0   # see if last row is all 0s
            else:
                return False    # if num vectors > dimension, can't be independent

        else:
            raise ValueError("Invalid input - Must be a vector or list of vectors.")

    def __add__(self, other):
        """
        Adding two vectors returns a vector with the respective components
        added together as expected. (does not affect this or other vector's
        components)

        :param other: the other vector to be added to current instance vector
        :type other: Vector
        :return: a vector with the resulting added components
        :rtype: Vector
        :raises: ValueError when vectors are not the same length
        """
        if len(self.comp) == len(other.comp):
            return Vector([x+y for x, y in zip(self.comp, other.comp)])
        else:
            raise ValueError("Invalid vectors - must be of same length.")

    def __sub__(self, other):
        """
       Subtracting two vectors returns a vector with the respective components
       subtracted. "current - other" is formatting. (this does not affect this
       or other vector's components)

       :param other: the other vector which is subtracting from the current vector
       :type other: Vector
       :return: a vector with the resulting subtracted components
       :rtype: Vector
       :raises: ValueError when vectors are not the same length
       """
        if len(self.comp) == len(other.comp):
            return Vector([x-y for x, y in zip(self.comp, other.comp)])
        else:
            raise ValueError("Invalid vectors - must be of same length.")

    def __mul__(self, other):
        """
        Multiplies the two vectors together; same functionality as calling the
        dot() function for dot product of current and other. Could also scale
        each component by a number

        :param other: the other vector
        :type other: Vector, integer, float
        :return: number value representing dot product of both vectors
        :rtype: int, float
        :raises: ValueError when vectors are not the same length
        """
        if isinstance(other, int) or isinstance(other, float):
            return Vector([x * other for x in self.comp])
        elif len(self.comp) == len(other.comp):
            return self.dot(other)
        else:
            raise ValueError("Invalid vectors - must be of same length.")

    def __eq__(self, other):
        """
        If two vectors have the same components, then they are equal. If the
        lists are not the same length, will always be False with no error thrown.

        :param other: other vector being tested for equality
        :type other: Vector
        :return: True or False based on equality
        :rtype: bool
        """
        return self.comp == other.comp  # compares lists

    def __pow__(self, power, modulo=None):
        """
        Allows you to raise each of the components of the current vector to a power.

        :param power: value to raise each component to
        :param modulo: optional parameter that applies the modulus operator to each result
        :type power: int, float
        :type modulo: int, float
        :return: a vector containing the appropriately scaled components
        :rtype: Vector
        """

        return Vector([pow(x, power) % modulo if modulo else pow(x, power) for x in self.comp])

    def __str__(self):
        """
        Converts vector to string by placing the components a, b , c, ... into arrow
        brackets, as such <a, b, c, ...>.  Empty arrays return <empty> for clarity.

        :return: a string detailing contents of vector with the format <a, b, c, ...> or <empty>
        :rtype: str
        """
        if self.length == 0:
            return "<empty>"
        vec = "<"
        for elem in self.comp:
            vec += str(elem) + ", "
        return vec[:-2] + ">"  # remove additional ", " and close

    def __len__(self):
        """
        :return: length of vector
        :rtype: int
        """
        return self.length

    def __getitem__(self, i):
        """
        Alternate for get(), allows you to reliably access components of vector.
        v = Vector([1,2]) v[0] -> 1

        :param i: index
        :type i: int
        :return: value at specified index in self.comp/vector
        :rtype: int, float
        """
        return self.get(i)

    def __setitem__(self, key, value):
        """
        Allows user to set value using index-based accessing.

        :param key:
        :param value:
        :return: item just inserted
        """
        return self.set(value, key)

class Matrix:
    """
    The Matrix class imitates the matrix concept from linear algebra and allows
    for different ways of dealing and interacting with matrices and vectors.

    class Matrix
    __init__(comp)          - takes in a list of components or a valid Vector
    resize(rows, cols)      - while preserving current elements or filling with 0's, changes current vector length
    set(comp, index=None)   - sets entire list at once or one specific index/value (tuple or array as (row, col))
    get(row=None,col=None)  - can get a specific row, column, or entire matrix composition (no args for matrix)
    zero()                  - replaces values in current matrix with all zeroes and returns it
    det()                   - takes the determinant of the current NxN matrix
    transpose()             - transposes the current mxn matrix to an nxm matrix (1st row becomes 1st col, etc.)
    row_echelon()           - returns the current matrix in row echelon form
    row_reduce()            - returns the current matrix to reduced row echelon form
    identity(n)             - static method that returns the nxn identity matrix
    combine(first, second)  - static method that combines two matrices by concatenation
    inverse()               - returns the inverse of current nxn matrix, or None if singular
    operator +              - returns sum of two matrices, component wise addition
    operator -              - returns difference of two matrices, component wise subtraction
    operator *              - matrix multiplication, matrix-vector product, scalar multiplication
    operator **             - returns original matrix with its components raised to power
    operator ==             - checks to see if internal lists are equal
    to string method        - format: "[row1\n row2\n row3\n ...]" and floats are shown as fractions
    len() method            - returns tuple formatted as (row, col)
    get and set [][]        - can get rows and specific values with [] or [][], and set specific values with [][]

    comp = matrix composition, list of lists where each list is a row
    rows = number of rows in matrix
    cols = number of columns in matrix
    """
    def __init__(self, comp=[]):
        """
        Initializes the matrix to the specified format. Default is an empty 0x0 matrix

        Note: It is up to the user to pass well-formed matrices, that is, two different
                rows cannot be different lengths, etc.

        :param comp: list of lists where each individual list represents a row,
                     similar to how numpy implements arrays; could also be a vector
        :type comp: list, Vector
        :return: none
        """

        if isinstance(comp, Vector):
            self.comp = [[row] for row in comp.comp]  # m x 1 Vector --> m rows
        else:
            self.comp = comp                          # list

        self.rows = len(self.comp)          # User should never be able to change instance
        if self.rows != 0:                  # variables directly, use an appropriate method
            self.cols = len(self.comp[0])
        else:
            self.cols = 0                   # if rows = 0, then no columns by default

    def resize(self, rows, cols):
        """
        Re-sizes the current matrix to the specified dimensions, rows x cols.
        Previous elements are left in place, if size is increased then new
        locations are filled with values of 0.

        :param rows: new row size
        :param cols: new column size
        :type rows: int, float
        :type cols: int, float
        :return: current matrix after resizing
        :rtype: Matrix
        """
        assert(rows >= 0 and cols >= 0)     # no negative dimensions allowed

        dist_rows = rows - self.rows
        dist_cols = cols - self.cols

        if dist_rows < 0:
            self.comp = self.comp[:dist_rows]
        else:
            for i in range(dist_rows):
                self.comp.append([0]*self.cols)  # update rows but don't have varying number of columns for each row

        if dist_cols < 0:                        # go through and shape columns now
            for i in range(rows):
                self.comp[i] = self.comp[i][:dist_cols]
        else:
            for i in range(rows):
                self.comp[i] += [0]*dist_cols

        self.rows = rows
        self.cols = cols
        return self

    def set(self, comp, index=None):
        """
        Set/change the current matrix. If index is not specified, then comp should
        be a list of lists detailing a new matrix. Otherwise, comp should be the
        integer value that goes in the specified index (row, column) tuple.

        :param comp: list of lists to replace matrix entirely, or single value
                     to replace a specific location in matrix
        :param index: optional tuple/list with (row, column) of value to be replaced
        :type comp: list of lists, int
        :type index: tuple, list
        :return: self, after edits are made
        :rtype: Matrix
        """

        if not index:
            assert(isinstance(comp, list))
            self.comp = comp
            self.rows = len(comp)
            if self.rows != 0:
                self.cols = len(comp[0])
            else:
                self.cols = 0
        else:
            assert(isinstance(comp, int))
            self.comp[index[0]][index[1]] = comp
        return self

    def get(self, row=None, col=None):
        """
        User can get rows, columns, the matrix comp list, or specific values
        in Matrix using this function and its optional parameters

        :param row: index of target row
        :param col: index of target col
        :type row: int
        :type col: int
        :return: element at specified row/col, or a row, or a col, or entire Matrix
        :rtype: int, list (row/col), List
        :raises: IndexError if row index or col index invalid
        """
        if row is not None and col is not None:  # value
            if 0 > row >= self.rows and 0 > col >= self.cols:
                raise IndexError("Row or column out of index bounds.")
            return self.comp[row][col]
        elif col is None and row is not None:    # row
            if 0 > row >= self.rows:
                raise IndexError("Row out of index bounds.")
            return self.comp[row]
        elif col is not None:                    # just col
            if 0 > col >= self.cols:
                raise IndexError("Col out of index bounds.")
            return [r[col] for r in self.comp]
        else:                                    # entire matrix
            return self.comp

    def zero(self):
        """
        Zeroes out the current matrix by replacing every element with a zero.

        :return: The current matrix, but updated to be the zero matrix.
        """
        self.comp = [[0]*self.cols for _ in range(self.rows)]
        return self

    def det(self):
        """
        Returns the determinant of an nxn matrix that is at least a 2x2. (recursive)

        :return: the determinant of the current matrix
        :rtype: int, float
        """

        if self.rows != self.cols:
            raise ValueError("Invalid matrix - only N x N matrices supported.")

        # base case -> 2 by 2
        if self.rows == 2 and self.cols == 2:  # ad - bc
            return self.comp[0][0] * self.comp[1][1] - self.comp[0][1] * self.comp[1][0]

        # going along top, along first row (not optimized to find best path)
        top_row = self.comp[0]
        determinant = 0
        for col_i in range(len(top_row)):
            # don't include in same row or column
            new_matrix = self.comp[1:]            # remove top row
            for r in range(len(new_matrix)):      # remove this column from each row
                new_matrix[r] = new_matrix[r][:col_i] + new_matrix[r][col_i + 1:]

            constant = top_row[col_i]
            if col_i % 2 == 1:
                constant *= -1  # every other constant is negative

            determinant += constant * Matrix(new_matrix).det()
        return determinant

    def transpose(self):
        """
        This function will return the transpose of the current matrix. (A -> A^T)
        "First row becomes first column, second row becomes second column, etc."

        :return: Transposed matrix
        :rtype: Matrix
        """
        return Matrix([[self.comp[r][c] for r in range(self.rows)] for c in range(self.cols)])

    @staticmethod
    def identity(n):
        """
        Static method for creating an identity matrix of dimension nxn.

        :param n: dimension of identity matrix
        :type n: int
        :return: identity matrix of size nxn
        :rtype: Matrix
        """
        return Matrix([[1 if i == j else 0 for j in range(n)] for i in range(n)])

    @staticmethod
    def combine(first, second):
        """
        Static method for concatenating two matrices, side by side.
        1 1   *combined  1  0    =      1 1 1 0
        2 2     with*    0  1    =      2 2 0 1

        Warning/Note: Matrices should have the same number of rows, otherwise
                        the minimum amount of rows will be present. (If first
                        has 3 rows and second has 5 rows, combined matrix has 3)

        :param first: first matrix
        :param second: second matrix
        :return: combined matrix, [[row1 + row2], ...]
        :rtype: Matrix
        """
        return Matrix([one + two for one, two in zip(first.comp, second.comp)])

    @staticmethod
    def _clean_matrix(new_matrix):
        """
        Not intended for client use. This method goes through matrix contents
        and reduces each row by the greatest common divisor of that row,
        multiplies row by -1 if leading pivot is negative, and turns floats
        into ints if no reduction occurs. self._clean_matrix or Matrix._clean_matrix

        :param new_matrix: matrix.comp, composition of matrix
        :type new_matrix: list
        :return: "cleaned" matrix comp
        :rtype: list
        """
        cols = len(new_matrix[0])
        for r, row in enumerate(new_matrix):
            gcf = row[0]
            for col in row[1:]:
                gcf = gcd(gcf, col)

            if gcf != 0:
                new_matrix[r] = row = [elem // gcf for elem in row]  # update row for next list comp

            c = 0
            while c < cols and row[c] == 0:
                c += 1
            if c < cols and row[c] < 0:
                new_matrix[r] = row = [-1*elem for elem in row]

            new_matrix[r] = [int(col) if int(col) == col else col for col in row]

        return new_matrix

    @staticmethod
    def _clear_pos(new_matrix, r, c, other_row):
        """
        Helper method for both row echelon functions.

        :param new_matrix: the matrix that will be updated by algorithmically
                            clearing one position in matrix
        :param r:  index of row to be changed
        :param c:   index of col to be changed
        :param other_row: index of other row being using in row operation
        :type r: int
        :type c: int
        :type other_row: list
        :return: matrix composition
        :rtype: list
        """
        above = new_matrix[r][c]
        const = new_matrix[other_row][c]  # row we will use

        # prioritize keeping numbers small / int division
        if const > above != 0 and const % above == 0:
            scale = const // above
            new_matrix[r] = [elem * scale for elem in new_matrix[r]]
        elif above >= const != 0 and const != 0 and above % const == 0:
            scale = above // const
            new_matrix[other_row] = [elem * scale for elem in new_matrix[other_row]]
        else:  # scale both
            new_matrix[r] = [elem * const for elem in new_matrix[r]]
            new_matrix[other_row] = [elem * above for elem in new_matrix[other_row]]
        new_matrix[r] = [other - curr for curr, other in
                         zip(new_matrix[r], new_matrix[other_row])]
        return new_matrix

    def row_echelon(self):
        """
        This function will row reduce the current matrix until it is in row echelon form.
        That is, until there is an upper triangular matrix. I've made a decent amount of
        optimizations in this function, but there definitely many others that could  be made.

        Note: This doesn't change the matrix internally, you will have to assign the
                return value to the your matrix variable if you want to change it.
              There is no guarantee that the matrix returned will contain only integers, may
                have floats.

        :return: row echelon form of current matrix
        :rtype: Matrix
        :return:
        """

        # adjust matrix so rows are in proper descending order / putting any pre-made pivots in place
        new_matrix = sorted(self.comp, reverse=True)
        pivot_row = 0

        for c in range(self.cols):
            new_matrix = sorted(new_matrix, reverse=True)   # swap any out of place rows
            if pivot_row < self.rows and new_matrix[pivot_row][c] != 0:
                for r in range(pivot_row+1, self.rows):     # use pivot row to clear other rows
                    while new_matrix[r][c] != 0:

                        new_matrix = self._clear_pos(new_matrix, r, c, pivot_row)

                pivot_row += 1
        new_matrix = self._clean_matrix(new_matrix)
        return Matrix(sorted(new_matrix, reverse=True))     # swap any out of place rows

    def row_reduce(self):
        """
        This function will row reduce the current matrix until it is in reduced row
        echelon form (RREF). The transpose of a matrix has the same RREF as original.

        Note: This doesn't change the matrix internally, you will have to assign the
                return value to the your matrix variable if you want to change it.

        :return: reduced row echelon form of current matrix
        :rtype: Matrix
        """

        new_matrix = self.row_echelon().comp     # get in row echelon form first

        pivots = {}     # store pivot indexes key-value for use later

        # store pivots as col : row pairs
        for r, row in enumerate(new_matrix):
            # identify pivot
            i = 0
            while i < self.cols and row[i] == 0:
                i += 1
            if i < self.cols:
                pivots[i] = r

        # apply only 0s above pivot (bottom part is done since already in row echelon form)
        offset = 0     # how far ahead the first pivot is (ex. may be zero cols before first pivot)
        for c in range(self.cols):
            if c in pivots:
                pivot_row = pivots[c]       # row the pivot is in
                for r in range(pivot_row):  # top part, don't loop past location of pivot
                    while new_matrix[r][c] != 0:    # stay in same column and fix parts above pivot
                        other_row = c-offset  # when no offset, col c can be cleared using row c since there are c zeros

                        new_matrix = self._clear_pos(new_matrix, r, c, other_row)
            else:
                offset += 1

        new_matrix = self._clean_matrix(new_matrix)  # this function also changes floats to perfect ints based on gcd

        # now, apply "each pivot is 1" rule, floats inevitable, but preserve as much ints as possible
        for r, row in enumerate(new_matrix):
            # identify pivot
            i = 0
            while i < self.cols and row[i] == 0:
                i += 1
            # divide row by proper amount to get a 1 on pivot
            if i < self.cols:
                pivot = row[i]
                new_matrix[r] = [elem // pivot if elem % pivot == 0 else elem / pivot for elem in row]
        return Matrix(sorted(new_matrix, reverse=True))  # ensure ordering is still valid

    def inverse(self):
        """
        Gets the inverse A^-1 of the current matrix A.

        :return: inverse matrix of current matrix, or None if not invertible (singular)
        :rtype: Matrix
        :raises: value error if current matrix is not nxn
        """
        n = self.cols
        identity = Matrix.identity(n)
        if self.rows != n:
            raise ValueError("Need an nxn matrix to calculate inverse.")
        # create combined matrix
        with_identity = Matrix.combine(self, identity).row_reduce()
        # if left side is identity, then right side is inverse
        if Matrix([row[:n] for row in with_identity.comp]) != identity:
            return None     # no inverse, singular
        else:
            return Matrix([row[-n:] for row in with_identity.comp])

    def __add__(self, other):
        """
        Adds two matrices and returns a matrix with the respective components
        added together as expected.

        :param other: the other matrix to be added to current instance matrix
        :type other: Matrix
        :return: a matrix with the resulting added components
        :rtype: Matrix
        :raises: ValueError when matrices do not have same dimensions
        """
        new_comp = []
        if self.rows == other.rows and self.cols == other.cols:
            for x, y in zip(self.comp, other.comp):
                new_comp.append([a + b for a, b in zip(x, y)])  # adding done in list comprehension
            return Matrix(new_comp)
        else:
            raise ValueError("Size mismatch, both matrices must have the same number of rows and columns.")

    def __sub__(self, other):
        """
       Subtracting two matrices returns a matrix with the respective components
       subtracted. "current - other" is formatting.

       :param other: the other matrix which is subtracting from the current matrix
       :type other: Matrix
       :return: a matrix with the resulting subtracted components
       :rtype: Matrix
       :raises: ValueError when matrices do not have same dimensions
       """
        new_comp = []
        if self.rows == other.rows and self.cols == other.cols:
            for x, y in zip(self.comp, other.comp):
                new_comp.append([a - b for a, b in zip(x, y)])  # subtracting done in list comprehension
            return Matrix(new_comp)
        else:
            raise ValueError("Size mismatch, both matrices must have the same number of rows and columns.")

    def __mul__(self, other):
        """
        Multiplies the two matrices together; aka Matrix Multiplication.
        Matrix-Vector product is also possible using the Vector class, though
        this method works for a mx1 matrix as well. Also configured to work with
        normal application of multiplying a scalar to a matrix.

        Notes: Approach is to take the dot product of each row of current matrix
                with each column of other matrix/vector. Since you typically write
                "Ax" where A is the matrix and x is the vector, this syntax should
                be adhered to when attempting matrix multiplication with these classes.

        :param other: the other matrix or vector, could also be an int or float for scaling
        :type other: Matrix, int, float
        :return: the resulting matrix
        :rtype: Matrix
        :raises: ValueError when there's a matrix multiplication size mismatch ([mxn]*[nxp]=[mxp])
        """
        new_matrix = []
        if isinstance(other, int) or isinstance(other, float):
            for row in self.comp:
                new_matrix.append([elem * other for elem in row])
            return Matrix(new_matrix)
        elif self.cols == other.rows:    # [m x n] * [n x p] = [m x p] i.e. [self.rows x other.cols] matrix
            other_cols = []
            for i in range(other.cols):  # extract columns from rows
                other_cols.append([row[i] if isinstance(other, Matrix) else row for row in other.comp])
            for row_me in self.comp:
                new_row = []
                for col_other in other_cols:
                    new_row.append(Vector(row_me) * Vector(col_other))   # Dot product of vectors
                new_matrix.append(new_row)
            return Vector([row[0] for row in new_matrix]) if other.cols == 1 else Matrix(new_matrix)
        else:
            raise ValueError("Size mismatch; [m x n] * [n x p] = [m x p] matrix")

    def __eq__(self, other):
        """
        If two matrices have the same components, then they are equal. If the
        lists are not the same length, will always be False with no error thrown.
        Have to compare each component due to necessity of using math.isclose()
        on floats in order to deal with floating point errors.

        :param other: other matrix being tested for equality
        :type other: Matrix
        :return: True or False based on equality
        :rtype: bool
        """
        if self.rows != other.rows or self.cols != other.cols:
            return False
        for my_row, other_row in zip(self, other):
            for my_val, other_val in zip(my_row, other_row):
                if not isclose(my_val, other_val):
                    return False

        return self.comp == other.comp  # compares lists

    def __pow__(self, power, modulo=None):
        """
        Allows you to raise a matrix to a power, that is, each of the
        components of the current matrix is raised to a power. Can use
        power 0 to fill the current matrix with all 1s.

        :param power: value to raise each component to
        :param modulo: optional parameter that applies the modulus operator to each result
        :type power: int, float
        :type modulo: int, float
        :return: a matrix containing the appropriately scaled components
        :rtype: Matrix
        """
        new_comp = []
        for row in self.comp:
            new_row = []
            for elem in row:
                if modulo:
                    elem = elem % modulo
                new_row.append(pow(elem, power))
            new_comp.append(new_row)
        return Matrix(new_comp)

    def __str__(self):
        """
        String representation of matrix is each row separated by new line
        characters. This is done so that when printed it resembles a normal
        matrix as closely as possible.

        :return: string representation of current matrix
        :rtype: str
        """

        # joins each row of matrix with a new line character and a space,
        # floats are converted to visual fractions, need to get rid of quotes around them
        return "[" + '\n '\
            .join([str([str(Fraction(elem).limit_denominator()) if isinstance(elem, float) else elem for elem in row])
                  .replace('\'', '') for row in self.comp])\
            + "]"

    def __len__(self):
        """
        :return: returns tuple formatted as (row, col)
        :rtype: tuple
        """
        return self.rows, self.cols

    def __getitem__(self, index):
        """
        Allows user to access internal self.comp without doing
        my_matrix.comp[i][j] and instead doing my_matrix[i][j]

        Note: the first [] calls this function, which returns row,
                that row is a list, which supports [] in the same way
                that this function does.

        :param index: index of row
        :type index: int
        :return: list or value for row or row+col value
        :rtype: list, value
        """
        return self.comp[index]

    def __setitem__(self, key, value):
        """
        Allows the user to set a value using brackets.

        Note: behavior undefined if user attempts to set a row.

        :param key: index of row to be changed
        :param value: value to be set
        :type key: int
        :type value: int
        :return: no return
        """

        self.set(value, key)

if __name__ == "__main__":
    test()


#=====================================================#
#
#    File:   zeta_zeroes_plot.py
#    Author: <NAME>
#    Date:   May 2021
#    Description:
#        Shows the position of the zeta zeroes along the critical strip
#
#=============================================================================#

from time import time
import numpy as np
from cmath import pi, sin
from scipy.special import gamma
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# returns the value of \frac{\zeta(s)}{\eta(s)}
# = \frac{ 2^{iy}}{2^{iy}-2^{1-x}}
def zeta_over_eta(x, y):
    A = -2**(1 - x)
    B = 2**(y * 1j)
    zeta_over_eta = B / (A + B)
    return zeta_over_eta

# approximates the eta function where
# r is the real part of the input (r\equiv\Re(s))
# and t is the imaginary part of the input (t\equiv\Im(s))
# \eta(r+it)=\sum_{n=1}^\infty\frac{(-1)^{n-1}}{n^r}[cos(y\phi_n)-isin(y\phi_n)]
def eta(x, y, limit=5*10**3):
    n = np.arange(1, limit)
    coefficient = (-1)**(n-1) / n**x
    phi_y = np.log(n) * y
    real_parts = coefficient * np.cos(phi_y)
    imaginary_parts = coefficient * np.sin(phi_y)

    real_summation = real_parts.sum()
    imaginary_summation = imaginary_parts.sum()
    eta = real_summation - imaginary_summation * 1j

    return eta

# returns value of \zeta(s)
# where s \equiv r + it ,r,t \in \mathbb{C}
# by computing \frac{\zeta(s)\eta(s)}
# for x = 1/2
def zeta(x, y=0):
    result = zeta_over_eta(x, y) * eta(x, y)
    return result

def is_sign_change(previous, number):
    # print(f'prev: {previous}, number {number}')
    if previous == 0:
        return False
    previous_sign = previous / abs(previous)
    number_sign = number / abs(number)
    return previous_sign != number_sign

def animate(i, x_vals, y_vals):
    PRECISION = 100
    previous_y_input = (i-1)/PRECISION
    current_y_input = i/PRECISION
    previous_zeta = zeta(1/2, previous_y_input)
    current_zeta = zeta(1/2, current_y_input)
    # print(current_zeta.real, current_zeta.imag)
    # print(y_vals)
    if is_sign_change(previous_zeta.real, current_zeta.real) and is_sign_change(previous_zeta.imag, current_zeta.imag):
        x_vals.append(1/2)
        y_vals.append(current_y_input)
        print(current_y_input)

    plt.cla()

    plt.scatter(x_vals, y_vals)

    plt.xlabel('Re')
    plt.ylabel('Im')
    plt.tight_layout()

def plot():
    plt.style.use('dark_background')

    x_vals = []
    y_vals = []

    animation = FuncAnimation(plt.gcf(), animate, fargs=(x_vals, y_vals,), interval=0)
    plt.xlabel('Re')
    plt.ylabel('Im')
    plt.tight_layout()
    plt.show()

if __name__=='__main__':
    start = time()
    plot()
    print(f'--- {time()-start} seconds ---')


import numpy as np
from scipy.ndimage import uniform_filter, gaussian_filter
from numpy.lib.arraypad import _as_pairs

_integer_types = (np.byte, np.ubyte,  # 8 bits
                      np.short, np.ushort,  # 16 bits
                      np.intc, np.uintc,  # 16 or 32 or 64 bits
                      np.int_, np.uint,  # 32 or 64 bits
                      np.longlong, np.ulonglong)  # 64 bits

_integer_ranges = {t: (np.iinfo(t).min, np.iinfo(t).max)
                   for t in _integer_types}

dtype_range = {np.bool_: (False, True),
               np.bool8: (False, True),
               np.float16: (-1, 1),
               np.float32: (-1, 1),
               np.float64: (-1, 1)}

dtype_range.update(_integer_ranges)

def _assert_compatible(im1, im2):
    """Raise an error if the shape and dtype do not match."""
    if not im1.shape == im2.shape:
        raise ValueError('Input images must have the same dimensions.')
    return

def _as_floats(im1, im2):
    """Promote im1, im2 to nearest appropriate floating point precision."""
    float_type = np.result_type(im1.dtype, im2.dtype, np.float32)
    im1 = np.asarray(im1, dtype=float_type)
    im2 = np.asarray(im2, dtype=float_type)
    return im1, im2

def crop(ar, crop_width, copy=False, order='K'):
    """Crop array `ar` by `crop_width` along each dimension.
    Parameters
    ----------
    ar : array-like of rank N
        Input array.
    crop_width : {sequence, int}
        Number of values to remove from the edges of each axis.
        ``((before_1, after_1),`` ... ``(before_N, after_N))`` specifies
        unique crop widths at the start and end of each axis.
        ``((before, after),)`` specifies a fixed start and end crop
        for every axis.
        ``(n,)`` or ``n`` for integer ``n`` is a shortcut for
        before = after = ``n`` for all axes.
    copy : bool, optional
        If `True`, ensure the returned array is a contiguous copy. Normally,
        a crop operation will return a discontiguous view of the underlying
        input array.
    order : {'C', 'F', 'A', 'K'}, optional
        If ``copy==True``, control the memory layout of the copy. See
        ``np.copy``.
    Returns
    -------
    cropped : array
        The cropped array. If ``copy=False`` (default), this is a sliced
        view of the input array.
    """
    ar = np.array(ar, copy=False)
    crops = _as_pairs(crop_width, ar.ndim, as_index=True)
    slices = tuple(slice(a, ar.shape[i] - b)
                   for i, (a, b) in enumerate(crops))
    if copy:
        cropped = np.array(ar[slices], order=order, copy=True)
    else:
        cropped = ar[slices]
    return cropped

def MSE(im1, im2):
    """Compute the mean-squared error between two images.
    Parameters
    ----------
    im1, im2 : ndarray
        Image.  Any dimensionality.
    Returns
    -------
    mse : float
        The mean-squared error (MSE) metric.
    """
    _assert_compatible(im1, im2)
    im1, im2 = _as_floats(im1, im2)
    return np.mean(np.square(im1 - im2), dtype=np.float64)

def PSNR(ref_img, pred_img):
    """
    Compute average PSNR score for a batch of pairs of images

    :param ref_img: ndarray of shape (batch_size, width, height, channels) with target noise-free image
    :param pred_img: ndarray of shape (batch_size, width, height, channels) with predicted denoised image
    :return: float, averaged over batch PSNR score
    """

    _assert_compatible(ref_img, pred_img)

    dmin, dmax = dtype_range[ref_img.dtype.type]
    true_min, true_max = np.min(ref_img), np.max(ref_img)
    if true_max > dmax or true_min < dmin:
        raise ValueError(
            "im_true has intensity values outside the range expected for "
            "its data type.  Please manually specify the data_range")
    if true_min >= 0:
        # most common case (255 for uint8, 1 for float)
        data_range = dmax
    else:
        data_range = dmax - dmin

    ref_img, pred_img = _as_floats(ref_img, pred_img)

    err = MSE(ref_img, pred_img)
    return 10 * np.log10((data_range ** 2) / err)

def SSIM(X, Y, win_size=None, gradient=False,
                 data_range=None, multichannel=False, gaussian_weights=False,
                 full=False, **kwargs):
    """Compute the mean structural similarity index between two images.
    Parameters
    ----------
    X, Y : ndarray
        Image. Any dimensionality.
    win_size : int or None
        The side-length of the sliding window used in comparison. Must be an
        odd value. If `gaussian_weights` is True, this is ignored and the
        window size will depend on `sigma`.
    gradient : bool, optional
        If True, also return the gradient with respect to Y.
    data_range : float, optional
        The data range of the input image (distance between minimum and
        maximum possible values). By default, this is estimated from the image
        data-type.
    multichannel : bool, optional
        If True, treat the last dimension of the array as channels. Similarity
        calculations are done independently for each channel then averaged.
    gaussian_weights : bool, optional
        If True, each patch has its mean and variance spatially weighted by a
        normalized Gaussian kernel of width sigma=1.5.
    full : bool, optional
        If True, also return the full structural similarity image.
    Other Parameters
    ----------------
    use_sample_covariance : bool
        If True, normalize covariances by N-1 rather than, N where N is the
        number of pixels within the sliding window.
    K1 : float
        Algorithm parameter, K1 (small constant, see [1]_).
    K2 : float
        Algorithm parameter, K2 (small constant, see [1]_).
    sigma : float
        Standard deviation for the Gaussian when `gaussian_weights` is True.
    Returns
    -------
    mssim : float
        The mean structural similarity over the image.
    grad : ndarray
        The gradient of the structural similarity index between X and Y [2]_.
        This is only returned if `gradient` is set to True.
    S : ndarray
        The full SSIM image.  This is only returned if `full` is set to True.
    Notes
    -----
    To match the implementation of Wang et. al. [1]_, set `gaussian_weights`
    to True, `sigma` to 1.5, and `use_sample_covariance` to False.
    References
    ----------
    .. [1] <NAME>., <NAME>., <NAME>., & <NAME>.
       (2004). Image quality assessment: From error visibility to
       structural similarity. IEEE Transactions on Image Processing,
       13, 600-612.
       https://ece.uwaterloo.ca/~z70wang/publications/ssim.pdf,
       :DOI:`10.1109/TIP.2003.819861`
    .. [2] <NAME>. (2009). Exact global histogram specification
       optimized for structural similarity. Optical Review, 16, 613-621.
       :arXiv:`0901.0065`
       :DOI:`10.1007/s10043-009-0119-z`
    """
    if not X.shape == Y.shape:
        raise ValueError('Input images must have the same dimensions.')

    if multichannel:
        # loop over channels
        args = dict(win_size=win_size,
                    gradient=gradient,
                    data_range=data_range,
                    multichannel=False,
                    gaussian_weights=gaussian_weights,
                    full=full)
        args.update(kwargs)
        nch = X.shape[-1]
        mssim = np.empty(nch)
        if gradient:
            G = np.empty(X.shape)
        if full:
            S = np.empty(X.shape)
        for ch in range(nch):
            ch_result = SSIM(X[..., ch], Y[..., ch], **args)
            if gradient and full:
                mssim[..., ch], G[..., ch], S[..., ch] = ch_result
            elif gradient:
                mssim[..., ch], G[..., ch] = ch_result
            elif full:
                mssim[..., ch], S[..., ch] = ch_result
            else:
                mssim[..., ch] = ch_result
        mssim = mssim.mean()
        if gradient and full:
            return mssim, G, S
        elif gradient:
            return mssim, G
        elif full:
            return mssim, S
        else:
            return mssim

    K1 = kwargs.pop('K1', 0.01)
    K2 = kwargs.pop('K2', 0.03)
    sigma = kwargs.pop('sigma', 1.5)
    if K1 < 0:
        raise ValueError("K1 must be positive")
    if K2 < 0:
        raise ValueError("K2 must be positive")
    if sigma < 0:
        raise ValueError("sigma must be positive")
    use_sample_covariance = kwargs.pop('use_sample_covariance', True)

    if gaussian_weights:
        # Set to give an 11-tap filter with the default sigma of 1.5 to match
        # Wang et. al. 2004.
        truncate = 3.5

    if win_size is None:
        if gaussian_weights:
            # set win_size used by crop to match the filter size
            r = int(truncate * sigma + 0.5)  # radius as in ndimage
            win_size = 2 * r + 1
        else:
            win_size = 7   # backwards compatibility

    if np.any((np.asarray(X.shape) - win_size) < 0):
        raise ValueError(
            "win_size exceeds image extent.  If the input is a multichannel "
            "(color) image, set multichannel=True.")

    if not (win_size % 2 == 1):
        raise ValueError('Window size must be odd.')

    if data_range is None:
        assert X.dtype == Y.dtype, \
            "Inputs have mismatched dtype. Setting data_range based on X.dtype."
        dmin, dmax = dtype_range[X.dtype.type]
        data_range = dmax - dmin

    ndim = X.ndim

    if gaussian_weights:
        filter_func = gaussian_filter
        filter_args = {'sigma': sigma, 'truncate': truncate}
    else:
        filter_func = uniform_filter
        filter_args = {'size': win_size}

    # ndimage filters need floating point data
    X = X.astype(np.float64)
    Y = Y.astype(np.float64)

    NP = win_size ** ndim

    # filter has already normalized by NP
    if use_sample_covariance:
        cov_norm = NP / (NP - 1)  # sample covariance
    else:
        cov_norm = 1.0  # population covariance to match Wang et. al. 2004

    # compute (weighted) means
    ux = filter_func(X, **filter_args)
    uy = filter_func(Y, **filter_args)

    # compute (weighted) variances and covariances
    uxx = filter_func(X * X, **filter_args)
    uyy = filter_func(Y * Y, **filter_args)
    uxy = filter_func(X * Y, **filter_args)
    vx = cov_norm * (uxx - ux * ux)
    vy = cov_norm * (uyy - uy * uy)
    vxy = cov_norm * (uxy - ux * uy)

    R = data_range
    C1 = (K1 * R) ** 2
    C2 = (K2 * R) ** 2

    A1, A2, B1, B2 = ((2 * ux * uy + C1,
                       2 * vxy + C2,
                       ux ** 2 + uy ** 2 + C1,
                       vx + vy + C2))
    D = B1 * B2
    S = (A1 * A2) / D

    # to avoid edge effects will ignore filter radius strip around edges
    pad = (win_size - 1) // 2

    # compute (weighted) mean of ssim
    mssim = crop(S, pad).mean()

    if gradient:
        # The following is Eqs. 7-8 of Avanaki 2009.
        grad = filter_func(A1 / D, **filter_args) * X
        grad += filter_func(-S / B2, **filter_args) * Y
        grad += filter_func((ux * (A2 - A1) - uy * (B2 - B1) * S) / D,
                            **filter_args)
        grad *= (2 / X.size)

        if full:
            return mssim, grad, S
        else:
            return mssim, grad
    else:
        if full:
            return mssim, S
        else:
            return mssim


import matplotlib
from hydroDL import kPath, utils
from hydroDL.app import waterQuality
from hydroDL.master import basins
from hydroDL.data import usgs, gageII, gridMET, ntn
from hydroDL.master import slurm
from hydroDL.post import axplot, figplot
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import scipy
import os
import json

dataName = 'rbWN5'
wqData = waterQuality.DataModelWQ(dataName)
siteNoLst = wqData.siteNoLst
label = 'QFP_C'
ep = 500
reTest = False
codeLst = sorted(usgs.newC)
nSite = len(siteNoLst)

hsLst = [16, 32, 64, 128, 256, 512]
corrMat = np.full([nSite, len(codeLst), len(hsLst)], np.nan)
rmseMat = np.full([nSite, len(codeLst), len(hsLst)], np.nan)

for k, hs in enumerate(hsLst):
    code = 'comb'
    trainSet = '{}-B10'.format('comb')
    testSet = '{}-A10'.format('comb')
    outName = '{}-{}-{}-{}-hs{}'.format(dataName, code, label, trainSet, hs)
    master = basins.loadMaster(outName)
    yP, ycP = basins.testModel(
        outName, testSet, wqData=wqData, ep=ep, reTest=reTest)
    ind = wqData.subset[testSet]
    info = wqData.info.iloc[ind].reset_index()
    siteNoTemp = info['siteNo'].unique()
    for iCode, code in enumerate(codeLst):
        ic = wqData.varC.index(code)
        if len(wqData.c.shape) == 3:
            p = yP[-1, :, master['varY'].index(code)]
            o = wqData.c[-1, ind, ic]
        elif len(wqData.c.shape) == 2:
            p = ycP[:, master['varYC'].index(code)]
            o = wqData.c[ind, ic]
        for siteNo in siteNoTemp:
            iS = siteNoLst.index(siteNo)
            indS = info[info['siteNo'] == siteNo].index.values
            rmse, corr = utils.stat.calErr(p[indS], o[indS])
            corrMat[iS, iCode, k] = corr
            rmseMat[iS, iCode, k] = rmse

# plot box
dirSel = os.path.join(kPath.dirData, 'USGS', 'inventory', 'siteSel')
with open(os.path.join(dirSel, 'dictRB_Y30N5.json')) as f:
    dictSite = json.load(f)
labLst1 = [usgs.codePdf.loc[code]['shortName'] +
           '\n'+code for code in codeLst]
labLst2 = hsLst
dataBox = list()
for k, code in enumerate(codeLst):
    siteNoCode = dictSite[code]
    indS = [siteNoLst.index(siteNo) for siteNo in siteNoCode]
    temp = list()
    for i in range(len(hsLst)):
        temp.append(corrMat[indS, k, i])
    dataBox.append(temp)
fig = figplot.boxPlot(dataBox, label1=labLst1, widths=0.5,
                      label2=labLst2, figsize=(12, 4), yRange=[0, 1])
fig.show()