#%%
import numpy as np

#%%
# runge-kutta (order four)
# 5.4.1.b
def func1(t, y):
    return 1 + (t - y) ** 2

def sol1(t):
    return t + 1 / (1 - t)

a = 2
b = 3
N = 10
alpha = 1

h = (b - a) / N
t = a
w = alpha
y = sol1(t)
error = abs(y - w)
print('{0:.1f} & {1:.7f} & {2:.7f} & {3} \\\\'.format(t, y, w, error))

for _ in range(N):
    K_1 = h * func1(t, w)
    K_2 = h * func1(t + h / 2, w + K_1 / 2)
    K_3 = h * func1(t + h / 2, w + K_2 / 2)
    K_4 = h * func1(t + h, w + K_3)
    w += (K_1 + 2*K_2 + 2*K_3 + K_4) / 6
    t += h
    y = sol1(t)
    error = abs(y - w)
    print('{0:.1f} & {1:.7f} & {2:.7f} & {3:.7f} \\\\'.format(t, y, w, error))

# 2.0 & 1.0000000 & 1.0000000 & 0.0 \\
# 2.1 & 1.1909091 & 1.1909088 & 0.0000003 \\
# 2.2 & 1.3666667 & 1.3666663 & 0.0000004 \\
# 2.3 & 1.5307692 & 1.5307688 & 0.0000004 \\
# 2.4 & 1.6857143 & 1.6857138 & 0.0000004 \\
# 2.5 & 1.8333333 & 1.8333329 & 0.0000004 \\
# 2.6 & 1.9750000 & 1.9749996 & 0.0000004 \\
# 2.7 & 2.1117647 & 2.1117643 & 0.0000004 \\
# 2.8 & 2.2444444 & 2.2444441 & 0.0000003 \\
# 2.9 & 2.3736842 & 2.3736839 & 0.0000003 \\
# 3.0 & 2.5000000 & 2.4999997 & 0.0000003 \\

#%%
# runge-kutta 4-4
# 5.6.1.c
def func2(t, y):
    return 1 + y / t

def sol2(t):
    return t * np.log(t) + 2 * t

a = 1
b = 2
N = 10
alpha = 2

h = (b - a) / N
t = a
w = alpha
y = sol2(t)
error = abs(y - w)
print('{0:.1f} & {1:.7f} & {2:.7f} & {3} \\\\'.format(t, y, w, error))

tLi = [t]
wLi = [w]
for _ in range(3):
    K_1 = h * func2(t, w)
    K_2 = h * func2(t + h / 2, w + K_1 / 2)
    K_3 = h * func2(t + h / 2, w + K_2 / 2)
    K_4 = h * func2(t + h, w + K_3)
    w += (K_1 + 2*K_2 + 2*K_3 + K_4) / 6
    t += h

    tLi.append(t)
    wLi.append(w)

    y = sol2(t)
    error = abs(y - w)
    print('{0:.1f} & {1:.7f} & {2:.7f} & {3:.7f} \\\\'.format(t, y, w, error))

for _ in range(3, N):
    t += h
    w = wLi[3] + h * (55 * func2(tLi[3], wLi[3]) - 59 * func2(tLi[2], wLi[2]) + 37 * func2(tLi[1], wLi[1]) - 9 * func2(tLi[0], wLi[0])) / 24
    w = wLi[3] + h * (9 * func2(t, w) + 19 * func2(tLi[3], wLi[3]) - 5 * func2(tLi[2], wLi[2]) + func2(tLi[1], wLi[1])) / 24
    y = sol2(t)
    error = abs(y - w)
    print('{0:.1f} & {1:.7f} & {2:.7f} & {3:.7f} \\\\'.format(t, y, w, error))

    tLi.pop(0)
    wLi.pop(0)
    tLi.append(t)
    wLi.append(w)

# 1.0 & 2.0000000 & 2.0000000 & 0.0 \\
# 1.1 & 2.3048412 & 2.3048409 & 0.0000003 \\
# 1.2 & 2.6187859 & 2.6187854 & 0.0000005 \\
# 1.3 & 2.9410735 & 2.9410729 & 0.0000007 \\
# 1.4 & 3.2710611 & 3.2710601 & 0.0000010 \\
# 1.5 & 3.6081977 & 3.6081963 & 0.0000014 \\
# 1.6 & 3.9520058 & 3.9520041 & 0.0000017 \\
# 1.7 & 4.3020680 & 4.3020661 & 0.0000020 \\
# 1.8 & 4.6580160 & 4.6580138 & 0.0000022 \\
# 1.9 & 5.0195224 & 5.0195199 & 0.0000025 \\
# 2.0 & 5.3862944 & 5.3862917 & 0.0000027 \\

#%%
# stiff
# 5.11.1.b
def func3(t, y):
    return -20 * (y - t ** 2) + 2 * t

def func3dy(t, y):
    return -20

def sol3(t):
    return t ** 2 + (1 / 3) * np.exp(-20 * t)

a = 0
b = 1
N = 10
alpha = 1 / 3
TOL = 1e-6
M = 10

h = (b - a) / N
t = a
w = alpha
y = sol3(t)
error = abs(y - w)
print('{0:.1f} & {1:.7f} & {2:.7f} & {3} \\\\'.format(t, y, w, error))

for _ in range(N):
    k_1 = w + h * func3(t, w) / 2
    w_0 = k_1
    j = 1
    FLAG = 0
    while FLAG == 0 and j <= M:
        w = w_0 - (w_0 - h * func3(t + h, w_0) / 2 - k_1) / (1 - h * func3dy(t+h, w_0) / 2)
        if abs(w - w_0) < TOL:
            FLAG = 1
        else:
            j += 1
            w_0 = w

    if j > M:
        print('The maxmum number of iterations exceeded')

    t += h
    y = sol3(t)
    error = abs(y - w)
    print('{0:.1f} & {1:.7f} & {2:.7f} & {3:.7f} \\\\'.format(t, y, w, error))

# 0.0 & 0.3333333 & 0.3333333 & 0.0 \\
# 0.1 & 0.0551118 & 0.0100000 & 0.0451118 \\
# 0.2 & 0.0461052 & 0.0400000 & 0.0061052 \\
# 0.3 & 0.0908263 & 0.0900000 & 0.0008263 \\
# 0.4 & 0.1601118 & 0.1600000 & 0.0001118 \\
# 0.5 & 0.2500151 & 0.2500000 & 0.0000151 \\
# 0.6 & 0.3600020 & 0.3600000 & 0.0000020 \\
# 0.7 & 0.4900003 & 0.4900000 & 0.0000003 \\
# 0.8 & 0.6400000 & 0.6400000 & 0.0000000 \\
# 0.9 & 0.8100000 & 0.8100000 & 0.0000000 \\
# 1.0 & 1.0000000 & 1.0000000 & 0.0000000 \\


import matplotlib.pyplot as plt
import numpy as np

import cvxpy as cp

def loss_fn(X, Y, beta):
    return cp.pnorm(cp.matmul(X, beta) - Y, p=2)**2

def regularizer(beta):
    return cp.pnorm(beta, p=2)**2

def objective_fn(X, Y, beta, lambd):
    return loss_fn(X, Y, beta) + lambd * regularizer(beta)

def mse(X, Y, beta):
    return (1.0 / X.shape[0]) * loss_fn(X, Y, beta).value

def generate_data(m: int = 1000, n: int = 30, sigma: int = 40):
    """Generates data for regression.

    To experiment with your own data, just replace the contents of this
    function with code that loads your dataset.

    Args
    ----
    m : int
        The number of examples.
    n : int
        The number of features per example.
    sigma : positive float
        The standard deviation of the additive noise.

    Returns
    -------
    X : np.array
        An array of featurized examples, shape (m, n), m the number of
        examples and n the number of features per example.

    Y : np.array
        An array of shape (m,) containing the observed labels for the
        examples.

    beta_star : np.array
        The true parameter. This is the quantity we are trying to
        estimate.
    """
    beta_star = np.random.randn(n)
    # Generate an ill-conditioned data matrix
    X = np.random.randn(m, n)
    U, _, V = np.linalg.svd(X)
    s = np.linspace(30, 1, min(m, n))
    S = np.zeros((m, n))
    S[:min(m, n), :min(m, n)] = np.diag(s)
    X = np.dot(U, np.dot(S, V))
    # Corrupt the observations with additive Gaussian noise
    Y = X.dot(beta_star) + np.random.normal(0, sigma, size=m)
    return X, Y, beta_star

def plot_train_test_errors(train_errors, test_errors, lambd_values):
    plt.plot(lambd_values, train_errors, label="Train error")
    plt.plot(lambd_values, test_errors, label="Test error")
    plt.xscale("log")
    plt.legend(loc="upper left")
    plt.xlabel(r"$\lambda$", fontsize=16)
    plt.title("Mean Squared Error (mSE)")
    plt.show()

def plot_regularization_path(lambd_values, beta_values):
    num_coeffs = len(beta_values[0])
    for i in range(num_coeffs):
        plt.plot(lambd_values, [wi[i] for wi in beta_values])
    plt.xlabel(r"$\lambda$", fontsize=16)
    plt.xscale("log")
    plt.title("Regularization Path")
    plt.show()

if __name__ == "__main__":
    m = 1000
    n = 30
    sigma = 4

    X, Y, beta_star = generate_data(m, n, sigma)
    X_train = X[:800, :]
    Y_train = Y[:800]
    X_test = X[800:, :]
    Y_test = Y[800:]

    beta = cp.Variable(n)
    lambd = cp.Parameter(nonneg=True)
    problem = cp.Problem(
                    cp.Minimize(objective_fn(X_train, Y_train, beta, lambd)))

    lambd_values = np.logspace(-2, 2, 50)
    train_errors = []
    test_errors = []
    beta_values = []
    for v in lambd_values:
        lambd.value = v
        problem.solve()
        train_errors.append(mse(X_train, Y_train, beta))
        test_errors.append(mse(X_test, Y_test, beta))
        beta_values.append(beta.value)

    plot_train_test_errors(train_errors, test_errors, lambd_values)
    plot_regularization_path(lambd_values, beta_values)


'''
Problem 44

Pentagonal numbers are generated by the formula, Pn=n(3n−1)/2. The first ten pentagonal numbers are:

1, 5, 12, 22, 35, 51, 70, 92, 117, 145, ...

It can be seen that P4 + P7 = 22 + 70 = 92 = P8. However, their difference, 70 − 22 = 48, is not pentagonal.

Find the pair of pentagonal numbers, Pj and Pk, for which their sum and difference are pentagonal and D = |Pk − Pj| is minimised; what is the value of D?

'''

import numpy as np

def pentagonal(a: int):

    pent = lambda x: x*(3*x-1)/2
    arr = np.arange(a)
    i = 0
    while i < len(arr):
        arr[i] = pent(i)
        i += 1

    return arr[arr != 0]

def find_pent(a: int):
    arr = pentagonal(a)
    for i in range(0, len(arr)):
        for j in range(0, len(arr)):
            if arr[i] + arr[j] in arr:
                if arr[j] - arr[i] in arr:
                    return arr[np.where(arr == arr[i] + arr[j])], arr[np.where(arr == arr[j] - arr[i])], arr[i], arr[j], arr[j]-arr[i]

def main():

    a = 10000
    print('Pj and Pk are: ', find_pent(a))
    #print(distinct_factors(factors(644),644))

if __name__=='__main__':
    main()


'''
    author: Virag Dosani (@repl.it/@VTech1699)
'''

import sympy
from sympy import *

def readData(f):
    with open(f) as f:
        data = []
        da = f.read().split("\n")
        for line in da:
            data.append(line.split("="))

    return data

def algebraicSimplify(data):
    x,y,z = symbols('x y z')
    simplfy = []
    for i in data:
        expr = eval(i[1])
        smpl = simplify(expr)
        simplfy.append([i[0],expr, smpl])

    return simplfy

def printData(data):
    print("Input Code: ")
    for i in data:
        print(i[0]," = ",i[1])

    print("\nCode after Algebraic Simplification: ")
    for i in data:
        print(i[0]," = ",i[2])

    return

if __name__ == '__main__':
    input = "input.txt"
    data = readData(input)
    output = algebraicSimplify(data)
    print("Algebraic Simplification")
    printData(output)

'''
OUTPUT:
Algebraic Simplification
Input Code:
a   =  sin(x)**2 + cos(x)**2
b   =  (x**3 + x**2 - x - 1)/(x**2 + 2*x + 1)
c   =  gamma(x)/gamma(x - 2)
d   =  (x**2 + y**2 + z**2)/(-x**2 - y**2 - z**2)

Code after Algebraic Simplification:
a   =  1
b   =  x - 1
c   =  (x - 2)*(x - 1)
d   =  -1
'''

Semester 6/SPCC/Code optimization - Algebraic simplification/main.py


from numpy import zeros, array_equal

from ubcs_auxiliary import traveling_salesman as ts

def test_euclidian_distance():
    x = (1,2,3)
    y = (2,3,4)
    assert (ts.euclidean_distance(x,y) == 3**0.5)
    x = (1,2,3)
    y = (2,4,4)
    assert (ts.euclidean_distance(x,y) == (1**2+2**2+1**2)**0.5)

def test_create_test_sequence():
    N = 20
    dim = 2
    arr,vary = ts.create_test_sequence(N,dim = dim)
    assert arr.shape == (N,dim)
    assert vary[0] == 0
    assert vary[-1] == 0
    assert vary.sum() == N-2

def test_total_distance():
    from numpy import array
    arr = array([[ 0.        ,  0.        ],
       [-3.16353381, -4.95307408],
       [-2.73465273, -3.9936309 ],
       [ 3.90539788, -0.03624822],
       [-3.26566328,  3.78071339],
       [ 4.05165855, -1.08678949],
       [ 2.27780317,  3.78714615],
       [ 4.76618918, -4.29571196],
       [ 0.89862702, -4.12542064],
       [ 0.        ,  0.        ]])
    vary = array([0., 1., 1., 1., 1., 1., 1., 1., 1., 0.])
    assert round(ts.total_distance(arr),7) == round(53.30737952083856,7)
    # rounded to 7 decimals due to precisiob in arr


# -*- coding: cp1252 -*-
# ------------------------- #
# Description:
# A collection of functions which implement special maths functions in my Masters dissertation.
# ------------------------- #
# Created by: Ashiv Dhondea, RRSG, UCT.
# Date created: 21 July 2016
# Edits: 22 July 2016: added fn_Create_Concatenated_Block_Diag_Matrix, originally from DynamicsFunctions.py
#        28 July 2016: implemented the function invSymQuadForm from the Tracker Component Library.
# ------------------------- #
# Import libraries
import numpy as np
from numpy import linalg
import scipy.linalg
# --------------------------------------------------------------------------- #
def schol(A): # Cholesky decomposition for PSD matrices.
    ## Emulates schol.m of ekfukf toolbox.
    ## Description from schol.m
    ##    %SCHOL  Cholesky factorization for positive semidefinite matrices
    ##% Syntax:
    ##%   [L,def] = schol(A)
    ## % In:
    ##%   A - Symmetric pos.semi.def matrix to be factorized
    ##%
    ##% Out:
    ##%   L   - Lower triangular matrix such that A=L*L' if def>=0.
    ##%   def - Value 1,0,-1 denoting that A was positive definite,
    ##%         positive semidefinite or negative definite, respectively.
    ## % Copyright (C) 2006 Simo Särkkä
    n = np.shape(A)[0];
    L = np.zeros((n,n),dtype=np.float64);
    definite = 1;

    for i in range(0,n):
        for j in range(0,i+1):
            s = A[i,j];
            for k in range (0,j):
                s = s - L[i,k]*L[j,k];
            if j < i :
                if L[j,j] > np.finfo(np.float64).eps:
                    L[i,j] = s/L[j,j];
                else:
                    L[i,j] = 0;
            else:
                if (s < - np.finfo(np.float64).eps ):
                    s = 0;
                    definite = -1;
                elif (s < np.finfo(np.float64).eps):
                    s = 0;
                    definite = min(0,definite);

                L[j,j] = np.sqrt(s);
    # if definite < 0, then negative definite
    # if definite == 0, then positive semidefinite
    # if definite == 1, then positive definite
    return L, definite
# ------------------------------------------------------------------- #
def fn_Create_Concatenated_Block_Diag_Matrix(R,stacklen):
    # fn_Create_Concatenated_Block_Diag_Matrix creates a block diagonal matrix of size (stacklen) x (stacklen)
    # whose diagonal blocks are copies of the matrix R.
##    L = [R];
##    for index in range (0,stacklen):
##        L.append(R);
##        ryn = scipy.linalg.block_diag(*L);
    ryn =np.kron(np.eye(stacklen+1),R); # Edit: 19/07/2016: probably better idea than using a for loop.
    return ryn
# ------------------------------------------------------------------------- #
def invSymQuadForm(x,M):
    """
    %%INVSYMQUADFORM Compute the quadratic form x'*inv(R)*x in a manner that
    %                should be more numerically stable than directly evaluating
    %                the matrix inverse, where R is a symmetric, positive
    %                definite matrix. The parameter M can either be the matrix
    %                R directly (R=M) (the default), or the lower-triangular
    %                square root of R (R=M*M'). The distance  can be computed
    %                for a single matrix and multiple vectors x. If x consists
    %                of vector differences and R is a covariance matrix, then
    %                this can be viewed as a numerically robust method of
    %                computing Mahalanobis distances.
    %
    %INPUTS: x       An mXN matrix of N vectors whose quadratic forms with the
    %                inverse of R are desired.
    %        M       The mXm real, symmetric, positive definite matrix R, or
    %                the square root of R, as specified by the following
    %                parameter.
    %        matType This optional parameter specified the type of matrix that
    %                M is. Possible values are
    %                0 (the default if omitted) M is the matrix R in the form
    %                  x'*inv(R)*x.
    %                1 M is the lower-triangular square root of the matrix R.
    %
    %OUTPUTS: dist   An NX1 vector of the values of x*inv(R)*x for every vector
    %                in the matrix x.
    %
    %As one can find in many textbooks, solvign A*x=b using matrix inversion is
    %generally a bad idea. This relates to the problem at hand, because one can
    %write
    %x'*inv(R)*x=x'inv(C*C')*x
    %           =x'*inv(C)'*inv(C)*x
    %where C is the lower-triangular Cholesky decomposition of R. Next, say
    %that
    %y=inv(C)*x.
    %Then, we are just computing y'*y.
    %What is a stable way to find y? Well, we can rewrite the equation as
    %C*y=x
    %Since C is lower-triangular, we can find x using forward substitution.
    %This should be the same as one of the many ways that Matlab solves the
    %equation C*y=x when one uses the \ operator. One can explicitely tell
    %Matlab that the matrix is lower triangular when using the linsolve
    %function, thus avoiding the need for loops or for Matlab to check the
    %structure of the matrix on its own.
    %
    %August 2014 David F. Crouse, Naval Research Laboratory, Washington D.C.
    %(UNCLASSIFIED) DISTRIBUTION STATEMENT A. Approved for public release.
    """
    ## 28 July 2016. Emulates the invSymQuadForm.m function from the Tracker Component Library. Removed matType to make things simpler.

    # implements x'M x
    C,definiteness = schol(M);
    y = np.linalg.solve(C,x);
    dist = np.dot(y.T,y);
    return dist


import fromage.utils.array_operations as ao
import numpy as np
from numpy.testing import assert_allclose
from pytest import approx

def test_distance(vec_100,vec_220):
    distance = ao.distance(vec_100,vec_220)
    assert distance == approx(2.236,rel=0.001)

def test_closest(vec_100,vec_220):
    origin = np.array([0.,0.,0.])
    close_point, dis = ao.closest(origin,[vec_100,vec_220])
    assert_allclose(close_point,vec_100)
    assert dis == approx(1.)

def test_dist_mat(h2o_dim_dist_arr):
    dis = ao.dist_mat(h2o_dim_dist_arr)
    assert dis[1][0] == approx(0.91,rel=0.1)

def test_find_largest(h2o_dim_dist_arr):
    largest_inds = ao.find_largest(h2o_dim_dist_arr,3)
    assert largest_inds == [(4, 0), (4, 2), (5, 0)]

# Angle tests
def test_angle_between_vectors(vec_100, vec_220):
    ang = ao.vec_angle(vec_100, vec_220)
    assert ang == approx(45.)

def test_angle_between_vectors_rad(vec_100, vec_220):
    ang = ao.vec_angle(vec_100, vec_220, degrees = False)
    assert ang == approx(45.*np.pi/180)

# Orthog test
def test_orthogonalise_sym():
    vectors = np.array([[0.1,0.,1.],
                        [-0.1,0.,1.]])

    o_vecs = ao.orthogonalise_sym(vectors)
    ovec_1 = o_vecs[0]/np.linalg.norm(o_vecs[0])
    ovec_2 = o_vecs[1]/np.linalg.norm(o_vecs[1])

    expected_1 = np.array([np.sqrt(2)/2,0.,np.sqrt(2)/2])
    expected_2 = np.array([-np.sqrt(2)/2,0.,np.sqrt(2)/2])

    assert_allclose(ovec_1,expected_1)
    assert_allclose(ovec_2,expected_2)

def test_orthogonalise_asym():
    vectors = np.array([[0.,0.,1.],
                        [0.,1.,1.]])

    o_vecs = ao.orthogonalise_asym(vectors)
    ovec_1 = o_vecs[0]/np.linalg.norm(o_vecs[0])
    ovec_2 = o_vecs[1]/np.linalg.norm(o_vecs[1])

    expected_1 = np.array([0.,0.,1.])
    expected_2 = np.array([0.,1.,0.])

    assert_allclose(ovec_1,expected_1)
    assert_allclose(ovec_2,expected_2,atol=1e-05)

def test_rmsd():
    ref = np.array([0.,0.,0.,0.,0.])
    alt = np.array([1.,-1.,-1.,-1.,1.])
    assert ao.rmsd(ref,alt) == 1.


"""Generalized Pauli matrices."""
import numpy as np

from toqito.matrices import shift
from toqito.matrices import clock

def gen_pauli(k_1: int, k_2: int, dim: int) -> np.ndarray:
    r"""
    Produce generalized Pauli operator [WikGenPaul]_.

    Generates a :code:`dim`-by-:code:`dim` unitary operator. More specifically,
    it is the operator :math:`X^k_1*Z^k_2`, where :math:`X` and :math:`Z` are
    the "shift" and "clock" operators that naturally generalize the Pauli X and
    Z operators. These matrices span the entire space of
    :code:`dim`-by-:code:`dim` matrices as :code:`k_1` and :code:`k_2` range
    from 0 to :code:`dim-1`, inclusive.

    Note that the generalized Pauli operators are also known by the name of
    "discrete Weyl operators". [WatrousLec6]_

    Examples
    ==========

    The generalized Pauli operator for :code:`k_1 = 1`, :code:`k_2 = 0` and
    :code:`dim = 2` is given as the standard Pauli-X matrix

    .. math::
        G_{1, 0, 2} = \begin{pmatrix}
                         0 & 1 \\
                         1 & 0
                      \end{pmatrix}.

    This can be obtained in :code:`toqito` as follows.

    >>> from toqito.matrices import gen_pauli
    >>> dim = 2
    >>> k_1 = 1
    >>> k_2 = 0
    >>> gen_pauli(k_1, k_2, dim)
    [[0.+0.j, 1.+0.j],
     [1.+0.j, 0.+0.j]])

    The generalized Pauli matrix :code:`k_1 = 1`, :code:`k_2 = 1`, and
    :code:`dim = 2` is given as the standard Pauli-Y matrix

    .. math::
        G_{1, 1, 2} = \begin{pmatrix}
                        0 & -1 \\
                        1 & 0
                      \end{pmatrix}.

    This can be obtained in :code:`toqito` as follows.`

    >>> from toqito.matrices import gen_pauli
    >>> dim = 2
    >>> k_1 = 1
    >>> k_2 = 1
    >>> gen_pauli(k_1, k_2, dim)
    [[ 0.+0.0000000e+00j, -1.+1.2246468e-16j],
     [ 1.+0.0000000e+00j,  0.+0.0000000e+00j]])

    References
    ==========
    .. [WikGenPaul] Wikipedia: Generalizations of Pauli matrices
        https://en.wikipedia.org/wiki/Generalizations_of_Pauli_matrices

    .. [WatrousLec6] Lecture 6: Further remarks on measurements and channels
        https://cs.uwaterloo.ca/~watrous/LectureNotes/CS766.Fall2011/06.pdf

    :param k_1: (a non-negative integer from 0 to :code:`dim-1` inclusive).
    :param k_2: (a non-negative integer from 0 to :code:`dim-1` inclusive).
    :param dim: (a positive integer indicating the dimension).
    :return: A generalized Pauli operator.
    """
    gen_pauli_x = shift(dim)
    gen_pauli_z = clock(dim)

    gen_pauli_w = np.linalg.matrix_power(gen_pauli_x, k_1) @ np.linalg.matrix_power(
        gen_pauli_z, k_2
    )

    return gen_pauli_w


import sys
import scipy.io as sio
from pprint import pprint
import matplotlib.pyplot as plt
import numpy as np

# h(X(i)) = theta_0 + theta_1 * X1(i)

def normalize(x):
   return (x - x.min())/(x.max() - x.min())

# theta = (X.T X)^(-1)(X.T)(Y)
def normal_equation(X, y):
    theta = []
    m = np.size(x)
    bias = np.ones((m, 1))
    X = np.reshape(X, (m, 1))
    X = np.append(bias, X, axis=1)
    theta = np.linalg.inv(X.T.dot(X)).dot(X.T).dot(y)
    return theta

x = np.array([5, 15, 25, 35, 45, 55])
X = normalize(x)
y = np.array([15, 11, 2, 8, 25, 32])

theta = normal_equation(X, y)
print("theta_0", theta[0], "\ntheta_1", theta[1])

y_pred = theta[0] + theta[1] * X
error = y - (y_pred)
print("Error:", error)

plt.figure(figsize=(8, 6))
plt.scatter(X, y, color = "r", marker = "o", s = 30)
plt.title("Predicted Regression Line", fontsize=16)
plt.plot(X, y_pred, color = "b")
plt.legend(['h(x)', '(x, y)'], loc='best')

plt.xlabel('X', fontsize=14)
plt.ylabel('Y', fontsize=14)
plt.show()


import numpy as np
'''
    Exponential kernel to measure the similarity betweeen two points.
'''
class GaussianKernel():
    def __init__(self, sig1=0.01, sig2=0.1):
        self.sig1 = sig1;
        self.sig2 = sig2;

    def dot(self, x1, x2, sig1=None, sig2=None):
        if (sig1 is None):
            sig1 = self.sig1;
        if (sig2 is None):
            sig2 = self.sig2;

        diff = np.linalg.norm(x1 - x2);
        return np.exp(-(diff ** 2) / (2 * sig1 * sig2));


#############################################################
# This is a class that collects statistical distributions.
# This distributions are the essential ones to be remembered.
#############################################################

import numpy as np
import random

from factory import *

class StatisticalDistributions(object):

	def __init__(self):
		super().__init__()

	def Exponential(self, x:float=0, l:float=0, vector:list=None):

		'''
		param x: this is the real number to be processed into an Exponential Distributions
		param l: this is the lambda value of an Exponential Distributions
		param vector: passing this list, these values will be indexes to get corresponding exponential values
		              This value should start from zero and has no upper limit (X belongs to REAL)
		'''

		if vector is None:
			return l * np.exp( -l * x )
		else:
			return [ self.Exponential(x, l, vector=None) for x in vector ]

	def Gaussian(self, x:float=0, mu:float=0, sigma:float=1, vector:list=None):

		'''
		param mu: this is the real mean of our gaussian distribution
		param sigma: this is the variance of our gaussian distribution
		param vector: passing this list, these values will be indexes to get corresponding exponential values
		'''

		if vector is None:
			return (1 / (sigma * (2 * np.pi) ** 0.5 )) * ( np.exp( -((x - mu) / (2 * sigma))**2 ) )
		else:
			return [ self.Gaussian(x=x / 100, mu=mu, sigma=sigma, vector=None) for x in vector ]

	def Poisson(self, x:int=0, l:float=0, vector:list=None):

		'''
		param l: this is the lambda value of a Poisson Distribution
		param vector: passing this list, these values will be indexes to get corresponding exponential values
		              This value should start from zero and has no upper limit (X belongs to REAL)
		'''

		if vector is None:
			return (np.exp(-l) * (l ** x)) / ( self._Factorial( x ) )
		else:
			return [ self.Poisson(x=x, l=l, vector=None) for x in vector ]

	def Binomial(self, n:int, k:int, p:float):

		'''
		param n: number of time the experiment is conducted
		param k: number of time the experiment is a success (k <= n)
		param p: probability a success occurs
		'''

		factory_    = StatisticalIndexes()
		combination = factory_.Combination(n, k)

		return combination * (p ** k) * ( (1 - p) ** (n - k) )


import math
import numpy as np
import scipy.io.wavfile
from scipy import signal
from scipy.ndimage.interpolation import map_coordinates
from imageio import imread
from skimage.color import rgb2gray

GRAYSCALE = 1

def DFT(signal):
    """
    Compute the one-dimensional discrete Fourier Transform
    :param signal: the original signal as array of float64 with shape (N,1)
    :return: DFT of signal as array of complex128 with shape (N,1)
    """
    N = signal.shape[0]
    if N == 0:
        return signal
    # create grid for DFT matrix
    x = np.arange(N)
    DFT_matrix = np.exp((-2 * math.pi * 1j * x.reshape(N, 1) * x) / N)
    return signal.T.dot(DFT_matrix).T

def IDFT(fourier_signal):
    """
    Compute the one-dimensional inverse discrete Fourier Transform
    :param fourier_signal: the fourier signal as array of complex128 with shape (N,1)
    :return: the original signal as array of float64 with shape (N,1)
    """
    N = fourier_signal.shape[0]
    if N == 0:
        return fourier_signal
    # create grid for IDFT matrix
    x = np.arange(N)
    IDFT_matrix = np.exp((2 * math.pi * 1j * x.reshape(N, 1) * x) / N) / N
    return np.real_if_close(fourier_signal.T.dot(IDFT_matrix).T).astype(np.complex128)

def DFT2(image):
    """
    Compute the 2-dimensional discrete Fourier Transform
    :param image: grayscale image of float64 with shape (M,N,1)
    :return: DFT of image as array of complex128 with shape (M,N,1)
    """
    # M, N = image.shape[:2]
    # image = image.astype(np.complex128)
    # for row in range(M):  # DFT over rows
    #     image[row] = DFT(image[row])
    #
    # for col in range(N):  # DFT over cols
    #     image[:, col] = DFT(image[:, col])
    #
    # return image

    # apply DFT matrix on all cols and then DFT matrix again on all rows
    # swapaxes is used to handle case of (M,N,1) where transpose doesn't work as needed
    return DFT(DFT(image).swapaxes(0, 1)).swapaxes(0, 1)

def IDFT2(fourier_image):
    """
    Compute the 2-dimensional inverse discrete Fourier Transform
    :param fourier_image: 2D array of complex128 with shape (M,N,1)
    :return: the original image as array of float64 with shape (M,N,1)
    """
    # M, N = fourier_image.shape[:2]
    # image = fourier_image.astype(np.complex128)
    # for row in range(M):  # DFT over rows
    #     image[row] = IDFT(image[row])
    #
    # for col in range(N):  # DFT over cols
    #     image[:, col] = IDFT(image[:, col])
    #
    # return np.real_if_close(image).astype(np.complex128)

    # apply IDFT matrix on all cols and then IDFT matrix again on all rows
    # swapaxes is used to handle case of (M,N,1) where transpose doesn't work as needed
    return np.real_if_close(IDFT(IDFT(fourier_image).swapaxes(0, 1))).swapaxes(0, 1).astype(np.complex128)

def change_rate(filename, ratio):
    """
    Changes the duration of an audio file by changing the sampling rate
    :param filename: string representing the path to a WAV file
    :param ratio: positive float64 representing the duration change
    """
    rate, data = scipy.io.wavfile.read(filename)
    scipy.io.wavfile.write('change_rate.wav', int(rate * ratio), data)

def change_samples(filename, ratio):
    """
    Changes the duration of an audio file by reducing the number of samples using Fourier
    :param filename: string representing the path to a WAV file
    :param ratio: positive float64 representing the duration change
    :return: 1D ndarray of dtype float64 representing the new sample points
    """
    rate, data = scipy.io.wavfile.read(filename)
    data = data.astype(np.float64)
    # data = data / np.iinfo(np.int16).max  # uncomment if want to save as normalized values for audio to make sense
    data = resize(data, ratio)
    scipy.io.wavfile.write('change_samples.wav', rate, data)
    return data

def resize_spectrogram(data, ratio):
    """
    Speeds up a WAV file, without changing the pitch, using spectrogram scaling
    :param data: 1D ndarray of dtype float64 representing the original sample points,
    :param ratio: positive float64 representing the rate change of the WAV file
    :return: new sample points according to ratio with the same datatype as data
    """
    spectrogram = stft(data)
    new_data = []
    for row in range(spectrogram.shape[0]):
        new_data.append(resize(spectrogram[row], ratio))
    new_data = np.vstack(new_data)
    new_data = istft(new_data)
    return new_data

def resize_vocoder(data, ratio):
    """
    Speeds up a WAV file by phase vocoding its spectrogram
    :param data: 1D ndarray of dtype float64 representing the original sample points,
    :param ratio: positive float64 representing the duration change
    :return: data rescaled according to ratio with the same data type as data
    """
    spectrogram = stft(data)
    new_data = phase_vocoder(spectrogram, ratio)
    return istft(new_data)

def resize(data, ratio):
    """
    Changes the number of samples by the given ratio
    :param data: 1D ndarray of dtype float64 or complex128 representing the original sample points
    :param ratio: positive float64 representing the duration change
    :return: 1D ndarray of the dtype of data representing the new sample points
    """
    if ratio == 1:
        return data
    data_type = data.dtype
    data = DFT(data)
    data = np.fft.fftshift(data)

    delta_samples = abs(data.size - math.floor(data.size / ratio))  # number of samples to add or remove
    if ratio > 1:  # speed up, remove samples
        data = data[math.floor(delta_samples / 2):data.size - math.ceil(delta_samples / 2)]  # clip data
    else:  # slow down, add padding samples
        left_pad = np.zeros(math.floor(delta_samples / 2))
        right_pad = np.zeros(math.ceil(delta_samples / 2))
        data = np.append(left_pad, np.append(data, right_pad))

    data = np.fft.ifftshift(data)
    data = IDFT(data)
    data = data.astype(data_type)
    return data

def conv_der(im):
    """
    Computes the magnitude of image derivatives using convolution
    :param im: grayscale image of type float64
    :return: magnitude of the derivative, with the same dtype and shape
    """
    conv_vec = np.array([[0.5, 0, -0.5]])
    dx = scipy.signal.convolve2d(im, conv_vec, mode='same')
    dy = scipy.signal.convolve2d(im, conv_vec.reshape(3, 1), mode='same')
    return np.sqrt(np.abs(dx) ** 2 + np.abs(dy) ** 2)

# noinspection DuplicatedCode
def fourier_der(im):
    """
    Computes the magnitude of image derivatives using Fourier transform
    :param im: grayscale image of type float64
    :return: magnitude of the derivative, with the same dtype and shape
    """
    rows, cols = im.shape
    fourier_transform = np.fft.fftshift(DFT2(im))

    # get dx by making [-N/2,N/2] coef vec
    x_coef = np.arange(int(-cols / 2), math.ceil(cols / 2)) * ((2 * np.pi * 1j) / cols)
    dx = fourier_transform * x_coef
    dx = IDFT2(np.fft.ifftshift(dx))

    # get dy by making [-M/2, M/2] coef vec
    y_coef = np.arange(int(-rows / 2), math.ceil(rows / 2)) * ((2 * np.pi * 1j) / rows)
    dy = fourier_transform * y_coef[:, None]  # make y_coef col vector
    dy = IDFT2(np.fft.ifftshift(dy))

    return np.sqrt(np.abs(dx) ** 2 + np.abs(dy) ** 2)

#  - - - - - - - - - - added functions from helper and prev ex- - - - - - - - - - - - - -
def read_image(filename, representation):
    """
    Reads an image file and converts it into a given representation
    :param filename: the filename of an image on disk
    :param representation: representation code, either 1 or 2 defining whether the output should be a
    grayscale image (1) or an RGB image (2)
    :return: the image represented by a matrix of type np.float64
    """
    im = imread(filename)
    if representation == GRAYSCALE and im.ndim == 3:
        im = rgb2gray(im)
        return im
    im = im.astype(np.float64)
    im /= 255
    return im

def stft(y, win_length=640, hop_length=160):
    fft_window = signal.windows.hann(win_length, False)

    # Window the time series.
    n_frames = 1 + (len(y) - win_length) // hop_length
    frames = [y[s:s + win_length] for s in np.arange(n_frames) * hop_length]

    stft_matrix = np.fft.fft(fft_window * frames, axis=1)
    return stft_matrix.T

def istft(stft_matrix, win_length=640, hop_length=160):
    n_frames = stft_matrix.shape[1]
    y_rec = np.zeros(win_length + hop_length * (n_frames - 1), dtype=np.float)
    ifft_window_sum = np.zeros_like(y_rec)

    ifft_window = signal.windows.hann(win_length, False)[:, np.newaxis]
    win_sq = ifft_window.squeeze() ** 2

    # invert the block and apply the window function
    ytmp = ifft_window * np.fft.ifft(stft_matrix, axis=0).real

    for frame in range(n_frames):
        frame_start = frame * hop_length
        frame_end = frame_start + win_length
        y_rec[frame_start: frame_end] += ytmp[:, frame]
        ifft_window_sum[frame_start: frame_end] += win_sq

    # Normalize by sum of squared window
    y_rec[ifft_window_sum > 0] /= ifft_window_sum[ifft_window_sum > 0]
    return y_rec

def phase_vocoder(spec, ratio):
    num_timesteps = int(spec.shape[1] / ratio)
    time_steps = np.arange(num_timesteps) * ratio

    # interpolate magnitude
    yy = np.meshgrid(np.arange(time_steps.size), np.arange(spec.shape[0]))[1]
    xx = np.zeros_like(yy)
    coordiantes = [yy, time_steps + xx]
    warped_spec = map_coordinates(np.abs(spec), coordiantes, mode='reflect', order=1).astype(np.complex)

    # phase vocoder
    # Phase accumulator; initialize to the first sample
    spec_angle = np.pad(np.angle(spec), [(0, 0), (0, 1)], mode='constant')
    phase_acc = spec_angle[:, 0]

    for (t, step) in enumerate(np.floor(time_steps).astype(np.int)):
        # Store to output array
        warped_spec[:, t] *= np.exp(1j * phase_acc)

        # Compute phase advance
        dphase = (spec_angle[:, step + 1] - spec_angle[:, step])

        # Wrap to -pi:pi range
        dphase = np.mod(dphase - np.pi, 2 * np.pi) - np.pi

        # Accumulate phase
        phase_acc += dphase

    return warped_spec


import numpy as np
import math
import engine.utils
from engine.utils.constants import *

def get_norm(vector):
    return sum([x**2 for x in vector])**0.5

def normalize(vector, fall_back=None):
    norm = get_norm(vector)
    if norm > 0:
        return np.array(vector) / norm
    elif fall_back is not None:
        return fall_back
    else:
        return np.zeros(len(vector))

def unit_vector(vector):
    """Returns the unit vector of the vector."""
    return vector / np.linalg.norm(vector)

def angle_between(v1, v2):
    """Finds angle between two vectors."""
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1., 1.))

def z_rotation(vector, theta):
    """Rotates a 3-D vector around z-axis"""
    R = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
    return np.dot(R, vector)


# -*- coding: utf-8 -*-
"""
Created on Sun Feb  2 11:47:27 2020

@author: Rapha
"""

import math
#import ctypes
import numpy as np
import OpenGL.GL as gl

class OrbitPolygon():

    def __init__(self, semi_major_axis, semi_minor_axis, focus, num_vertices):

        if(num_vertices < 10):
            num_vertices = 10

        self.__numVertices = num_vertices
        self.__semiMajorAxis = semi_major_axis
        self.__semiMinorAxis = semi_minor_axis
        self.__focus = focus

        self.__vertices = np.empty(self.__numVertices * 4, dtype=np.float32)
        self.__calculateVertices()

    def __calculateVertices(self):

        inc_theta = 2 * math.pi / self.__numVertices
        theta = 0

        vertices_pos = 0

        for i in range(self.__numVertices):

            x, y = self.getPositionAt(theta)

            self.__vertices[vertices_pos] = x
            self.__vertices[vertices_pos + 1] = y
            self.__vertices[vertices_pos + 2] = 0.0
            self.__vertices[vertices_pos + 3] = 1.0

            vertices_pos += 4
            theta += inc_theta

    def getPositionAt(self, angle):

        x = self.__semiMajorAxis * np.sin(angle) + self.__focus
        y = self.__semiMinorAxis * np.cos(angle)

        return (x, y)

    def getVertexPositions(self):

        return self.__vertices


import numpy as np
import matplotlib.pyplot as plt

f1 = 5 # Hertz
w1 = 2*np.pi*f1 # angular frequency
f2 = 50 # Hertz
w2 = 2*np.pi*f2 # angular frequency

# Create an array of time points
N = 500
t = np.linspace(0,5,N)

# Create signals
pure = np.sin(w1*t) # pure sine wave
noise = np.sin(w2*t) + np.random.normal(0, 0.2, N) # Normally distributed random noise
 # 0   = mean of the normal distribution
 # 0.2 = standard deviation of the normal distribution
 # N   = number of elements in the array
signal = pure + noise # noisy sine wave

#################################################
### Code your recursive low pass filter below ###
#################################################
###  Use test run to visualize your solution  ###
#################################################
### Then submit it when your done to grade it ###
#################################################

# Recursive low-pass filter
alpha = 0.2 # alpha closer to zero means stronger smoothing effect
          # be sure to try differnt values and observe its effects in the plot below
filtered = np.zeros(N) # initialize array to store values
for i in range(N):
    if i > 1:
        # Fill out the filtered[i] variable with the code for your low pass filer
        filtered[i] = filtered[i-1] + alpha*(signal[i] - filtered[i-1])

fig = plt.figure()
ax1 = fig.add_subplot(311)
ax1.plot(t, pure)
ax1.set_xlabel('Time, (sec)')
ax1.set_ylabel('Amplitude, (m)')
ax1.set_title('5 Hz, Pure Signal')

ax2 = fig.add_subplot(312)
ax2.plot(t, signal)
ax2.set_xlabel('Time, (sec)')
ax2.set_ylabel('Amplitude, (m)')
ax2.set_title('Noisy Signal')

ax3 = fig.add_subplot(313)
ax3.plot(t, filtered)
ax3.set_xlabel('Time, (sec)')
ax3.set_ylabel('Amplitude, (m)')
ax3.set_title('Filtered Signal')

plt.tight_layout()
plt.show()


#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
    Copyright 2018 Bryan Webb

    Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
    associated documentation files (the "Software"), to deal in the Software without restriction,
    including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
    and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
    subject to the following conditions:

        The above copyright notice and this permission notice shall be included in all copies
        or substantial portions of the Software.

        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
        INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE
        AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
        DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

"""
Adapted from:

https://glowingpython.blogspot.com/2012/03/linear-regression-with-numpy.html
"""

import numpy as np

def linear_least_squares_fit_xy_lists(x_list, y_list):
    """
    Calculate the least squares fit of corresponding lists of X and Y values
    as in:

        y = mx + b

    :param x_list: list of x values
    :param y_list: list of y values to be LS Fit to the list of x values

    :return:    (slope, y-intercept) list
    """
    assert isinstance(x_list, list) and 2 <= len(x_list)
    assert isinstance(y_list, list) and 2 <= len(y_list)
    assert len(x_list) == len(y_list)
    assert all([isinstance(v, float) for v in x_list])
    assert all([isinstance(v, float) for v in y_list])

    try:
        x_matrix = np.array([x_list, np.ones(len(x_list))])
        mb_list = np.linalg.lstsq(x_matrix.T, y_list)[0]
    except:
        mb_list = [None for i in range(len(x_list))]

    return mb_list

def linear_least_squares_fit_tup_list(v_list):
    """
    Calculate the least squares fit of a list of (X, Y) tuples, as in:

        y = mx + b

    :param v_list:  a list of (x, y) tuples, where x and y are floats

    :return:    (slope, y-intercept) tuple
    """
    assert isinstance(v_list, list) and 2 <=len(v_list)
    assert all([isinstance(v, (list, tuple)) and 2 == len(v) for v in v_list])

    x_list = [v[0] for v in v_list]
    y_list = [v[1] for v in v_list]

    assert all([isinstance(v, float) for v in x_list])
    assert all([isinstance(v, float) for v in y_list])

    mb_list = linear_least_squares_fit_xy_lists(x_list, y_list)

    return mb_list

if "__main__" == __name__:
    # Nominal test case
    x_list = []
    y_list = []
    tup_list = []

    for i in range(10):
        x = float(i)
        for j in range(i):
            y = float(j)
            x_list.append(x)
            y_list.append(y)
            tup = (x, y)
            tup_list.append(tup)

    print linear_least_squares_fit_xy_lists(x_list, y_list)
    print linear_least_squares_fit_tup_list(tup_list)

    # Reasonable result
    tup_list = [(0.0, 0.0), (0.0, 1.0), (1.0, 0.0), (1.0, 1.0)]
    x_list = [v[0] for v in tup_list]
    y_list = [v[1] for v in tup_list]

    print linear_least_squares_fit_xy_lists(x_list, y_list)
    print linear_least_squares_fit_tup_list(tup_list)

    # This test case should have yielded infinite slope m
    # If the data is set so that it always yields an ~ 45deg angle, this slope of 0.0 will never happen
    x_list = []
    y_list = []
    tup_list = []

    for i in range(10):
        x = 0.0
        y = float(i)
        x_list.append(x)
        y_list.append(y)
        tup = (x, y)
        tup_list.append(tup)

    print linear_least_squares_fit_xy_lists(x_list, y_list)
    print linear_least_squares_fit_tup_list(tup_list)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Dec  2 15:37:10 2020

Copyright 2020 by Hadrien Montanelli.
"""
# %% Imports.

# Standard library imports:
import matplotlib.pyplot as plt
import numpy as np
import time

# Chebpy imports:
from chebpy.cheb import chebpts, coeffs2vals, multmat, spconvert, vals2coeffs

# %% Multiplication g = f*h on [-1,1].

# Functions:
f = lambda x: np.cos(x)*np.exp(x)
h = lambda x: np.sin(x)*x**2
gex = lambda x: f(x)*h(x)

# Grid:
n = 10000
x = chebpts(n)

# Compute coeffs of f:
F = vals2coeffs(f(x))

# Multiplication by h matrix in coefficient space:
start = time.time()
M = multmat(n, h)
end = time.time()
print(f'Time   (setup): {end-start:.5f}s')
plt.figure()
plt.spy(M)

# Multiply:
start = time.time()
G = M @ F
end = time.time()
print(f'Time (product): {end-start:.5f}s')

# Convert to value space:
g = coeffs2vals(G)

# Error:
error = np.max(np.abs(g - gex(x)))/np.max(np.abs(gex(x)))
print(f'Error  (L-inf): {error:.2e}')

# %% Multiplication g = f*h on [0,2].

# Functions:
f = lambda x: np.cos(x)*np.exp(x)
h = lambda x: np.sin(x)*x**2
gex = lambda x: f(x)*h(x)

# Grid:
n = 10000
x = chebpts(n, [0, 2])

# Compute coeffs of f:
F = vals2coeffs(f(x))

# Multiplication by h matrix in coefficient space:
start = time.time()
M = multmat(n, h, [0, 2])
end = time.time()
print(f'Time   (setup): {end-start:.5f}s')
plt.figure()
plt.spy(M)

# Multiply:
start = time.time()
G = M @ F
end = time.time()
print(f'Time (product): {end-start:.5f}s')

# Convert to value space:
g = coeffs2vals(G)

# Error:
error = np.max(np.abs(g - gex(x)))/np.max(np.abs(gex(x)))
print(f'Error  (L-inf): {error:.2e}')

# %% Multiplication g = f*h in the C^{(lam)} basis.

# Basis lambda:
lam = 2

# Functions:
f = lambda x: np.cos(x)*np.exp(x)
h = lambda x: np.sin(x)*x**2
gex = lambda x: f(x)*h(x)

# Grid:
n = 400
x = chebpts(n)

# Compute coeffs of f and convert to the C^{(lam)} basis:
F = vals2coeffs(f(x))
for k in range(lam):
    S = spconvert(n, k)
    F = S @ F

# Multiplication by h matrix in in the C^{(lam)} basis:
start = time.time()
M = multmat(n, h, [-1, 1], lam)
end = time.time()
print(f'Time   (setup): {end-start:.5f}s')
plt.figure()
plt.spy(M)

# Multiply in the C^{(lam)} basis:
start = time.time()
G = M @ F
end = time.time()
print(f'Time (product): {end-start:.5f}s')

# Exact coefficients in the C^{(lam)} basis:
Gex = vals2coeffs(gex(x))
for k in range(lam):
    S = spconvert(n, k)
    Gex = S @ Gex

# Error:
error = np.max(np.abs(G - Gex))/np.max(np.abs(Gex))
print(f'Error  (L-inf): {error:.2e}')


import numpy as np
from math import floor, cos, radians, degrees, acos

class ColorConverter():

    @staticmethod
    def rgb_to_gray(rgb):
        return np.dot(rgb[...,:3], [0.2989, 0.5870, 0.1140])

    @staticmethod
    def rgb_to_gray_via_weighted_average(r, g, b):
        return (r * 0.2989) + (g * 0.5870) + (b * 0.1140)

    @staticmethod
    def rgb_to_gray_via_simple_average(r, g, b):
        return (r + g + b) / 3

    @staticmethod
    def calculate_hsv_hue(r, g, b):
        max_value = max(r, g, b)
        min_value = min(r, g, b)
        diff = max_value - min_value
        if max_value == min_value:
            hue = 0
        elif max_value == r:
            hue = (60 * ((g - b) / diff) + 360) % 360
        elif max_value == g:
            hue = (60 * ((b - r) / diff) + 120) % 360

        return hue

    @staticmethod
    def calculate_hsv_saturation(r, g, b):
        max_value = max(r, g, b)
        min_value = min(r, g, b)
        diff = max_value - min_value

        if max_value == 0:
            saturation = 0
        else:
            saturation = diff / max_value

        return saturation

    @staticmethod
    def rgb_to_hsv(r, g, b):
        '''
        Convert a RGB image to HSV image.
        Input: R, G, B values are [0, 255].
        Output: H value is [0, 360]. S, V values are [0, 1].
        '''
        r, g, b = r / 255.0, g / 255.0, b / 255.0

        h = ColorConverter.calculate_hsv_hue(r, g, b)
        s = ColorConverter.calculate_hsv_saturation(r, g, b)
        v = max(r, g, b)
        return int(h), s, v

    @staticmethod
    def hsv_to_rgb(h, s, v):
        '''
        Convert a HSV image to RGB image.
        Input: H value is [0, 360]. S, V values are [0, 1].
        Output: R, G, B values are [0, 255].
        '''
        h = float(h)
        s = float(s)
        v = float(v)
        h60 = h / 60.0
        h60f = floor(h60)
        hi = int(h60f) % 6
        f = h60 - h60f
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        r, g, b = 0, 0, 0
        if hi == 0:
            r, g, b = v, t, p
        elif hi == 1:
            r, g, b = q, v, p
        elif hi == 2:
            r, g, b = p, v, t
        elif hi == 3:
            r, g, b = p, q, v
        elif hi == 4:
            r, g, b = t, p, v
        elif hi == 5:
            r, g, b = v, p, q
        return ColorConverter.normalize_rgb(r, g, b)

    @staticmethod
    def rgb_to_hsi(r, g, b):
        '''
        Convert a RGB image to HSI image.
        Input: R, G, B values are [0, 255].
        Output: H value is [0, 360]. S, I values are [0, 1].
        '''
        r, g, b = ColorConverter.normalize_to_zero_one(r, g, b)

        # Hue
        numerator = 0.5 * ((r - g) + (r - b))
        denominator = ((r - g) ** 2 + ((r - b) * (g - b))) ** 0.5
        # A small number was added in the donominator to avoid dividing by 0
        theta = degrees(acos(numerator / (denominator + 0.000001)))
        if b <= g:
            h = theta
        else:
            h = 360 - theta

        # Saturation
        s = 1 - ((3 / (r + g + b)) * min([r, g, b]))

        # Intensity
        i = (r + g + b) / 3

        return np.round(h, 2), np.round(s, 2), np.round(i, 2)

    @staticmethod
    def hsi_to_rgb(h, s, i):
        '''
        Convert a HSI image to RGB image.
        Input: H value is [0, 360]. S, I values are [0, 1].
        Output: R, G, B values are [0, 255].
        '''
        r, g, b = [0, 0, 0]

        # RG sector (0 degrees <= hue < 120 degrees)
        if 0 <= h < 120:
            b = i * (1 - s)
            r = i * (1 + ((s * cos(radians(h))) / (cos(radians(60 - h)))))
            g = 3 * i - (r + b)

        # GB sector (120 degrees <= hue < 240 degrees)
        elif 120 <= h < 240:
            h = h - 120
            r = i * (1 - s)
            g = i * (1 + ((s * cos(radians(h))) / (cos(radians(60 - h)))))
            b = 3 * i - (r + g)

        # BR sector (240 degrees <= hue <= 360 degrees)
        elif 240 <= h <= 360:
            h = h - 240
            g = i * (1 - s)
            b = i * (1 + ((s * cos(radians(h))) / (cos(radians(60 - h)))))
            r = 3 * i - (b + g)

        return ColorConverter.normalize_rgb(r, g, b)

    @staticmethod
    def normalize_to_zero_one(r, g, b):
        '''
        Convert RGB values from [0, 255] to [0, 1] range.
        '''
        r = float(r) / 255.0
        g = float(g) / 255.0
        b = float(b) / 255.0

        return r, g, b

    @staticmethod
    def normalize_rgb(r, g, b):
        '''
        Convert RGB values from [0, 1] to [0, 255] range.
        '''
        r = int(r * 255)
        g = int(g * 255)
        b = int(b * 255)

        return r, g, b


from numpy import matrix, linalg

def gcd(a,b):
    a,b = (a,b) if a>b else (b,a)
    rem = a%b
    return gcd(b, rem) if rem != 0 else b

def pulverizer(a,b):
    x,y = (a,b) if a>b else (b,a)
    q = x//y
    rem = x%y
    print (rem,q)
    return pulverizer(y, rem) if rem != 0 else y

def bezout(a,b):
    """Finding Bezout's Identity by Extended Euclid's method"""
    px,py = (a,b)
    a,b = (a,b) if a>b else (b,a)
    x,y = a,b
    q = a//b
    rem = a%b
    s = t = 0
    s_2, s_1 = [1, 0]
    t_2, t_1 = [0, 1]

    while rem != 0:
        #print (rem,q)
        q, rem = a//b, a%b
        a,b = (b,rem)
        s = s_2 - q*s_1
        t = t_2 - q*t_1
        s_2, s_1 = s_1, s
        t_2, t_1 = t_1, t
    assert abs(a*s) == abs(y)
    assert abs(a*t) == abs(x)
    s_2, t_2 = (s_2, t_2) if px>py else (t_2, s_2)
    assert (s_2*px + t_2*py) == a
    #print a, b
    #print s, t
    #print (a*s), (a*t)
    #print x,y
    #print px,py
    return a, s_2, t_2

def multiplicative_inverse(a,n):
    t = 0
    newt = 1
    r = n
    newr = a
    while newr != 0:
        quotient = r // newr
        (t, newt) = (newt, t - quotient * newt)
        (r, newr) = (newr, r - quotient * newr)
    if r > 1: return "a is not invertible"
    if t < 0: t = t + n
    return t

#REWORK - FIX CONCEPT!
def pulverizer_bezout_matrix(a,b):
    """Pulverizing to find Bezout's Identity by Matrix Method"""
    mat_list = []
    a,b = (a,b) if a>b else (b,a)
    x,y = a,b
    q = a//b
    rem = a%b

    while rem != 0:
        #return pulverizer(y, rem) if rem != 0 else y
        #print (rem,q)
        mat_list.append(matrix([[q, 1], [1,0]]))
        a,b = (b,rem) if b>rem else (rem,b)
        q, rem = a//b, a%b

    #print a, b, rem  #b - previous remainder - r[N-1] or r0 when rem == 0
    mat_list.append(matrix([[q, 1], [1,0]]))
    rem_mat = matrix([[b],[rem]])
    m = reduce(lambda m1,m2: (m1*m2), mat_list)
    #print mat_list
    #print rem_mat
    #print m * rem_mat
    mInv = linalg.inv(m)
    #print mInv
    #print  mInv * matrix([[x],[y]])
    assert all([i for i in (mInv * matrix([[x],[y]]) == rem_mat)])
    #print mInv[0,0], mInv[0,1]
    assert long(mInv[0,0]*x + mInv[0,1]*y) == b
    return (mInv[0,0], mInv[0,1], mInv[0,0]*x + mInv[0,1]*y, b)

def euler_tot(n):
    """Euler's Totient Function"""
    return len([i for i in range(1, n) if gcd(n, i)==1])


import B3_VanillaOptions_MC_BlackScholes as B3MC
import B3_VanillaOptionsinBlackScholesWorld as B3BS
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import norm

def delta_call_BSM(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    d1 = (np.log(spot/strike) + (rate - dividend + 0.5*vol**2)*(maturity))/(vol*np.sqrt(maturity))
    Nd1 = norm.cdf(d1)
    return np.exp(-dividend*maturity)*Nd1

def _dN(x):
    return np.exp((-x**2)/2)/(np.sqrt(2*np.pi))

def gamma_call_BSM(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    d1 = (np.log(spot/strike) + (rate - dividend + 0.5*vol**2)*(maturity))/(vol*np.sqrt(maturity))
    dN = _dN(d1)
    return (dN/(spot*vol*np.sqrt(maturity)))*np.exp(-dividend*maturity)

def vega_call_BSM(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    d1 = (np.log(spot/strike) + (rate - dividend + 0.5*vol**2)*(maturity))/(vol*np.sqrt(maturity))
    dN = _dN(d1)
    return np.exp(-dividend*maturity)*spot*np.sqrt(maturity)*dN

def theta_call_BSM(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    d1 = (np.log(spot/strike) + (rate - dividend + 0.5*vol**2)*(maturity))/(vol*np.sqrt(maturity))
    d2 = d1 - vol*np.sqrt(maturity)
    dN = _dN(d1)
    return -np.exp(-dividend*maturity)*spot*dN*vol/(2*np.sqrt(maturity)) + \
           dividend*np.exp(-dividend*maturity)*spot*norm.cdf(d1) - \
           rate*strike*np.exp(-rate*maturity)*norm.cdf(d2)

def rho_call_BSM(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    d1 = (np.log(spot/strike) + (rate - dividend + 0.5*vol**2)*(maturity))/(vol*np.sqrt(maturity))
    d2 = d1 - vol*np.sqrt(maturity)
    return strike*maturity*np.exp(-rate*maturity)*norm.cdf(d2)

def delta_call_finite_difference(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    e = 0.0001  ##this works better than using 0.0001*spot
    call_up = B3BS.europeanCallOptionPrice(spot + e, strike, maturity, rate, dividend, vol)
    call = B3BS.europeanCallOptionPrice(spot, strike, maturity, rate, dividend, vol)
    return (call_up-call)/e

def gamma_call_finite_difference(spot,strike,maturity,rate,vol,dividend=None):
    if dividend is None:
        dividend = 0
    e = 0.0001
    call_up = B3BS.europeanCallOptionPrice(spot + e, strike, maturity, rate, dividend, vol)
    call_dn = B3BS.europeanCallOptionPrice(spot - e, strike, maturity, rate, dividend, vol)
    call = B3BS.europeanCallOptionPrice(spot, strike, maturity, rate, dividend, vol)
    return (call_up - 2*call + call_dn)/e**2

def vega_call_finite_difference(spot, strike, maturity, rate, vol, dividend=None):
    if dividend is None:
        dividend = 0
    e = 0.0001  ##this works better than using 0.0001*spot
    call_up = B3BS.europeanCallOptionPrice(spot , strike, maturity, rate, dividend, vol + e)
    call = B3BS.europeanCallOptionPrice(spot, strike, maturity, rate, dividend, vol)
    return (call_up-call)/e

def rho_call_finite_difference(spot, strike, maturity, rate, vol, dividend=None):
    if dividend is None:
        dividend = 0
    e = 0.0001  ##this works better than using 0.0001*spot
    call_up = B3BS.europeanCallOptionPrice(spot , strike, maturity, rate + e, dividend, vol)
    call = B3BS.europeanCallOptionPrice(spot, strike, maturity, rate, dividend, vol)
    return (call_up-call)/e

def theta_call_finite_difference(spot, strike, maturity, rate, vol, dividend=None):
    if dividend is None:
        dividend = 0
    e = 0.0001  ##this works better than using 0.0001*spot
    call_up = B3BS.europeanCallOptionPrice(spot , strike, maturity + e, rate, dividend, vol)
    call = B3BS.europeanCallOptionPrice(spot, strike, maturity, rate, dividend, vol)
    return (call_up-call)/e


import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig
from tools import get_img

class Image_Filters:
    """
    A class simply to hold all image filters for more elegant code
    """

    def __init__(self):
        self.img_store = {}
        return

    # complex filters
    Sobel_0_3 = np.array([[-1, 0, 1],
                         [ -2, 0, 2],
                         [ -1, 0, 1]])
    Sobel_45_3 = np.array([[0, 1, 2],
                          [-1, 0, 1],
                          [-2,-1, 0]])
    Sobel_0_5 = np.array([[1, 2, 3, 2, 1],
                         [ 2, 3, 5, 3, 2],
                         [ 0, 0, 0, 0, 0],
                         [-2,-3,-5,-3,-2],
                         [-1,-2,-3,-2,-1]])
    Sobel_45_5 = np.array([[0, 2, 3, 2, 1],
                          [-2, 0, 3, 5, 2],
                          [-3,-3, 0, 3, 3],
                          [-2,-5,-3, 0, 2],
                          [-1,-2,-3,-2, 0]])
    Sobel_90_3 = Sobel_0_3.T
    Sobel_135_3 = Sobel_45_3.T
    Sobel_90_5 = Sobel_0_5.T
    Sobel_135_5 = Sobel_45_5.T
    Laws_L5 = np.array([1,4,6,4,1])
    Laws_E5 = np.array([-1,-2,0,2,1])
    Laws_S5 = np.array([-1,0,2,0,-1])
    Laws_W5 = np.array([-1,2,0,-2,1])
    Laws_R5 = np.array([1,-4,6,-4,1])

    def avg(self, size):
        return np.ones([size, size]) / (size*size)

    def laplacian_8(self, size):
        assert size%2==1, 'size must be odd'
        lap = np.ones([size, size])
        i = int(size/2)
        lap[i, i] = 1 - size*size
        return lap

    def gabor_fn(self, sigma, theta, Lambda, gamma=1., psi=0.):
        """
        Code from https://en.wikipedia.org/wiki/Gabor_filter#cite_note-8
        sigma: std of the gaussian
        theta: angle of the sinasoid in 2d
        Lambda: sinasoid wavelength
        psi: phase ofset of the sinasoid
        gamma: elpisive factor in the gaussian
        """
        sigma_x = sigma
        sigma_y = float(sigma) / gamma

        # Bounding box
        nstds = 3 # Number of standard deviation sigma
        xmax = max(abs(nstds * sigma_x * np.cos(theta)),
                   abs(nstds * sigma_y * np.sin(theta)))
        xmax = np.ceil(max(1, xmax))
        ymax = max(abs(nstds * sigma_x * np.sin(theta)),
                   abs(nstds * sigma_y * np.cos(theta)))
        ymax = np.ceil(max(1, ymax))
        xmin = -xmax
        ymin = -ymax
        (y, x) = np.meshgrid(np.arange(ymin, ymax + 1),
                             np.arange(xmin, xmax + 1))

        # Rotation
        x_theta = x * np.cos(theta) + y * np.sin(theta)
        y_theta = -x * np.sin(theta) + y * np.cos(theta)

        gauss = np.exp(-(x_theta**2 / sigma_x**2 + y_theta**2 / sigma_y**2) / 2)
        trig = np.cos(2 * np.pi / Lambda * x_theta + psi)
        return gauss*trig

    def gabor_bank(self, n_theta, sigma, Lambda, gamma=1., psi=0.,
                   verbose=False):

        angs = np.linspace(0, 180, n_theta)
        filt_bank = [self.gabor_fn(sigma, np.deg2rad(ang), Lambda, gamma, psi)
                     for ang in angs]

        if verbose:
            x = 4
            y = int(n_theta/x) + 1
            fig, axs = plt.subplots(y, x, figsize=[16, 16])
            for filt, ax, ang in zip(filt_bank,
                                     axs.ravel()[:n_theta],
                                     angs):
                ax.imshow(filt, cmap='gray')
                ax.set(title=ang)
                ax.axis('off')

            # remove unused axis
            for ax in axs.ravel()[n_theta:]:
                fig.delaxes(ax)

        return filt_bank

    def apply_bank(self, filt_bank, grey_img, normed=True, verbose=False):
        output = np.zeros_like(grey_img)
        for filt in filt_bank:
            output += sig.convolve2d(grey_img, filt, mode='same')

        if normed:
            output -= output.min()
            output /= output.max()

        if verbose:
            fig, ax = plt.subplots(figsize=[22, 11])
            col = ax.imshow(output)
            plt.colorbar(col)

        return output

    def mask_outine(self, mask, size=3):
        """
        Take a 2d mask and use a laplacian convolution to find the segment
        outlines for plotting. Option decides if all directions are to be
        included 'full' or just horizontal and vertical ones 'edge'
        """
        conv = sig.convolve2d(mask, self.laplacian_8(size), mode='valid')
        conv = np.pad(conv, 1, 'edge') # ignore edges
        return 1. - np.isclose(conv, 0) # not zero grad (i.e. an edge)

    def sobel_edge_detection(self, grey_img, verbose=False):
        """
        Sobel edge detection convolution in 4 diretions returning a 3d array
        with angles (0, 45, 90, 135) on the final axis.
        """
        assert len(grey_img.shape) == 2, 'grey_img must be 2d'

        zeros = np.zeros_like(grey_img)
        edge_det = np.dstack([zeros, zeros, zeros, zeros])

        for i, filt in enumerate([self.Sobel_0_5, self.Sobel_45_5,
                                  self.Sobel_90_5, self.Sobel_135_5]):
            edge_det[:, :, i] = sig.convolve2d(grey_img, filt, mode='same')

        if verbose:
            for i in range(4):
                plt.figure(figsize=[22,22])
                plt.imshow(edge_det[:,:,i])
                plt.title('angle ' + str(i*45) + ' deg')

        return edge_det

    def rgba(self, mask, color='r', opaqueness=1):
        "Take a 2d mask and return a 4d rgba mask for imshow overlaying"

        # create the transparent mask
        zeros = np.zeros_like(mask)
        rgba = np.dstack([zeros, zeros, zeros, mask*opaqueness])

        # set the correct color channel
        i = ['r', 'g', 'b'].index(color)
        rgba[:, :, i] = mask

        return rgba

    def laws_texture_filters(self, size):

        filts = []
        arrs = [self.Laws_L5, self.Laws_E5, self.Laws_S5,
                self.Laws_W5, self.Laws_R5]
        arrs = [np.repeat(arr, size) for arr in arrs]

        for arr1 in arrs:
            for arr2 in arrs:
                filts.append(np.outer(arr1.T, arr2.T))

        return filts

    def laws_texture_features(self, grey_img, size, verbose=False):

        assert len(grey_img.shape) == 2, 'image must be grey'

        filts = self.laws_texture_filters(size)
        n_filts = len(filts)

        output = np.empty([*grey_img.shape, n_filts])

        for i, filt in enumerate(filts):
            output[:, :, i] = sig.convolve2d(grey_img, filt, mode='same')
            output[:, :, i] -= output[:, :, i].mean()
            output[:, :, i] -= output[:, :, i].std()

            if verbose:
                fig, ax = plt.subplots(figsize=[22,11])
                col = ax.imshow(output[:, :, i])
                plt.colorbar(col)

        return output

if __name__ == '__main__':

    # gabor filter example
    IF_obj = Image_Filters()
    bank = IF_obj.gabor_bank(13, 5, 20, gamma=1., verbose=True)

    img = get_img('/content/images/TX1_white_cropped.tif')
    grey_img = img.mean(axis=2)

    IF_obj.apply_bank(bank, grey_img, verbose=True)

    # sobel_edge_detection example
    IF_obj = Image_Filters()

    img = get_img('/content/images/TX1_white_cropped.tif')
    grey_img = img.mean(axis=2)

    IF_obj.sobel_edge_detection(grey_img, verbose=True)


import numpy as np

class SimpleLinearRegression1:

    def __init__(self):
        self.a_ = None
        self.b_ = None

    def fit(self, x_train, y_train):
        x_mean = np.mean(x_train)
        y_mean = np.mean(y_train)

        num = 0.0
        d = 0.0
        for i, j in zip(x_train, y_train):
            num += (i - x_mean) * (j - y_mean)
            d += (i - x_mean) ** 2

        self.a_ = num / d
        self.b_ = y_mean - self.a_ * x_mean
        return self

    def predict(self, x_predict):
        """
        :param x_predict: np.array()
        :return:
        """
        return np.array([self._predict(x) for x in x_predict])

    def _predict(self, x_single):
        return self.a_ * x_single + self.b_

    def __repr__(self):
        return "SimpleLinearRegression1()"

class SimpleLinearRegression2:

    def __init__(self):
        self.a_ = None
        self.b_ = None

    def fit(self, x_train, y_train):
        x_mean = np.mean(x_train)
        y_mean = np.mean(y_train)

        num = (x_train - x_mean).dot(y_train - y_mean)
        d = (x_train - x_mean).dot(x_train - x_mean)

        self.a_ = num / d
        self.b_ = y_mean - self.a_ * x_mean
        return self

    def predict(self, x_predict):
        """
        :param x_predict: np.array()
        :return:
        """
        return np.array([self._predict(x) for x in x_predict])

    def _predict(self, x_single):
        return self.a_ * x_single + self.b_

    def __repr__(self):
        return "SimpleLinearRegression2()"

def test():
    x_train = np.array([1, 2, 3, 4, 5])
    y_train = np.array([1, 3, 2, 3, 5])
    simple1 = SimpleLinearRegression1()
    simple1.fit(x_train, y_train)
    print(simple1.predict(np.array([6, 7, 10])))

    simple2 = SimpleLinearRegression2()
    simple2.fit(x_train, y_train)
    print(simple2.predict(np.array([6, 7, 10])))


# -*- coding: utf-8 -*-
"""
Created on Mon Dec 11 20:02:37 2017

@author: tzave
"""

import numpy as np
import solver

def determinant(A):
    rows = A.shape[0]
    det = 1.  # Det of a triangular matrix is equal
    for i in range(rows):  # to the product of the
        det = det * A[i, i]  # diagonal entries
    return det

A = np.array([[1., 1., 2.], [1., -1., 0.], [-2., 2., 4.]])
B = np.array([7., 1., 2.])
x = np.zeros(3)

result1 = solver.GaussElimination(A, B, 'partial')
det1 = determinant(result1['E'])
print('Partial Pivot Determinant is: ', det1)
print('Partial Pivot Solution is: ', result1['x'])

result2 = solver.GaussElimination(A, B, 'complete')
det2 = determinant(result2['E'])
print('Complete Pivot Determinant is: ', det2)
print('Complete Pivot Solution is: ', result2['x'])


# There exists exactly one Pythagorean triplet for which a + b + c = 1000.
# Find the product abc.

from sympy.solvers import solve
from sympy import Symbol
import math

def is_int(x):
    return int(x) == x

def find_pythag_triplet():
    c = Symbol('c')
    a, b, soln = 0, 0, [0]
    for i in range(4):
        for j in range(4):
            a = 2 ** i * 5 ** j
            r = 1000 // a
            soln = solve(a ** 2 + (r * (500 - c)) ** 2 - c ** 2, c)
            if soln[0].is_Integer:
                b = r * (500 - int(soln[0]))
                break
    return a * b * soln[0], (a, b, soln[0])

def solve_test():
    x = Symbol('x')
    soln = solve(x + 1, x)
    print(soln[0].is_Integer)
    return

def find_pythag_triplet_optimized(s=1000):
    """
    Find the Pythagorean triplet (a, b, c) that satisfies a + b + c = s

    :param s: a + b + c
    :return: (a * b * c, (a, b, c))
    """
    s2 = s // 2
    lim = math.ceil(math.sqrt(s2)) - 1
    for m in range(2, lim + 1):
        if s2 % m == 0:
            s_m = s2 // m
            while s_m % 2 == 0:
                s_m //= 2
            k = m + 1
            if m % 2 == 1:
                k = m + 2
            while k < 2 * m and k <= s_m:
                if s_m % k == 0 and math.gcd(k, m) == 1:
                    d = s2 // (k * m)
                    n = k - m
                    a = d * (m ** 2 - n ** 2)
                    b = 2 * d * m * n
                    c = d * (m ** 2 + n ** 2)
                    return a * b * c, (a, b, c)
                k += 2
    return None

if __name__ == "__main__":
    ans = find_pythag_triplet_optimized()[0]
    a, b, c = find_pythag_triplet_optimized()[1]
    print(f"abc = {ans}")
    print(f"a = {a}, b = {b}, c = {c}")
    # solve_test()


import numpy as np
from numba import njit
from numba import types
from numba.typed import Dict

@njit
def mixing_matrix(edges, in_deg, out_deg, mapping):

    n_degrees = len(mapping)
    M = np.zeros((n_degrees, n_degrees))

    for i in range(edges.shape[1]):
        u = edges[0][i]
        v = edges[1][i]
        du = out_deg[u]
        dv = in_deg[v]
        x = mapping[du]
        y = mapping[dv]
        M[x, y] += 1
        M[y, x] += 1
    return M

def degree_mixing_matrix(adj_matrix, normalize=True):
    out_deg = adj_matrix.sum(1).A1.astype('int64')
    in_deg = adj_matrix.sum(0).A1.astype('int64')
    deg_set = np.union1d(out_deg, in_deg)

    mapping = Dict.empty(
        key_type=types.int64,
        value_type=types.int64,
    )
    for x, y in enumerate(deg_set):
        mapping[y] = x

    edges = np.asarray(adj_matrix.nonzero())

    M = mixing_matrix(edges, in_deg, out_deg, mapping)

    if normalize:
        M /= M.sum()

    return M, mapping

def numeric_ac(M, mapping):
    # M is a numpy matrix or array
    # numeric assortativity coefficient, pearsonr
    x = y = np.array(list(mapping.keys()), dtype='float')
    a = M.sum(0)
    b = M.sum(1)
    # D(x) = E(x^2) - [E(x)]^2
    vara = (a * x**2).sum() - ((a * x).sum())**2
    varb = (b * x**2).sum() - ((b * x).sum())**2
    xy = np.outer(x, y)
    ab = np.outer(a, b)

    return (xy * (M - ab)).sum() / vara

def degree_assortativity_coefficient(adj_matrix):
    """Compute degree assortativity of graph.

    Assortativity measures the similarity of connections
    in the graph with respect to the node degree.

    Parameters
    ----------
    adj_matrix: Scipy sparse adjacency matrix representing a graph

    Returns
    -------
    r : float
        Assortativity of graph by degree.

    Notes
    --------
    This is a faster implementation of
        networkx.degree_assortativity_coefficient

    See Also
    --------
    graphgallery.degree_mixing_matrix
    networkx.degree_assortativity_coefficient

    References
    ----------
    .. [1] M. E. J. Newman, Mixing patterns in networks,
    Physical Review E, 67 026126, 2003
    .. [2] Foster, J.G., Foster, D.V., Grassberger, P. & Paczuski, M.
    Edge direction and the structure of networks, PNAS 107, 10815-20 (2010).
    """
    M, mapping = degree_mixing_matrix(adj_matrix)
    return numeric_ac(M, mapping)


from sympy import gcd

__all__ = ('phi', 'partition_sequences')

def phi(n):
    """ Euler's totient function."""
    assert n >= 0, 'Negative integer.'

    out = 0
    for i in range(1, n + 1):
        if gcd(n, i) == 1:
            out += 1

    return out

def partition_sequences(k):
    """ Generates a set of partition_sequences of size :math:`k`, i.e.
    sequences :math:`(n_i)` such that :math:`\sum_{i=1}^k i n_i = k`."""
    assert k >= 0, 'Negative integer.'

    def f(k, c, i):

        if k == 0:
            yield [0] * c

        elif i == k:
            yield [1] + [0] * (c - 1)

        elif i < k:
            for n in range(0, k // i + 1):
                for ps in f(k - (i * n), c - 1, i + 1):
                    yield [n] + ps

    return f(k, k, 1)