# Modelling spread of disease
# https://www.maa.org/press/periodicals/loci/joma/the-sir-model-for-spread-of-disease-the-differential-equation-model
import random
import numpy as np
import matplotlib.pyplot as plt
from random import uniform

def w_choice(seq):
    """
    Takes an input on the form of [('A', 30), ('B', 40), ('C', 30)] and returns
    one of the items based on the probality from the second value in each tuple.
    Eg. B will get chosen 40% of the time while A and C will get chosen 30% of
    the time.
    """
    total_prob = sum(item[1] for item in seq)
    chosen = random.uniform(0, total_prob)
    cumulative = 0
    for item, probality in seq:
        cumulative += probality
        if cumulative > chosen:
            return item

class Disease:
    """
    Uses the SIR model to model the spread of a certain disease.

    At each time t, s(t) + i(t) + r(t) = 1
    """

    def __init__(self, N, b, k, susceptible, infected, recovered):
        """
        params:
            N = The total population
            b = The average _FIXED_ number of contacts per day per person.
            k = The average _FIXED_ fraction of infected individuals that will recover per day.
            susceptible = The initial amount of susceptible individuals
            infected = The initial amount of infected individuals
            recovered = The initial amount of recovered individuals
        """
        self.time = 0
        self._N = N
        self._b = b
        self._k = k
        self._susceptible = susceptible
        self._infected = infected
        self._recovered = recovered

    def N(self):
        return self._N

    def vaccinate(self, amount):
        """
        Vaccinates the specified amount of the population.
        """
        self.susceptible(self.susceptible() - amount)

    def susceptible(self, v=None):
        """
        The number of susceptible individuals.
        S = S(t)

        No one is added to this group due to births and immigrations begin ignored.
        An individual will leave this group once infected.
        """
        if v is None:
            return self._susceptible
        self._susceptible = v

    def susceptible_dt(self):
        """
        The rate of change for susceptible individuals.
        dS/dt = -b*s(t)*I(t)
        """
        return -1 * self._b * self.susceptible_fraction() * self.infected()

    def infected(self, v=None):
        """
        The number of infected individuals.
        I = I(t)

        Each infected individual infects b * s(t) new individuals.
        """
        if v is None:
            return self._infected
        self._infected = v

    def recovered(self, v=None):
        """
        The number of recovered individuals.
        R = R(t)
        """
        if v is None:
            return self._recovered
        self._recovered = v

    def susceptible_fraction(self):
        """
        The susceptible fraction of the population.

        s(t)
        """
        return self.fraction(self.susceptible())

    def susceptible_fraction_dt(self):
        """
        The rate of change for the susceptible fraction.
        ds/dt = -b*s(t)*i(t)
        """
        return -1 * self._b * self.susceptible_fraction() * self.infected_fraction()

    def recovered_fraction(self):
        """
        The recovered fraction of the population.

        r(t)
        """
        return self.fraction(self.recovered())

    def recovered_fraction_dt(self):
        """
        The rate of change for the recovered fraction of the population.
        dr/dt = k*i(t)
        """
        return self._k * self.infected_fraction()

    def infected_fraction(self):
        """
        The infected fraction of the population.

        i(t)
        """
        return self.fraction(self.infected())

    def infected_fraction_dt(self):
        """
        The rate of change for the infected fraction of the population.
        di/dt = b * s(t) * i(t) - k * i(t)
        di/dt = -ds/dt - dr/dt
        """
        return -1 * (self.susceptible_fraction_dt() + self.recovered_fraction_dt())

    def fraction(self, nr):
        """
        Returns the fraction of the total population.
        nr / N
        """
        return nr / self._N

    def number(self, fraction):
        """
        Returns the number of the total population
        fraction * N
        """
        return fraction * self._N

    def step(self):
        """
        Steps forward one unit in time.
        """
        i_dt = self.number(self.infected_fraction_dt())
        r_dt = self.number(self.recovered_fraction_dt())
        s_dt = self.number(self.susceptible_fraction_dt())
        self.susceptible(self.susceptible() + s_dt)
        self.recovered(self.recovered() + r_dt)
        self.infected(self.infected() + i_dt)

class City:
    """
    City.
    """

    def __init__(self, name, position, N, b, k, susceptible, infected, recovered):
        """
        params:
            name = The name of the city.
            population = The population of the city
            position = The position of the city
        """
        self._position = position
        self._name = name
        self._disease = Disease(N=N, b=b, k=k, susceptible=susceptible, infected=infected, recovered=recovered)
        self.Y_susceptible = []
        self.Y_infected = []
        self.Y_recovered = []

    def population(self):
        """
        Returns the population of the city.
        """
        return self.disease().N()

    def position(self):
        """
        Returns the position of the city.
        """
        return self.position

    def disease(self):
        return self._disease

    def vaccinate(self, n):
        """
        Vaccinates n people of the population.
        """
        self.disease().vaccinate(n)

    def name(self):
        return self._name

    def step(self, t, _print=False):
        """
        Steps forward one step in time.
        params:
            t = The current timestamp (Used when printing).
            _print = If output should be printed to the console
        """
        d = self.disease()
        self.log(d, t, _print)
        d.step()

    def log(self, d, t, _print):
        """
        Logs the current state of the disease for later plotting and/or dumping
        to the console.
        """
        self.Y_susceptible.append(d.susceptible_fraction())
        self.Y_infected.append(d.infected_fraction())
        self.Y_recovered.append(d.recovered_fraction())
        if _print:
            print("S("+str(t)+")", "%0.2f" % d.susceptible(), "\ts("+str(t)+")=", "%0.2f" % d.susceptible_fraction())
            print("I("+str(t)+")", "%0.2f" % d.infected(), "\ti("+str(t)+")=", "%0.2f" % d.infected_fraction())
            print("R("+str(t)+")", "%0.2f" % d.recovered(), "\tr("+str(t)+")=", "%0.2f" % d.recovered_fraction())
            print()

    def plot(self):
        """
        Plots the progress of the diesease for this City.
        """
        _ys, = plt.plot(self.Y_susceptible, label='s(t)')
        _yi, = plt.plot(self.Y_infected, label="i(t)")
        _yr, = plt.plot(self.Y_recovered, label="r(t)")
        plt.legend(handles=[_ys, _yi, _yr])
        plt.title(self.name())
        plt.show()

    def plot_get_infected(self):
        _yi, = plt.plot(self.Y_infected, label="i(t) - "+str(self.name()))
        return _yi

def example_graph():
    """
    Example graph that is the same as the one found at the link:
    https://www.maa.org/press/periodicals/loci/joma/the-sir-model-for-spread-of-disease-the-differential-equation-model
    """
    c = City(name="Sweden", position=(0,0), N=7900000, b=1/2, k=1/3, susceptible=7900000, infected=10, recovered=0)

    for t in range(140):
        c.step(t)

    c.plot()

def example_cities():
    b = 0.99
    k = 1/10
    cities = [
        City(name="Stockholm", position=(0,0),      N=1400000, b=b, k=k, susceptible=1400000, infected=1, recovered=0),
        City(name="Göteborg",  position=(-50, -50), N=1000000, b=b, k=k, susceptible=1000000, infected=0,  recovered=0),
        City(name="Umeå",      position=(-4, 100),  N=100000,  b=b, k=k, susceptible=100000,  infected=0,  recovered=0),
        City(name="Örebro",    position=(-10, 5),   N=350000,  b=b, k=k, susceptible=350000,  infected=0,  recovered=0)
    ]

    print(cities)
    for c in cities:
        print(cities[:].remove(c))

    for t in range(140):
        for c in cities:
            c.step(t)

    plt.legend(handles=[c.plot_get_infected() for c in cities])
    plt.show()

example_cities()


import numpy as np

class TensorDP:

    def __init__(self,
                 gamma=1.0,
                 error_tol=1e-5):
        self.gamma = gamma
        self.error_tol = error_tol

        # Following attributes will be set after call "set_env()"

        self.env = None  # environment
        self.policy = None  # policy
        self.ns = None  # Num. states
        self.na = None  # Num. actions
        self.P = None  # Transition tensor
        self.R = None  # Reward tensor

    def set_env(self, env, policy=None):
        self.env = env
        if policy is None:
            self.policy = np.ones([env.nS, env.nA]) / env.nA

        self.ns = env.nS
        self.na = env.nA
        self.P = env.P_tensor  # Rank 3 tensor [num. actions x num. states x num. states]
        self.R = env.R_tensor  # Rank 2 tensor [num. actions x num. states]

        print("Tensor DP agent initialized")
        print("Environment spec:  Num. state = {} | Num. actions = {} ".format(env.nS, env.nA))

    def reset_policy(self):
        self.policy = np.ones([self.ns, self.na]) / self.na

    def set_policy(self, policy):
        assert self.policy.shape == policy.shape
        self.policy = policy

    def get_r_pi(self, policy):
        r_pi = (policy * self.R).sum(axis=-1)  # [num. states x 1]
        return r_pi

    def get_p_pi(self, policy):
        p_pi = np.einsum("na,anm->nm", policy, self.P)  # [num. states x num. states]
        return p_pi

    def policy_evaluation(self, policy=None, v_init=None):
        """
        :param policy: policy to evaluate (optional)
        :param v_init: initial value 'guesstimation' (optional)
        :param steps: steps of bellman expectation backup (optional)
        if none, repeat the backup until converge.

        :return: v_pi: value function of the input policy
        """
        if policy is None:
            policy = self.policy

        r_pi = self.get_r_pi(policy)  # [num. states x 1]
        p_pi = self.get_p_pi(policy)  # [num. states x num. states]

        if v_init is None:
            v_old = np.zeros(self.ns)
        else:
            v_old = v_init

        while True:
            # perform bellman expectation back
            v_new = r_pi + self.gamma * np.matmul(p_pi, v_old)

            # check convergence
            bellman_error = np.linalg.norm(v_new - v_old)
            if bellman_error <= self.error_tol:
                break
            else:
                v_old = v_new

        return v_new

    def policy_improvement(self, policy=None, v_pi=None):
        if policy is None:
            policy = self.policy

        if v_pi is None:
            v_pi = self.policy_evaluation(policy)

        # Compute Q_pi(s,a) from V_pi(s)
        r_pi = self.get_r_pi(policy)
        q_pi = r_pi + self.P.dot(v_pi)  # q_pi = [num.action x num states]

        # Greedy improvement
        policy_improved = np.zeros_like(policy)
        policy_improved[np.arange(q_pi.shape[1]), q_pi.argmax(axis=0)] = 1
        return policy_improved

    def policy_iteration(self, policy=None):
        if policy is None:
            pi_old = self.policy
        else:
            pi_old = policy

        info = dict()
        info['v'] = list()
        info['pi'] = list()
        info['converge'] = None

        steps = 0
        converged = False
        while True:
            v_old = self.policy_evaluation(pi_old)
            pi_improved = self.policy_improvement(pi_old, v_old)
            steps += 1

            info['v'].append(v_old)
            info['pi'].append(pi_old)

            # check convergence
            policy_gap = np.linalg.norm(pi_improved - pi_old)

            if policy_gap <= self.error_tol:
                if not converged:  # record the first moment of within error tolerance.
                    info['converge'] = steps
                break
            else:
                pi_old = pi_improved
        return info

    def value_iteration(self, v_init=None, compute_pi=False):
        """
        :param v_init: (np.array) initial value 'guesstimation' (optional)
        :param compute_pi: (bool) compute policy during VI
        :return: v_opt: the optimal value function
        """

        if v_init is not None:
            v_old = v_init
        else:
            v_old = np.zeros(self.ns)

        info = dict()
        info['v'] = list()
        info['pi'] = list()
        info['converge'] = None

        steps = 0
        converged = False

        while True:
            # Bellman optimality backup
            v_improved = (self.R.T + self.gamma * self.P.dot(v_old)).max(axis=0)
            info['v'].append(v_improved)

            if compute_pi:
                # compute policy from v
                # 1) Compute v -> q
                q_pi = (self.R.T + self.gamma * self.P.dot(v_improved))

                # 2) Construct greedy policy
                pi = np.zeros_like(self.policy)
                pi[np.arange(q_pi.shape[1]), q_pi.argmax(axis=0)] = 1
                info['pi'].append(pi)

            steps += 1

            # check convergence
            policy_gap = np.linalg.norm(v_improved - v_old)

            if policy_gap <= self.error_tol:
                if not converged:  # record the first moment of within error tolerance.
                    info['converge'] = steps
                break
            else:
                v_old = v_improved
        return info


"""

http://scipp.ucsc.edu/~haber/ph116A/Rotation2.pdf
https://math.stackexchange.com/questions/3501114/understanding-a-consequence-of-eulers-rotation-theorem
https://en.wikipedia.org/wiki/Euler%27s_rotation_theorem
https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions
https://en.wikipedia.org/wiki/Euler_angles
https://en.wikipedia.org/wiki/Euclidean_vector
https://en.wikipedia.org/wiki/Direction_cosine
https://d3cw3dd2w32x2b.cloudfront.net/wp-content/uploads/2012/07/euler-angles1.pdf
https://www.mathworks.com/help/aerotbx/ug/angle2dcm.html
https://en.wikipedia.org/wiki/Rodrigues'_rotation_formula
https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation
https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation

"""

import math
import numpy as np
import sympy
import sympy.abc
from scipy.stats import special_ortho_group

class Quat:
    def __init__(self, x, y, z, w):
        self.x = x
        self.y = y
        self.z = z
        self.w = w

    def __str__(self):
        return sympy.pretty([self.x, self.y, self.z, self.w])

    @staticmethod
    def from_axis_angle(axis, angle):
        s = sympy.sin(angle/2)
        c = sympy.cos(angle/2)
        q = Quat(axis[0]*s, axis[1]*s, axis[2]*s, c)
        return q

    def rotation_matrix(self):
        x = self.x
        y = self.y
        z = self.z
        w = self.w
        x2 = x*x
        y2 = y*y
        z2 = z*z
        xy = x*y
        xz = x*z
        yz = y*z
        wx = w*x
        wy = w*y
        wz = w*z
        s = 1
        return sympy.Matrix([[1 - 2*s*(y2+z2), 2*s*(xy-wz), 2*s*(xz+wy)],
                             [2*s*(xy+wz), 1 - 2*s*(x2+z2), 2*s*(yz-wx)],
                             [2*s*(xz-wy), 2*s*(yz+wx), 1 - 2*s*(x2+y2)]])

"""

The Euler rotation theorem states that for any rotation around an axis in 3D
(represented by a rotation matrix) there will be an axis that is unchanged by the rotation

Mathematically, this means the following:
R - rotation matrix
v - a vector that represents an axis of rotation
There is such a vector v where
R*v = R*R*v = R*R*R*v ...

To find that vector, we need to find the eigenvalues/eigenvectors of R
Since R is a NxN matrix where N is the dimension, we are going to get
N eigenvalues with their magnitude to be 1.

One of these eigenvalue will be (1+0i), and that corresponding eigenvector
will be the fixed axis of rotation because its eigenvalue is 1

ew = 1
R^n * ev = ew^n * ev  = ev

Generalizing this, for even dimensional space (2, 4, 6 ...),
there is no fixed axis of rotation because of the fundamental
theorem of algebra says that the complex roots always comes in pair,
so there is no eigenvalue of (1+0i) for those

The corollary is that in 3D, any number of rotations can be represented
by one rotation around a fixed axis

"""

def test_rotation_theorem(n):
    R = special_ortho_group.rvs(n)

    print("Rotation Matrix {}x{}".format(n, n))
    print(R)

    print("Eigenvalues and Eigenvectors")
    ew, ev = np.linalg.eig(R)
    vf = None
    for i in range(len(ew)):
        if abs(np.real(ew[i])-1) < 1e-6:
            vf = ev[:, i]
        print(ew[i])
        print(ev[:, i])
        print()

    print("Fixed axis of rotation")
    print(vf)
    if vf is not None:
        print(np.dot(R, vf[:]))
    print()

def make_matrix_from_angles(angle, axis, sym):
    if not sym:
        c = math.cos(angle)
        s = math.sin(angle)
    else:
        c = sympy.cos(angle)
        s = sympy.sin(angle)
    if axis == 'X':
        M = [[1, 0, 0], [0, c, -s], [0, s, c]]
    elif axis == 'Y':
        M = [[c, 0, s], [0, 1, 0], [-s, 0, c]]
    else:
        M = [[c, -s, 0], [s, c, 0], [0, 0, 1]]

    if sym:
        M = sympy.simplify(sympy.Matrix(M))

    return M

def extract_matrix_angle(M, conv):
    if conv == 'XYZ':
        t1 = math.atan2(M[1][2], M[2][2])
        c2 = math.sqrt(M[0][0]**2 + M[0][1]**2)
        t2 = math.atan2(-M[0][2], c2)
        s1 = math.sin(t1)
        c1 = math.cos(t1)
        t3 = math.atan2(s1*M[2][0] - c1*M[1][0], c1*M[1][1] - s1*M[2][1])
        t1, t2, t3 = -t1, -t2, -t3
    elif conv == 'ZYX':
        t1 = math.atan2(M[2][1], M[2][2])
        t2 = math.asin(-M[2][0])
        t3 = math.atan2(M[1][0], M[0][0])
    else:
        assert 0
    return (t1, t2, t3)

def angle2dcm(yaw, pitch, roll, conv='ZYX'):
    M = None
    if conv == 'RPY' or conv == 'ZYX':
        Z = make_matrix_from_angles(yaw, 'Z', False)
        Y = make_matrix_from_angles(pitch, 'Y', False)
        X = make_matrix_from_angles(roll, 'X', False)
        M = np.dot(Z, np.dot(Y, X))
    else:
        assert 0

    return np.linalg.inv(M)

def vector2dcm(E, N):
    n = E.shape[0]
    if n == 2:
        e1 = E[:, 0]
        e2 = E[:, 1]
        n1 = N[:, 0]
        n2 = N[:, 1]
        return np.array([
            [np.dot(n1, e1), np.dot(n1, e2)],
            [np.dot(n2, e1), np.dot(n2, e2)],
        ])
    elif n == 3:
        e1 = E[:, 0]
        e2 = E[:, 1]
        e3 = E[:, 2]
        n1 = N[:, 0]
        n2 = N[:, 1]
        n3 = N[:, 2]
        return np.array([
            [np.dot(n1, e1), np.dot(n1, e2), np.dot(n1, e3)],
            [np.dot(n2, e1), np.dot(n2, e2), np.dot(n2, e3)],
            [np.dot(n3, e1), np.dot(n3, e2), np.dot(n3, e3)],
        ])
    assert 0

def rodrigues_rotation_matrix(axis, angle):
    a1 = axis[0]
    a2 = axis[1]
    a3 = axis[2]
    K = sympy.Matrix([[0, -a3, a2], [a3, 0, -a1], [-a2, a1, 0]])
    I = sympy.eye(3)
    return sympy.trigsimp(sympy.expand(I + K*sympy.sin(angle) + np.dot(K, K)*(1-sympy.cos(angle))))

"""

A frame of reference can be constructed from 3 angles conventionally denoted as

(alpha, beta, gamma) or
(psi, theta, phi)

There are 12 different sequence of rotation axes one can make with 3 angles
and are classified into 2 categories
Proper Euler angles (z-x-z, x-y-x, y-z-y, z-y-z, x-z-x, y-x-y)
Tait–Bryan angles   (x-y-z, y-z-x, z-x-y, x-z-y, z-y-x, y-x-z)

These represent rotation of an axis around the angle alpha, beta, or gamma
and we multiply all of those matrices generated together to get the concatenated matrix

These are active rotations, to find the passive rotations, transpose the matrix (which is the inverse)

"""

def test_matrix_angle_conventions():
    print("Matrix Angle Conventions")

    X1 = make_matrix_from_angles(sympy.abc.a, 'X', True)
    Y1 = make_matrix_from_angles(sympy.abc.a, 'Y', True)
    Z1 = make_matrix_from_angles(sympy.abc.a, 'Z', True)
    X2 = make_matrix_from_angles(sympy.abc.b, 'X', True)
    Y2 = make_matrix_from_angles(sympy.abc.b, 'Y', True)
    Z2 = make_matrix_from_angles(sympy.abc.b, 'Z', True)
    X3 = make_matrix_from_angles(sympy.abc.c, 'X', True)
    Y3 = make_matrix_from_angles(sympy.abc.c, 'Y', True)
    Z3 = make_matrix_from_angles(sympy.abc.c, 'Z', True)

    print()
    print("Proper Euler Angles")
    print()
    print("ZXZ")
    sympy.pprint(np.dot(X1, np.dot(Z2, X3)))

    print("XYX")
    sympy.pprint(np.dot(X1, np.dot(Y2, X3)))

    print("YXY")
    sympy.pprint(np.dot(Y1, np.dot(X2, Y3)))

    print("YZY")
    sympy.pprint(np.dot(Y1, np.dot(Z2, Y3)))

    print("ZYZ")
    sympy.pprint(np.dot(Z1, np.dot(Y2, Z3)))

    print("ZXZ")
    sympy.pprint(np.dot(Z1, np.dot(X2, Z3)))

    print()
    print("Tait-Bryan Angles")
    print()
    print("XZY")
    sympy.pprint(np.dot(X1, np.dot(Z2, Y3)))

    print("XYZ")
    sympy.pprint(np.dot(X1, np.dot(Y2, Z3)))

    print("YXZ")
    sympy.pprint(np.dot(Y1, np.dot(X2, Z3)))

    print("YZX")
    sympy.pprint(np.dot(Y1, np.dot(Z2, X3)))

    print("ZYX")
    sympy.pprint(np.dot(Z1, np.dot(Y2, X3)))

    print("ZXY")
    sympy.pprint(np.dot(Z1, np.dot(X2, Y3)))

    print()

"""

Given a rotation matrix, we want to extract
the 3 angles that make it up, the angles
returned depends on what convention we used
generally we return the angles in the [-pi, pi] range

"""

def test_matrix_angle_extraction():
    print("Matrix Angle Extraction")
    R1 = make_matrix_from_angles(0.3, 'X', False)
    R2 = make_matrix_from_angles(-0.5, 'Y', False)
    R3 = make_matrix_from_angles(0.6, 'Z', False)
    print(extract_matrix_angle(np.dot(R1, np.dot(R2, R3)), 'XYZ'))

    R1 = make_matrix_from_angles(-0.2, 'X', False)
    R2 = make_matrix_from_angles(-1.4, 'Y', False)
    R3 = make_matrix_from_angles(-math.pi, 'Z', False)
    print(extract_matrix_angle(np.dot(R1, np.dot(R2, R3)), 'XYZ'))

    R1 = make_matrix_from_angles(2.1, 'X', False)
    R2 = make_matrix_from_angles(-1.2, 'Y', False)
    R3 = make_matrix_from_angles(0.5, 'Z', False)
    print(extract_matrix_angle(np.dot(R1, np.dot(R2, R3)), 'XYZ'))

    R1 = make_matrix_from_angles(1.1, 'Z', False)
    R2 = make_matrix_from_angles(0.5, 'Y', False)
    R3 = make_matrix_from_angles(0.3, 'X', False)
    print(extract_matrix_angle(np.dot(R1, np.dot(R2, R3)), 'ZYX'))

    R1 = make_matrix_from_angles(-0.8, 'Z', False)
    R2 = make_matrix_from_angles(-1.2, 'Y', False)
    R3 = make_matrix_from_angles(1.3, 'X', False)
    print(extract_matrix_angle(np.dot(R1, np.dot(R2, R3)), 'ZYX'))

    print()

"""

A direction cosine matrix is making a rotation matrix
out of a given set of angles commonly referred to as (yaw, pitch, roll)

The direction cosine matrix name comes from constructing the cosine out of a set of vectors,
which when multiplied together can give us a rotation matrix

Given a vector P and Q, we want to find the angle between them:
dot(P, Q) = cos(angle) assuming P/Q are normalized which they usually are

Some common conventions:
RPY is used in airplanes which is the ZYX convention

Passive rotation is used to specify the rotation matrix
so need the inverse of the standard construction of the rotation matrix

"""

def test_angle2dcm():
    print("Direction Cosine Matrix from Angles")
    yaw = 0.7854
    pitch = 0.1
    roll = 0
    dcm = angle2dcm(yaw, pitch, roll)
    print(dcm)

    print(angle2dcm(0, 0, 0))

"""

If we are not given a set of angles to construct a direction cosine matrix
but two sets of orthogonal unit basis vectors
(e1, e2)     (n1, n2)     2D
(e1, e2, e3) (n1, n2, n3) 3D

We want to construct a matrix that transforms vectors from basis E to N.
We can construct a direction cosine matrix A as follows:
A_ij = dot(n_i, e_j)

This matrix that we constructed from this procedure is equivalent to the following transformation:

E = orientation matrix of basis vectors e_i
N = orientation matrix of basis vectors n_i
A = N^-1*E

"""

def test_vector2dcm(n):
    E = special_ortho_group.rvs(n)
    N = special_ortho_group.rvs(n)
    I = np.eye(n)

    EN = vector2dcm(E, N)
    NE = vector2dcm(N, E)

    print("Direction Cosine Matrix from Vectors")

    print("The basis matrix")
    print(E)
    print()
    print(N)
    print()

    # Transforming to the same basis gives us the identity matrix as the transformation matrix
    print("Transforming to the same Basis")
    print(vector2dcm(E, E))
    print(vector2dcm(N, N))
    print()

    # The basis transformation E -> N is the inverse of N <- E and vice versa
    # Since this is a orthonormal matrix, they are transpose of each others
    print("From <-> To")
    print(EN)
    print(np.dot(np.linalg.inv(N), E))
    print()
    print(NE)
    print(np.dot(np.linalg.inv(E), N))
    print()

    # Transforming from the standard basis I -> N gives us the transformation N^-1
    print("Identity -> To")
    print(np.linalg.inv(E))
    print(vector2dcm(I, E))
    print()
    print(np.linalg.inv(N))
    print(vector2dcm(I, N))
    print()

"""

If we want to rotate around an arbitrary axis instead of a product of rotations around the standard axes,
we can use the rodrigues rotation formula.

The matrix form of the formula can be derived using the exponential mapping

R = exp(K*theta)
K = w x v where w is the unit vector and v is the vector to be rotated
K is written out as a cross product matrix

We can expand the matrix exponential as a taylor series
I + theta*K + (theta*K)^2/2! + (theta*K)^3/3! ...
Because K is skew-symmetric and the sum of squares of the above diagonal is 1, the characteristic polynomial
implies that K^3 = -K, K^4 = -K^2 ...
So we can group these together and simplify to get

R = I + sin(theta)*K + (1-cos(theta))*K^2

"""

def test_rodrigues_formula():
    print()
    print("Rodrigues Rotation Formula")
    # gives the standard euler rotation matrix around X, Y, Z axes
    print("\nX Axis\n")
    sympy.pprint(rodrigues_rotation_matrix([1, 0, 0], sympy.abc.theta))
    print("\nY Axis\n")
    sympy.pprint(rodrigues_rotation_matrix([0, 1, 0], sympy.abc.theta))
    print("\nZ Axis\n")
    sympy.pprint(rodrigues_rotation_matrix([0, 0, 1], sympy.abc.theta))

    # general formula for rotation around arbitrary axis
    print("\nArbitrary Axis\n")
    sympy.pprint(rodrigues_rotation_matrix(
        [sympy.abc.x, sympy.abc.y, sympy.abc.z], sympy.abc.theta))

"""

Quaternions can be used as an alternative to Euler angles, using standard quaternion construction formulas
we can show that it gives us the same formula as the rotation matrices generated from Euler angles

"""

def test_quaternion():
    # gives standard euler rotation matrix around X, Y, Z axes
    X = Quat.from_axis_angle([1, 0, 0], sympy.abc.theta)
    Y = Quat.from_axis_angle([0, 1, 0], sympy.abc.theta)
    Z = Quat.from_axis_angle([0, 0, 1], sympy.abc.theta)
    R = Quat.from_axis_angle(
        [sympy.abc.x, sympy.abc.y, sympy.abc.z], sympy.abc.theta)
    print("\nQuaternion\n")
    print("\nX Axis\n")
    sympy.pprint(sympy.simplify(X.rotation_matrix()))
    print("\nY Axis\n")
    sympy.pprint(sympy.simplify(Y.rotation_matrix()))
    print("\nZ Axis\n")
    sympy.pprint(sympy.simplify(Z.rotation_matrix()))

    # this gives the same formula as the rodrigues formula
    print("\nArbitrary Axis\n")
    s = R.rotation_matrix()
    s = s.subs(sympy.sin(sympy.abc.theta/2)**2,
               (1-sympy.cos(sympy.abc.theta))/2)
    s = sympy.expand(sympy.simplify(s))
    sympy.pprint(s)

sympy.init_printing()

for i in range(2, 6):
    test_rotation_theorem(i)
test_matrix_angle_conventions()
test_matrix_angle_extraction()
test_angle2dcm()
test_vector2dcm(2)
test_vector2dcm(3)
test_rodrigues_formula()
test_quaternion()


"""Known matrices related to physics"""

from sympy import Matrix, I

def msigma(i):
    """Returns a Pauli matrix sigma_i. i=1,2,3
    See also:
    http://en.wikipedia.org/wiki/Pauli_matrices
    """
    if i==1:
        mat=( (
            (0, 1),
            (1, 0)
            ) )
    elif i==2:
        mat=( (
            (0, -I),
            (I, 0)
            ) )
    elif i==3:
        mat=( (
            (1, 0),
            (0, -1)
            ) )
    else:
        raise IndexError("Invalid Pauli index")
    return Matrix(mat)

def mgamma(mu,lower=False):
    """Returns a Dirac gamma matrix gamma^mu in the standard
    (Dirac) representation.

    If you want gamma_mu, use gamma(mu, True).

    We use a convention:

    gamma^5 = I * gamma^0 * gamma^1 * gamma^2 * gamma^3
    gamma_5 = I * gamma_0 * gamma_1 * gamma_2 * gamma_3 = - gamma^5

    See also:

    http://en.wikipedia.org/wiki/Gamma_matrices

    """
    if not mu in [0,1,2,3,5]:
        raise IndexError("Invalid Dirac index")
    if mu == 0:
        mat = (
                (1,0,0,0),
                (0,1,0,0),
                (0,0,-1,0),
                (0,0,0,-1)
                )
    elif mu == 1:
        mat = (
                (0,0,0,1),
                (0,0,1,0),
                (0,-1,0,0),
                (-1,0,0,0)
                )
    elif mu == 2:
        mat = (
                (0,0,0,-I),
                (0,0,I,0),
                (0,I,0,0),
                (-I,0,0,0)
                )
    elif mu == 3:
        mat = (
                (0,0,1,0),
                (0,0,0,-1),
                (-1,0,0,0),
                (0,1,0,0)
                )
    elif mu == 5:
        mat = (
                (0,0,1,0),
                (0,0,0,1),
                (1,0,0,0),
                (0,1,0,0)
                )
    m= Matrix(mat)
    if lower:
        if mu in [1,2,3,5]:
            m = - m
    return m

#Minkowski tensor using the convention (+,-,-,-) used in the Quantum Field
#Theory
minkowski_tensor = Matrix( (
    (1,0,0,0),
    (0,-1,0,0),
    (0,0,-1,0),
    (0,0,0,-1)
    ))


import matplotlib.pyplot as plt
from scipy.stats import poisson
from functools import reduce
from numpy.polynomial import Polynomial

def compose(fn, x, n):
    return reduce(lambda res, func: func(res), [fn] * n, x)

def equal_probability():
    number_of_surnames = 10
    number_of_surviving_sons = 5
    number_of_generations = 10

    fn = Polynomial([1 / (number_of_surviving_sons + 1)] * (number_of_surviving_sons + 1))

    plot_graph(fn,
               number_of_surnames,
               number_of_generations,
               "Equal Probability with Number of Surnames = {0} and Number of Surviving Sons = {1}".format(
                   number_of_surnames, number_of_surviving_sons)
               )

def varying_likelyhood():
    number_of_surnames = 10
    number_of_surviving_sons = 5
    number_of_generations = 10

    chance_of_i_surviving_sons = [1 / (number_of_surviving_sons + i) for i in range(1, number_of_surviving_sons + 1)]
    chance_of_0_surviving_sons = 1 - sum(chance_of_i_surviving_sons)

    fn = Polynomial([chance_of_0_surviving_sons] + chance_of_i_surviving_sons)

    plot_graph(fn,
               number_of_surnames,
               number_of_generations,
               "Varying Likelyhood with Number of Surnames = {0} and Number of Surviving Sons = {1}".format(
                   number_of_surnames, number_of_surviving_sons)
               )

def poisson_distribution():
    number_of_surnames = 10
    number_of_surviving_sons = 5
    mean_number_of_surviving_sons = 1
    number_of_generations = 10

    chance_of_i_surviving_sons = [
        poisson.pmf(i, mean_number_of_surviving_sons) for i in range(0, number_of_surviving_sons + 1)
    ]

    fn = Polynomial(chance_of_i_surviving_sons)

    plot_graph(fn,
               number_of_surnames,
               number_of_generations,
               "Poisson Distribution with Number of Surnames = {0} and Number of Surviving Sons = {1}".format(
                   number_of_surnames, number_of_surviving_sons)
               )

def plot_graph(fn, number_of_surnames, number_of_generations, title):
    generating_functions = [compose(fn, 0, generation) for generation in range(0, number_of_generations)]

    points = [number_of_surnames - number_of_surnames * func for func in generating_functions]

    plt.plot(points)
    plt.xlabel('Generation')
    plt.ylabel('Number of Surnames')
    plt.title(title)
    plt.axis([0, number_of_generations, 0, number_of_surnames])
    plt.show()

if __name__ == "__main__":
    equal_probability()
    varying_likelyhood()
    poisson_distribution()


#!/usr/bin/env python
'''
| Filename    : polyinterp.py
| Description : Polynomial interpolation of data points.
| Author      : Pushpendre Rastogi
| Created     : Tue Sep  6 00:10:47 2016 (-0400)
| Last-Updated: Tue Sep  6 00:11:35 2016 (-0400)
|           By: Pushpendre Rastogi
|     Update #: 1
Copied from http://hsakamoto.com/docs/numerical/ch3/301.html
September 6 2016
'''
import numpy as np
import matplotlib.pyplot as plt

def PI(x_data, y_data):
    """ return the polynomially interpolated function
    for a given set of data points by Newton's method
    """
    a = coeff(x_data, y_data)
    n = len(x_data)

    def f(xi):
        """ polynomially interpolated function
        """
        v = a[n - 1]  # initial value
        for k in range(1, n):
            # compute backwardly
            v = a[-k - 1] + (xi - x_data[-k - 1]) * v
        return v
    return f

def coeff(x_data, y_data):
    """ return the coefficients for Newton's polynomial
    """
    x = np.array(x_data)
    y = np.array(y_data)
    m = len(y)
    a = y.copy()  # initial value (a[0] = y[0])
    for k in range(m - 1):
        for i in range(k, m - 1):
            # compute divided difference of k-th order for all i >= k
            a[i + 1] = (a[i + 1] - a[k]) / (x[i + 1] - x[k])
    return a  # coefficient


import math
import numpy as np
import matplotlib.pyplot as plt
import timeit
"""
Constants for movement, subject to change.
May also make hor_offset a non-constant to allow matplotlib to
"animate" the wave moving from left to right.
"""
AMPLITUDE = 140
WAVELENGTH = 840
HOR_OFFSET = 0
VER_OFFSET = 0

SEG_LEN = 280
HEAD_LEN = 270
NUM_JOINTS = 6

def sine(x : float) -> float:
    """
    Trig Sine Function in the form f(x) = (a * sin((2pi/wv)*x + c)) + d
    :param x: angle in radians
    :return: sin value given the constants
    """

    return (AMPLITUDE * np.sin((2 * np.pi / WAVELENGTH) * x + HOR_OFFSET)) + VER_OFFSET

def find_joint_coords(sine_func) -> list:
    """
    Function using euclidean distance formula and sine function to solve the
    positions of each joint according to the given sine function. Assuming the
    snake is oriented with its head on the left hand side at x = 0, and all "tail"
    to the right of the head in positive x-axis.
    :param: sine function callable
    :return: array of tuples in format [(0, y_0), (x_1, y_1)...]
        representing position of each joint. y_0 representing the leftmost tip
        of the head of the snake, which will be x = 0.
    """

    result = [[0.0], [sine_func(0)]]
    # result = [[0.0], [0.0]]

    for i in range(0, NUM_JOINTS):
        guessed_dist = 0.0
        min_x = 0.0 + result[0][i]
        max_x = SEG_LEN + result[0][i]
        guessed_x = 0.0
        # essentially estimate the x coordinate of the next joint in the sequence
        # getting down to +/- 0.5mm in precision to the ideal distance.
        while not ((guessed_dist <= SEG_LEN + 0.5) and (guessed_dist >= SEG_LEN - 0.5)
                   or (max_x - min_x <= 1.0)):
            guessed_x = ((max_x - min_x) / 3.0) + min_x
            # guessed_x = ((max_x + min_x) / 2.0)
            guessed_dist = math.sqrt((guessed_x - result[0][i])**2
                                     + (sine_func(guessed_x) - result[1][i])**2)
            if guessed_dist < SEG_LEN:
                min_x = guessed_x
            else:
                max_x = guessed_x
        result[0].append(guessed_x)
        result[1].append(sine_func(guessed_x))

    return result

if __name__ == "__main__":
    found_coords = find_joint_coords(sine)
    print(found_coords)
    print("Time taken: " + str(timeit.timeit(lambda: find_joint_coords(sine), number=1)))

    plt.plot(found_coords[0], found_coords[1])

    sine_x = range(0, WAVELENGTH * 2)
    sine_y = []
    for i in sine_x:
        sine_y.append(sine(i))
    plt.plot(sine_x, sine_y)
    plt.xlim([0, 4 * AMPLITUDE])
    plt.ylim([-2 * AMPLITUDE, 2 * AMPLITUDE])
    plt.show()


import numpy as np
import OpenCV as opencv
# sigmoid function

def nonlin(x,deriv=False):
  if(deriv==True):
    return x*(1-x)
    return 1/(1+np.exp(-x))

# input dataset

X = np.array([  [0,0,1],[0,1,1],[1,0,1],[1,1,1] ])
# output dataset
y = np.array([[0,0,1,1]]).T

# seed random numbers to make calculation

# deterministic (just a good practice)

np.random.seed(1)
# initialize weights randomly with mean 0

syn0 = 2*np.random.random((3,1)) - 1

for iter in xrange(10000):
# forward propagation

  l0 = X

  l1 = nonlin(np.dot(l0,syn0))
  # how much did we miss?

  l1_error = y - l1
  # multiply how much we missed by the

  # slope of the sigmoid at the values in l1

  l1_delta = l1_error * nonlin(l1,True)

  # update weights

  syn0 += np.dot(l0.T,l1_delta)

print("Output After Training:")

print(l1)


import numpy as np
import scipy.stats as ss

def dnorm(x, mu, sigma):
    sigma += np.eye(sigma.shape[0]) * 1e-8
    return ss.multivariate_normal.logpdf(x, mu, sigma)

class GMM_EM:
    """
    GMM by EM.

    Methods:
        fit(data, max_iter, threshold): Fit the model to data.
        predict(x): Predict cluster labels for x.
    """

    def __init__(self, n_clusters):
        """
        Constructor Methods:

        Args:
            n_clusters(int): number of clusters.
        """

        self.n_clusters = n_clusters

        self.pi = None
        self.mus = [None] * self.n_clusters
        self.sigmas = [None] * self.n_clusters

    def fit(self, data, max_iter=200, threshold=1e-8):
        """
        Fit the model to data.

        Args:
            data: Array-like, shape (n_samples, n_dim)
            max_iter: maximum number of EM steps.
            threshold: threshold to step iteration.
        """

        assert data.ndim == 2
        n_data = data.shape[0]

        # Initialize
        z = self._initialization(data)
        self.pi = np.array([np.mean(z == cluster) for cluster in range(self.n_clusters)])
        self.mus = [np.mean(data[z == cluster, :], axis=0) for cluster in range(self.n_clusters)]
        self.sigmas = [np.cov(data[z == cluster, :].T) for cluster in range(self.n_clusters)]

        old_ll = 0
        for iter in range(max_iter):

            # E-step
            log_p = np.array([dnorm(data, mu=self.mus[cluster], sigma=self.sigmas[cluster]) + np.log(self.pi[cluster])
                          for cluster in range(self.n_clusters)])
            max_p = np.max(log_p, axis=0)
            sum_p = np.log(np.sum(np.exp(log_p - max_p), axis=0)) + max_p
            log_p -= sum_p
            p = np.exp(log_p)

            # M-step
            self.pi = np.sum(p, axis=1) / n_data
            for cluster in range(self.n_clusters):
                effective_size = np.sum(p, axis=1)[cluster]
                self.mus[cluster] = np.sum(p[cluster].reshape(-1, 1) * data, axis=0) / effective_size
                self.sigmas[cluster] = ((data - self.mus[cluster]).T * p[cluster]) @ (data - self.mus[cluster]) / effective_size

            # Calculate (negative) log_likelihood
            new_ll = -np.sum(sum_p)
            if abs(new_ll-old_ll) <= threshold:
                break
            else:
                old_ll = new_ll

        return self

    def predict(self, x):
        """
        Predict cluster labels for x.

        Args:
            x: Array-like, shape (n_samples, n_dim)
        Return:
            Array-like, shape (n_samples, )
        """
        log_prob = [dnorm(x, self.mus[cluster], self.sigmas[cluster]) + np.log(self.pi[cluster])
                     for cluster in range(self.n_clusters)]
        log_prob = np.vstack(log_prob)
        z = np.argmax(log_prob, axis=0)
        return z

    def _initialization(self, data, max_iter=50):
        """
        Initialization by K-Means.
        """
        means = data[np.random.choice(data.shape[0], self.n_clusters, replace=False)]    # pick random samples as center
        z = np.zeros(data.shape[0])

        for iter in range(max_iter):
            dist = [np.sum((data - means[cluster]) ** 2, axis=1) for cluster in range(self.n_clusters)]
            dist = np.vstack(dist)
            z = np.argmin(dist, axis=0)
            means = [np.mean(data[z == cluster], axis=0) for cluster in range(self.n_clusters)]

        return z


import numpy as np
from numpy import pi as pi
from numpy import cos as cos
from numpy import sin as sin
from numpy import linalg
from matplotlib import pyplot as plt

# input units and values
units = input('units: ')

sigX = float(input("sigma-x: "))
sigY = float(input("sigma-y: "))
tauXY = float(input("tau-XY: "))

#calculate principal normal stresses and absolute maximum shear stress
Cauchy = np.matrix([[sigX,tauXY],[tauXY,sigY]])

sig_prncpl = linalg.eigvals(Cauchy)
sig_prncpl = sorted(sig_prncpl,reverse=True)
tau_abs = abs(sig_prncpl[0]-sig_prncpl[1])/2

#print value outputs
print("\nPrincipal normal stresses:\n" +
	str(round(sig_prncpl[0],2)) + units + "\n" +
	str(round(sig_prncpl[1],2)) + units + "\n")
print("Absolute maximum shear stress:\n" +
	str(round(tau_abs,2))  + units + "\n")

#calculations for plot of Mohr's circle
#center of circle, corresponding to average in-plane normal stresses
sig_avg = (sig_prncpl[0] + sig_prncpl[1])/2

#initialize theta
theta = np.linspace(0,2*pi,256)

#create plot
plt.figure(figsize=(10,10),dpi=80)
plt.subplot(1,1,1)

#plot Mohr's circle
X = tau_abs*cos(theta)+sig_avg
Y = tau_abs*sin(theta)
plt.plot(X,Y,color="blue")

#notable points on graph
#absolute maximum shear stress
plt.plot(sig_avg,tau_abs,marker='o',markersize=6,color="black",
	label="Absolute maximum shear stress\n" + "(" + str(round(tau_abs,2)) + units + ")")
#principal normal stresses
#maximum
plt.plot(sig_prncpl[0],0,marker='o',markersize=6,color="black",
	label="Principal normal stresses\n" + "(" + str(round(sig_prncpl[0],2)) + units +
	", " + str(round(sig_prncpl[1],2)) + units + ")")
#minimum
plt.plot(sig_prncpl[1],0,marker='o',markersize=6,color="black")
#current stress state
#X-axis
plt.plot(sigX,tauXY,marker='o',markersize=4,color="red",
	label="Current stress state\n" + "(" + str(round(sigX,2)) + units +
	", " + str(round(sigY,2)) + units + ")")
#Y-axis
plt.plot(sigY,-tauXY,marker='o',markersize=6,color="red")
plt.grid(True)

#format axes and grid
ax = plt.gca()  # gca stands for 'get current axis'
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.spines['bottom'].set_position(('data',0))
ax.yaxis.set_ticks_position('left')
ax.spines['left'].set_position(('data',sig_avg))
plt.gca().invert_yaxis()
plt.title("Mohr's Circles")
plt.legend(loc='upper left')
plt.show()


import math
import copy
import numpy as np
import prettymatrix
from prettytable import PrettyTable
from sympy import *
# Métodos diretos
## Método de Cholesky
def Determinantes(A, ordem):
    k = 1
    tempA = [[]]
    while k <= ordem:
        for i in range(k):
            for j in range(k):
                tempA[i].append(A[i][j])
        if(round(np.linalg.det(tempA), 8) == 0): return (k,round(np.linalg.det(tempA), 8))
        if k < ordem:
            for i in range(k):
                for j in range(k):
                    tempA[i].pop()
            tempA.append([])
        k+=1
    return (0,round(np.linalg.det(tempA), 8))

def Transposta(A, ordem):
    k = ordem
    mT = []
    for c in range(k):
        mT.append([])
    for i in range(k):
        for j in range(k):
            mT[i].append( A[j][i] )
    soma = 0
    for i in range(k):
        for j in range(k):
            if(mT[i][j] != A[i][j]):
                soma += 1
    if soma == 0:
        return 1

def valoresG(A, G, ordem):
    k = ordem
    for i in range(k):
        for j in range(k):
            if(i == 0 and j == 0):
                G[0][0] = A[0][0] ** (1/2)
            elif(i == j):
                soma = 0
                for c in range(0, i):
                    soma += G[i][c]**2
                G[i][j] = (A[i][j] - soma)**(1/2)
            elif(j == 0 and i >= 1):
                G[i][0] = A[i][0]/G[0][0]
            elif(j == 1 and i >= 2):
                G[i][1] = (A[i][1] - G[i][0]*G[1][0])/G[1][1]
                G[i][1] = round(G[i][1], 8)
            elif(1 <= j and j < i):
                soma = 0
                for c in range(i):
                    soma += G[i][c]*G[j][c]
                G[i][j] = (A[i][j] - soma)/G[j][j]

    G = np.asarray(G)
    return G

def sistLinear(G, B, ordem):
    y = symbols('y:'+str(ordem))
    mY = []
    for i in range(len(y)):
        mY.append(y[i])
    D = np.linalg.det(G)
    tempG = G.copy()
    for j in range(ordem):
        for i in range(ordem):
            tempG[i][j] = B[i]
        tempD = np.linalg.det(tempG)
        tempG = G.copy()
        mY[j] = round(tempD/D, 8)
    mTemp = []
    for i in range(len(mY)):
        mTemp.append([mY[i]])
    mY = mTemp.copy()
    mY = np.asarray(mY)
    return mY

def Cholesky(A, B):
    ordem = np.shape(A)
    if ordem[0] == 0: return print("Isso não é uma matriz")
    if ordem[0] == ordem[1]:
        det = Determinantes(A, ordem[0])
        if( det[1] <= 0):
            print("Determinante <= 0 para k = " + str(det[0]))
            return
        trans = Transposta(A, ordem[0])
        if( not trans ):
            return print("Matriz não é igual sua transposta")

        mG = np.zeros((ordem[0],ordem[1]))
        mG = valoresG(A, mG, ordem[0])
        print(prettymatrix.matrix_to_string(mG, name='G = '))

        mY = sistLinear(mG, B, ordem[0])
        print(prettymatrix.matrix_to_string(mY, name='Y = '))

        transpG = np.zeros((ordem[0],ordem[1]))
        for i in range(ordem[0]):
            for j in range(ordem[0]):
                transpG[i][j] = mG[j][i]

        transpG = np.asarray(transpG)

        mX = sistLinear(transpG, mY, ordem[0])
        print(prettymatrix.matrix_to_string(mX, name='X = '))

# A = [[4,2,-4],
#      [2,10,4],
#      [-4,4,9]]
# B = [0,
#      6,
#      5]
# Cholesky(A,B)

## Método LU
def valoresLU(L, U, A, ordem):
    for i in range(ordem):
        for j in range(ordem):
            if i == 0:
                U[0][j] = A[0][j]
            elif i == 1 and j >= 1:
                U[1][j] =  A[1][j] - L[1][0]*U[0][j]
            elif i <= j:
                soma = 0
                for c in range(i):
                    soma += L[i][c]*U[c][j]
                U[i][j] = A[i][j] - soma

        for j in range(ordem):
            if i == 0 and j >= 1:
                L[j][0] = A[j][0]/U[0][0]
            elif i == 1 and j >= 2:
                L[j][1] = (A[j][1] - L[j][0]*U[0][1])/U[1][1]
            elif j > i:
                soma = 0
                for c in range(i):
                    soma += L[j][c]*U[c][i]
                L[j][i] = (A[j][i] - soma)/U[i][i]

    return L, U

def LU(A,B):
    ordem = np.shape(A)
    if ordem[0] == 0: return print("Isso não é uma matriz")
    if ordem[0] != ordem[1]: return print("Esta matriz não é quadrada")
    det = Determinantes(A, ordem[0]-1)
    if( det[1] == 0):
        return print("Determinante = 0 para k = " + str(det[0]))

    L = np.eye(ordem[0])
    U = np.zeros((ordem[0],ordem[1]))

    mLU = valoresLU(L, U, A, ordem[0])

    mY = sistLinear(L, B, ordem[0])
    mX = sistLinear(U, mY, ordem[0])

    print(prettymatrix.matrices_to_string(L, U, names=['L','U']))
    print(prettymatrix.matrix_to_string(mY, name='Y = '))
    print(prettymatrix.matrix_to_string(mX, name='X = '))

# A = [[5,2,1],[-1,4,2],[2,-3,10]]
# B = [-12, 20, 3]
# LU(A,B)

## Método de Eliminação de Gauss
def pivotParcial(A,B,MA,MB,ordem,k):
    if A[k][k] == 0:
        for i in range(k+1,ordem):
            if A[i][k] != 0:
                temp1 = copy.copy(A[i])
                temp2 = copy.copy(B[i])
                temp3 = copy.copy(MA[i])
                temp4 = copy.copy(MB[i])
                A[i] = copy.copy(A[k])
                B[i] = copy.copy(B[k])
                MA[i] = copy.copy(MA[k])
                MB[i] = copy.copy(MB[k])
                A[k] = copy.copy(temp1)
                B[k] = copy.copy(temp2)
                MA[k] = copy.copy(temp3)
                MB[k] = copy.copy(temp4)
                break
    save = 0
    maior = -math.inf
    for i in range(ordem):
        if abs(A[i][k]) > maior and i >= k:
            maior = copy.copy(abs(A[i][k]))
            save = copy.copy(i)
    temp1 = copy.copy(A[save])
    temp2 = copy.copy(B[save])
    temp3 = copy.copy(MA[save])
    temp4 = copy.copy(MB[save])
    A[save] = copy.copy(A[k])
    B[save] = copy.copy(B[k])
    MA[save] = copy.copy(MA[k])
    MB[save] = copy.copy(MB[k])
    A[k] = copy.copy(temp1)
    B[k] = copy.copy(temp2)
    MA[k] = copy.copy(temp3)
    MB[k] = copy.copy(temp4)

def MA(A, B, ordem):
    MA = np.zeros((ordem,ordem))
    MB = np.zeros(ordem)

    for i in range(ordem):
        MB[i] = B[i]
        for j in range(ordem):
            MA[i][j] = A[i][j]

    mAum = np.zeros((ordem,ordem+1))
    for k in range(ordem-1):
        pivotParcial(A,B,MA,MB,ordem,k)
        for i in range(k+1,ordem):
            MB[i] = round(B[i] - (B[k] * A[i][k])/A[k][k], 8)
            for j in range(ordem):
                MA[i][j] = round(A[i][j] - (A[k][j]*A[i][k])/A[k][k], 8)

        for i in range(ordem):
            for j in range(ordem+1):
                if j == ordem:
                    mAum[i][j] = MB[i]
                else:
                    mAum[i][j] = MA[i][j]
        mAum = np.asarray(mAum)
        print(prettymatrix.matrix_to_string(mAum, name='Matriz aumentada = '))

        for i in range(ordem):
            B[i] = MB[i]
            for j in range(ordem):
                A[i][j] = MA[i][j]
    return MA, MB

def Gauss(A, B):
    ordem = np.shape(A)
    if ordem[0] == 0: return print("Isso não é uma matriz")
    if ordem[0] != ordem[1]: return print("Esta matriz não é quadrada")
    mA, mB = MA(A, B, ordem[0])
    mX = sistLinear(mA, mB, ordem[0])
    print(prettymatrix.matrix_to_string(mX, name='X = '))

A = [[2,2,-1],
     [3,3,1],
     [1,-1,5]]
B = [3, 7, 5]
Gauss(A,B)

# Métodos iterativos
## Norma Linha

def norma_linha(A):
    ordem = np.shape(A)
    if ordem[0] == 0 or ordem[1] == 0: return print("Isso não é uma matriz")

    somaMaior = 0
    for i in range(ordem[0]):
        soma = 0
        for j in range(ordem[0]):
            soma += abs(A[i][j])
        if soma > somaMaior:
            somaMaior = soma
        print("Soma dos elementos da linha ", i+1, " = ", round(soma, 8))

    return somaMaior

def norma_coluna(A):
    ordem = np.shape(A)
    if ordem[0] == 0 or ordem[1] == 0: return print("Isso não é uma matriz")

    somaMaior = 0
    for j in range(ordem[0]):
        soma = 0
        for i in range(ordem[0]):
            soma += abs(A[i][j])
        if soma > somaMaior:
            somaMaior = soma
        print("Soma dos elementos da coluna ", j+1, " = ", soma)

    return somaMaior

def norma_euclidiana(A):
    ordem = np.shape(A)
    if ordem[0] == 0 or ordem[1] == 0: return print("Isso não é uma matriz")
    soma = 0
    for i in range(ordem[0]):
        for j in range(ordem[0]):
            soma += A[i][j]**2

    soma = soma ** (1/2)
    return soma

def norma_inf(A):
    somaMaior = 0
    for i in range(len(A)):
        if abs(A[i]) > somaMaior:
            somaMaior = abs(A[i])

    return somaMaior

# A = [[3,-5,7],
#      [1,-2,4],
#      [-8,1,-7]]

# print("||A||_inf = ", norma_linha(A))
# print("||A||_1 = ", norma_coluna(A))
# print("||A||_E = ", norma_euclidiana(A))

## Método Gauss Jacobi
def condicaoEDD(A, ordem):
    for i in range(ordem):
        for j in range(ordem):
            if A[i][j] > A[i][i]:
              return false

def condicaoF(A, ordem, tipo):
    F = np.zeros((ordem,ordem))
    for i in range(ordem):
        for j in range(ordem):
            if i != j:
                F[i][j] = -(A[i][j]/A[i][i])
    print(prettymatrix.matrix_to_string(F, name='Matriz F = '))
    print("Cálculo da soma em valor absoluto de cada linha da matriz F:")
    if norma_linha(F) < 1:
        return true
    elif tipo == 1:
        print("Cálculo da soma em valor absoluto de cada coluna da matriz F:")
        if norma_coluna(F) < 1:
            return true
    else:
        return false

def GaussJacobi(A, B, X0, e):
    flag = 0
    ordem = np.shape(A)
    if ordem[0] == 0 or ordem[1] == 0: return print("Isso não é uma matriz")
    if condicaoEDD(A,ordem[0]) == false:
        flag += 1
        print("A matriz não é E.D.D.")
    if condicaoF(A, ordem[0], 1) == false:
        flag += 1
    if flag != 2:
        xk = []
        x_x = []
        xk.append(X0)
        x_x.append('-')
        controle = 0
        end_condition = 0
        while not end_condition:
            array = []
            for i in range(len(X0)):
                array2 = []
                temp = 0
                for j in range(len(X0)):
                    if i != j:
                        temp += A[i][j]*X0[j]
                array.append(round((1/A[i][i]) * (B[i] - temp),8))

            for i in range(len(array)):
                array2.append(array[i]-xk[controle][i])

            xk.append(array)
            x_x.append(norma_inf(array2)/norma_inf(array))
            X0 = array
            controle += 1
            if norma_inf(array2)/norma_inf(array) < e:
                end_condition = 1

        Table = PrettyTable(["k", "xk", "||x^(k+1) - x^(k)||/||x^(k+1)||"])
        for k in range(0, len(xk)):
            Table.add_row([k, xk[k], x_x[k]])

        print(Table)
    else:
        print("A matriz não contém os requisitos quando feita a matriz F")

# A = [[10,2,1],[1,5,1],[2,3,10]]
# B = [14, 11, 8]
# X0 = [0,0,0]
# GaussJacobi(A,B,X0,0.01)

## Método Gauss Seidel
def Sanssenfeld(A, ordem):
    B = []
    maior = -math.inf
    for i in range(ordem):
        soma = 0
        for j in range(i):
            if i == 0: continue
            soma += abs(A[i][j]/A[i][i]) * B[j]
        temp = copy.copy(soma)
        soma = 0
        for j in range(i+1,ordem):
            soma += abs(A[i][j]/A[i][i])
        print("\u03B2", i, " = ", round(temp+soma,8))
        B.append(temp+soma)
        if B[len(B)-1] > maior:
            maior = B[len(B)-1]
    if maior < 1:
        return true
    else:
        return false

def GaussSeidel(A, B, X0, e):
    flag = 0
    ordem = np.shape(A)
    if ordem[0] == 0 or ordem[1] == 0: return print("Isso não é uma matriz")
    if condicaoEDD(A,ordem[0]) == false:
        flag += 1
        print("Matriz não é E.D.D.")
    if condicaoF(A, ordem[0], 2) == false:
        flag += 1
        print("Não atende aos requisitos quando transformada na matriz F")
    print("Critério de Sanssenfeld:")
    if Sanssenfeld(A, ordem[0]) == false:
        flag += 1
        print("A matriz não atende ao critério de Sanssenfeld")
    if flag != 3:
        xk = []
        x_x = []
        xk.append(X0)
        x_x.append('-')
        controle = 0
        end_condition = 0
        while not end_condition:
            array = []
            array2 = []
            for i in range(len(X0)):
                temp = 0
                for j in range(len(X0)):
                    if i != j:
                        if j < i:
                            temp += A[i][j]*array[j]
                        else:
                            temp += A[i][j]*X0[j]
                array.append(round((1/A[i][i]) * (B[i] - temp),8))

            for i in range(len(array)):
                array2.append(array[i]-xk[controle][i])

            xk.append(array)
            x_x.append(norma_inf(array2)/norma_inf(array))
            X0 = array
            controle += 1
            if norma_inf(array2)/norma_inf(array) < e:
                end_condition = 1

        Table = PrettyTable(["k", "xk", "||x^(k+1) - x^(k)||/||x^(k+1)||"])
        for k in range(0, len(xk)):
            Table.add_row([k, xk[k], x_x[k]])

        print(Table)

    else:
        return print("A matriz não atende aos requisitos do método")

# A = [[5,1,1],[3,4,1],[3,3,6]]
# B = [5, 6, 0]
# X0 = [0,0,0]
# GaussSeidel(A, B, X0, 0.01)


from sklearn.linear_model import LinearRegression
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn

# change the coin_type as you want. Support: btc, eth, abc

coin_type = 'btc'

Df = pd.read_csv( 'data/' + coin_type + '_price.csv',
    parse_dates=True,
    index_col=0)

Df = Df[['Price']]
Df = Df.dropna()
Df['S_3'] = Df['Price'].shift(1).rolling(window=3).mean()
Df['S_9'] = Df['Price'].shift(1).rolling(window=9).mean()
Df = Df.dropna()
x = Df[['S_3', 'S_9']]
x.head()
y = Df['Price']
y.head()

t = .8
t = int(t * len(Df))

x_train = x[:t]
y_train = y[:t]

x_test = x[t:]
y_test = y[t:]

linear = LinearRegression().fit(x_train, y_train)

print("Price =", round(linear.coef_[0], 2),
      "* 3 Days Moving Average", round(linear.coef_[1], 2),
      "* 9 Days Moving Average +", round(linear.intercept_, 2))

predicted_price = linear.predict(x_test)

predicted_price = pd.DataFrame(
    predicted_price,
    index=y_test.index,
    columns=['price'])

predicted_price.plot(figsize=(10, 5))

y_test.plot()

plt.legend(['predicted_price', 'actual_price'])

plt.ylabel(coin_type.upper() + "Price")

r2_score = linear.score(x[t:], y[t:]) * 100

print('\nThe heigher, the better (0 - 100)')
print(float("{0:.2f}".format(r2_score)))

plt.show()


import logging
logger = logging.getLogger(__name__)

import itertools
import random

from collections import OrderedDict

import functools

from fractions import Fraction

from math import comb

import numpy as np
import pandas as pd

class ProbabilityPredictor:

  @staticmethod
  @functools.lru_cache(maxsize=256)
  def probExactlyXDice(numInstances, numDice, numDieFaces):
    """
    # Probability exactly <numInstances> occurrences of a dice value are
    # shown when <numDice> are thrown, each die having <numDieFaces> sides
    """

    # probability die value will be shown
    probOccur = Fraction(1, numDieFaces)

    # probability exactly <numInstances> of this die value appearing is:
    #  -total number of ways this can occur given the number of dice we have
    # *
    #  -probability this value appears <numInstances> times
    # *
    #  -probability this value DOESN'T appear the remaining times
    #
    return comb(numDice, numInstances) * (probOccur ** numInstances) * ((1-probOccur) ** (numDice-numInstances))

  @staticmethod
  @functools.lru_cache(maxsize=256)
  def probAtLeastXDice(numInstances, numDice, numDieFaces):
    """
    # Probability at least <numInstances> occurrences of a dice value are
    # shown when <numDice> are thrown, each die having <numDieFaces> sides
    """

    # add up the probabilities that:
    #  -exactly <numInstances> values are shown
    # +
    #  -exactly <numInstances>-1 values are shown
    # +
    #  -exactly <numInstances>-2 values are shown
    # +
    # ....
    return sum([ProbabilityPredictor.probExactlyXDice(i, numDice, numDieFaces)
                  for i in range(numInstances,numDice+1)])

  @staticmethod
  def DEPrangeOfAKind(numDice, kind, numDieFaces=6):
    """
    # Probability of getting 1-<numOccurrences> of <kind> dice values
    #  -e.g., when rolling 5 dice, probability of getting 1, 2, 3, 4 or 5 "4"s
    #
    #
    """

    # current number of <kind> we have
    currentNum = diceCount.get(kind, 0)

    # calculate the chance of getting X of <kind>
    #  -e.g., for kind=3 with 5 dice, probability of getting 1x3, 2x3, 3x3, 4x3, 5x3
    for numOccurrences in range(1, numDice + 1):

      # if we already have this number of <kind> then probability is 1.0
      if currentNum >= numOccurrences:
        prob[numOccurrences] = Fraction(1)

      # else, calculate the probability of rolling the number of <kind> needed
      # with the dice we have available
      else:

        # if we can hold the current <kind> dice we have, then we only need to
        # roll for the remaining number of <kind> we need
        numKindNeeded    = numOccurrences - currentNum
        numAvailableDice = numDice - currentNum

        # probability of getting the number of <kind> we need, given the dice we have
        prob[numOccurrences] = ProbabilityPredictor.probAtLeastXDice(numKindNeeded,
                                                                     numAvailableDice,
                                                                     numDieFaces)

    return prob

  @staticmethod
  def ofAKind(kind, diceValues, turnsRemaining, numDieFaces=6, canHold=True):
    """  """

    numDice = len(diceValues)

    diceCount = pd.Series(diceValues).value_counts()

    # current number of <kind> we have
    currentNum = diceCount.get(kind, 0)

    """
    prob 5-of-a-kind of 1s:
    -prob(getting 5 | got exactly 4 before) + prob(getting 5 | got exactly 3 before) +

    p(5|0) + p(4|1) + p(3|2) + p(2|3) + p(1|4) + p(5)

    """

    probSum = 0
    for numOccurrences in range(numDice):

      pA = ProbabilityPredictor.probAtLeastXDice(numDice-numOccurrences, numDice-numOccurrences, numDieFaces=6)
      pB = ProbabilityPredictor.probExactlyXDice(numOccurrences, numDice, numDieFaces=6)
      print("{}: {} * {}".format(numOccurrences, pA, pB))
      probSum += pA * pB

    return probSum

    prob = {}

    # prob getting X of a kind next turn
    #prob[1] = 1.0 if currentNum >= 1 else ProbabilityPredictor.probAtLeastXDice(1, numDice, numDieFaces)

    # calculate the chance of getting X of <kind>
    #  -e.g., for kind=3 with 5 dice, probability of getting 1x3, 2x3, 3x3, 4x3, 5x3
    for numOccurrences in range(1, numDice+1):

      # if we already have this number of <kind> then probability is 1.0
      if currentNum >= numOccurrences:
        prob[numOccurrences] = Fraction(1)

      # else, calculate the probability of rolling the number of <kind> needed
      # with the dice we have available
      else:

        # if we can hold the current <kind> dice we have, then we only need to
        # roll for the remaining number of <kind> we need
        numKindNeeded    = numOccurrences - currentNum if canHold else numOccurrences
        numAvailableDice = numDice - currentNum        if canHold else numDice

        # probability of getting the number of <kind> we need, given the dice we have
        prob[numOccurrences] = ProbabilityPredictor.probAtLeastXDice(numKindNeeded,
                                                                     numAvailableDice,
                                                                     numDieFaces)

    return prob

    # probability for each turn
    probs = {}

    for iTurn in range(1, turnsRemaining+1):

      # prob getting X of a kind this turn

      probs[iTurn] = {}

      for numOccurences in range(1, numDice + 1):
        probs[iTurn][numOccurences] = ProbabilityPredictor.probAtLeastXDice(numOccurences, numDice, numDieFaces)

  @staticmethod
  def OLDofAKind(kind, diceValues, numDieFaces=6, numTurns=3, canHold=True):
    """  """

    numDice = len(diceValues)

    # probability for each turn
    probs = {}

    for iTurn in range(numTurns):

      probs[iTurn] = {}

      for numOccurences in range(1, numDice+1):
        probs[iTurn][numOccurences] = ProbabilityPredictor.probAtLeastXDice(numOccurences, numDice, numDieFaces)

    #probKind = Fraction(1,6)

    #diceCounts = pd.Series(diceValues).value_counts()

    #currNum = diceCounts.get(kind, 0)
    #diceRemaining = numDice - currNum

    return probs


import numpy as np
import math
import copy
import scipy.stats

def log_poisson_likelihood(y, y_obs):
    """
    Calculate the naural logarithm of the Poisson likelihood for
    modeled number of detections y and the observed number of
    detections y_obs. The Poisson PMF is

    pmf(y_obs | y) = (y^(y_obs) * e^(-y))/ y_obs!

    Parameters
    ----------
    y : numpy array
        The population counts, also called lambda.
    y_obs : numpy array
        The observed counts, also called k.

    Returns
    -------
    log_likelihood : float
        The sum of of the Poisson log-likelihoods.

    Notes
    ------
    This code can be condesned down to
    log_likelihood = np.sum(scipy.stats.poisson.logpmf(k=y_obs, mu=y))
    but I chose to be explicit for the learning opportunity.
    """
    assert len(y) == len(y_obs), ('The modeled and observed count'
                                ' lengths must be identical.')

    # Drop 0 y values because they diverge. The probability of
    # any y_obs other than 0, given y is 0 is 0.
    y_ind = np.where(np.array(y) != 0)[0]
    y = y[y_ind]
    y_obs_copy = copy.copy(np.array(y_obs))
    y_obs_copy = y_obs_copy[y_ind]

    # Take the factorial and apply Stirling's approximation
    # for y_obs > 100.
    log_factorial = [math.log(math.factorial(y_i)) if y_i < 100
                    else y_i*np.log(y_i) - y_i for y_i in y_obs_copy]
    # Analytic derivation of the Poisson log-likelihood for a product
    # of Poisson random variables
    log_likelihood = np.sum(y_obs_copy*np.log(y) - y - log_factorial)
    return log_likelihood

def log_normal_likelihood(y_obs, p):
    """
    Calculate the naural logarithm of the Gaussian (normal) likelihood
    for the observation y_obs and Gaussian parameter array p. p[0] is
    the offset, and p[1] is the standard deviation. The norm PDF is

    pdf(x|p) = exp(-((x-p[0])/sqrt(2)*p[1])**2)/(sqrt(2*pi)*p[1])

    Parameters
    ----------
    x : numpy array
        The observation.
    p : numpy array
        An array of length 2 with the norm offset and standard deviation.

    Returns
    -------
    log_likelihood : float
        The sum of of the norm log-likelihoods.
    """
    log_pdf = scipy.stats.norm.logpdf(y_obs, loc=p[0], scale=p[1])
    log_likelihood = np.nansum(log_pdf)
    return log_likelihood


import matplotlib.pyplot as plt
import numpy as np
from math import *

x = np.array([i for i in range(-7, 7)])
y = np.array([sin(i) for i in x])

# Sine Function
# f = sin(x)
plt.plot(x, y, label='sin(x)')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title('The Graph of Trigonometric Functions')
plt.show()

# Cosine Function
y = np.array([cos(i) for i in x])
plt.plot(x, y, label='cos(x)')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()

yx = lambda x: pow(sin(x),2) + pow(cos(x),2)
y = np.array([yx(i) for i in x])
plt.plot(x, y, label='sin(x)**2 + cos(x)**2')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()

# Tangent Function
y = np.array([tan(i) for i in x])
plt.plot(x, y, label='tan(x)')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()


"""Change coordinates"""
from numpy import arccos, arctan2, hypot, sqrt

def to_polar(coo_x, coo_y):
    """
    r, θ = to_polar(x, y)

    Change Cartesian coordinates to Polar coordinates.

    Parameters
    ----------
    x : array_like
        first Cartesian coordinate
    y : array_like
        seconde Cartesian coordinate

    Returns
    -------
    r : ndarray
        radial coordinate
    θ : ndarray
        angular coordinate
    """

    return (hypot(coo_x, coo_y), arctan2(coo_y, coo_x))

def to_spheric(coo_x, coo_y, coo_z):
    """
    r, θ, φ = to_spheric(x, y, z)

    Change Cartesian coordinates to Spherical coordinates.

    Parameters
    ----------
    x : array_like
        first Cartesian coordinate
    y : array_like
        seconde Cartesian coordinate
    z : array_like
        third Cartesian coordinate

    Returns
    -------
    r : ndarray
        radial distance
    θ : ndarray
        azimuthal angle
    φ : ndarray
        polar angle
    """

    coo_r = sqrt(coo_x * coo_x + coo_y * coo_y + coo_z * coo_z)
    return (coo_r, arctan2(coo_y, coo_x), arccos(coo_z / coo_r))


import numpy as np

def ExpRegression(xval, yval):
    """Takes an array of x-values, and y-values and perform a origin exponential regression, with least square method. Returns the function with corresponding constants, and the constants.
        One notes that f(x) = a* e^(-x/b) and thus returns the constants on the form (a,b)."""
    if not isinstance((xval, yval), (np.generic, np.ndarray)):
        if isinstance((xval, yval), (list, tuple)):
            xval, yval = np.array(xval), np.array(yval)
        else:
            raise TypeError("[ExpRegression]: Need a iterable as input.")
    if len(xval) != len(yval):
        raise KeyError("[ExpRegression]: Input xval and yval are not of the same size")
    try:
        A_Matrix = np.array([xval, -np.ones(len(xval))]).T
        A_Matrix_Transposed = A_Matrix.T
        A_AT_Inverse = np.linalg.inv(A_Matrix_Transposed.dot(A_Matrix))
        #Computation to find constants.
        Corrected_yval = np.log(yval)
        ylist = A_Matrix_Transposed.dot(Corrected_yval)
        Constants = A_AT_Inverse.dot(ylist)
        a = np.exp(Constants[1])
        b = Constants[0]
        Expfunc = lambda x: b[0] * np.exp(x/a[0])
        if np.isnan(Constants[0]) or np.isnan(Constants[1]):
            raise TypeError("[ExpRegression]: Performing edited version")
        return Expfunc, (b[0], a[0])
    except Exception as e:
        try:
            print("[ExpRegression]: Performing edited version")
            A_Matrix = np.array([xval, -np.ones(len(xval))]).T
            A_Matrix_Transposed = A_Matrix.T
            A_AT_Inverse = np.linalg.inv(A_Matrix_Transposed.dot(A_Matrix))
            #Computation to find constants.
            ylist1 = np.array([[arg + 1] for arg in yval])
            Corrected_yval = np.log(ylist1)
            ylist = A_Matrix_Transposed.dot(Corrected_yval)
            Constants = A_AT_Inverse.dot(ylist)
            a = np.exp(Constants[1])
            b = Constants[0]
            Expfunc = lambda x: b[0] * np.exp(x/a[0])
            return Expfunc, (b[0],a[0])
        except Exception as e:
            raise e
"""
xlist = (1,2,3,4,5,6,7,8,9)
ylist = (0,0,0,0,1,2,4,8,16)
print(ExpRegression(xlist, ylist))
"""


import numpy as np

def polyfit(x, y, degree):
    '''
    Description: returns the results of the numpy polyfit function
                    polynomial regression
    Parameters: x is input variable (a list of indepedent values)
                y is output variable (a list of dependent values)
                degree order of polynomial to fit x to y
    '''
    results = {}

    coeffs = np.polyfit(x, y, degree)

     # Polynomial Coefficients
    results['polynomial'] = coeffs.tolist()

    # r-squared
    p = np.poly1d(coeffs)
    # fit values, and mean
    yhat = p(x)                      # or [p(z) for z in x]
    ybar = np.sum(y)/len(y)          # or sum(y)/len(y)
    ssreg = np.sum((yhat-ybar)**2)   # or sum([ (yihat - ybar)**2 for yihat in yhat])
    sstot = np.sum((y - ybar)**2)    # or sum([ (yi - ybar)**2 for yi in y])
    results['determination'] = ssreg / sstot

    return results

def regress(i, out, deg):
    '''
    Description: sets up a polynomial regress w/ lists and labels for visualization
    Parameters: x is input variable (a list of indepedent values)
                y is output variable (a list of dependent values)
                degree order of polynomial to fit x to y
    '''

    results = polyfit(i, out, deg)

    txt = 'r_sq: %.4f' % results['determination']
    p = np.poly1d(results['polynomial'])
    x = np.arange(min(i), max(i), 0.01)
    y = p(x)

    return txt,p,x,y


import numpy as np

from tensor import Tensor
from graph import Node

class Operation(Node):
	base_name:str=None
	output:Tensor = None
	inputs:list = []
	def __init__(self):
		super().__init__()
		self.inputs=[]
		self.output=None

	def forward(self):
		raise NotImplementedError
	def backward(self, gradOutput: Tensor):
		raise NotImplementedError

class Negative(Operation):
	base_name="Negative"
	def __init__(self, A: Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)

	def forward(self):
		A = self.inputs[0]
		self.output = -A.output

class Add(Operation):
	base_name = "Add"
	def __init__(self, A: Node,B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name
		self.inputs.append(A)
		self.inputs.append(B)

	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		self.output = A.output + B.output

class Substract(Operation):
	base_name = "Substract"
	def __init__(self, A: Node,B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)
		self.inputs.append(B)

	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		self.output = A.output - B.output

class Divide(Operation):
	base_name = "Divide"
	def __init__(self, A: Node,B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)
		self.inputs.append(B)

	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		self.output = A.output/B.output

class Multiply(Operation):
	base_name = "Multiply"
	def __init__(self, A: Node,B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name
		self.inputs.append(A)
		self.inputs.append(B)
	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		self.output = A.output * B.output

class Sum(Operation):
	base_name = "Sum"
	def __init__(self, A: Node,axis:int = -1, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)
		self.axis = axis

	def forward(self):
		A = self.inputs[0]
		self.output = np.sum(A.output,self.axis)

class MatMul(Operation):
	base_name = "MatMul"
	def __init__(self, A: Node, B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)
		self.inputs.append(B)

	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		self.output = np.matmul( A.output , B.output )

class MatMulNaive(Operation):
	base_name = "MatMulNaive"
	def __init__(self, A: Node,B:Node, Name:str=None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)
		self.inputs.append(B)

	def forward(self):
		A = self.inputs[0]
		B = self.inputs[1]
		cshape=(A.shape[0],B.shape[1])
		C=Tensor([x for x in range(np.prod(cshape))]).reshape(cshape)
		for i in range(0, A.shape[0]):
			for j in range(0, B.shape[1]):
				C[i,j]=0
				for k in range(0, A.shape[1]):
					C[i,j]+= A[i,k]*B[k,j]
		self.output = C

class sin(Operation):
	base_name = "Sin"

	def __init__(self, A: Node, Name: str = None):
		super().__init__()
		if Name is not None:
			self.name = Name

		self.inputs.append(A)

	def forward(self):
		A = self.inputs[0]
		self.output = np.sin(A.output)


import numpy
import torch

def hermgauss(n, dtype=torch.float32):
    x, w = numpy.polynomial.hermite.hermgauss(n)
    x, w = torch.as_tensor(x, dtype=dtype), torch.as_tensor(w, dtype=dtype)
    return x, w

def mvhermgauss(H, D, dtype=torch.float32):
    """
    Return the evaluation locations 'xn', and weights 'wn' for a multivariate
    Gauss-Hermite quadrature.

    The outputs can be used to approximate the following type of integral:
    int exp(-x)*f(x) dx ~ sum_i w[i,:]*f(x[i,:])

    :param H: Number of Gauss-Hermite evaluation points.
    :param D: Number of input dimensions. Needs to be known at call-time.
    :return: eval_locations 'x' ((H**D)xD), weights 'w' (H**D)
    """
    gh_x, gh_w = hermgauss(H)
    x = numpy.array(numpy.meshgrid(*(D*(gh_x,))))
    w = numpy.array(numpy.meshgrid(*(D*(gh_w,)))).prod(1)
    x, w = torch.as_tensor(x, dtype=dtype), torch.as_tensor(w, dtype=dtype)
    return x, w

def mvnquad(f, means, covs, H, Din, Dout=()):
    """
    Computes N Gaussian expectation integrals of a single function 'f'
    using Gauss-Hermite quadrature.
    :param f: integrand function. Takes one input of shape ?xD.
    :param means: NxD
    :param covs: NxDxD
    :param H: Number of Gauss-Hermite evaluation points.
    :param Din: Number of input dimensions. Needs to be known at call-time.
    :param Dout: Number of output dimensions. Defaults to (). Dout is assumed
    to leave out the item index, i.e. f actually maps (?xD)->(?x*Dout).
    :return: quadratures (N,*Dout)
    """
    xn, wn = mvhermgauss(H, Din)
    N = means.size(0)

    # transform points based on Gaussian parameters
    Xt = []
    for c in covs:
        chol_cov = torch.potrf(c, upper=False) # DxD each
        Xt.append(torch.matmul(chol_cov, xn.t()))
    Xt = torch.stack(Xt, dim=0) # NxDx(H**D)
    X = 2.0 ** 0.5 * Xt + means.unsqueeze(2)  # NxDx(H**D)
    Xr = X.permute(2, 0, 1).view(-1, Din)  # (H**D*N)xD

    # perform quadrature
    fX = f(Xr).view(*((H ** Din, N,) + Dout))
    wr = (wn * float(numpi.pi) ** (-Din * 0.5)).view(*((-1,) + (1,) * (1 + len(Dout))))
    return (fX * wr).sum(0)


import numpy as np
import pandas as pd

def numerical_summary(series: pd.Series) -> dict:
    """

    Args:
        series: series to summarize

    Returns:

    """
    aggregates = [
        "mean",
        "std",
        "var",
        "max",
        "min",
        "median",
        "kurt",
        "skew",
        "sum",
        "mad",
    ]
    summary = series.agg(aggregates).to_dict()

    quantiles = [0.05, 0.25, 0.5, 0.75, 0.95]
    for percentile, value in series.quantile(quantiles).to_dict().items():
        summary["quantile_{:d}".format(int(percentile * 100))] = value
    summary["iqr"] = summary["quantile_75"] - summary["quantile_25"]

    summary["range"] = summary["max"] - summary["min"]
    summary["cv"] = summary["std"] / summary["mean"] if summary["mean"] else np.NaN

    # TODO: only calculations for histogram, not the plotting
    # summary['image'] = plotting.histogram(series)
    return summary


import numpy as np
import matplotlib.pyplot as plt

def ImgCompress(img, method, level = 1):
    """
    Image compression using different compression method.
    This function compress image to a user defined compression level.

    SVD Method:
        SVD (Singular Value Decomposition) on RGB image.
        Note: Actual image memory size is not changed from this function.
              It does not handle the Jpeg compression algorithm

    Resize:
        Resize image to different size using interval pixel selection method.

    Parameters
    ----------
    img : numpy.array
        Input image for compression.
    method : string
        Compression methods: Resize, SVD
    level : int
        Level of compression: 1-High, 2-Med, 3-Low

    Returns
    -------
    numpy.array
        The compressed image.

    Examples
    --------
    >>> import matplotlib.pyplot as plt
    >>> from imghelp.ImgCompress import ImgCompress
    >>> image = plt.imread('../test_img/ubc.jpeg')
    >>> plt.imshow(image) #show the image
    >>> img_comp = ImgCompress(image, method='resize', level=1)
    >>> plt.imshow(img_comp) #show the compressed (resized) image
    """
    #Function input checks:
    assert type(img) == np.ndarray, "Invalid image type, expecting numpy.ndarray."
    assert np.min(img.shape[:2]) > 2 and len(img.shape) == 3 , "Input image has invalid dimension. Expecting a 3d array (RGB)."
    assert level in (1,2,3), "Invalid level, expecting 1, 2, 3"
    assert method.lower() in ('svd','resize'), "Invalid method, expecting 'SVD' or 'Resize'."

    #Calculate image size:
    h, w, c = img.shape

    #SVD image compression method:
    if method.lower() == "svd":
        level_scale = [0.05, 0.3, 0.90]
        img_comp = np.zeros(img.shape, dtype="uint8")
        for i in range(c):
            U, S, D = np.linalg.svd(img[:,:,i])
            #x = int((h * w * level_scale[level-1]) / (h + w + 1))
            #x = (np.cumsum(S)/np.sum(S) < level_scale[level-1]).sum()
            x = max(int(len(S) * level_scale[level-1]), 2)
            C = U[:, :x] @ np.diag(S[:x]) @ D[:x, :]
            img_comp[:,:,i] = C.astype("uint8")
        return img_comp

    #Resize method:
    elif method.lower() == "resize":
        level_scale = [0.3, 0.5, 0.70]
        h_new = int(h * level_scale[level-1])
        w_new = int(w * level_scale[level-1])
        img_comp = np.zeros([w_new, h_new, c], dtype="uint8")
        for s in range(c):
            for i in range(w_new):
                for j in range(h_new):
                    img_comp[i, j, s] = img[int(i / level_scale[level-1]), int(j / level_scale[level-1]), s]

        return img_comp

    else:
        pass


import pandas as pd
import sys
import numpy as np

filename = sys.argv[1]
#filename = 'outlier.csv'

def iqr_row_removal(filename):

    data = pd.read_csv(filename)
    nrows, ncols = data.shape

    data2 = [[val for val in data[col]] for col in data]
    data2.insert(0,list(data.keys()))

    for num in range(1, ncols+1):
        if(type(data2[num][0]) != int and type(data2[num][0]) != float):
            data = data.drop(data2[0][num-1], axis=1)
    ncols = data.shape[1]

    row_hash = [0] * nrows
    for i in range(ncols):
        col = data.iloc[:,i]

        q25, q75 = np.percentile(col, [25,75])
        iqr = q75 - q25

        lower_bound = q25 - (iqr * 1.5)
        upper_bound = q75 + (iqr * 1.5)
        for row in range(nrows):
            if (data.iloc[row,i] < lower_bound or data.iloc[row,i] > upper_bound):
                row_hash[row] = 1

    inds = []
    for row in range(nrows):
        if(row_hash[row]):
            inds.append(row)

    data = data.drop(inds, axis = 0)
    nrows2 = data.shape[0]
    print('No. of rows removed by IQR method: ',nrows-nrows2)
    return data

def zscore_row_removal(filename, threshold = 2):

    data = pd.read_csv(filename)
    nrows, ncols = data.shape

    data2 = [[val for val in data[col]] for col in data]
    data2.insert(0,list(data.keys()))

    for num in range(1, ncols+1):
        if(type(data2[num][0]) != int and type(data2[num][0]) != float):
            data = data.drop(data2[0][num-1], axis=1)
    ncols = data.shape[1]

    row_hash = [0] * nrows
    for i in range(ncols):
        col = data.iloc[:,i]

        mean1 = np.mean(col)
        std_dev = np.std(col)
        for row in range(nrows):
            z_score = (data.iloc[row,i] - mean1)/std_dev

            if abs(z_score) > threshold:
                row_hash[row] = 1

    inds = []
    for row in range(nrows):
        if(row_hash[row]):
            inds.append(row)

    data = data.drop(inds, axis = 0)
    nrows2 = data.shape[0]
    print('No. of rows removed by z-score method: ',nrows-nrows2)
    return data

print(iqr_row_removal(filename))
print(zscore_row_removal(filename))


import numpy as np
import os
import fnmatch
import scipy.linalg as sp

# PCA
def makeTwoDim(monsters):
    i,j,k=monsters.shape
    tmp=np.zeros((i*j,k))
    for x in range(k):
        vec=np.reshape(monsters[:,:,x],(i*j),order='F')
        tmp[:,x]=vec
    return tmp

def calcMean(vec):
    i,j=vec.shape
    meanVec=np.zeros(i)
    for x in range(j):
        meanVec+=vec[:,x]/float(j)
    return meanVec

def calcCov(vec):
    i,j=vec.shape
    meanVec=calcMean(vec)
    covMatrix=np.zeros((i,i))
    for x in range(j):
        t1=((vec[:,x]-meanVec).T).reshape(i,1)
        t2=(vec[:,x]-meanVec).reshape(1,i)
        covMatrix+=np.dot(t1,t2)/float(j)
    return covMatrix

def sortEig(eVal, eVec):
    i,j=eVec.shape
    ePair=list(range(j))
    for y in range(j):
        ePair[y]=[(np.abs(eVal[y]), eVec[:,y])]
        ePair[y].sort()
        ePair[y].reverse()
    return ePair

def pairEig(eVal, eVec):
    i,j=eVec.shape
    ePair=list(range(j))
    for y in range(j):
        ePair[y]=[(np.abs(eVal[y]), eVec[:,y])]
    return ePair

def makeTransformMatrix(ePair,pcA,pcB):
    tmp=ePair[0][0][1]
    i=tmp.shape
    i=i[0]
    vec1=ePair[pcA][0][1]
    vec2=ePair[pcB][0][1]

    transform=np.zeros((i,2))
    transform[:,0]=vec1.reshape(i).real
    transform[:,1]=vec2.reshape(i).real
    return transform

def plotTanProj(monsters,pcA,pcB):
    i, j, k = monsters.shape
    twoDim=makeTwoDim(monsters)
    mShape=calcMean(twoDim)
    covMatrix=calcCov(twoDim)
    if k > i * j:  # limit results returned if sample number is less than observations
        self.val, self.vec = sp.eigh(covMatrix)
    else:
        self.val, self.vec = sp.eigh(covMatrix, eigvals=(i * j - k, i * j - 1))
    eigVal=eigVal[::-1]
    eigVec=eigVec[:, ::-1]
    eigPair=pairEig(eigVal,eigVec)
    transform=makeTransformMatrix(eigPair,pcA,pcB)
    transform=np.transpose(transform)
    coords=np.dot(transform,twoDim)
    tmp=np.column_stack((coords[0,:],coords[1,:]))
    return tmp

def plotTanProj(monsters,eigSort,pcA,pcB):
    twoDim=makeTwoDim(monsters)
    transform=makeTransformMatrix(eigSort,pcA,pcB)
    transform=np.transpose(transform)
    coords=np.dot(transform,twoDim)
    tmp=np.column_stack((coords[0,:],coords[1,:]))
    return tmp

############## GPA original
#center shape by removing the mean of each column
def centerShape(shape):
    shape=shape-shape.mean(axis=0)
    return shape

#scale shape by divinding by frobinius norm
def scaleShape(shape):
    shape=shape/np.linalg.norm(shape)
    return shape

#align one shape to a reference shape, soley by rotation
def alignShape(refShape,shape):
    """
    Align the shape to the mean.
    """
    u,s,v=sp.svd(np.dot(np.transpose(refShape),shape), full_matrices=True)
    rotationMatrix=np.dot(np.transpose(v), np.transpose(u))
    shape=np.dot(shape,rotationMatrix)
    return shape

def meanShape(monsters):
    return monsters.mean(axis=2)

def procDist(monsters,mshape):
    i,j,k=monsters.shape
    procDists=np.zeros(k)
    for x in range(k):
        tmp=monsters[:,:,x]-mshape
        procDists[x]=np.linalg.norm(tmp,'fro')
    return procDists

################# GPA update
def runGPA(allLandmarkSets):
  i,j,k=allLandmarkSets.shape
  for index in range(k):
    landmarkSet=allLandmarkSets[:,:,index]
    tempSet = applyCenterScale(landmarkSet)
    allLandmarkSets[:,:,index]=tempSet
  allLandmarkSets = procrustesAlign(allLandmarkSets[:,:,0],allLandmarkSets)
  initialMeanShape=meanShape(allLandmarkSets)
  initialMeanShape = scaleShape(initialMeanShape)
  diff=1
  tries=0
  while diff>0.0001 and tries<5:
    allLandmarkSets = procrustesAlign(initialMeanShape,allLandmarkSets)
    currentMeanShape=meanShape(allLandmarkSets)
    currentMeanShape= scaleShape(currentMeanShape)
    diff=np.linalg.norm(initialMeanShape-currentMeanShape)
    initialMeanShape=currentMeanShape
    tries=tries+1
  return allLandmarkSets, currentMeanShape

def procrustesAlign(mean, allLandmarkSets):
  mean = scaleShape(mean)
  i,j,k=allLandmarkSets.shape
  for index in range(k):
    allLandmarkSets[:,:,index] = alignShape(mean, allLandmarkSets[:,:,index])
  return allLandmarkSets

def applyCenterScale(landmarkSet):
  landmarkSet=centerShape(landmarkSet)
  landmarkSet=scaleShape(landmarkSet)
  return landmarkSet

def runGPANoScale(allLandmarkSets):
  i,j,k=allLandmarkSets.shape
  for index in range(k):
    landmarkSet=allLandmarkSets[:,:,index]
    tempSet = applyCenter(landmarkSet)
    allLandmarkSets[:,:,index]=tempSet
  allLandmarkSets = procrustesAlignNoScale(allLandmarkSets[:,:,0],allLandmarkSets)
  initialMeanShape=meanShape(allLandmarkSets)
  diff=1
  tries=0
  while diff>0.0001 and tries<5:
    allLandmarkSets = procrustesAlignNoScale(initialMeanShape,allLandmarkSets)
    currentMeanShape=meanShape(allLandmarkSets)
    diff=np.linalg.norm(initialMeanShape-currentMeanShape)
    initialMeanShape=currentMeanShape
    tries=tries+1
  return allLandmarkSets, currentMeanShape

def procrustesAlignNoScale(mean, allLandmarkSets):
  i,j,k=allLandmarkSets.shape
  for index in range(k):
    allLandmarkSets[:,:,index] = alignShape(mean, allLandmarkSets[:,:,index])
  return allLandmarkSets

def applyCenter(landmarkSet):
  landmarkSet=centerShape(landmarkSet)
  return landmarkSet


# -*- coding: utf-8 -*-
"""
Module containing different OWA operators.

@author: Javier Fumanal Idocin (UPNA).

To suggest changes or submit new code please use the github page.
"""

import numpy as np
# =============================================================================
#   ~OWAs
# =============================================================================

def owa(X, weights):
    '''

    :param X:
    :param weights:
    :return:
    '''
    X_sorted = np.sort(X) #Order decreciente

    return np.sum(X_sorted * weights[::-1])

def OWA_generic(X, a, b, axis=0, keepdims=True):
    '''
    OWA operator that generates the weight vector using a quantifier
    function determined by a and b. (Check std_quantifier() understand this process)

    :param X: data to aggregate.
    :param a: quantifier parameter 1
    :param b: quantifier parameter 2
    :param axis: axis to reduce.
    :param keepdims: if true, the shape will have the same length.
    :return: matrix with the aggregated axis.
    '''
    X_sorted = -np.sort(-X, axis = axis)
    w = generate_owa_weights(X.shape[axis], lambda x: std_quantifier(x, a=a, b=b))
    X_agg  = np.apply_along_axis(lambda a: np.dot(a, w), axis, X_sorted)

    if keepdims:
        X_agg = np.expand_dims(X_agg, axis=axis)

    return X_agg

def OWA1(X, axis=0, keepdims=True):
    '''
    Performs the OWA operation using a=0.1 and b=0.5 in the quantifier
    function that generates the weights.

    :param X: data to aggregate.
    :param axis: axis to reduce.
    :param keepdims: if true, the shape will have the same length.
    :return: matrix with the aggregated axis.
    '''
    return OWA_generic(X, a=0.1, b=0.5, axis=axis, keepdims=keepdims)

def OWA2(X, axis=0, keepdims=True):
    '''
    Performs the OWA operation using a=0.5 and b=1 in the quantifier
    function that generates the weights.

    :param X: data to aggregate.
    :param axis: axis to reduce.
    :param keepdims: if true, the shape will have the same length.
    :return: matrix with the aggregated axis.
    '''
    return OWA_generic(X, a=0.5, b=1, axis=axis, keepdims=keepdims)

def OWA3(X, axis=0, keepdims=True):
    '''
    Performs the OWA operation using a=0.3 and b=0.8 in the quantifier
    function that generates the weights.

    :param X: data to aggregate.
    :param axis: axis to reduce.
    :param keepdims: if true, the shape will have the same length.
    :return: matrix with the aggregated axis.
    '''
    return OWA_generic(X, a=0.3, b=0.8, axis=axis, keepdims=keepdims)

def generate_owa_weights(n, quantifier):
    '''
    Quantifier function that generates a vector of weights using a quantifier
    function.

    :param quantifier: quantifier function.
    :return: a vector of weights.
    '''
    weights = np.zeros((n))
    for i in range(n):
        ri = i + 1
        weights[i] = quantifier(ri / n) - quantifier((ri - 1) / n)

    return weights

def std_quantifier(x, a=0.0, b=1.0):
    '''
    Standar quantifier, given two parameters: a,b.
    It follows the expression:
        Q = (x-a)/(b-a)
        return 0, if Q < a
        return 1, if Q < b
        return Q, otherwise

    :param a: a real number between 0 and 1.
    :param b: a real number between 0 and 1. Bigger than a.
    :return: the quantifier result.
    '''
    Q = (x-a) / (b - a)
    if x < a:
        return 0
    elif x > b:
        return 1
    else:
        return Q


# Standardize data (0 mean, 1 stdev)
from sklearn.preprocessing import StandardScaler
from pandas import read_csv
from numpy import set_printoptions
filename = 'pima-indians-diabetes.data.csv'
names = ['preg', 'plas', 'pres', 'skin', 'test', 'mass', 'pedi', 'age', 'class']
dataframe = read_csv(filename, names=names)
array = dataframe.values
# separate array into input and output components
X = array[:,0:8]
Y = array[:,8]
scaler = StandardScaler().fit(X)
rescaledX = scaler.transform(X)
# summarize transformed data
set_printoptions(precision=3)
print(rescaledX[0:5,:])

07 Prepare Your Data For Machine Learning/standardize_data.py


import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(1, 100, 100)
y1 = x * np.log2(x)
y2 = x * np.sqrt(x)

plt.plot(x, y1)
plt.plot(x, y2)

plt.grid(True)
plt.savefig('../../img/question_4_plots/h.png')


import numpy as np
from numba import njit, types
from numba.extending import overload, register_jitable
from numba.core.errors import TypingError

import scipy.linalg

@register_jitable
def _oneD_norm_2(a):
    # re-usable implementation of the 2-norm
    val = np.abs(a)
    return np.sqrt(np.sum(val * val))

@overload(scipy.linalg.norm)
def jit_norm(a, ord=None):
    if isinstance(ord, types.Optional):
        ord = ord.type
    # Reject non integer, floating-point or None types for ord
    if not isinstance(ord, (types.Integer, types.Float, types.NoneType)):
        raise TypingError("'ord' must be either integer or floating-point")
    # Reject non-ndarray types
    if not isinstance(a, types.Array):
        raise TypingError("Only accepts NumPy ndarray")
    # Reject ndarrays with non integer or floating-point dtype
    if not isinstance(a.dtype, (types.Integer, types.Float)):
        raise TypingError("Only integer and floating point types accepted")
    # Reject ndarrays with unsupported dimensionality
    if not (0 <= a.ndim <= 2):
        raise TypingError('3D and beyond are not allowed')
    # Implementation for scalars/0d-arrays
    elif a.ndim == 0:
        return a.item()
    # Implementation for vectors
    elif a.ndim == 1:
        def _oneD_norm_x(a, ord=None):
            if ord == 2 or ord is None:
                return _oneD_norm_2(a)
            elif ord == np.inf:
                return np.max(np.abs(a))
            elif ord == -np.inf:
                return np.min(np.abs(a))
            elif ord == 0:
                return np.sum(a != 0)
            elif ord == 1:
                return np.sum(np.abs(a))
            else:
                return np.sum(np.abs(a)**ord)**(1. / ord)
        return _oneD_norm_x
    # Implementation for matrices
    elif a.ndim == 2:
        def _two_D_norm_2(a, ord=None):
            return _oneD_norm_2(a.ravel())
        return _two_D_norm_2

if __name__ == "__main__":
    @njit
    def use(a, ord=None):
        # simple test function to check that the overload works
        return scipy.linalg.norm(a, ord)

    # spot check for vectors
    a = np.arange(10)
    print(use(a))
    print(scipy.linalg.norm(a))

    # spot check for matrices
    b = np.arange(9).reshape((3, 3))
    print(use(b))
    print(scipy.linalg.norm(b))


"""Test gisin."""
import numpy as np

from toqito.states import gisin

def test_gisin_valid():
    """Standard Gisin state."""
    expected_res = np.array(
        [
            [1 / 4, 0, 0, 0],
            [0, 0.35403671, -0.22732436, 0],
            [0, -0.22732436, 0.14596329, 0],
            [0, 0, 0, 1 / 4],
        ]
    )

    res = gisin(0.5, 1)

    bool_mat = np.isclose(res, expected_res)
    np.testing.assert_equal(np.all(bool_mat), True)

def test_gisin_invalid():
    """Invalid Gisin state parameters."""
    with np.testing.assert_raises(ValueError):
        gisin(5, 1)

if __name__ == "__main__":
    np.testing.run_module_suite()