import sys
import numpy as np
import time
import logging
from matplotlib import pyplot as plt
import seaborn as sns

kmax = int(sys.argv[1])
r = int(sys.argv[2])
amin = sys.argv[3]
amax = sys.argv[4]

def plotResults(k, tNaive, tList, tStrassen):
    x = list(range(1, k+1))
    a = sns.lineplot(x, tNaive, label="Naive")
    a = sns.lineplot(x, tList, label="List Comprehesion")
    a = sns.lineplot(x, tStrassen, label="Strassen")
    a.set(xlabel='Tamanho da entrada 2^x', ylabel='Tempo de execução médio (s)')
    a.legend()
    plt.show()

def MultiplyMatrix(A, B):
    """Multiply two matrices using naive algorithm"""
    C = np.zeros((len(A), len(B)), dtype=int)

    for i in range(len(A)):
        for j in range(len(B)):
            for k in range(len(B)):
                C[i][j] += A[i][k] * B[k][j]
    return C

def addMatrix(A,B):
    """Add two matrices"""
    return [[A[i][j] + B[i][j] for j in range(len(A))] for i in range(len(B))]

def subMatrix(A,B):
    """subtraction two matrices"""
    return [[A[i][j] - B[i][j] for j in range(len(A))] for i in range(len(B))]

def Strassen(A, B):
    """Strassen Algorithm to squared matrices"""
    if len(A) <= 2:
        return  matrixMultiplicationListComprehesion(A, B)
    else:
        n = int(len(A)/2)

        a11 = [[A[i][j] for j in range(len(A)) if j < n  ] for i in range(len(A)) if i < n]
        a12 = [[A[i][j] for j in range(len(A)) if j >= n ] for i in range(len(A)) if i < n]
        a21 = [[A[i][j] for j in range(len(A)) if j < n  ] for i in range(len(A)) if i >= n]
        a22 = [[A[i][j] for j in range(len(A)) if j >= n ] for i in range(len(A)) if i >= n]

        b11 = [[B[i][j] for j in range(len(B)) if j < n  ] for i in range(len(B)) if i < n]
        b12 = [[B[i][j] for j in range(len(B)) if j >= n ] for i in range(len(B)) if i < n]
        b21 = [[B[i][j] for j in range(len(B)) if j < n  ] for i in range(len(B)) if i >= n]
        b22 = [[B[i][j] for j in range(len(B)) if j >= n ] for i in range(len(B)) if i >= n]

        m1 = Strassen(addMatrix(a11,a22), addMatrix(b11,b22))
        m2 = Strassen(addMatrix(a21,a22), b11)
        m3 = Strassen(a11, subMatrix(b12, b22))
        m4 = Strassen(a22, subMatrix(b21, b11))
        m5 = Strassen(addMatrix(a11, a12), b22)
        m6 = Strassen(subMatrix(a21, a11), addMatrix(b11, b12))
        m7 = Strassen(subMatrix(a12, a22), addMatrix(b21, b22))

        c11 = addMatrix(addMatrix(m1, m4), subMatrix(m7, m5))
        c12 = addMatrix(m3, m5)
        c21 = addMatrix(m2, m4)
        c22 = addMatrix(subMatrix(m1, m2), addMatrix(m3, m6))

        C = np.concatenate((np.concatenate((c11, c12), axis=1), np.concatenate((c21, c22), axis=1)), axis = 0 )

        return C

def matrixMultiplicationListComprehesion(A, B):
    """Multiply two squared matrices using list comprehesion"""
    return [[sum([x*y for (x, y) in zip(row, col)]) for col in zip(*B)] for row in A]

if __name__ == "__main__":

    tNaive = [0] * kmax
    tList = [0] * kmax
    tStrassen = [0] * kmax

    format = "%(asctime)s: %(message)s"
    logging.basicConfig(format=format, level=logging.INFO, datefmt="%Y-%m-%d %H:%M:%S")

    for n in range(1, kmax+1):
        for _ in range(r):
            A = np.random.randint(low = amin, high = amax, size = (2**n,2**n))
            B = np.random.randint(low = amin, high = amax, size = (2**n,2**n))

            start_time = time.time()
            C = MultiplyMatrix(A,B)
            end_time = time.time()
            tNaive[n-1] += (end_time - start_time)

            start_time = time.time()
            D = matrixMultiplicationListComprehesion(A,B)
            end_time = time.time()
            tList[n-1] += (end_time - start_time)

            start_time = time.time()
            E = Strassen(A,B)
            end_time = time.time()
            tStrassen[n-1] += (end_time - start_time)

        logging.debug("2^%d: {Naive: %f}, {List: %f}, {Strassen: %f}", n, tNaive[n-1], tList[n-1], tStrassen[n-1])
        logging.info("2^%d: Average {Naive: %f}, {List: %f}, {Strassen: %f}", n, tNaive[n-1]/r, tList[n-1]/r, tStrassen[n-1]/r)

    plotResults(kmax, [x/r for x in tNaive], [x/r for x in tList], [x/r for x in tStrassen])


import numpy as np

def NW_corner_rule(a, b):
    """
    North-West Corner rule for discrete optimal transport:
        Finding the vertex of the polytope of the feasible set U(a, b),
        where the minimum is attained.

    Reference: P44, Sec 3.4.2

    @input a, b: np.ndarray, of shape (n, ) and (m, ), the source histogram and the target histogram
    @return P: np.ndarray, of shape (n, m), one of the vertex of the polytope, P[i, j]: the flow from a[i] to b[j]

    >>> a = np.array([0.2, 0.5, 0.3])
    >>> b = np.array([0.5, 0.1, 0.4])
    >>> NW_corner_rule(a, b)
    array([[0.2, 0. , 0. ],
           [0.3, 0.1, 0.1],
           [0. , 0. , 0.3]])
    """
    P = np.zeros((a.shape[0], b.shape[0]))
    i = j = 0
    r = a[0]
    c = b[0]
    while i < a.shape[0] and j < b.shape[0]:
        P[i, j] = min(r, c)
        r -= P[i, j]
        c -= P[i, j]
        if r == 0:
            i += 1
            if i < a.shape[0]:
                r = a[i]
        if c == 0:
            j += 1
            if j < b.shape[0]:
                c = b[j]
    return P

def compute_NW_solutions(a, b, sigma1, sigma2):
    """
    compute an arbitrary NW solution by permutation

    Reference: Page 45, Section 3.4.2

    @param a, b: the source histogram and the target histogram
    @param sigma1, sigma2: the permutation vector for a and b respectively
    @return P_permutation: the permutated NW solution by
        P_permutation = sigma1_inv(sigma2_inv(P, column), row)
        where P is the transport matrix of sigma1(a) and sigma2(b)

    >>> a = np.array([0.2, 0.5, 0.3])
    >>> b = np.array([0.5, 0.1, 0.4])
    >>> sigma1 = np.array([2, 0, 1])
    >>> sigma2 = np.array([2, 1, 0])
    >>> compute_NW_solutions(a, b, sigma1, sigma2)
    array([[0., 0.1, 0.1],
        [0.5, 0., 0.],
        [0., 0., 0.3]])
    """
    a = a[sigma1]
    b = b[sigma2]
    P = NW_corner_rule(a, b)
    print(P)
    sigma1_inv = np.argsort(sigma1)
    sigma2_inv = np.argsort(sigma2)
    P = P[:, sigma2_inv]
    P = P[sigma1_inv, :]
    return P

if __name__ == "__main__":
    import doctest
    doctest.testmod()


"""
Exercise 3
Find the area enclosed by two curves between two points.
"""

from sympy import Integral, Symbol, SympifyError, sympify

def find_area(f1x, f2x, var, a, b):
    return Integral(f1x - f2x, (var, a, b)).doit()

def validate_function(function):

    try:
        return sympify(function)
    except SympifyError:
        print("Invalid function entered")
        sys.exit(1)

    return None

if __name__ == "__main__":
    f1 = input("Enter the upper function in one variable: ")
    f1 = validate_function(f1)

    f2 = input("Enter the lower upper function in one variable: ")
    f2 = validate_function(f2)

    variable = input("Enter the variable: ")
    variable = Symbol(variable)

    lower_bound = float(input("Enter the lower bound of the enclosed region: "))
    upper_bound = float(input("Enter the upper bound of the enclosed region: "))

    print(
        "Area enclosed by {} and {} is {} ".format(
            f1, f2, find_area(f1, f2, variable, lower_bound, upper_bound)
        )
    )


import numpy as np

def calc_weights(x, y):
    """
    Calculating weights in least square

    Parameters
    ----------
    x: numpy.ndarray or shape (n_samples, n_features)
        the data to be finding the weights for features.

    y: numpy.ndarray or shape (n_samples, )
        the target data of x.

    Returns
    -------
    output: numpy.ndarray or shape (n_features, )
        the weights of x
    """

    x = np.array(x)
    y = np.array(y)

    mean_x = x.mean(axis=0)
    mean_y = y.mean(axis=0)

    distance_x = x - mean_x
    distance_y = y - mean_y

    distance_xy = distance_x * distance_y.reshape((distance_y.shape[0], 1))

    x_square = (x - mean_x) ** 2

    numerator = np.sum(distance_xy, axis=0)
    denumerator = np.sum(x_square, axis=0)

    output = numerator / denumerator

    return output

def calc_intercept(X, y, slope):
    """
    Calculating intercept in least square

    Parameters
    ----------
    x: numpy.ndarray or shape (n_samples, n_features)
        the data to be finding the weights for features.

    y: numpy.ndarray or shape (n_samples, )
        the target data of x.

    slope: numpy.ndarray or shape (n_features, )
        the weighted of x

    Returns
    -------
    output: int
        the result of calculating intercept in least square

    """

    X = np.array(X)
    y = np.array(y)

    mean_y = y.mean(axis=0)
    mean_x = X.mean(axis=0)

    weighted_mean_x = slope * mean_x

    summing_weighted_mean_x = np.sum(weighted_mean_x)

    output = mean_y - summing_weighted_mean_x

    return output

def calc_with_svd(a, b):
    """
    Finding x with svd in equation ax = b.
    where a and b are known.

    Parameters
    ----------
    a: numpy.ndarray or shape (n, f)
        The independent variables to be calculate with svd

    b: numpy.ndarray or shape (n, )
        The dependent variable of a

    Returns
    -------
    x: numpy.ndarray or shape (f, )
        Solved to find x of ax = b equation

    """
    # change to numpy.array
    a = np.array(a)
    b = np.array(b)

    # raises the error
    if a.shape[0] != b.shape[0]:
        raise ValueError("the rows of a and b supposed to be same")
    elif len(b.shape) != 1:
        raise ValueError("for now the shape of b supposed to be (n, ")

    # solving least square with svd
    x, residuals, rank, s = np.linalg.lstsq(a, b, rcond=None)

    return x


import numpy as np
from numpy import ndarray

def linear_score_function(n: int) -> ndarray:
    """
    With the linear score function the "points" awarded scale linearly from first place
    through last place. For example, improving from 2nd to 1st place has the same sized
    benefit as improving from 5th to 4th place.

    :param n: number of players
    :return: array of the points to assign to each place (summing to 1)
    """
    return np.array([(n - p) / (n * (n - 1) / 2) for p in range(1, n + 1)])

def create_exponential_score_function(base: float):
    """
    With an exponential score function with base > 1, more points are awarded to the top
    finishers and the point distribution is flatter at the bottom. For example, improving
    from 2nd to 1st place is more valuable than improving from 5th place to 4th place. A
    larger base value means the scores will be more weighted towards the top finishers.

    :param base: base for teh exponential score function (> 1)
    :return: a function that takes parameter n for number of players and returns an array
    of the points to assign to each place (summing to 1)
    """
    return lambda n: _exponential_score_template(n, base)

def _exponential_score_template(n: int, base: float) -> ndarray:
    if base < 1:
        raise ValueError("base must be >= 1")
    if base == 1:
        return linear_score_function(n)  # it converges to this as base -> 1

    out = np.array([base ** (n - p) - 1 for p in range(1, n + 1)])
    return out / sum(out)


# # basicMLpy.decomposition module
import numpy as np
from .utils import z_normalize

def SVD(A,compute_uv:bool=True):
    """
    Performs the SVD decomposition of an array.
    Inputs:
        A: ndarray
            input the array on which to perform the SVD.
        compute_uv: bool, default = True
            choose whether to calculate the matrices of singular vectors.

    Returns:
        U: array (if compute_uv = True)
            outputs the array of right singular vectors.
        sigma: array
            outputs the array of singular values in decreasing order.
        V: array (if compute_uv = True)
            outputs the array of left singular vectors.
    """
    eig_vals,V = np.linalg.eig(A.T @ A)
    indx = np.argsort(-1 * eig_vals)
    eig_vals = eig_vals[indx]
    V = V[:,indx]
    sigma = np.sqrt(eig_vals)[::-1]

    if compute_uv:
        sigma_mat = np.diag(sigma)
        U = A @ V @ np.linalg.pinv(sigma_mat)

        return U,sigma,V
    else:
        return sigma

def PCA(A,n_components:int=None,normalize:bool=True):
    """
    Performs the PCA algorithm on an array.
    Inputs:
        A: ndarray
            input the array on which to perform the PCA. assumar A has the following format: rows = samples, columns = features, pass A.T otherwise.
        n_components: int, default = None
            number of principal components to return. if None, will return all components.
        normalize: bool, default = True
            choose wheter to z-normalize the data before applying the PCA algorithm.

    Returns:
        squared_singular_values: array_like
            returns the calculated PCA scores for each component.
        principal_components: array
            returns the array of principal component direction vectors.
    """
    if(normalize == True):
        A_norm = z_normalize(A)
    else:
        if(n_components != None):
            _,singular_values,principal_components = SVD(A.T @ A)
            return np.square(singular_values[:n_components]),principal_components[:n_components]
        else:
            _,singular_values,principal_components = SVD(A.T @ A)
            return np.square(singular_values),principal_components


import pandas as pd
from statsmodels.tsa.stattools import adfuller

def adfuller_test(series, signif=0.05, name='', verbose=False):
    """Perform ADFuller to test for Stationarity of given series"""
    stationary = -1
    r = adfuller(series, autolag='AIC')
    output = {'test_statistic':round(r[0], 4),
              'pvalue':round(r[1], 4),
              'n_lags':round(r[2], 4),
              'n_obs':r[3]}
    p_value = output['pvalue']
    test_statistic = output['test_statistic']
    critical_value = r[4]['5%']
#    def adjust(val, length= 6): return str(val).ljust(length)

    # Print Summary
    print(f'\nAugmented Dickey-Fuller Test on "{name}"', "\n", '-'*47)
    print(f'Null Hypothesis: Data has unit root. Non-Stationary.')
    print(f'Significance Level\t= {signif}')
    print(f'Test Statistic\t= {test_statistic}')
    print(f'No. Lags Chosen\t= {output["n_lags"]}')

    for key,val in r[4].items():
        print(f'Critical value {key}\t = {round(val, 3)}')

    if (p_value <= signif) and (test_statistic <= critical_value):
        print(f"=> P-Value\t = {p_value}. Rejecting Null Hypothesis.")
        print(f"=> Series is Stationary.")
        stationary = 1
    else:
        print(f"=> P-Value\t = {p_value}. Weak evidence to reject the Null Hypothesis.")
        print(f"=> Series is Non-Stationary.")
        stationary = 0
    print(f'-'*47)
    return stationary

sample = pd.read_csv('sample.csv', header='infer')
st_flag0 = adfuller_test(sample['volume'], name='SalesVolume')
st_flag1 = adfuller_test(sample['volume'].diff().dropna(), name='SalesVolume_Diff')


import numpy as np
import matplotlib .pyplot as plt

N = 50

X = np.linspace(0, 10, N)
Y = 0.5 * X + np.random.randn(N)
# outlier
Y[-1] += 30
Y[-2] += 30

plt.scatter(X, Y)
plt.show()

X = np.vstack([np.ones(N), X]).T

w_ml = np.linalg.solve(X.T.dot(X), X.T.dot(Y))
y_hat_ml = X.dot(w_ml)
plt.scatter(X[:, 1], Y)
plt.plot(X[:, 1], y_hat_ml)
plt.show()

l2 = 1000
w_map = np.linalg.solve(l2 * np.eye(2) + X.T.dot(X), X.T.dot(Y))
y_hat_map = X.dot(w_map)
plt.scatter(X[:, 1], Y)
plt.plot(X[:, 1], y_hat_ml, label='maxmimum likelihood')
plt.plot(X[:, 1], y_hat_map, label='map')
plt.legend()
plt.show()


"""
reference : https://en.wikipedia.org/wiki/PageRank
"""
import numpy as np

def pr(adjacency_arr, N_iter = 20, d = 0.85):
    M = np.zeros(adjacency_arr.shape)
    for m in range(M.shape[0]):
        M[:, m] = adjacency_arr[:, m] / adjacency_arr[:, m].sum()

    e = np.ones(adjacency_arr.shape[0])*1.0 / adjacency_arr.shape[0]
    v0 = e.copy()
    for i in range(N_iter):
        vn = np.dot(M, v0)*d + e*(1-d)
        v0 = vn
    return vn

if __name__ == "__main__":
	adjacency_arr = np.array([[0, 0, 0, 0, 1],
	              [0.5, 0, 0, 0, 0],
	              [0.5, 0, 0, 0, 0],
	              [0, 1, 0.5, 0, 0],
	              [0, 0, 0.5, 1, 0]])
	print("ground truth\n", [[0.25419178],
	       [0.13803151],
	       [0.13803151],
	       [0.20599017],
	       [0.26375504]])
	vn = pr(adjacency_arr, N_iter = 100, d = 0.85)
	print("result\n", vn)


"""
Given an array of points where points[i] = [xi, yi]
represents a point on the X-Y plane and an integer k,
return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is
the Euclidean distance (i.e., √(x1 - x2)2 + (y1 - y2)2).

You may return the answer in any order.
The answer is guaranteed to be unique (except for the order that it is in).

- Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest k = 1 points from the origin, so the answer is just [[-2,2]].

- Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

Source from: LeetCode

# Results:
Runtime: 2236 ms, faster than 5.01% of Python3 online submissions for K Closest Points to Origin.
Memory Usage: 50.1 MB, less than 9.07% of Python3 online submissions for K Closest Points to Origin.

Reference:
- Good solution: https://leetcode.com/problems/k-closest-points-to-origin/discuss/1647325/Python3-ONE-LINER-Explained
"""
# import math
from scipy.spatial import distance

class Solution:
    def kClosest(self, points, k):

        dis_list = []
        for i, point in enumerate(points):
            # distance = math.dist([point[0], point[1]], [0, 0])
            eu_distance = distance.euclidean((point[0], point[1]), (0, 0))
            dis_list.append([eu_distance, i])

        dis_list.sort()
        dis_list = dis_list[:k]

        for j in range(len(dis_list)):
            dis_list[j] = points[dis_list[j][1]]

        return dis_list

if __name__ == "__main__":
    Solution = Solution()

    points = [[1,3],[-2,2]]
    k = 1

    result = Solution.kClosest(points, k)
    print(f"My ans is: {result}\n")

    points = [[3,3],[5,-1],[-2,4]]
    k = 2

    result = Solution.kClosest(points, k)
    print(f"My ans is: {result}\n")


import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import pandas as pd

'''

last task from basic.

'''

import random
def flip():
    if random.random() < 0.5:
        return 0 # heads
    else:
        return 1 # tails

def sample():
    count = 0
    outcome = flip()
    while outcome == 0:
        count += 1 # equivalent to count = count + 1
        outcome = flip()
    return count

samples = np.fromiter((sample() for _ in range(100000)), dtype=int)
plt.hist(samples, bins=20)
plt.show()

'''

task: create a beautiful plot of x against x^2 for 0 <= x < 10. Use
numpy.linspace to find values of x and numpy.power to compute the
square of x. Save this figure to disk.

'''
plt.style.use('ggplot')
plt.rcParams['figure.figsize'] = (6, 6)
plt.rcParams['figure.dpi'] = 200
samples = 100
x = np.linspace(0, 10, samples)
y = np.power(x, 2)
plt.plot(x, y)
plt.xlabel('x')
plt.ylabel('y')
plt.xlim(0, 10)
plt.ylim(0, 100)
plt.title('dat plot')
plt.savefig('dat_plot.png', dpi='figure')
plt.show()

'''

task: Use scipy.optimize.curve_fit to fit a curve of the form
 "y=a+bx", where a and b are coefficients to this data:
x = [1, 2, 3, 4]
y = [13.46688604, 18.58857204, 20.51804834, 27.6660338]

First, plot the given x and y data. Second, fit parameters (a, b) to
the data, third plot the given x and y together with the fit
curve. Use np.linspace to find values of x.

'''

from scipy.optimize import curve_fit
x = [1, 2, 3, 4]
y = [13.46688604, 18.58857204, 20.51804834, 27.6660338]
f = lambda x, a, b: a + b*x
p, _ = curve_fit(f, x, y)
print(p)
plt.figure()
plt.plot(x, y)

t = np.linspace(1, 4, 100)
plt.plot(t, [f(x, *p) for x in t])
plt.show()

'''

task: use pandas.read_csv to read the csv file "scatter.csv", print it
to terminal, and then plot its "x" column against its "y" column using
pyplot.

'''

# code used to generate the data
# n = 100
# x = np.linspace(0, 1, n)
# y = np.log2(x) + np.random.randn(n)
# df = pd.DataFrame({'x': x, 'y': y})
# df.to_csv('scatter.csv', index=False)

# read and plot it
df = pd.read_csv('scatter.csv')
plt.scatter(df['x'], df['y'])
plt.show()


import numpy as np
from scipy import integrate

conturs_areas = np.array([194135, 136366, 79745, 38335, 18450, 9635, 3895])
x = np.array([0,25,50,75,100,125,150])

vol_traps = integrate.trapz(conturs_areas, x)
vol_simps = integrate.simps(conturs_areas, x)

print('The trapezoidal rule returns a volume of {:.0f} cubic meters'.format(vol_traps))
print('The composite Simpson rule returns a volume of {:.0f} cubic meters'.format(vol_simps))

'''
Output:
The trapezoidal rule returns a volume of 9538650 cubic meters
The composite Simpson rule returns a volume of 9431367 cubic meters
'''


from math import *
import operator as op
import numpy as np

dim  = 512
k = 8
q = 12289
bound_CC = 34000    # Bound n*mu+beta to use in Cramer-Chernoff Ineq.
n = 2 * dim / 2      # Number of samples in Cramer-Chernoff Ineq.

check_each_y = True
# Check each distribution phi_y (slower)
#experimentally, they are always equal probably due to symmetries

subg_param = sqrt(k/2)
tau = 17.6

subg_value = sqrt(bound_CC) * subg_param * tau
p_subg = exp(-tau**2/2 )

print "#### "
print "#### Correctness check :"
print "#### "

print "target bound (before +e'')=", subg_value
print "target bound =", subg_value + 2*k
print "(sqrt(2)/2.) * 3q/4 =", (sqrt(2)/2.) * 3*q/4.
print  "Correct ?", subg_value + 2*k < (sqrt(2)/2.) * 3*q/4.
assert subg_value + 2*k < (sqrt(2)/2.) * 3*q/4.

print
print "#### "
print "#### Lemma 6 Verification :"
print "#### "

def rot(v):
	return [-v[1],v[0]]

if check_each_y:
	y_list = [
		[ 1, 1],
		[ 1,-1],
		[-1, 1],
		[-1,-1]
	]
else:
	y_list = [[ 1, 1]]

### Construct the binomial law
supp = xrange(-k,k+1)

def binomial(n, r):
    r = min(r, n-r)
    if r == 0: return 1
    numer = reduce(op.mul, xrange(n, n-r, -1))
    denom = reduce(op.mul, xrange(1, r+1))
    return numer//denom

def pdf_binom(x):
	return binomial(2*k,x+k) / 2.**(2*k)

# Union bound on each application of Chernoff-Cramer
Union_CC = 0.

for y in y_list:
	print "y=",y
	### Construct a 4*4 matrix associated to y
	v1 = y
	v2 = rot(v1)
	mat = np.matrix([v1,v2])

	### Initialize a table to store the pdf of \varphi_y
	### the range of the above distribution
	ima = xrange(4 * (k*4)**2 )
	varphi_y_Table = [0 for i in ima]

	### Brute force computation of the pdf of || x*y ||^2
	s = 0.
	for x1 in supp:
		for x2 in supp:
			x = np.array([x1,x2])
			mx = x * mat
			s = mx[0,0]**2 + mx[0,1]**2
			p = pdf_binom(x1) * pdf_binom(x2)
			varphi_y_Table[s] += p                   ### The result is always a multiple of 4

	avg = sum([i*varphi_y_Table[i] for i in ima])
	print "varphi_y has average ",avg

	def Moment(t):
		return sum([varphi_y_Table[i] * exp(t*(i - avg)) for i in range(len(ima))])

	def ChernoffCramer(n,t,bound):
		beta = bound - n*avg
		p =  exp(- t*beta + n*log(Moment(t)))
		return p

	t_CC = 0.0060
	p_CC = ChernoffCramer(n,t_CC,bound_CC)

	print "Chernoff-Cramer Bound", p_CC, " = 2^", log(p_CC)/log(2)

	if check_each_y:
		Union_CC += p_CC

else:
	Union_CC = 4*p_CC

print
print "#### "
print "#### Conclusion of Lemma 6 :"
print "#### "

print "with parameter k=",k, "and parameter n = ",n,"(v has dimension 4n)"
print "The bound ||v||_2^2 <=", bound_CC, " fails with probability less than 2^", log(Union_CC)/log(2)

print "#### "
print "#### Corollary 1 :"
print "#### "

print "subgaussian parameter sigma = ", subg_param
print "tailcut parameter tau = ", tau
print "tailcut value ||v||* sigma * tau = ", subg_value

print "sub-gaussian tail-bound (Lemma 5): 2^", log(p_subg) /log(2)
print "(failure proba per bit-agreement and per y)"
print
print "#### "
print "#### Conclusion of Corollary 1 :"
print "#### "
print "Partial bound ||(es'-e's )_i||_1 <=", subg_value

print "Final bound ||(es'-e's +e'')_i||_1 <=", subg_value + 2*k
print "fails with probability at most 2^",
print log(Union_CC + 4 * 256 * p_subg)/log(2)


import numpy

def feature_normalize(x_array):
    """
    Normalizes the features in x_array returns a normalized version of x_array where the mean value of each
    feature is 0 and the standard deviation is 1. This is often a good preprocessing step to do when
    working with learning algorithms.

    Parameters
    ----------
    x_array : array_like
        An dataset which is a (m x n) matrix, where m is the number of examples,
        and n is the number of dimensions for each example.

    Returns
    -------
    x_norm : array_like
        The normalized input dataset.

    mu : array_like
        A vector of size n corresponding to the mean for each dimension across all examples.

    sigma : array_like
        A vector of size n corresponding to the standard deviations for each dimension across
        all examples.
    """
    mu = numpy.mean(x_array, axis=0)
    sigma = numpy.std(x_array, axis=0, ddof=1)

    x_norm = (x_array - mu)/sigma

    return x_norm, mu, sigma


import numpy as np

### Functions for you to fill in ###

def polynomial_kernel(X, Y, c, p):
    """
        Compute the polynomial kernel between two matrices X and Y::
            K(x, y) = (<x, y> + c)^p
        for each pair of rows x in X and y in Y.

        Args:
            X - (n, d) NumPy array (n datapoints each with d features)
            Y - (m, d) NumPy array (m datapoints each with d features)
            c - a coefficient to trade off high-order and low-order terms (scalar)
            p - the degree of the polynomial kernel

        Returns:
            kernel_matrix - (n, m) Numpy array containing the kernel matrix
    """
    kernel_matrix = (np.inner(X, Y) + c) ** p
    return kernel_matrix

def rbf_kernel(X, Y, gamma):
    """
        Compute the Gaussian RBF kernel between two matrices X and Y::
            K(x, y) = exp(-gamma ||x-y||^2)
        for each pair of rows x in X and y in Y.

        Args:
            X - (n, d) NumPy array (n datapoints each with d features)
            Y - (m, d) NumPy array (m datapoints each with d features)
            gamma - the gamma parameter of gaussian function (scalar)

        Returns:
            kernel_matrix - (n, m) Numpy array containing the kernel matrix
    """
    XT_X = np.mat([np.matmul(row, row.T) for row in X]).T
    YT_Y = np.mat([np.matmul(row, row.T) for row in Y]).T

    XTX_matrix = np.repeat(XT_X, Y.shape[0], axis=1)
    YTY_matrix = np.repeat(YT_Y, X.shape[0], axis=1).T

    distance_matrix = np.asarray(XTX_matrix + YTY_matrix - 2 * (np.matmul(X, Y.T)), dtype='float64')
    kernel_matrix = np.exp(-gamma * distance_matrix)
    return kernel_matrix


"""
Problem 4
"""

import numpy as np

s_inf = np.pi/4

def term(n):
    return (-1) ** n / (2 * n + 1)

def basic_sum(tol):
    n = 0
    val = 0
    while True:
        val += term(n)
        if np.abs(s_inf - val) < tol:
            return val, n
        n += 1

def aitkin_sum(tol):
    n = 0
    while True:
        sn2 = np.sum(term(np.arange(0, n+3)))
        sn1 = sn2 - term(n+2)
        sn0 = sn1 - term(n+1)
        s_aitkin = sn0 - (sn1 - sn0)**2 / (sn2 - 2*sn1 + sn0)
        if np.abs(s_inf - s_aitkin) < tol:
            return s_aitkin, n
        n += 1

if __name__ == '__main__':
    print(basic_sum(1e-7))
    print(aitkin_sum(1e-7))


"""This module creates the contour plots for the bivariate normal distribution."""
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import multivariate_normal

from fig_config import OUTPUT_DIR

plt.style.use("resources/grmpy.mplstyle")

y_min, y_max = -4, 4

x = np.linspace(y_min, y_max, 100)
y = np.linspace(y_min, y_max, 100)
X, Y = np.meshgrid(x, y)
pos = np.dstack((X, Y))

mean = np.tile(0.2, 2)
cov = np.identity(2) * 2

rv = multivariate_normal(mean, cov)
rv.pdf(pos)

ax = plt.figure().add_subplot(111)

levels = np.linspace(0.1, 1.0, 10, endpoint=True)
cns = plt.contourf(X, Y, rv.pdf(pos) / np.max(rv.pdf(pos)), levels=levels)

ax.set_ylabel("$Y_1$")
ax.set_xlabel("$Y_0$")
ax.set_ylim(y_min, y_max)
ax.set_xlim(y_min, y_max)

ax.set_yticklabels([])
ax.set_xticklabels([])

plt.plot(
    np.arange(y_min, y_max), np.arange(y_min, y_max), color="black", linestyle="--"
)

plt.colorbar(cns)

ax.text(3.3, 3.3, r"$45^o$")

plt.savefig(OUTPUT_DIR + "/fig-distribution-joint-potential.png", dpi=300)

# This plot shows the joint distribution of surplus and benefits.
x = np.linspace(y_min, y_max, 100)
y = np.linspace(y_min, y_max, 100)
X, Y = np.meshgrid(x, y)
pos = np.dstack((X, Y))

mean = np.tile(0.2, 2)
cov = np.identity(2) * 2

rv = multivariate_normal(mean, cov)
rv.pdf(pos)

ax = plt.figure().add_subplot(111)

levels = np.linspace(0.1, 1.0, 10, endpoint=True)
cns = plt.contourf(X, Y, rv.pdf(pos) / np.max(rv.pdf(pos)), levels=levels)

ax.set_ylabel("$B$")
ax.set_xlabel("$S$")
ax.set_ylim(y_min, y_max)
ax.set_xlim(y_min, y_max)

ax.set_yticklabels(["", "", "", "", 0])
ax.set_xticklabels(["", "", "", "", 0])

plt.colorbar(cns)

ax.axvline(x=0, color="black", linestyle="--")
ax.axhline(y=0, color="black", linestyle="--")

ax.axvline(x=-0.5, color="lightgray", linestyle="--")

ax.text(x=-3.5, y=3.5, s="I")

ax.text(x=3.5, y=3.5, s="II")

ax.text(x=3.5, y=-3.5, s="III")

ax.text(x=-3.5, y=-3.5, s="IV")

plt.savefig(OUTPUT_DIR + "/fig-distribution-joint-surplus.png", dpi=300)


"""

Nombre del archivo: mandelbrot.py
Fecha de creación: 19/09/2019
Fecha de modificación: 20/09/2019
Autores: Bryan Steven Biojó R.     1529879-2711
         Joel Alexander Ramírez N. 1528879-2711

"""

import numpy # Librería para computación científica.
import matplotlib.pyplot as plt # Librería para graficar.
from numba import jit # Compilador de JIT que mejora la rapidez del código.

# ******************************************* FUNCIÓN 1 *******************************************

@jit
def mandelbrot1(Re, Im, maxIteraciones): # Primera función. Pinta z^2+C
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 2) + c

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 1 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 1, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-1, 1, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot1(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'gray', interpolation = 'bilinear', extent = [-2, 1, -1, 1])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 2 *******************************************

@jit
def mandelbrot2(Re, Im, maxIteraciones): # Segunda función. Pinta z^3+C
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 3) + c

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 2 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 1, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-1.5, 1.5, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot2(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'hot', interpolation = 'bilinear', extent = [-2, 1, -1.5, 1.5])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 3 *******************************************

@jit
def mandelbrot3(Re, Im, maxIteraciones): # Tecera función. Pinta z^5+C
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 5) + c

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 3 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 1, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-1, 1, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot3(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'winter', interpolation = 'bilinear', extent = [-2, 1, -1, 1])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 4 *******************************************

@jit
def mandelbrot4(Re, Im, maxIteraciones): # Cuarta función. Pinta z^2+(1/C)
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 2) + (1 / c)

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 4 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 4, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-3, 3, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot4(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'Wistia', interpolation = 'bilinear', extent = [-2, 4, -3, 3])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 5 *******************************************

@jit
def mandelbrot5(Re, Im, maxIteraciones): # Quinta función. Pinta z^7+(1/C)
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 7) + (1 / c)

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 5 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-2, 2, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot5(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'autumn', interpolation = 'bilinear', extent = [-2, 2, -2, 2])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 6 *******************************************

@jit
def mandelbrot6(Re, Im, maxIteraciones): # Sexta función. Pinta z^2+(1/C^2)
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 2) + (1 / (c * c))

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 6 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-3, 1.5, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot6(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'bone', interpolation = 'bilinear', extent = [-2, 2, -3, 1.5])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 7 *******************************************

@jit
def mandelbrot7(Re, Im, maxIteraciones): # Séptima función. Pinta z^2+C+(1/C)
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 2) + c + (1 / c)

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 7 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-3, 2, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot7(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'magma', interpolation = 'bilinear', extent = [-2, 2, -3, 2])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 8 *******************************************

@jit
def mandelbrot8(Re, Im, maxIteraciones): # Octava función. Pinta [(z+C^2-1)/C^2]^2
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(((z + (c * c) - 1) / (c * c)), 2)

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 8 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-2, 2, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot8(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'binary', interpolation = 'bilinear', extent = [-2, 2, -2, 2])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 9 *******************************************

@jit
def mandelbrot9(Re, Im, maxIteraciones): # Novena función. Pinta [z^2+(1/C) interseccción z^2+C (o eso parece).
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 2) + (1 / c)
        z = pow(z, 2) + c

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 9 *******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-2, 2, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot9(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'Greens', interpolation = 'bilinear', extent = [-2, 2, -2, 2])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()

# ******************************************* FUNCIÓN 10 ******************************************

@jit
def mandelbrot10(Re, Im, maxIteraciones): # Décima función. Pinta z^3-z+C
    z = 0.0j
    c = complex(Re, Im)

    for i in range(maxIteraciones):
        z = pow(z, 3) - z + c

        if((pow(z.real, 2) + pow(z.imag, 2)) >= 4):
            return i

    return maxIteraciones

# ******************************************* GRÁFICA 10 ******************************************

filas = 2000
columnas = 2000

resultado = numpy.zeros([filas, columnas])
for fila_index, Re in enumerate(numpy.linspace(-2, 2, num = filas)):
    for columna_index, Im in enumerate(numpy.linspace(-1, 1, num = columnas)):
        resultado[fila_index, columna_index] = mandelbrot10(Re, Im, 100)

plt.figure(dpi = 300)
plt.imshow(resultado.T, cmap = 'YlOrBr', interpolation = 'bilinear', extent = [-2, 1, -1, 1])
plt.xlabel('Eje Real')
plt.ylabel('Eje Imaginario')
plt.show()


import numpy as np
import matplotlib.pyplot as plt

# plotting Rosenbrock funtion
x, y = np.meshgrid(np.linspace(-50, 50, 200),np.linspace(-500, 1000, 200))
z = (1-x)**2 + 100*((y-x**2)**2)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap = 'jet', alpha=0.7)

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
ax.scatter(1, 1, marker='o', color='red')
plt.tight_layout()
plt.show()

# plot the contour
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contour(x, y, z, 200, cmap = 'jet')

ax.scatter(1, 1, marker='o', color='red')

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
plt.tight_layout()
plt.show()

# using autodiff to calculate the gradient and perform gradient descent
from autograd import grad
import autograd.numpy as np
import numpy

def rosenbrock(x):
    z = (1-x[0])**2 + 100*((x[1]-x[0]**2)**2)
    return z
# testing how to calculate the gradient!
gradient = grad(rosenbrock)
# it works!
gradient(np.array([3., 5.]))

# now doing regular gradient descent
w = np.array([2., 3.])
gamma = 10e-4
epochs = 1000
trace_x = np.zeros(1 + epochs)
trace_y = np.zeros(1 + epochs)
trace_z = np.zeros(1 + epochs)
trace_x[0] = w[0]
trace_y[0] = w[1]
trace_z[0] = rosenbrock(w)

for e in range(epochs):
  grad_eval = gradient(w)
  w = w - gamma*grad_eval
  trace_x[e+1] = w[0]
  trace_y[e+1] = w[1]
  trace_z[e+1] = rosenbrock(w)
  #print(w)
print(trace_z)

# plotting the trace
import numpy as nmp
import matplotlib.pyplot as plt

# plotting Rosenbrock funtion
x, y = nmp.meshgrid(nmp.linspace(-50, 50, 200),nmp.linspace(-500, 1000, 200))
z = (1-x)**2 + 100*((y-x**2)**2)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap = 'jet', alpha=0.7)

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
ax.scatter(trace_x, trace_y, trace_z, marker='o', color='red')
plt.tight_layout()
plt.show()

# plot the contour
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contour(x, y, z, 200, cmap = 'jet')
ax.scatter(trace_x, trace_y, marker='o', color='red')

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
plt.tight_layout()
plt.show()

# now doing gradient descent with momentum
# I used some code from http://people.duke.edu/~ccc14/sta-663-2018/notebooks/S09G_Gradient_Descent_Optimization.html
beta = 0.9
#w = np.array([2., 3.])
w = np.array([-2., 0.])
v=0
trace_x = np.zeros(1 + epochs)
trace_y = np.zeros(1 + epochs)
trace_z = np.zeros(1 + epochs)
trace_x[0] = w[0]
trace_y[0] = w[1]
trace_z[0] = rosenbrock(w)
gamma = 10e-3
for i in range(epochs):
  v = beta*v + (1-beta)*gradient(w)
  vc = v/(1+beta**(i+1))
  w = w - gamma * vc
  trace_x[i+1] = w[0]
  trace_y[i+1] = w[1]
  trace_z[i+1] = rosenbrock(w)
  #print(w)
print(trace_z)
print(trace_x)
print(trace_y)

# plotting the trace
import numpy as nmp
import matplotlib.pyplot as plt

# plotting Rosenbrock funtion
x, y = nmp.meshgrid(nmp.linspace(-5, 5, 200),nmp.linspace(-5, 5, 200))
z = (1-x)**2 + 100*((y-x**2)**2)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap = 'jet', alpha=0.7)
ax.view_init(elev=10, azim=60)

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
#plt.grid(color='silver', linestyle='-', linewidth=0.2)
ax.scatter(trace_x, trace_y, trace_z, marker='x', color='red', s=10)
plt.tight_layout()
plt.show()

# plot the contour
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.contour(x, y, z, 200, cmap = 'jet')
ax.plot(trace_x, trace_y, color='blue')
ax.scatter(trace_x, trace_y, marker='x', color='red', s=5)
ax.scatter(1, 1, marker='o', color='magenta', s=30)

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
plt.tight_layout()
plt.show

# plot the contour
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.scatter(trace_x, trace_y, marker='x', color='red', s=5)

plt.xlabel("x")
plt.ylabel("y")
plt.title("Rosenbrock Function", fontsize='small')
plt.grid(color='silver', linestyle='-', linewidth=0.2)
plt.tight_layout()
plt.show


# Linear System Solvers
#
# Available Methods:
# - LU (Own Implementation)
# - LU_sparse (Own Implementation)
# - LU_compiled (Own Implementation)
# - LU_scipy_sparse (Scipy Implementation)
# - QR (Own Implementation)
# - QR_compiled (Own Implementation)
# - QR_scipy (Scipy Implementation)
# - scipy (Scipy Implementation)
# - jacobi (Own Implementation)
# - jacobi_compiled (Own Implementation)
# - jacobi_sparse (Own Implementation)
# - gauss_seidel (Own Implementation)
# - gauss_seidel_compiled (Own Implementation)
# - gauss_seidel_sparse (Own Implementation)
# - sor (Own Implementation)
# - sor_compiled (Own Implementation)
# - sor_sparse (Own Implementation)
#
# Assignment 1 for ME4233
# Author: Giovani Hidalgo Ceotto
# Prof: Mengqi Zhang

import numpy as np
from scipy.sparse import dia_matrix, csr_matrix, csc_matrix, lil_matrix, identity, tril, triu
from scipy.sparse.linalg import splu
from scipy.sparse.linalg import inv as sparse_inv
from scipy.linalg import solve, inv, qr
from numba import jit

# Wrapper for all solvers
def solve_algebraic_system(A, b, method, initial_guess=0, true_solution=0, tol=1e-6, w=1):
    """ Solves the algebraic system A*x = b, for the vector x.

        Parameters
        ----------
        A : array
            Matrix A corresponding to the system A*x = b. Must be in sparse
            format.
        b : array
            Column vector b corresponding to the system A*x = b.
        method : string
            Mehtod to be used. Available methods are:
                - LU (Own Implementation)
                - LU_sparse (Own Implementation)
                - LU_compiled (Own Implementation)
                - LU_scipy_sparse (Scipy Implementation)
                - QR (Own Implementation)
                - QR_compiled (Own Implementation)
                - QR_scipy (Scipy Implementation)
                - scipy (Scipy Implementation)
                - jacobi (Own Implementation)
                - jacobi_compiled (Own Implementation)
                - jacobi_sparse (Own Implementation)
                - gauss_seidel (Own Implementation)
                - gauss_seidel_compiled (Own Implementation)
                - gauss_seidel_sparse (Own Implementation)
                - sor (Own Implementation)
                - sor_compiled (Own Implementation)
                - sor_sparse (Own Implementation)
        initial_guess : array, optional
            Initial solution used for iterative solvers.
        true_solution : array, optional
            If given, iterative solver residuals will be calculated considering
            this.
        tol : float, optional
            Convergence tolerance for iterative solvers.
        w : float, optiona
            Relaxation parameter to be used with the SOR methods.

        Returns
        -------
        x : array
            Solution vector corresponding to A*x = b.
        residue : array
            Residue array given by iterative solver.
    """
    if method == "LU":
        L_prime, U = LU_factorization(A.toarray())
        return solve_upper_triangular_system(U, L_prime.dot(b))
    elif method == "LU_sparse":
        L_prime, U = LU_factorization_sparse(A)
        return solve_upper_triangular_system(U.toarray(), L_prime.dot(b))
    elif method == "LU_compiled":
        L, U = LU_factorization_compiled(A.toarray())
        y = solve_lower_triangular_system(L, b)
        return solve_upper_triangular_system(U, y)
    elif method == "LU_scipy_sparse":
        LU = splu(csc_matrix(A))
        return LU.solve(b)
    elif method == "QR":
        Q, R = QR_factorization(A.toarray())
        return solve_upper_triangular_system(R, np.dot(Q.T, b))
    elif method == "QR_compiled":
        Q, R = QR_factorization_compiled(A.toarray())
        return solve_upper_triangular_system(R, np.dot(Q.T, b))
    elif method == "QR_scipy":
        Q, R = qr(A.toarray())
        return solve_upper_triangular_system(R, np.dot(Q.T, b))
    elif method == "jacobi":
        return jacobi(A.toarray(), b, initial_guess, true_solution, tol)
    elif method == "jacobi_compiled":
        return jacobi_compiled(A.toarray(), b, initial_guess, true_solution, tol)
    elif method == "jacobi_sparse":
        return jacobi_sparse(A, b, initial_guess, true_solution, tol)
    elif method == "gauss_seidel":
        return gauss_seidel(A.toarray(), b, initial_guess, true_solution, tol)
    elif method == "gauss_seidel_compiled":
        return gauss_seidel_compiled(A.toarray(), b, initial_guess, true_solution, tol)
    elif method == "gauss_seidel_sparse":
        return gauss_seidel_sparse(A, b, initial_guess, true_solution, tol)
    elif method == "sor":
        return sor(A.toarray(), b, initial_guess, true_solution, tol, w)
    elif method == "sor_compiled":
        return sor_compiled(A.toarray(), b, initial_guess, true_solution, tol, w)
    elif method == "sor_sparse":
        return sor_sparse(A, b, initial_guess, true_solution, tol, w)
    elif method == "scipy":
        return solve(A.toarray(), b)

# Direct triangular system Solvers
def solve_upper_triangular_system(U, b):
    """ Solves the linear algebraic equation U*x = b, where b is a column
    vector, x is the unknown column vector and U is an upper triangular
    matrix.

    Parameters
    ----------
    U : array
        Square upper triangular matrix.
    b : array
        Column vector with the same number of lines as U.

    Returns
    -------
    x : array
        Column vector with the same number of lines as U which solves U*x=b.
    """
    # Retrieve matrix dimension
    n, m = U.shape

    # Initialize solution vector x:
    x = np.zeros(n)

    # Iterate line by line, beginning from the last one.
    for i in range(n-1, -1, -1):
        x[i] = (b[i] - sum([U[i, j]*x[j] for j in range(i+1, n)]))/U[i, i]

    # Return the solution vector x
    return x

def solve_lower_triangular_system(L, b):
    """ Solves the linear algebraic equation L*y = b, where b is a column
    vector, y is the unknown column vector and L is a lower triangular
    matrix.

    Parameters
    ----------
    L : array
        Square lower triangular matrix.
    b : array
        Column vector with the same number of lines as L.

    Returns
    -------
    y : array
        Column vector with the same number of lines as L which solves L*y=b.
    """
    # Retrieve matrix dimension
    n, m = L.shape

    # Initialize solution vector y
    y = np.zeros(n)

    # Iterate line by line, beginning from the first one.
    for i in range(n):
        y[i] = (b[i] - sum([L[i, j]*y[j] for j in range(i)]))/L[i, i]

    # Return the solution vector y
    return y

# Factorization methods
def LU_factorization(A, return_L=False):
    """ Decomposes the given A matrix into its LU factors, that is A=LU,
    where L is a lower triangular matrix, with diagonal entries equal to 1 and
    U is an upper triangular matrix. By default, it returns U and L^(-1).

    Parameters
    ----------
    A : array
        Matrix to be decomposed into A=LU. A must be a square matrix with no
        diagonal entries equal to 0.
    return_L : bool, optional
        If True, returns U and L. If false, which it is by default, return U
        and L^(-1).

    Returns
    -------
    L : array
        Lower triangular matrix with the same shape as A, equal to L^(-1) if
        return_L is False. If return_L is true, equal to L.
    U : array
        Upper triangular matrix with the same shape as A.

    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Initialize L prime (L^(-1)) and U
    L_prime = np.eye(n, n, dtype=np.float64)
    U = A.copy()

    # Iterate column by column, except for the last column
    for j in range(m - 1):
        # Create temporary L_1
        L_temp = np.eye(n, n, dtype=np.float64)
        L_temp[(j+1):, j] = U[(j+1):, j]/(-U[j, j])

        # Multiply L_temp by current version of U to get new U
        U = L_temp.dot(U)

        # Multiply L_temp by current version of L_prime to get new L_prime
        L_prime = L_temp.dot(L_prime)

    # Return L and U
    if return_L:
        return inv(L_prime), U
    else:
        return L_prime, U

def LU_factorization_sparse(A, return_L=False):
    """ Decomposes the given A matrix into its LU factors, that is A=LU,
    where L is a lower triangular matrix, with diagonal entries equal to 1 and
    U is an upper triangular matrix. By default, it returns U and L^(-1). It
    makes use of sparse properties to optimize speed.

    Parameters
    ----------
    A : array
        Matrix to be decomposed into A=LU. A must be a square matrix with no
        diagonal entries equal to 0.
    return_L : bool, optional
        If True, returns U and L. If false, which it is by default, return U
        and L^(-1).

    Returns
    -------
    L : array
        Lower triangular matrix with the same shape as A, equal to L^(-1) if
        return_L is False. If return_L is true, equal to L.
    U : array
        Upper triangular matrix with the same shape as A.

    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Initialize L prime (L^(-1)) U
    L_prime = identity(n, np.float64, 'csr')
    U = csc_matrix(A)
    L_temp1 = identity(n, np.float64, 'lil')

    # Iterate column by column, except for the last column
    for j in range(m - 1):
        # Create temporary Ls
        L_temp1[(j+1):, j] = U[(j+1):, j]/(-U[j, j])
        L_temp2 = csr_matrix(L_temp1)

        # Roll back L_temp1
        L_temp1[(j+1):, j] = 0

        # Multiply L_temp by current version of U to get new U
        U = L_temp2.dot(U)

        # Multiply L_temp by current version of L_prime to get new L_prime
        L_prime = L_temp2.dot(L_prime)

    # Return L and U
    if return_L:
        return Inv(L_prime), U
    else:
        return L_prime, U

@jit(nopython=True)
def LU_factorization_compiled(A):
    """ Decomposes the given A matrix into its LU factors, that is A=LU,
    where L is a lower triangular matrix, with diagonal entries equal to 1 and
    U is an upper triangular matrix. The function is compiled the first time it
    is ran by numba.

    Parameters
    ----------
    A : array
        Matrix to be decomposed into A=LU. A must be a square matrix with no
        diagonal entries equal to 0 and in numpy format.
    return_L : bool, optional
        If True, returns U and L. If false, which it is by default, return U
        and L^(-1).

    Returns
    -------
    L : array
        Lower triangular matrix with the same shape as A.
    U : array
        Upper triangular matrix with the same shape as A.

    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Initialize L and U
    L = np.eye(n, n)
    U = np.eye(n, n)

    # Create first column of L and first row of U
    L[:, 0] = A[:, 0]/A[0, 0]
    U[0, :] = A[0, :]

    # Iterate line by line, except for the first and last ones
    for i in range(1, n):
        #  Complete ith row of U and ith column of L
        for j in range(i, n):
            sum1 = 0
            for k in range(i):
                sum1 += L[i, k]*U[k, j]
            U[i, j] = A[i, j] - sum1
            sum2 = 0
            for k in range(i):
                sum2 += L[j, k]*U[k, i]
            L[j, i] = (A[j, i] - sum2)/U[i, i]

    # Return L and U
    return L, U

def QR_factorization(A):
    """Decomposes the given A matrix into its QR factors, that is A=QR,
    where Q is a orthogonal matrix and R is an upper triangular matrix.

    Parameters
    ----------
    A : array
        Matrix to be decomposed into A=QR. A must be a square matrix with no
        diagonal entries equal to 0.

    Returns
    -------
    Q : array
        Orthogonal matrix with the same shape as A.
    R : array
        Upper triangular matrix with the same shape as A.
    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Initialize Q and R matrix
    Q = A.copy()
    R = np.eye(n, n)

    # Compute columns of Q matrix
    for j in range(n):
        # Subtract projections on previous columns
        temp1 = Q[:, j]
        temp2 = temp1 - sum([np.dot(temp1.T, Q[:, k])*Q[:, k] for k in range(j)])
        # Normalize result
        Q[:, j] = temp2/np.linalg.norm(temp2)

    # Compute lines of R matrix
    for i in range(n):
        R[i, i:] = [np.dot(A[:, j], Q[:, i]) for j in range(i, n)]

    # Return Q and R
    return Q, R

@jit(nopython=True)
def QR_factorization_compiled(A):
    """Decomposes the given A matrix into its QR factors, that is A=QR,
    where Q is a orthogonal matrix and R is an upper triangular matrix.

    Parameters
    ----------
    A : array
        Matrix to be decomposed into A=QR. A must be a square matrix with no
        diagonal entries equal to 0.

    Returns
    -------
    Q : array
        Orthogonal matrix with the same shape as A.
    R : array
        Upper triangular matrix with the same shape as A.
    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Initialize Q and R matrix
    Q = A.copy()
    R = np.eye(n, n)

    # Compute columns of Q matrix
    for j in range(n):
        # Subtract projections on previous columns
        current_column = Q[:, j].copy()
        for k in range(j):
            reference_column = Q[:, k]
            # Project current column on reference column
            proj = 0
            for l in range(n):
                proj += current_column[l] * reference_column[l]
            # Subtract projection from current column
            current_column -= proj*reference_column
        # Calculate norm
        norm = 0.0
        for k in range(n):
            norm += current_column[k]**2
        norm = (norm)**0.5
        # Normalize result
        Q[:, j] = current_column/norm

    # Compute lines of R matrix
    for i in range(n):
        for j in range(i, n):
            # Scalar projection of jth column from A into ith column from Q
            proj = 0
            A_column = A[:, j]
            Q_column = Q[:, i]
            for k in range(n):
                proj += A_column[k] * Q_column[k]
            R[i, j] = proj

    # Return Q and R
    return Q, R

# Iterative Solvers
def jacobi(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a jacobi iterative solver.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    print("Initializing Jacobi Solver", end='\r')

    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()
    u_new = u.copy()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u_new[i] = (1/A[i, i]) * (b[i] - sum_factor)

        # Update u
        u = u_new.copy()

        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    print("\nJacobi solver converged after ", len(res_array), " iterations.")
    return u, res_array

@jit(nopython=True)
def jacobi_compiled(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a jacobi iterative solver.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()
    u_new = u.copy()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u_new[i] = (1/A[i, i]) * (b[i] - sum_factor)

        # Update u
        u = u_new.copy()

        # Update residue
        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    return u, res_array

def jacobi_sparse(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a jacobi iterative solver.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b. Should be in sparse
        format.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # print("Initializing Jacobi Sparse Solver", end='\r')
    # Retrieve matrix dimension
    n, m = A.shape

    # Define D by extracting main diagonal from A
    main_diagonal = A.diagonal()
    D = dia_matrix(([main_diagonal], [0]), shape=(n, n), dtype=np.float64)

    # Compute R
    R = A - D

    # Compute D^-1
    D_inverse = dia_matrix(([1/main_diagonal], [0]), shape=(n, n), dtype=np.float64)

    # Compute (D^-1)*b
    D_inverse_dot_b =  D_inverse.dot(b.flatten())

    # Compute -(D^-1)*R
    D_inverse_dot_R = D_inverse.dot(R)

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        # print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        # Update u
        u = D_inverse_dot_b - D_inverse_dot_R.dot(u)

        # Update residue
        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    # print("\nJacobi solver converged after ", len(res_array), " iterations.")
    return u, res_array

def gauss_seidel(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel iterative solver.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    print("Initializing Gauss-Seidel Solver", end='\r')

    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u[i] = (1/A[i, i]) * (b[i] - sum_factor)

        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    print("\nGauss-Seidel solver converged after ", len(res_array), " iterations.")
    return u, res_array

@jit(nopython=True)
def gauss_seidel_compiled(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel iterative solver.
    This function takes advantage of just in time compilation to speed up
    evaluations.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u[i] = (1/A[i, i]) * (b[i] - sum_factor)

        # Update residue
        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    return u, res_array

def gauss_seidel_sparse(A, b, initial_guess, true_solution, tol):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel iterative solver.
    This function makes use of the sparsity of the A matrix to speed up
    function evaluation.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b. Should be in sparse
        format.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # print("Initializing Gauss-Seidel Sparse Solver", end='\r')
    # Retrieve matrix dimension
    n, m = A.shape

    # Compute L by extracting lower triangular matrix from A
    L = tril(A)

    # Compute U by extracting upper triangular matrix from A
    U = triu(A, k=1)

    # Compute L^-1
    L_inverse = sparse_inv(csc_matrix(L))

    # Compute (L^-1)*b
    L_inverse_dot_b =  L_inverse.dot(b.flatten())

    # Compute -(L^-1)*U
    L_inverse_dot_U = L_inverse.dot(U)

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        # print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        # Update u
        u = L_inverse_dot_b - L_inverse_dot_U.dot(u)

        # Update residue
        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    # print("\nGauss-Seidel solver converged after ", len(res_array), " iterations.")
    return u, res_array

def sor(A, b, initial_guess, true_solution, tol, w):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel Successive
    Over Relaxation iterative solver.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.
    w : float
        Relaxation parameter.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    print("Initializing Gauss-Seidel SOR Solver", end='\r')

    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u[i] = (1-w)*u[i] + w*(1/A[i, i]) * (b[i] - sum_factor)

        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    print("\nGauss-Seidel Successive Over Relaxation solver converged after ", len(res_array), " iterations.")
    return u, res_array

@jit(nopython=True)
def sor_compiled(A, b, initial_guess, true_solution, tol, w):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel Successive
    Over Relaxation iterative solver. This function takes advantage of just
    in time compilation to speed up evaluations.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.
    w : float
        Relaxation parameter.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # Retrieve matrix dimension
    n, m = A.shape

    # Flatten b
    b = b.flatten()

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        for i in range(n):
            sum_factor = 0
            for j in range(i):
                sum_factor += A[i, j] * u[j]

            for j in range(i + 1, n):
                sum_factor += A[i, j] * u[j]

            u[i] = (1-w)*u[i] + w*(1/A[i, i]) * (b[i] - sum_factor)

        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    return u, res_array

def sor_sparse(A, b, initial_guess, true_solution, tol, w):
    """ Solves the algebraic system A*x=b using a Gauss-Seidel Successive
    Over Relaxation iterative solver. This function makes use of the sparsity
    of the A matrix to speed up function evaluation.

    Parameters
    ----------
    A : array
        Matrix A corresponding to the system A*x = b. Should be in sparse
        format.
    b : array
        Column vector b corresponding to the system A*x = b.
    initial_guess : array, optional
        Initial solution.
    true_solution : array, optional
        Used for residual calculation
    tol : float, optional
        Convergence tolerance.
    w : float
        Relaxation parameter.

    Returns
    -------
    x : array
        Solution vector corresponding to A*x = b.
    residue : array
        Residue array.
    """
    # print("Initializing Gauss-Seidel Successive Over Relaxation Sparse Solver", end='\r')
    # Retrieve matrix dimension
    n, m = A.shape

    # Compute D by extracting main diagonal from A
    main_diagonal = A.diagonal()
    D = dia_matrix(([main_diagonal], [0]), shape=(n, n), dtype=np.float64)

    # Compute L by extracting lower triangular matrix from A
    L = tril(A, k=-1)

    # Compute U by extracting upper triangular matrix from A
    U = triu(A, k=1)

    # Compute (D + w*L)^-1
    D_plus_w_L_inverse = sparse_inv(csc_matrix(D + w*L))

    # Compute w*b
    w_b =  w*b.flatten()

    # Compute w*U + (w-1)*D
    w_U_plus_w_minus_1_D = w*U + (w-1)*D

    # Compute ((D + w*L)^-1)*(w*U + (w-1)*D)
    D_plus_w_L_inverse_dot_w_U_plus_w_minus_1_D = D_plus_w_L_inverse.dot(w_U_plus_w_minus_1_D)

    # Compute ((D + w*L)^-1)*w*b0
    D_plus_w_L_inverse_dot_w_b = D_plus_w_L_inverse.dot(w_b)

    # Initialize u
    u = initial_guess.flatten()
    u_true = true_solution.flatten()

    # Initialize residue
    res = np.linalg.norm(u - true_solution)
    res_array = [res]

    # Iterate
    while res > tol:
        # print("Current Iteration and Residue: {:06d} - {:05.4E}".format(len(res_array), res), end='\r')

        # Update u
        u = D_plus_w_L_inverse_dot_w_b - D_plus_w_L_inverse_dot_w_U_plus_w_minus_1_D.dot(u)

        # Update residue
        res = np.linalg.norm(u - true_solution)
        res_array += [res]

    # print("\nGauss-Seidel Successive Over Relaxation solver converged after ", len(res_array), " iterations.")
    return u, res_array


"""
Static methods that all take a float parameter meant to be between 0 and 1 (but can be any), and return a float between
0 and 1. Meant to be used as fill color value in order to have custom shadings. Some return a list of 3 floats
for rgb values.
"""
import math
from typing import Callable

from scipy.stats import norm

# import numpy as np

def _clamp(x: float, lower_limit: float = 0.0, upper_limit: float = 1.0) -> float:
    if x > upper_limit:
        return upper_limit
    if x < lower_limit:
        return lower_limit
    return x

def zero(ratio: float) -> float:
    return 0.0

def _power(exponent: float) -> Callable[[float], float]:
    return lambda x: _clamp(x ** exponent)

def identity(ratio: float) -> float:
    return _power(1.0)(ratio)

def square(ratio: float) -> float:
    return _power(2)(ratio)

def cube(ratio: float) -> float:
    return _power(3)(ratio)

def sinus(ratio: float) -> float:
    return _clamp(math.sin(2 * ratio / math.pi))

def _gauss(mean: float = 0.5, sd: float = 0.3) -> Callable[[float], float]:
    return lambda x: norm.cdf(x, mean, sd)

def gauss(ratio: float) -> float:
    return _gauss()(ratio)

def gauss_heavy(ratio: float) -> float:
    return _gauss(0.8, 0.4)(ratio)

def _logistic_curve(midpoint: float = 0.5, steepness: float = 1.0, max_value: float = 1.0) -> Callable[[float], float]:
    return lambda x: max_value / (1 + math.e ** (-steepness * (x - midpoint)))

def logistic_curve(ratio: float) -> float:
    return _logistic_curve(0.5, 5.0, 1.0)(ratio)

def smooth_step(ratio: float, edge_1: float = 0.0, edge_2: float = 1.0) -> float:
    x = _clamp((ratio - edge_1) / (edge_2 - edge_1))
    return x * x * (3 - 2 * x)

def smoother_step(ratio: float, edge_1: float = 0.0, edge_2: float = 1.0) -> float:
    x = _clamp((ratio - edge_1) / (edge_2 - edge_1))
    return x * x * x * (x * (x * 6 - 15) + 10)

def general_smooth_step(N, x):
    x = _clamp(x, 0.0, 1.0)
    result = 0

    for i in range(N):
        result += pascal_triangle(-N - 1, i) * pascal_triangle(2 * N + 1, N - i) * math.pow(x, N + i + 1)
    return result

def pascal_triangle(a, b):
    result = 1
    for i in range(b):
        result *= (a - i) / (i + 1)
    return result


from scipy.integrate import quad
import numpy as np

import matplotlib.pyplot as plt

#-----------------------------------
#INTEGRATION EXAMPLE
#----------------------------------

#ANTIDEVIATIVE OF x^2--> x^3/3
def integrand(x):
    return x**2

print(1/3-0,quad(integrand, 0, 1))

#-----------------------------------
#NORMAL DISTIBUTION ENTROPY
#----------------------------------

s=1; u=0
def p(x):
	return np.exp(-0.5*((x-u)/s)**2)/(s*np.sqrt(2*3.1415))

def I(x): return -p(x)*np.log(p(x))

S=quad(I, -20, 20)
print(S, np.log(2*3.1415*np.exp(1)*s)) #,np.exp(1))

#-----------------------------------
#KL DIVERGENCE AND CROSS ENTROPY
#----------------------------------

u1=0; s1=1
u2=3; s2=1

#INTEGRATION LIMITS
x1=min(u1,u2)-4*max(s1,s2)
x2=max(u1,u2)+4*max(s1,s2); print(x1,x2)

#DEFINE TWO NORMAL DISTRIBUTINO
def p1(x): return np.exp(-0.5*((x-u1)/s1)**2)/(s1*np.sqrt(2*3.1415))
def p2(x): return np.exp(-0.5*((x-u2)/s2)**2)/(s2*np.sqrt(2*3.1415))

def I_S1(x): return -p1(x)*np.log(p1(x))  		#P1 ENTROPY
def I_S2(x): return -p2(x)*np.log(p2(x))  		#P2 ENTROPY
def I_CE(x): return -p1(x)*np.log(p2(x))  		#CROSS ENTROPY
def I_KL(x): return -p1(x)*np.log(p2(x)/p1(x))  #D_KL

def plot(KLD,S1,S2,CE):
	fig = plt.figure(figsize=(20,12))
	ax = fig.add_subplot(111)
	x=np.linspace(x1,x2,500)
	plt.plot(x, p1(x), label="p(x)", linewidth=4)
	plt.plot(x, p2(x), label="q(x)", linewidth=4)
	plt.plot(x, I_S1(x), label="I_S1", linewidth=4)
	# plt.plot(x, I_S2(x), label="I_S2", linewidth=4)
	plt.plot(x, I_CE(x), label="I_CE", linewidth=4)
	plt.plot(x, I_KL(x), label="S1,S2,CE,KLD="+str(S1)+' '+str(S2)+' '+str(CE)+' '+str(KLD), linewidth=4)
	plt.legend(loc="best")
	ax.set_xlabel('x')
	ax.set_ylabel('probablity density function')
	plt.show()

S1=round(quad(I_S1, x1, x2)[0],3)
S2=round(quad(I_S1, x1, x2)[0],3)
KLD=round(quad(I_KL, x1, x2)[0],3)
CE=round(quad(I_CE, x1, x2)[0],3)
print(CE,S1+KLD)
plot(KLD,S1,S2,CE)

exit()

#-----------------------------------
#KL DIVERGENCE: MOVIE
#----------------------------------

fig = plt.figure(figsize=(20,12))
ax = fig.add_subplot(111)

for var in [15,1.42,3,4,5,6,7,8,9]:
	u1=0; s1=1
	u2=0; s2=var

	#INTEGRATION LIMITS
	x1=min(u1,u2)-4*max(s1,s2)
	x2=max(u1,u2)+4*max(s1,s2); print(x1,x2)

	#DEFINE TWO NORMAL DISTRIBUTINO
	def p1(x): return np.exp(-0.5*((x-u1)/s1)**2)/(s1*np.sqrt(2*3.1415))
	def p2(x): return np.exp(-0.5*((x-u2)/s2)**2)/(s2*np.sqrt(2*3.1415))

	#INTEGRAND
	def I(x): return -p1(x)*np.log(p2(x)/p1(x))

	def plot(KLD):

		x=np.linspace(x1,x2,500)
		ax.clear()

		ax.plot(x, p1(x), label="p(x)", linewidth=4)
		ax.plot(x, p2(x), label="q(x)", linewidth=4)
		ax.plot(x, I(x), label="p(x)log(p(x)/q(x)): KLD="+str(KLD), linewidth=4)
		ax.legend(loc="best")
		ax.set_xlabel('x')
		ax.set_ylabel('probablity density function')
		plt.pause(1)

	KLD=quad(I, x1, x2)[0]
	plot(KLD)

plt.show()


import numpy as np

def func1(x):
    # function f(x)= 10 + sum_i(-x_i^2)
    # for 2D: f(x)= 10 - x1^2 - x2^2
    return 10 + np.sum(-1*np.power(x, 2), axis=1)

def mc_integrate(func, a, b, dim, n = 1000):
    # Monte Carlo integration of given function over domain from a to b (for each parameter)
    # dim: dimensions of function

    x_list = np.random.uniform(a, b, (n, dim))
    print(x_list)
    y = func(x_list)
    print(y)

    y_mean =  y.sum()/len(y)
    domain = np.power(b-a, dim)

    integ = domain * y_mean

    return integ

# Examples
print("For f(x)= 10 - x1\u00b2 - x2\u00b2, integrated from -2 to 2 (for all x's)")
print(f"Monte Carlo solution for : {mc_integrate(func1, -2, 2, 2, 1000000): .3f}")
print(f"Analytical solution: 117.333")

print("For f(x)= 10 - x1\u00b2 - x2\u00b2 - x3\u00b2, integrated from -2 to 2 (for all x's)")
print(f"Monte Carlo solution: {mc_integrate(func1, -2, 2, 3, 1000000): .3f}")
print(f"Analytical solution: 384.000")


import numpy as np

def heuristic_aux(x, y, x_g, y_g, name):
  """Admissible heuristics

  Parameters
  ----------
  x, y: int
    state coordinate
  x_g, y_g: int
    goal coordinate
  name: str
    heuristic function name

  Return
  ------
  Float:
    Corresponding heuristic value for a single coordinate
  """
  if name == 'tiles-out':
    return x != x_g or y != y_g
  elif name == 'manhattan':
    return abs(x-x_g) + abs(y-y_g)
  elif name == 'euclidean':
    return np.sqrt((x-x_g) ** 2 + (y-y_g) ** 2)
  elif name == 'uniform-cost':
    return 0

def heuristic(goal, state, name='euclidean'):
  """Heuristic function

  Parameters
  ----------
  goal: list
    Puzzle state we want to achieve
  state: list
    Current Puzzle state
  name: str (Default: 'euclidean')
    Name of heuristic function to use ∈ ['euclidean', 'manhattan', 'tiles-out', 'uniform-cost']

  Returns
  -------
  One of
    'Tiles-out': float
      Number of tiles in the wrong position
    'Manhattan': float
      Sum of the horizontal and vertical distances between,
      current position and desired position, i.e. ∑ |state - goal|
    'Euclidean': float
      Sum of the distance between tiles
      ∑ √(x - x_g)² + (y - y_g)²
  """
  size = int(np.sqrt(len(goal)))

  coord_goal  = np.zeros((size ** 2, 2)) # in index i coordinates of tile i in goal
  coord_state = np.zeros((size ** 2, 2)) # in index i coordinates of tile i in state
  for x in range (size):
    for y in range (size):
      coord_goal[ goal[ x * size + y]] = [x, y]
      coord_state[state[x * size + y]] = [x, y]

  return np.sum([heuristic_aux(*coord_state[i], *coord_goal[i], name) for i in range(size ** 2)])


from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import math
import scipy.stats as stats

from .. import samp

def samp_size(s_p, width, alpha):
    """
    Input: s_p (sample proportion), width, alpha
    Output: Estimated sample size
    """
    z_cv = stats.norm.ppf(1 - alpha / 2)
    return (z_cv * math.sqrt(s_p * (1 - s_p)) / width) ** 2

def con_level(s_p, n, alpha, show=True, Wilson=False, N=False, correction=True):
    """
    Caution: np̂ > 5 and n(1 - p̂) > 5
    Input: s_p, n, alpha, show=True, Wilson=False, N=False, correction=True
    Output: {"lat": lat, "lcl": lcl, "ucl": ucl, "z_cv": z_cv}

    Note:
        z_cv = stats.norm.ppf(1 - alpha / 2)
        lat = z_cv * math.sqrt(s_p * (1 - s_p)/n)
        or for Wilson: lat = z_cv * math.sqrt(s_p * (1 - s_p)/(n + 4))

    Just in case:
    if there is no need for correction, but is is corrected
    go through 'lat' and do:
        lcl = s_p - lat
        ucl = s_p + lat
    """
    con_coef = 1 - alpha
    z_cv = stats.norm.ppf(1 - alpha / 2)
    if not Wilson and not samp.check5(n, s_p):
        print('Not satisfying np̂ > 5 and n(1 - p̂) > 5...')
    if Wilson:
        # make sure that you have arrange s_p to (x + 2) / (n + 4)
        lat = z_cv * math.sqrt(s_p * (1 - s_p)/(n + 4))
    else:
        lat = z_cv * math.sqrt(s_p * (1 - s_p)/n)

    lcl = s_p - lat
    ucl = s_p + lat

    if N:
        if n / N > 0.5 and correction:
            print("Corrected...")
            fpcf = math.sqrt((N - n)/(N - 1))
            lcl = s_p - lat * fpcf
            ucl = s_p + lat * fpcf
        elif correction:
            print("Corrected...")
            fpcf = math.sqrt((N - n)/(N - 1))
            lcl = s_p - lat * fpcf
            ucl = s_p + lat * fpcf
        if lcl < 0:
            lcl = 0
        if ucl < 0:
            ucl = 0
        result = f"""{con_coef * 100:.1f}% Confidence Interval: N [{lcl:.4f}, {ucl:.4f}] = [{N * lcl:.4f}, {N * ucl:.4f}]
p̂: {s_p:.4f}
Sample Size: {n}
z_cv (Critical value): {z_cv:.4f}
    """
    else:
        if lcl < 0:
            lcl = 0
        if ucl < 0:
            ucl = 0
        result = f"""{con_coef * 100:.1f}% Confidence Interval: [{lcl:.4f}, {ucl:.4f}]
p̂: {s_p:.4f}
Sample Size: {n}
z_cv (Critical value): {z_cv:.4f}
    """
    if show:
        print(result)
    return {"lat": lat, "lcl": lcl, "ucl": ucl, "z_cv": z_cv}

def rejection_region_method(s_p, h0_p, nsize, alpha, option='left', precision=4, show=True, ignore=False):
    """
    Input: s_p, h0_p, nsize, alpha, option='left', precision=4, show=True, ignore=False
    Output:
        if opt == 't':
            return p_l, p_u
        else:
            return p_c
    """
    opt = option.lower()[0]
    if not samp.check5(nsize, h0_p):
        print('Not satisfying np_0 > 5 and n(1 - p_0) > 5...')
    if opt == 't':
        option = 'Two-Tail Test'
        z_cv = stats.norm.ppf(1 - alpha/2)
        p_u = h0_p + z_cv * math.sqrt(h0_p * (1 - h0_p)/nsize)
        p_l = h0_p - z_cv * math.sqrt(h0_p * (1 - h0_p)/nsize)
        flag = s_p < p_l or s_p > p_u
        if not ignore:
            result = f'''======= The Rejection Region Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z (Critical value) = {z_cv:.{precision}f}

Using {option}:
p̂ =  {s_p:.{precision}f}
p_l (Lower bound for the critical value) = {p_l:.{precision}f}
p_u (Upper bound for the critical value) = {p_u:.{precision}f}
Reject H_0 → {flag}
            '''
        else:
            result = f'''======= The Rejection Region Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z (Critical value) = {z_cv:.{precision}f}

Using {option}:
p_l (Lower bound for the critical value) = {p_l:.{precision}f}
p_u (Upper bound for the critical value) = {p_u:.{precision}f}
            '''

    else:
        z_cv = stats.norm.ppf(1 - alpha)
        if opt == 'l':
            # left tail
            option = 'One-Tail Test (left tail)'
            p_c = h0_p - z_cv * math.sqrt(h0_p * (1 - h0_p)/nsize)
            flag = s_p < p_c
        elif opt == 'r':
            option = 'One-Tail Test (right tail)'
            p_c = h0_p + z_cv * math.sqrt(h0_p * (1 - h0_p)/nsize)
            flag = s_p > p_c
        if not ignore:
            result = f'''======= The Rejection Region Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z (Critical value) = {z_cv:.{precision}f}

Using {option}:
p̂ =  {s_p:.{precision}f}
p_c (Critical value) = {p_c:.{precision}f}
Reject H_0 → {flag}
            '''
        else:
            result = f'''======= The Rejection Region Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z (Critical value) = {z_cv:.{precision}f}

Using {option}:
p_c (Critical value) = {p_c:.{precision}f}
            '''

    if show:
        print(result)

    if opt == 't':
        return p_l, p_u
    else:
        return p_c

def testing_statistic_method(s_p, h0_p, nsize, alpha, option='left', precision=4, ignore=False):
    """
    Input: s_p, h0_p, nsize, alpha, option='left', precision=4, ignore=False
    Output: z_stats, z_cv
    """
    opt = option.lower()[0]
    z_stats = (s_p - h0_p)/math.sqrt(h0_p * (1 - h0_p)/nsize)
    if not samp.check5(nsize, h0_p):
        print('Not satisfying np_0 > 5 and n(1 - p_0) > 5...')

    if opt == 't':
        z_cv = stats.norm.ppf(1 - alpha / 2)
        option = 'Two-Tail Test'
        flag = z_stats < -z_cv or z_stats > z_cv

        if not ignore:
            result = f'''======= Testing Statistic Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z_cv (Critical value) = {z_cv:.{precision}f}

Using {option}:
z_stats (Observed value) =  {z_stats:.{precision}f}
-z_cv (Lower bound for the critical value) = {-z_cv:.{precision}f}
z_cv (Upper bound for the critical value) = {z_cv:.{precision}f}
Reject H_0 → {flag}
            '''
        else:
            result = f'''======= Testing Statistic Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z_cv (Critical value) = {z_cv:.{precision}f}

Using {option}:
-z_cv (Lower bound for the critical value) = {-z_cv:.{precision}f}
z_cv (Upper bound for the critical value) = {z_cv:.{precision}f}
            '''

    else:
        z_cv = stats.norm.ppf(1 - alpha)
        if opt == 'l':
            # left tail
            option = 'One-Tail Test (left tail)'
            z_cv = -z_cv
            flag = z_stats < z_cv
        elif opt == 'r':
            option = 'One-Tail Test (right tail)'
            flag = z_stats > z_cv

        if not ignore:
            result = f'''======= Testing Statistic Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z_cv (Critical value) = {z_cv:.{precision}f}

Using {option}:
z_stats (Observed value) =  {z_stats:.{precision}f}
z_cv (Critical value) = {z_cv:.{precision}f}
Reject H_0 → {flag}
            '''

        else:
            result = f'''======= Testing Statistic Method =======
Significant Level (alpha) = {alpha:.{precision}f}
z_cv (Critical value) = {z_cv:.{precision}f}

Using {option}:
z_cv (Critical value) = {z_cv:.{precision}f}
            '''

    print(result)

    return z_stats, z_cv

def inter_p_value(p_value):
    # interpretation
    if p_value >= 0 and p_value < 0.01:
        inter_p = 'Overwhelming Evidence'
    elif p_value >= 0.01 and p_value < 0.05:
        inter_p = 'Strong Evidence'
    elif p_value >= 0.05 and p_value < 0.1:
        inter_p = 'Weak Evidence'
    elif p_value >= .1:
        inter_p = 'No Evidence'
    return inter_p

def p_value_method(s_p, h0_p, nsize, alpha, option='left', precision=4):
    """
    Input: s_p, h0_p, nsize, alpha, option='left', precision=4
    Output: z_cv, z_stats, p_value
    """
    opt = option.lower()[0]
    z_stats = (s_p - h0_p)/math.sqrt(h0_p * (1 - h0_p)/nsize)
    if not samp.check5(nsize, h0_p):
        print('Not satisfying np_0 > 5 and n(1 - p_0) > 5...')
    if opt == 't':
        # two-tail test
        option = 'Two-Tail Test'
        if s_p > h0_p:
            p_value = stats.norm.sf(z_stats) * 2
        else:
            p_value = stats.norm.cdf(z_stats) * 2

        z_cv = stats.norm.ppf(1 - alpha/2)
        flag = p_value < alpha
        sub_result = f'''Using {option}:
Difference = {s_p - h0_p}
z_cv (Critical value) = {-z_cv:.{precision}f}, {z_cv:.{precision}f}
z_stats (Observed value) = {z_stats:.{precision}f}
p-value = {p_value:.{precision}f} ({inter_p_value(p_value)})
Reject H_0 → {flag}
        '''
    else:
        if opt == 'l':
            option = 'One-Tail Test (left tail)'
            p_value = stats.norm.cdf(z_stats)
            z_cv = -stats.norm.ppf(1 - alpha)
        elif opt == 'r':
            option = 'One-Tail Test (right tail)'
            p_value = stats.norm.sf(z_stats)
            z_cv = stats.norm.ppf(1 - alpha)
        flag = p_value < alpha
        sub_result = f'''Using {option}:
Difference = {s_p - h0_p}
z_cv (Critical value) = {z_cv:.{precision}f}
z_stats (Observed value) = {z_stats:.{precision}f}
p-value = {p_value:.{precision}f} ({inter_p_value(p_value)})
Reject H_0 → {flag}
        '''

    result = f"""======= p-value Method =======
p̂ = {s_p:.{precision}f}
Number of Observation = {nsize:.{precision}f}
Hypothesized Proportion (H0 Mean) = {h0_p:.{precision}f}
Significant Level (alpha) = {alpha:.{precision}f}

""" + sub_result

    print(result)

    return z_cv, z_stats, p_value


import matplotlib.pyplot as plt
import numpy as np

x = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
y = [0.99, 2.3, 2.9, 4.01, 4.85, 5.80, 7.2]
size = 7

plt.plot(x, y, 'go')
plt.ylabel('some numbers')

# Algorithm goes here
xsum = 0
x2sum = 0
ysum = 0
xysum = 0
slope = 0
intercept = 0

for i in range(size):
    xsum = xsum + x[i]
    ysum = ysum + y[i]
    x2sum = x2sum + x[i]**2
    xysum = xysum + x[i] * y[i]

slope = (size * xysum - xsum * ysum) / (size * x2sum - xsum * xsum)
intercept = (x2sum * ysum - xsum * xysum) / (x2sum * size - xsum * xsum)
y_fit = []

for i in range(size):
    y_fit.append(slope * x[i] + intercept)

plt.plot(x, y_fit, 'r')

plt.show()

#t = np.arange(0., 5., 0.2)
#plt.plot(t, t, 'r--', t, t**2, 'bs', t, t**3, 'g^')
#plt.show()


import numpy as np

def CorBrownian(mu, E, sampleSize):
    # Algorithm generates samples of increments from a correlated Brownian motion with a given mean and Variance-Covariance matrix (E).
    # The algorithm uses the fact that if you have n independent brownian motions, the samples given by "mu+ C*Z" are distributed as N(mu,E), where mu is the vector of means and C is the square root of the Variance-Covariance matrix.
    # For calculating the square root of the VarCovar matrix, the Cholesky decomposition is implemented.
    # Arguments:
    #     mu        = Array with n elements containing the mean of each BM
    #     E         = n x n numpy array with the Variance-Covariance matrix
    #    sampleSize = integer representing the number of samples
    #
    # Returns:
    #     sampleSize x n numpy array containing sampled increments for the correlated Brownian motion
    #
    # Note:
    #  The algorithm is not optimized for speed and no testing of inputs is implemented. If this would be usefull to you, let us know and we can extend the code.
    #
    # Example of use:
    # import numpy as np
    # mu = [1,2]
    # VarCovar = np.matrix('1,0.8; 0.8,3')
    # sampleSize = 5
    # out = CorBrownian(mu,VarCovar, sampleSize)
    # > [ 2.83211068  4.50021193]
    #   [ 0.26392619  1.56450446]
    #   [-0.25928109  0.97167124]
    #   [ 1.52038489  1.76274556]]
    #

    def Cholesky(X):
        # Cholesky–Banachiewicz algorithm decomposes a Hermitian matrix into a product of a lower triangular matrix and its conjugate transpose.
        # Arguemnts:
        #    X = n x n ndarray representing a Hermitian matrix that the user wants to decompose
        # Returns:
        #    n x n ndarray lower triangular matrix such that the matrix product between it and its conjugate transpose returns X
        #
        # More info on: https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky.E2.80.93Banachiewicz_and_Cholesky.E2.80.93Crout_algorithms

        L = np.zeros_like(X)
        n = X.shape[0]

        for i in range(0, n):
            for j in range(0, i+1):
                sum = 0
                for k in range(0, j):
                    sum = sum+ L[i,k]*L[j,k]
                if (i==j):
                    L[i,j] = np.sqrt(X[i,i]-sum)
                else:
                    L[i,j] = 1.0/L[j,j] * (X[i,j]-sum)
        return L

    dim = E.shape[0]                                         # Guess the number of Brownian motions (dimension) from the size of the Var-Covar matrix
    Z = np.random.default_rng().normal(0, 1, (sampleSize, dim)) # Generate independent increments of a simpleSize dimensional Brownian motion
    Y = np.zeros((sampleSize, dim))                          # Predefine the final output
    L = Cholesky(E)                                          # Calculate the square root of the Var-Covar matrix

    for iSample in range(sampleSize): # For each sample, calculate mu + L*Z
        Y[iSample] = np.transpose(mu) +  L @ np.transpose(Z[iSample])
    return Y


# PROJECT: Cross-Entropy Script
# PROGRAMMER: Carlos Mertens - Udacity Student
#				Udacity Instructors (Machine Learning Engineer Nanodegree)
# DATE CREATED: (DD/MM/YY) 10/01/2019
# REVISED DATE: (DD/MM/YY)
# PURPOSE: To show the algorithm to implement Cross-Entropy from the scratch.
# USAGE: ...

import numpy as np

# Write a function that takes as input two lists Y, P,
# and returns the float corresponding to their cross-entropy.
def cross_entropy(y, p):
    y = np.float_(y)
    p = np.float_(p)
    result = -np.sum(y * np.log(p) + (1 - y) * np.log(1 - p))
    return result

y = [1,0,1,1]
p = [0.4,0.6,0.1,0.5]

# Call function to calculate the cross-entropy
my_cross_entropy = cross_entropy(y, p)

print(my_cross_entropy)


import numpy as np

def mse(e):
    return 1/2 * np.mean(e**2)

def rmse(e):
    return np.sqrt(2 * mse(e))

def mae(e):
    return np.mean(np.abs(e))

def sigmoid(z):
    z = np.clip(z, -1000, 1000)
    return 1.0 / (1 + np.exp(-z))

def compute_mse(y, tx, w):
    """Compute the loss by MSE (Mean Square Error)."""
    e = y - tx.dot(w)
    return mse(e)

def compute_rmse(y, tx, w):
    """Compute the loss by RMSE (Root Mean Square Error)."""
    e = y - tx.dot(w)
    return rmse(e)

def compute_mae(y, tx, w):
    """Compute the loss by MAE (Mean Absolute Error)."""
    e = y - tx.dot(w)
    return mae(e)

def compute_log_likelihood_error(y, tx, w):
    """Compute the loss of the log-likelihood cost function."""
    tx_dot_w = tx.dot(w)
    return np.sum(np.log(1 + np.exp(tx_dot_w))) - y.dot(tx_dot_w)


import numpy as np

def gaussian_elimination_with_pivot(m):

    # forward elimination
    n = len(m)
    print(n)
    for k in range(n):
        max = -1e100
        for r in range(k, n):
            if max < abs(m[r][k]):
                max_row = r
                max = abs(m[r][k])
        m[k], m[max_row] = m[max_row], m[k]
        print(m)
        for i in range(k+ 1, n):
            m[i] = [m[i][j] - m[k][j] * m[i][k] / m[k][k] for j in range(n + 1)]

    if m[n - 1][n - 1] == 0: raise ValueError('No unique solution')

    print(m)

    # backward substitution
    x = [0] * n
    for i in range(n - 1, -1, -1):
        s = sum(m[i][j] * x[j] for j in range(i, n))
        x[i] = (m[i][n] - s) / m[i][i]
    return x

'''
# shorter way to pivot but cannot run in trinket
def pivot(m, n, i):
  max_row = max(range(i, n), key=lambda r: abs(m[r][i]))
  m[i], m[max_row] = m[max_row], m[i]
'''

def pivot(m, n, i):
    max = -1e100
    for r in range(i, n):
        if max < abs(m[r][i]):
            max_row = r
            max = abs(m[r][i])
    m[i], m[max_row] = m[max_row], m[i]

if __name__ == '__main__':
    # m = [[0,-2,6,-10], [-1,3,-6,5], [4,-12,8,12]]
    # m = [[1,-1,3,2], [3,-3,1,-1], [1,1,0,3]]
    m = [[0.5, 1.1, 3.1, 6], [5, 0.96, 6.5, 0.96], [2, 4.5, 0.36, 0.02]]
    print(gaussian_elimination_with_pivot(m))
    print(len(m))

    """
    m = [[4,4,0,400], [-1,4,2,400], [0,-2,4,400]]   # aj Montri p80  [50, 50, 125]
    print(gaussian_elimination_with_pivot(m))
    """