import numpy as np

def fitness_function_float(function_number):

    if function_number == 1:
        return lambda chromosome: -(np.abs(chromosome[0]) + np.cos(chromosome[0]))
    elif function_number == 2:
        return lambda chromosome: -(np.abs(chromosome[0]) + np.sin(chromosome[0]))
    elif function_number == 3:
        return lambda chromosome: -(chromosome ** 2).sum()
    elif function_number == 4:
        return lambda chromosome: -np.sum(
            np.abs(chromosome) - 10 * np.cos(np.sqrt(np.abs(10 * chromosome)))
        )
    elif function_number == 5:
        return lambda chromosome: -(chromosome[0] ** 2 + chromosome[0]) * np.cos(
            chromosome[0]
        )
    elif function_number == 6:
        return lambda chromosome: -(
            chromosome[0] * np.sin(4 * chromosome[0])
            + 1.1 * chromosome[1] * np.sin(2 * chromosome[1])
        )
    elif function_number == 7:  # Bohachevsky, bound in [-100,100], opt at (0,0)
        return lambda chromosome: -(
            0.7
            + chromosome[0] ** 2
            + 2.0 * chromosome[1] ** 2
            - 0.3 * np.cos(3 * np.pi * chromosome[0])
            - 0.4 * np.cos(4 * np.pi * chromosome[1])
        )

    elif function_number == 8:  # Goldstein, bounds in [-2,2], opt at (0,-1)
        a = 1 + (chromosome[0] + chromosome[1] + 1) ** 2 * (
            19
            - 14 * chromosome[0]
            + 3 * chromosome[0] ** 2
            - 14 * chromosome[1]
            + 6 * chromosome[0] * chromosome[1]
            + 3 * chromosome[1] ** 2
        )
        b = 30 + (2 * chromosome[0] - 3 * chromosome[1]) ** 2 * (
            18
            - 32 * chromosome[0]
            + 12 * chromosome[0] ** 2
            + 48 * chromosome[1]
            - 36 * chromosome[0] * chromosome[1]
            + 27 * chromosome[1] ** 2
        )
        return lambda chromosome: -a * b

    elif function_number == 9:  # Hartmann3, bounds in [0,1], opt at (0.11,0.55,0.85)
        a = np.array(
            [[3.0, 10.0, 30.0], [0.1, 10.0, 35.0], [3.0, 10.0, 30.0], [0.1, 10.0, 35.0]]
        )
        c = np.array([1.0, 1.2, 3.0, 3.2])
        p = np.array(
            [
                [0.36890, 0.11700, 0.26730],
                [0.46990, 0.43870, 0.74700],
                [0.10910, 0.87320, 0.55470],
                [0.03815, 0.57430, 0.88280],
            ]
        )
        return lambda chromosome: np.sum(
            c * np.exp(-np.sum(a * (chromosome[:] - p) ** 2, -1)), -1
        )

    elif (
        function_number == 10
    ):  # Hartmann6, bounds in [0,1]. opt at (0.2, 0.15, 0.47, 0.27, 0.31, 0.65)
        a = np.array(
            [
                [10.0, 3.0, 17.0, 3.5, 1.7, 8.0],
                [0.05, 10.0, 17.0, 0.1, 8.0, 14.0],
                [3.0, 3.5, 1.7, 10.0, 17.0, 8.0],
                [17.0, 8.0, 0.05, 10.0, 0.1, 14.0],
            ]
        )
        c = np.array([1.0, 1.2, 3.0, 3.2])
        p = np.array(
            [
                [0.1312, 0.1696, 0.5569, 0.0124, 0.8283, 0.5886],
                [0.2329, 0.4135, 0.8307, 0.3736, 0.1004, 0.9991],
                [0.2348, 0.1451, 0.3522, 0.2883, 0.3047, 0.6650],
                [0.4047, 0.8828, 0.8732, 0.5743, 0.1091, 0.0381],
            ]
        )
        return (
            lambda chromosome: (
                np.sum(c * np.exp(-np.sum(a * (chromosome[:] - p) ** 2, -1)), -1) + 2.58
            )
            / 1.94
        )


from sympy.functions import exp
from symplyphysics import (
    symbols, Eq, pretty, solve, Quantity, units, S,
    Probability, validate_input, expr_to_quantity, convert_to
)

# Description
## The resonance escape probability, symbolized by p, is the probability that a neutron will be
## slowed to thermal energy and escape resonance capture. This probability is defined as the ratio
## of the number of neutrons that reach thermal energies to the number of fast neutrons that slow down.

# Conditions
## - Homogeneous reactor.
## - Weak fast absorptions.
## - Very predominant absorber.

## Law: p ≈ e^(-(Na * Ieff) / (ξ * Σs_mod))
## Where:
## e - exponent.
## Na - atomic number density of the absorber.
##    See [atomic number density](./chemistry/atomic_number_density_from_material_density_atomic_weight.py) implementation.
## ξ - average lethargy gain per scattering event. Lethargy is a dimensionless logarithm of the ratio of the energy of
##    source neutrons to the energy of neutrons after a collision.
## Ieff - effective resonance integral. In practical situations, this integral strongly depends on the geometry of the unit cell.
##    It is too complex to calculate Ieff. Please refer to the tables and specific energy levels for proper data.
## Σs_mod - moderator's macroscopic scattering cross-section.
##   See [macroscopic cross-section](./macroscopic_cross_section_from_free_mean_path.py) implementation.
## p - resonance escape probability

absorber_atomic_number_density = symbols('absorber_atomic_number_density')
effective_resonance_integral = symbols('effective_resonance_integral')
average_lethargy_change = symbols('average_lethargy_change')
macroscopic_scattering_cross_section_moderator = symbols('macroscopic_scattering_cross_section_moderator')
resonance_escape_probability = symbols('resonance_escape_probability')

law = Eq(resonance_escape_probability,
    exp(-1 * (absorber_atomic_number_density * effective_resonance_integral) /
    (average_lethargy_change * macroscopic_scattering_cross_section_moderator)))

def print():
    return pretty(law, use_unicode=False)

@validate_input(
    absorber_atomic_number_density_=(1 / units.length ** 3),
    effective_resonance_integral_=(units.length ** 2),
    macroscopic_scattering_cross_section_moderator_=(1 / units.length))
def calculate_resonance_escape_probability(
    absorber_atomic_number_density_: Quantity,
    effective_resonance_integral_: Quantity,
    average_lethargy_change_: float,
    macroscopic_scattering_cross_section_moderator_: Quantity) -> Probability:

    result_factor_expr = solve(law, resonance_escape_probability, dict=True)[0][resonance_escape_probability]
    result_expr = result_factor_expr.subs({
        absorber_atomic_number_density: absorber_atomic_number_density_,
        effective_resonance_integral: effective_resonance_integral_,
        average_lethargy_change: average_lethargy_change_,
        macroscopic_scattering_cross_section_moderator: macroscopic_scattering_cross_section_moderator_})
    result_factor = expr_to_quantity(result_expr, 'resonance_escape_factor')
    return Probability(convert_to(result_factor, S.One).n())


'''
This program plot compton scatteing

the equation is:

             h
      Δλ = ----- (1 - cos(θ))
             mc

'''

import sys,os
sys.path.insert(1,os.path.join(sys.path[0],'..'))
from constants.constants import h,m_e,c

import matplotlib.pyplot as plt
import numpy as np

theta = np.linspace(0,np.pi,100)
lambda_ = h / (m_e *c) * (1 - np.cos(theta))

plt.plot(theta , lambda_,label = r'$\Delta\lambda = \frac{h}{m_0c}\left(1-\cos\theta\right)$')
plt.axvline(np.pi/2 , ymax = 0.5 , color='r' , linestyle = ':')
plt.xticks([0 , np.pi/2 , theta[-1]] , [0 , r'$\pi/2$' , r'$\pi$'])
plt.yticks([0 , lambda_[-1]] , [0 , r'$\frac{2h}{m_0c}$'])
plt.xlabel(r'$\theta\;\longrightarrow$')
plt.ylabel(r'$\Delta\lambda\;\longrightarrow$')
plt.legend()
plt.show()


# Comparison of high-level programming languages
# This is Python (requires version 3.6+)

# Function
def power(x, a):
    return x**a

# Lambda function
power = lambda x, a: x**a

# Variable assignment
x = power(2, 0.5)

# Standard output
print(x)

# Formatted output
print("{:5.2f}".format(x))  # Since 3.4
print(f"{x:5.2f}")  # Since 3.6
print("%5.2f" % x)  # 'Old style'

# Definite loop
for i in range(5):
    print(i)

# Indefinite loop
i = 0
while i < 5:
    print(i)
    i += 1

# List
symbols = ['H', 'He', 'Li']
print(symbols)

# Dictionary
elements = {
    'H': 1,
    'He': 2,
    'Li': 3
}
print(elements)

# Iterating over collections
for symbol, n in elements.items():
    print(f"{symbol}: {n}")

# Array literals
import numpy as np

# Vector (1d)
x = np.array([1, 2, 3])

# Matrix (2d)
A = np.array([[1, 2],
              [3, 4],
              [5, 6]])
print(x)
print(A)

# 3d array (3x2x3)
Z = np.ones((3, 2, 3))
print(Z)

# Array dimensions
print(x.shape)
print(A.shape)
print(Z.shape)

# Indexing arrays
print(x[1])
print(A[1,1])
print(A[1])

# Slicing arrays
print(x[1:])
print(A[:,1])

# Concatenation of vector and 2-d array
C = np.hstack([x.reshape(-1, 1), A])
# or
C = np.column_stack([x, A])
print(C)

# Array broadcasting
print(1 - A)
print(A + x.reshape(-1, 1))

# Linear Algebra
A = np.array([[0.8, 0], [0, 1]])
K = np.array([[0.942],[1.074]])
C = np.array([[0.2, 1]])

# Matrix inverse
# Since Python 3.5
X = np.linalg.inv(np.eye(2) - A + K @ C)
# Prior to 3.5
X = np.linalg.inv(np.eye(2) - A + K.dot(C))

print(X)

# Symbolic math
from sympy import symbols, solve, diff
a, b, c, x = symbols('a b c x')
y = a*x**2 + b*x + c
dydx = diff(y, x)
x_sol = solve(dydx, x)
y_sol = y.subs(x, x_sol[0])
print(dydx)
print(x_sol)
print(y_sol)

# Vectorizable function
def lorenz(y, sigma, beta, rho):
    return (sigma * (y[1] - y[0]),
            y[0] * (rho - y[2]) - y[1],
            y[0] * y[1] - beta * y[2])

# Parameter values
beta = 8 / 3
sigma = 10
rho = 28

# Test1 - Single function call
y0 = (-8, 8, 27)
assert lorenz(y0, sigma, beta, rho) == (160, -16, -136)

# Test 2 - Single function calls with for loop
def make_function_calls(f, y, n, *args, **kwargs):
    dy = np.empty_like(y)
    for i in range(n):
        dy[:, i] = f(y[:, i], *args, **kwargs)
    return dy

y = np.repeat(np.array((-8, 8, 27)).reshape(-1,1), 10, axis=1)
dy = make_function_calls(lorenz, y, 10, sigma=sigma, beta=beta, rho=rho)
assert dy.shape == y.shape
assert np.array_equiv(dy.T, (160, -16, -136))

# Test 3 - vectorized function call
y = np.repeat(np.array((-8, 8, 27)).reshape(-1,1), 10, axis=1)
dy = lorenz(y, sigma, beta, rho)
dy = np.vstack(dy)
assert dy.shape == y.shape
assert np.array_equiv(dy.T, (160, -16, -136))

# Test 4 - Numpy apply method (slowest, not vectorized)
dy = np.apply_along_axis(lorenz, 0, y, sigma=sigma, beta=beta, rho=rho)
assert dy.shape == y.shape
assert np.array_equiv(dy.T, (160, -16, -136))

# Test 5 - Numpy vectorized function
vlorenz = np.vectorize(lorenz, excluded={"sigma", "beta", "rho"})

# Speed test results
# Test 1 - Single function call: 504 ns ± 6.08 ns
# Test 2 - Single function calls with for loop: 37.2 ms ± 302 µs
# Test 3 - Vectorized function calls: 59.0 µs ± 381 ns
# Test 4 - Numpy apply method 54.7 ms ± 511 µs

# Timing code (use in IPython or Jupyter notebook):
# n = 10000
# y = np.random.randn(3, n)
# y0 = y[:, 0]
# def make_function_calls(f, y, n, *args, **kwargs):
#     dy = np.empty_like(y)
#     for i in range(n):
#         dy[:, i] = f(y[:, i], *args, **kwargs)
#     return dy
# %timeit dy = lorenz(y0, sigma, beta, rho)
# %timeit dy = make_function_calls(lorenz, y, n, sigma=sigma,
#                                  beta=beta, rho=rho)
# %timeit dy = lorenz(y, sigma, beta, rho)
# %timeit dy = np.apply_along_axis(lorenz, 0, y, sigma=sigma,
#                                  beta=beta, rho=rho)

# Lorenz system with forcing
def lorenz(y, u, sigma=10, beta=8/3, rho=28):
    return (sigma * (y[1] - y[0]) + u,
            y[0] * (rho - y[2]) - y[1],
            y[0] * y[1] - beta * y[2])


# Sampling from a discrete distribution

import numpy as np
import numpy.random as npr
import matplotlib.pyplot as plt
import bisect

def normalize(pmf):
    return pmf/(np.sum(np.abs(pmf)))

# This sampling function samples from a discrete distribution with probability mass function "pmf"
# and support coordinates specified in "support" by mapping a uniform random variable "u" onto the
# cumulative distribution function "cdf" and finding the associated support coordinate index
def sample(pmf, support=None):
    if support is None:
        support = np.arange(pmf.size)
    # Compute the cdf
    cdf = np.cumsum(pmf)
    # Sample from the uniform distribution
    u = npr.rand()
    # Find first index where cdf exceeds the uniform sample
    idx = bisect.bisect_left(cdf, u)
    # Map the index onto the support and return the sample
    return support[idx]

# This sampling function does the same thing as sample(), but uses the NumPy built-in function numpy.random.choice()
def sample_choice(pmf, support=None):
    return npr.choice(support, p=pmf)

# Number of points in the pmf support
n = 50
n2 = int(n/2)
# Define a support
support = np.hstack([np.linspace(-10, 0, n2), np.linspace(0+10/n2, 30, n-n2)])
# Define a pmf
a = np.linspace(0, 6, n)
pmf = normalize(np.abs(0.1*a**2 + np.sin(a)))
# Compute the pdf
d = np.diff(support)
ds = (np.hstack([d, d[-1]]) + np.hstack([d[0], d]))/2
pdf = pmf/ds

# Number of Monte Carlo samples
N = 100000
# Take Monte Carlo samples to get an empirical estimate of the pmf
samples = np.zeros(N)
for i in range(N):
    samples[i] = sample(pmf, support)

# Plot the pdf and empirical estimate
plt.close('all')
plt.style.use('../conlab.mplstyle')
plt.step(support, pdf, where='post', color='C0', label='True pdf')
plt.hist(samples, bins=support, density=True, align='mid', color='k', alpha=0.4, label='Monte Carlo pdf')
plt.legend()
plt.tight_layout()
plt.show()


import numpy as np

############################################################################
################### THIS HELPER IS ONLY FOR CONV WITH BN ###################
############################################################################
############### Note that: There are no parameters b anymore. ##############

def nn_f(X, W):
	# X : (m, n_prev)
	# W : (n, n_prev)
	Z = X @ W.T
	cache = (X, W)
	return Z, cache

def nn_b(dZ, cache):

	# dZ : (m, n)
	# cache : X, W
	X, W = cache
	m, _ = dZ.shape

	dW = dZ.T @ X / m
	dX = dZ @ W

	return dX, dW

def zero_pad(X, pad):

	#  X  : shape(m, n_H, n_W, n_C)
	# pad : scalar

	return np.pad(X,((0,0),(pad,pad),(pad,pad),(0,0)),'constant',constant_values=(0,0))

def conv_f(X, W, params):

	# X -> input,  shape : (m, in_H, in_W, in_C)
	# W -> filter, shape : (filters, f, f, in_C)
	# params: dict{'s':s, 'pad':pad}

	# Return:
	# Z -> conv output, shape : (m, out_H, out_W, filters)
	# cache -> Z, W (cache them for backward)

	filters, f, f, in_C = W.shape
	m, in_H, in_W, in_C = X.shape
	s, pad = params['s'], params['pad']

	out_H = int((in_H+2*pad-f)/s)+1
	out_W = int((in_W+2*pad-f)/s)+1

	X_pad = zero_pad(X, pad) if pad != 0 else X
	Z = np.zeros((m, out_H, out_W, filters), dtype=np.float32)

	for i in range(m):
		for h in range(out_H):
			for w in range(out_W):
				v_s, v_e, h_s, h_e = h*s, h*s+f, w*s, w*s+f
				x_slice = X_pad[i, v_s:v_e, h_s:h_e, :]
				Z[i,h,w,:] = ((W * x_slice).reshape(filters, -1).sum(1)).astype(np.float32)

	cache = (X, W, params)
	return Z, cache

def conv_b(dZ, cache):

	# dZ is the derivative
	# cache (X, W, params)

	X, W, params = cache
	m, n_H, n_W, _ = dZ.shape
	filters, f, f, _ = W.shape
	s, pad = params['s'], params['pad']

	dX = np.zeros(X.shape, dtype=np.float32)
	dW = np.zeros(W.shape, dtype=np.float32)

	for i in range(m):
		for h in range(n_H):
			for w in range(n_W):
				v_s, v_e, h_s, h_e = h*s, h*s+f, w*s, w*s+f
				x_slice = X[i, v_s:v_e, h_s:h_e, :]
				dW += (x_slice * dZ[i,h,w,:][:, None, None, None]).astype(np.float32)
				dX[i, v_s:v_e, h_s:h_e, :] += np.sum(W * dZ[i,h,w,:][:, None, None, None], axis=0).astype(np.float32)

	return dX, dW/m

def pool_f(X, params, mode=0):
	# NOTE THAT: POOLING LAYER DOESN'T HAVE WEIGHTS
	# Input parameters:
	# X -> input shape:(m, in_H, in_W, in_C)
	# params -> dict{f, s, out_C}
	# mode -> max=0, avg=1, default:0

	# Return:
	# Z -> output shape:(m, out_H, out_W, out_C)
	# cache -> X, params, mode (sava them for backwards)

	m, in_H, in_W, in_C = X.shape
	f, s = params['f'], params['s']

	out_H = int((in_H-f)/s)+1
	out_W = int((in_W-f)/s)+1

	Z = np.zeros((m, out_H, out_W, in_C), dtype=np.float32)
	pool = np.max if mode == 0 else np.mean

	for i in range(m):
		x_single = X[i]
		for h in range(out_H):
			for w in range(out_W):
				v_s, v_e, h_s, h_e = h*s, h*s+f, w*s, w*s+f
				x_slice = x_single[v_s:v_e, h_s:h_e, :]
				Z[i,h,w,:] = pool(x_slice.reshape(-1, in_C), axis=0)

	cache = (X, params, mode)
	return Z, cache

def pool_b(dZ, cache):

	# mode: 0->max, 1->avg
	X, params, mode = cache
	m, n_H, n_W, n_C = dZ.shape
	f, s = params['f'], params['s']
	dX = np.zeros(X.shape)
	pool = max_pool_b if mode == 0 else avg_pool_b

	for i in range(m):
		for h in range(n_H):
			for w in range(n_W):
				v_s, v_e, h_s, h_e = h*s, h*s+f, w*s, w*s+f
				dz_block = dZ[i, v_s, h_s, :]
				x_block = X[i, v_s:v_e, h_s:h_e, :]
				dX[i, v_s:v_e, h_s:h_e, :] = pool(dz_block, x_block)
	return dX

def max_pool_b(dz, x):
	# dz : (n_c,)
	# x  : (f,f,n_c)
	return (x == x.max(axis=(0,1))).astype(np.int16) * dz

def avg_pool_b(dz, x):
	# dz : (n_c,)
	# x  : (f,f,n_c)
	return np.full(x.shape, dz/(x.shape[0]**2))

def relu(Z):
	# Z : any shape
	return np.maximum(0, Z)

def relu_b(dA, Z):
	dZ = np.array(dA, copy=True)
	dZ[ Z<=0 ] = 0
	return dZ

def sigmoid(Z):
	# Z : any shape
	return 1.0 / (1.0 + np.exp(-Z))

def sigmoid_b(dA, Z):
	return dA * sigmoid(Z) * (1-sigmoid(Z))

#########################################################
################ 2018.04.10 / 2018.04.20 ################
#########################################################

def bn_conv_f(X, gamma, beta, epsilon=1e-6):

	#  batch_norm
	#  between: (conv,pool) and (activation)
	#
	#     X : (m, n_H, n_W, n_C)
	#  beta : (n_C,)
	# gamma : (n_C,)
	#     Y : (m, n_H, n_W, n_C) <- output

	# NOTE THAT: ***********************************************************************************
	# We learn a pair of parameters gamma(k) and beta(k) per feature map, rather than per activation
	# **********************************************************************************************

	m, n_H, n_W, n_C = X.shape
	# dof : degree of freedom
	dof_mu = m*n_H*n_W
	dof_s = (m-1)*n_H*n_W

	mu = X.sum(axis=(0,1,2)) / dof_mu
	s = np.power(X-mu, 2).sum(axis=(0,1,2)) / dof_s
	X_norm = ((X-mu) / np.sqrt(s+epsilon)).astype(np.float32)
	Y = gamma*X_norm + beta

	cache = (mu, s, X, X_norm, gamma, beta, epsilon)
	return Y, cache

def bn_conv_b(dY, cache):

	# dY : (m, n_H, n_W, n_C)
	# cache : the forward cache
	m, n_H, n_W, n_C = dY.shape
	mu, s, X, X_norm, gamma, beta, epsilon = cache

	# dof : degree of freedom
	dof = m*n_H*n_W
	dof_s = (m-1)*n_H*n_W

	# calculate the derivation
	dX_norm = dY*gamma
	dbeta = dY.sum(axis=(0,1,2)) / dof
	dgamma = ((dY * X_norm).sum(axis=(0,1,2)) / dof).astype(np.float32)
	ds = -(dX_norm*(X-mu)*np.power(s+epsilon, -1.5)).astype(np.float32).sum(axis=(0,1,2)) / (2*dof_s)
	dmu = -(dX_norm*np.power(s+epsilon, -0.5).astype(np.float32) + 2*ds*(X-mu) / dof_s).sum(axis=(0,1,2)) / dof
	dX = dX_norm*np.power(s+epsilon, -0.5).astype(np.float32) + 2*ds*(X-mu) / dof_s + dmu / dof

	return dX, dgamma, dbeta

def bn_nn_f(X, gamma, beta, epsilon=1e-6):
	# X : (m, n)
	m, n = X.shape
	mu = X.sum(axis=0) / m
	s = np.power(X-mu, 2).sum(axis=0) / (m-1)
	X_norm = (X-mu) / np.sqrt(s+epsilon)
	Z = X_norm*gamma + beta

	cache = (mu, s, X, X_norm, gamma, beta, epsilon)
	return Z, cache

def bn_nn_b(dY, cache):

	#    dY : (m, n)
	# cache : (mu, s, X, X_norm, gamma, beta, epsilon)
	mu, s, X, X_norm, gamma, beta, epsilon = cache
	m, n = dY.shape

	# calculate the derivation
	dX_norm = dY*gamma
	dbeta = dY.sum(axis=(0)) / m
	dgamma = (dY*X_norm).sum(axis=(0)) / m
	ds = -(dX_norm*(X-mu)*np.power(s+epsilon, -1.5)).astype(np.float32).sum(axis=(0)) / (2*(m-1))
	dmu = -(dX_norm*np.power(s+epsilon, -0.5).astype(np.float32) + 2*ds*(X-mu)/(m-1)).sum(axis=(0)) / m
	dX = dX_norm*np.power(s+epsilon, -0.5).astype(np.float32) + 2*ds*(X-mu)/(m-1) + dmu/m

	return dX, dgamma, dbeta

def adam(params, grads, hps, t, lr, beta1=0.9, beta2=0.999, epsilon=1e-8):
	# combained momentun and RMSprob together

	# params : W,  gamma,  beta
	#    hps : Vdw, Vdgamma, Vdbeta, Sdw, Sdgamma, Sdbeta.
	W, gamma, beta = params
	keys = ['VdW', 'Vdgamma', 'Vdbeta', 'SdW', 'Sdgamma', 'Sdbeta']
	VdW, Vdgamma, Vdbeta, SdW, Sdgamma, Sdbeta = [hps[key] for key in keys]

	for l in W.keys():
		corr = 1-np.power(beta1, t)
		# update W
		VdW[l] = beta1*VdW[l] + (1-beta1)*grads[l]['dW']
		SdW[l] = beta2*SdW[l] + (1-beta2)*(grads[l]['dW']**2).astype(np.float32)
		grads[l]['dW'] -= lr*(VdW[l] / corr) / (np.sqrt(SdW[l] / corr)+epsilon)

		# update gamma
		Vdgamma[l] = beta1*Vdgamma[l] + (1-beta1)*grads[l]['dgamma']
		Sdgamma[l] = beta2*Sdgamma[l] + (1-beta2)*(grads[l]['dgamma']**2).astype(np.float32)
		grads[l]['dgamma'] -= lr*(Vdgamma[l] / corr) / (np.sqrt(Sdgamma[l] / corr)+epsilon)

		# update beta
		Vdbeta[l] = beta1*Vdbeta[l] + (1-beta1)*grads[l]['dbeta']
		Sdbeta[l] = beta2*Sdbeta[l] + (1-beta2)*(grads[l]['dbeta']**2).astype(np.float32)
		grads[l]['dbeta'] -= lr*(Vdbeta[l] / corr) / (np.sqrt(Sdbeta[l] / corr)+epsilon)


Concept
==============

$$
\newcommand{\sumN}{\sum_{n = 1}^N}
\newcommand{\sumn}{\sum_n}
\newcommand{\bx}{\mathbf{x}}
\newcommand{\bbeta}{\boldsymbol{\beta}}
\newcommand{\btheta}{\boldsymbol{\theta}}
\newcommand{\bbetahat}{\boldsymbol{\hat{\beta}}}
\newcommand{\bthetahat}{\boldsymbol{\hat{\theta}}}
\newcommand{\dadb}[2]{\frac{\partial #1}{\partial #2}}
\newcommand{\by}{\mathbf{y}}
\newcommand{\bX}{\mathbf{X}}
\newcommand{\bphi}{\boldsymbol{\phi}}
\newcommand{\bPhi}{\boldsymbol{\Phi}}
$$

## Model Structure

*Linear regression* is a relatively simple method that is extremely widely-used. It is also a great stepping stone for more sophisticated methods, making it a natural algorithm to study first.

In linear regression, the target variable $y$ is assumed to follow a linear function of one or more predictor variables, $x_1, \dots, x_D$, plus some random error. Specifically, we assume the model for the $n^\text{th}$ observation in our sample is of the form

$$
y_n = \beta_0 + \beta_1 x_{n1} + \dots + \beta_Dx_{nD} + \epsilon_n.
$$

Here $\beta_0$ is the intercept term, $\beta_1$ through $\beta_D$ are the coefficients on our feature variables, and $\epsilon$ is an error term that represents the difference between the true $y$ value and the linear function of the predictors. Note that the terms with an $n$ in the subscript differ between observations while the terms without (namely the $\beta\text{s}$) do not.

The math behind linear regression often becomes easier when we use vectors to represent our predictors and coefficients. Let's define $\bx_n$ and $\bbeta$ as follows:

$$
\begin{align}
\bx_n &= \begin{pmatrix} 1 & x_{n1} & \dots & x_{nD} \end{pmatrix}^\top \\
\bbeta &= \begin{pmatrix} \beta_0 & \beta_1 & \dots & \beta_D \end{pmatrix}^\top.
\end{align}
$$

Note that $\bx_n$ includes a leading 1, corresponding to the intercept term $\beta_0$. Using these definitions, we can equivalently express $y_n$ as

$$
y_n = \bbeta^\top \bx_n + \epsilon_n.
$$

Below is an example of a dataset designed for linear regression. The input variable is generated randomly and the target variable is generated as a linear combination of that input variable plus an error term.

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# generate data
np.random.seed(123)
N = 20
beta0 = -4
beta1 = 2
x = np.random.randn(N)
e = np.random.randn(N)
y = beta0 + beta1*x + e
true_x = np.linspace(min(x), max(x), 100)
true_y = beta0 + beta1*true_x

# plot
fig, ax = plt.subplots()
sns.scatterplot(x, y, s = 40, label = 'Data')
sns.lineplot(true_x, true_y, color = 'red', label = 'True Model')
ax.set_xlabel('x', fontsize = 14)
ax.set_title(fr"$y = {beta0} + ${beta1}$x + \epsilon$", fontsize = 16)
ax.set_ylabel('y', fontsize=14, rotation=0, labelpad=10)
ax.legend(loc = 4)
sns.despine()

## Parameter Estimation

The previous section covers the entire structure we assume our data follows in linear regression. The machine learning task is then to estimate the parameters in $\bbeta$. These estimates are represented by $\hat{\beta}_0, \dots, \hat{\beta}_D$ or $\bbetahat$. The estimates give us *fitted values* for our target variable, represented by $\hat{y}_n$.

This task can be accomplished in two ways which, though slightly different conceptually, are identical mathematically. The first approach is through the lens of *minimizing loss*. A common practice in machine learning is to choose a loss function that defines how well a model with a given set of parameter estimates the observed data. The most common loss function for linear regression is squared error loss. This says the *loss* of our model is proportional to the sum of squared differences between the true $y_n$ values and the fitted values, $\hat{y}_n$. We then *fit* the model by finding the estimates $\bbetahat$ that minimize this loss function. This approach is covered in the subsection {doc}`s1/loss_minimization`.

The second approach is through the lens of *maximizing likelihood*. Another common practice in machine learning is to model the target as a random variable whose distribution depends on one or more parameters, and then find the parameters that maximize its likelihood. Under this approach, we will represent the target with $Y_n$ since we are treating it as a random variable. The most common model for $Y_n$ in linear regression is a Normal random variable with mean $E(Y_n) = \bbeta^\top \bx_n$. That is, we assume

$$
Y_n|\bx_n \sim \mathcal{N}(\bbeta^\top \bx_n, \sigma^2),
$$

and we find the values of $\bbetahat$ to maximize the likelihood. This approach is covered in subsection {doc}`s1/likelihood_maximization`.

Once we've estimated $\bbeta$, our model is *fit* and we can make predictions. The below graph is the same as the one above but includes our estimated line-of-best-fit, obtained by calculating $\hat{\beta}_0$ and $\hat{\beta}_1$.

# generate data
np.random.seed(123)
N = 20
beta0 = -4
beta1 = 2
x = np.random.randn(N)
e = np.random.randn(N)
y = beta0 + beta1*x + e
true_x = np.linspace(min(x), max(x), 100)
true_y = beta0 + beta1*true_x

# estimate model
beta1_hat = sum((x - np.mean(x))*(y - np.mean(y)))/sum((x - np.mean(x))**2)
beta0_hat = np.mean(y) - beta1_hat*np.mean(x)
fit_y = beta0_hat + beta1_hat*true_x

# plot
fig, ax = plt.subplots()
sns.scatterplot(x, y, s = 40, label = 'Data')
sns.lineplot(true_x, true_y, color = 'red', label = 'True Model')
sns.lineplot(true_x, fit_y, color = 'purple', label = 'Estimated Model')
ax.set_xlabel('x', fontsize = 14)
ax.set_title(fr"Linear Regression for $y = {beta0} + ${beta1}$x + \epsilon$", fontsize = 16)
ax.set_ylabel('y', fontsize=14, rotation=0, labelpad=10)
ax.legend(loc = 4)
sns.despine()

## Extensions of Ordinary Linear Regression

There are many important extensions to linear regression which make the model more flexible. Those include {doc}`/content/c2/s1/regularized`—which balances the bias-variance tradeoff for high-dimensional regression models—{doc}`/content/c2/s1/bayesian`—which allows for prior distributions on the coefficients—and {doc}`/content/c2/s1/GLMs`—which introduce non-linearity to regression models. These extensions are discussed in the next chapter.

```{toctree}
:hidden:
:titlesonly:

s1/loss_minimization
s1/likelihood_maximization
```


"""
Minimal Polynomials of Linear Recurrence Sequences

AUTHORS:

- William Stein
"""
# ****************************************************************************
#       Copyright (C) 2005 William Stein <wstein@gmail.com>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  https://www.gnu.org/licenses/
# ****************************************************************************

import sage.rings.rational_field

def berlekamp_massey(a):
    r"""
    Use the Berlekamp-Massey algorithm to find the minimal polynomial
    of a linear recurrence sequence `a`.

    The minimal polynomial of a linear recurrence `\{a_r\}` is
    by definition the unique monic polynomial `g`, such that if
    `\{a_r\}` satisfies a linear recurrence
    `a_{j+k} + b_{j-1} a_{j-1+k} + \cdots + b_0 a_k=0`
    (for all `k\geq 0`), then `g` divides the
    polynomial `x^j + \sum_{i=0}^{j-1} b_i x^i`.

    INPUT:

    -  ``a`` -- a list of even length of elements of a field (or domain)

    OUTPUT:

    the minimal polynomial of the sequence, as a polynomial over the
    field in which the entries of `a` live

    .. WARNING::

         The result is only guaranteed to be correct on the full
         sequence if there exists a linear recurrence of length less
         than half the length of `a`.

    EXAMPLES::

        sage: from sage.matrix.berlekamp_massey import berlekamp_massey
        sage: berlekamp_massey([1,2,1,2,1,2])
        x^2 - 1
        sage: berlekamp_massey([GF(7)(1),19,1,19])
        x^2 + 6
        sage: berlekamp_massey([2,2,1,2,1,191,393,132])
        x^4 - 36727/11711*x^3 + 34213/5019*x^2 + 7024942/35133*x - 335813/1673
        sage: berlekamp_massey(prime_range(2,38))
        x^6 - 14/9*x^5 - 7/9*x^4 + 157/54*x^3 - 25/27*x^2 - 73/18*x + 37/9

    TESTS::

        sage: berlekamp_massey("banana")
        Traceback (most recent call last):
        ...
        TypeError: argument must be a list or tuple
        sage: berlekamp_massey([1,2,5])
        Traceback (most recent call last):
        ...
        ValueError: argument must have an even number of terms
    """
    if not isinstance(a, (list, tuple)):
        raise TypeError("argument must be a list or tuple")
    if len(a) % 2:
        raise ValueError("argument must have an even number of terms")

    M = len(a) // 2

    try:
        K = a[0].parent().fraction_field()
    except AttributeError:
        K = sage.rings.rational_field.RationalField()
    R = K['x']
    x = R.gen()

    f = {-1: R(a), 0: x**(2 * M)}
    s = {-1: 1, 0: 0}
    j = 0
    while f[j].degree() >= M:
        j += 1
        qj, f[j] = f[j - 2].quo_rem(f[j - 1])
        s[j] = s[j - 2] - qj * s[j - 1]
    t = s[j].reverse()
    return ~(t[t.degree()]) * t  # make monic  (~ is inverse in python)


# pylint: disable=invalid-name

"""
Module for obtaining quadrature points and weights.

Three types of gaussian quadrature are supported: normal Gaussian,
Radau quadrature, and Lobatto quadrature.  The first method does not
include either endpoint of integration, Radau quadrature includes one
endpoint of the integration range, and Lobatto quadrature includes both
endpoints.::

    import iad.quadrature

    n=8
    x, w = iad.quadrature.gauss(n)
    print(" i        x         weight")
    for i,x in enumerate(xi):
        print("%2d   %+.12f %+.12f" % (i, x[i], w[i]))

    x, w = iad.quadrature.radau(n)
    print(" i        x         weight")
    for i,x in enumerate(xi):
        print("%2d   %+.12f %+.12f" % (i, x[i], w[i]))

    x, w = iad.quadrature.lobatto(n)
    print(" i        x         weight")
    for i,x in enumerate(xi):
        print("%2d   %+.12f %+.12f" % (i, x[i], w[i]))

"""

import functools
import scipy.special
import scipy.optimize
import numpy as np

__all__ = ('gauss',
           'radau',
           'lobatto',
           )

def _gauss_func(n, x):
    """Zeroes of this function are the Gaussian quadrature points."""
    return scipy.special.legendre(n)(x)

def _radau_func(n, x):
    """Zeros of this function are the Radau quadrature points."""
    return (scipy.special.eval_legendre(n - 1, x) + scipy.special.eval_legendre(n, x)) / (1 + x)

def _lobatto_func(n, x):
    """Zeros of this function are the Lobatto quadrature points."""
    return scipy.special.legendre(n - 1).deriv(1)(x)

def gauss(n, a=-1, b=1):
    """
    Return abscissas and weights for Gaussian quadrature.

    The definite integral ranges from a to b.  The default
    interval is -1 to 1.  The quadrature approximation is
    just the sum of w_i f(x_i).  Neither a nor b is included
    in the list of quadrature abscissas.

    The result should be exact when integrating any polynomial
    of degree 2n-1 or less.

    If -a=b, then abscissas will be symmetric about the origin

    Args:
        n: number of quadrature points
        a: lower limit of integral
        b: upper limit of integral
    Returns:
        x: array of abscissas of length n
        w: array of weights of length n
    """
    x, w, _ = scipy.special.roots_legendre(n, mu=True)

    # scale for desired interval
    x *= 0.5 * (a - b)
    x += 0.5 * (a + b)
    w *= 0.5 * (b - a)

    return np.flip(x), np.flip(w)

def radau(n, a=-1, b=1):
    """
    Return abscissas and weights for Radau quadrature.

    The definite integral ranges from a to b.  The default
    interval is -1 to 1.  The quadrature approximation is
    just the sum of w_i f(x_i).  The upper endpoint b is include
    in the list of quadrature abscissas.

    The result should be exact when integrating any polynomial
    of degree 2n-2 or less.

    Args:
        n: number of quadrature points
        a: lower limit of integral
        b: upper limit of integral
    Returns:
        x: array of abscissas of length n
        w: array of weights of length n
    """
    x = np.zeros(n)
    w = np.zeros(n)
    x[0] = -1

    # the roots of P_{n} bracket the roots of P_{n}':
    brackets = scipy.special.roots_legendre(n)[0]

    f = functools.partial(_radau_func, n)
    for i in range(n - 1):
        x[i + 1] = scipy.optimize.brentq(f, brackets[i], brackets[i + 1])

    pp = scipy.special.legendre(n - 1).deriv(1)
    w[0] = 2 / n**2
    w[1:] = 1 / pp(x[1:])**2 / (1 - x[1:])

    # scale for desired interval
    x *= 0.5 * (a - b)
    x += 0.5 * (b + a)
    w *= 0.5 * (b - a)

    return np.flip(x), np.flip(w)

def lobatto(n, a=-1, b=1):
    """
    Return abscissas and weights for Lobatto quadrature.

    The definite integral ranges from a to b.  The default
    interval is -1 to 1.  The quadrature approximation is
    just the sum of w_i f(x_i).  Both endpoints a and b are include
    in the list of quadrature abscissas.

    The result should be exact when integrating any polynomial
    of degree 2n-3 or less.

    If -a=b, then abscissas will be symmetric about the origin

    Args:
        n: number of quadrature points
        a: lower limit of integral
        b: upper limit of integral
    Returns:
        x: array of abscissas of length n
        w: array of weights of length n
    """
    x = np.zeros(n)
    w = np.full(n, 2 / n / (n - 1))
    x[0] = -1
    x[-1] = 1

    # The roots of P_{n-1} bracket the roots of P_{n-1}':
    brackets = scipy.special.roots_legendre(n - 1)[0]

    f = functools.partial(_lobatto_func, n)
    for i in range(n - 2):
        x[i + 1] = scipy.optimize.brentq(f, brackets[i], brackets[i + 1])

    pp = scipy.special.legendre(n - 1)(x)
    w[1:-1] = w[1:-1] / pp[1:-1]**2

    # scale for desired interval
    x *= 0.5 * (a - b)
    x += 0.5 * (a + b)
    w *= 0.5 * (b - a)

    return np.flip(x), np.flip(w)


#!/usr/bin/env
# -*- coding: utf-8 -*-
# Módulo Jacobi: Métodos de Jacobi para resolver Sistemas de Equações

from pprint import pprint

import numpy as np

import construtor

# Calcula o vetor solução pelo método de Jacobi.
# Entradas: matriz, vetor de termos independentes, número máximo de iterações, tolerância
# Retorno: vetor solução
def v_sol(m, v, n_max, tol, prin):
    # Cria o chute inicial
    x = np.zeros(len(m[0]))

    # Cria o vetor d com os elementos da diagonal principal de m
    d = np.diag(m)
    # Cria o vetor r com os elementos de m e diagonal principal nula
    r = m - np.diagflat(d)

    # Executa a iteração
    i = 0
    while i < n_max:
        i += 1
        x_old = x.copy()
        x = (v - np.dot(r, x)) / d
        # Verifica a tolerância
        # O método np.allclose comapara todos os elementos dos dois vetores e retorna True caso a norma < tol.
        if np.allclose(x, x_old, atol=tol, rtol=tol):
            break
    # Retorna o vetor solução
    if prin:
        print('\nForam feitas ', i, ' iterações.')
    return x.tolist()

# Calcula o vetor solução, para a matriz de uma equação, pelo método de Jacobi.
# Entradas: q(x), r(x), malha de pontos, passo, número de pontos, y(a), y(b), número máximo de iterações, tolerância
# Retorno: vetor solução
def v_sol_mh(q, r, x, h, n, a_, b_, n_max, tol, prin=True):
    # Calcula a matriz e o vetor de termos independentes
    m_h = construtor.matriz(q, x, h, n)
    v_h = construtor.vetor(r, x, h, n, a_, b_)

    # Calcula e retorna o vetor solução
    return v_sol(m_h, v_h, n_max, tol, prin)

if __name__ == "__main__":
    # TESTE 1
    m = np.array([[2.0, 1.0], [5.0, 7.0]])
    v = np.array([11.0, 13.0])
    guess = np.array([1.0, 1.0])

    sol = v_sol(m, v, 25, 0.0001)

    print('m:')
    pprint(m)

    print('v:')
    pprint(v)

    print('x:')
    pprint(sol)


import numpy as np

class OrdinaryLeastSquares:
    """Ordinary Least Squares Linear Regression.

    Parameters:
    -----------
    None.
    """
    def __init__(self):
        self.weights = None

    def fit(self, X, y):
        if len(y.shape) == 1:
            y = y.reshape(-1, 1)
        self._fit(X, y)

    @staticmethod
    def _pad(X):
        return np.pad(X, ((0, 0), (1, 0)), 'constant', constant_values=1)

    def _fit(self, X, y, pad=True):
        if pad:
            X = self._pad(X)
        X_pseudo_inverse = np.linalg.pinv(X)  # Moore-Penrose pseudo-inverse for more efficiency
        self.weights = np.dot(X_pseudo_inverse, y)

    def predict(self, X):
        if self.weights is None:
            raise ValueError('no weights, run fit(X, y) to fit model before prediction')
        return self._predict(X)

    def _predict(self, X, pad=True):
        if pad:
            y_pred = np.dot(self._pad(X), self.weights)
        else:
            y_pred = np.dot(X, self.weights)
        return y_pred

    def score(self, X, y):
        if len(y.shape) == 1:
            y = y.reshape(-1, 1)
        y_pred = self.predict(X)
        mse = np.square(y_pred - y).mean()
        return mse

class OrdinaryLeastSquaresGD(OrdinaryLeastSquares):
    """Ordinary Least Squares Linear Regression with Gradient Descent.

    Parameters:
    -----------
    learning_rate: float
        The rate at which to update learned weights. Also known as alpha.
    n_iter: int
        Number of iterations to run in order to learn weights.
    early_stop: boolean
        Allows to stop the learning process if loss has plateaued.
    """
    def __init__(self, learning_rate=0.05, n_iter=1000, early_stop=True):
        super().__init__()
        self.learning_rate = learning_rate
        self.n_iter = n_iter
        self.early_stop = early_stop
        self.loss = []

    def _fit(self, X, y, pad=True):
        if pad:
            X = self._pad(X)
        self.weights = np.zeros((X.shape[1], 1))
        for iter in range(self.n_iter):
            # print('iter %d/%d' % (iter, self.n_iter))
            self._step(X, y)
            curr_loss = self.score(X[:, 1:], y)
            self.loss.append(curr_loss)
            if self.early_stop and (len(self.loss) >= 3) and (self.loss[-2] < self.loss[-1]):
                # print('early stop!')
                break

    def _step(self, X, y):
        n = X.shape[0]
        y_est = np.dot(X, self.weights)
        error = y_est - y
        grad = np.dot(X.T, error) / n
        self.weights = self.weights - (self.learning_rate * grad)


from sympy import symbols, sin, cos
from shenfun import *

# Use sympy to compute manufactured solution
x, y, z = symbols("x,y,z")
ue = (cos(4*x) + sin(2*y) + sin(4*z))*(1-x**2)
fe = ue.diff(x, 2) + ue.diff(y, 2) + ue.diff(z, 2)

C0 = FunctionSpace(32, 'Chebyshev', bc=(0, 0))
F1 = FunctionSpace(32, 'Fourier', dtype='D')
F2 = FunctionSpace(32, 'Fourier', dtype='d')
T = TensorProductSpace(comm, (C0, F1, F2))
u = TrialFunction(T)
v = TestFunction(T)

# Assemble left and right hand
f_hat = inner(v, Array(T, buffer=fe))
A = inner(v, div(grad(u)))

# Solve
solver = chebyshev.la.Helmholtz(*A)
u_hat = Function(T)
u_hat = solver(f_hat, u_hat)
uj = u_hat.backward()
assert np.linalg.norm(u_hat.backward()-Array(T, buffer=ue)) < 1e-12
print(u_hat.shape)


import numpy as np

def perceptron(data,b=0.1,a=1.0,K=200):
    (N,L)=data.shape;
    D=L-1;
    labs=np.unique(data[:,L-1]);
    C=labs.size;
    w = np.zeros((L,C));
    for k in range(1,K+1):
        E=0;
        for n in range(N):
             xn=np.concatenate(([1],data[n,:D]));
             cn=np.where(labs==data[n,L-1])[0][0];
             er=0;
             g=np.dot(w[:,cn],xn);
             for c in range(C):
                 if c != cn and np.dot(w[:,c],xn) + b > g:
                     w[:,c] = w[:,c] - a*xn;
                     er=1;
             if er==1:
                w[:,cn] = w[:,cn] + a*xn;
                E=E+1;
        if E==0:
            break;
    return w,E,k;


import numpy as np

def get_primes(n):
    arr = np.ones(n + 1, dtype=np.int8)
    arr[0], arr[1] = 0, 0
    for i in range(2, n + 1):
        if arr[i]:
            arr[i * i:n + 1:i] = 0
    return arr.nonzero()[0]

def count_of_occurrences(n, p):
    """
    Return Number of times 'p' occurs in  n!

    for eg: f(30 , 7) = 4
    """
    cnt = 0
    while n:
        n //= p
        cnt += n
    return cnt

def p231(n, r):
    primes = get_primes(n)

    res = 0

    for p in primes:
        res += p*(count_of_occurrences(n, p) -
                  count_of_occurrences(r, p) - count_of_occurrences(n-r, p))

    return res

if __name__ == "__main__":
    print(p231(20000000,5000000))


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#  ch2Python.py. Finding the best fitting line.
#  This is demonstration code, so it is transparent but inefficient.

import numpy as np
import matplotlib.pyplot as plt

x = [1.00,1.25,1.50,1.75,2.00,2.25,2.50,2.75,3.00,3.25,3.50,3.75,4.00]
y = [3.34,4.97,4.15,5.40,5.21,4.56,3.69,5.86,4.58,6.94,5.57,5.62,6.87]

# convert data to vectors.
x = np.array(x)
y = np.array(y)
xmean = np.mean(x)
ymean = np.mean(y)
# Find zero mean versions of x and y.
xzm = x - xmean
yzm = y - ymean

numerator = np.sum(xzm * yzm)
denominator = np.sum(xzm * xzm)
print("numerator = %0.3f" %numerator) # 8.693
print("denominator = %0.3f" %denominator) # 11.375

# Find slope b1.
# b1 = numerator/denominator # 0.764
b1 = np.sum(xzm * yzm) / ( np.sum(xzm * xzm) ) # 0.764

# Find intercept b0.
b0 = ymean - b1*xmean # 3.225
print('slope b1 = %6.3f\nintercept b0 = %6.3f' % (b1, b0))

# Draw best fitting line.
xx = np.array([-0.3, 3, 5])
yline = b1 * xx + b0
fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(x, y, "ko", label="Data")
ax.plot(xx, yline, "k-",label="Best fitting line (LSE)")
ax.legend(loc="best")
plt.grid()

###############################
# END OF FILE.
###############################


import numpy as np
def sigmoid(x):  return 1 / (1 + np.exp(-x))
def dsigmoid(y): return y * (1.0 - y)
def tanh(x):     return np.tanh(x)
def dtanh(y):    return 1 - y*y


from sympy import *
import numpy as np

libceed = True

def div(v):
    d  = diff(v[0],x)
    d += diff(v[1],y)
    d += diff(v[2],z)
    return d

def curl(v):
    c = []
    c.append(  diff(v[2],y)-diff(v[1],z) )
    c.append(-(diff(v[2],x)-diff(v[0],z)))
    c.append(  diff(v[1],x)-diff(v[0],y) )
    return np.asarray(c)

x,y,z = symbols('x y z')
a0,b0,c0,d0,a1,b1,c1,d1,a2,b2,c2,d2 = symbols("a0 b0 c0 d0 a1 b1 c1 d1 a2 b2 c2 d2")
r0,r1,r2,r3,s0,s1,s2,s3,t0,t1,t2,t3 = symbols("r0 r1 r2 r3 s0 s1 s2 s3 t0 t1 t2 t3")

BDDF1 = np.asarray([a0+b0*x+c0*y+d0*z,
                    a1+b1*x+c1*y+d1*z,
                    a2+b2*x+c2*y+d2*z])
BDDF1 += r0*curl([0       ,0       ,x*y*z   ])
BDDF1 += r1*curl([0       ,0       ,x*y**2  ])
BDDF1 += r2*curl([0       ,0       ,x**2*z  ])
BDDF1 += r3*curl([0       ,0       ,x**2*y*z])
BDDF1 += s0*curl([x*y*z   ,0       ,0       ])
BDDF1 += s1*curl([y*z**2  ,0       ,0       ])
BDDF1 += s2*curl([x*y**2  ,0       ,0       ])
BDDF1 += s3*curl([x*y**2*z,0       ,0       ])
BDDF1 += t0*curl([0       ,x*y*z   ,0       ])
BDDF1 += t1*curl([0       ,x**2*z  ,0       ])
BDDF1 += t2*curl([0       ,y*z**2  ,0       ])
BDDF1 += t3*curl([0       ,x*y*z**2,0       ])

"""
     local numbering of Hex

     5--------7           z
   / |      / |            |
 6 --|---- 8  |            |
 |   |     |  |            /----y
 |   |     |  |           /
 |   1 -------3         x
 | /       | /
 2 --------4
"""
# normals
nl = [0.0,-1.0, 0.0]
nr = [0.0, 1.0, 0.0]
nbt = [0.0, 0.0, -1.0]
nt = [0.0, 0.0, 1.0]
nf = [1.0, 0.0, 0.0]
nbk = [-1.0, 0.0, 0.0]
# nodes
n1 = [-1., -1., -1.]
n2 = [1., -1., -1]
n4 = [1., 1., -1.]
n3 = [-1., 1., -1.]
n5 = [-1., -1., 1.]
n6 = [1., -1., 1]
n8 = [1., 1., 1.]
n7 = [-1., 1., 1.]

nodes = [n1, n2, n3, n4, n5, n6, n7, n8, n2, n1, n6, n5, n4, n3, n8, n7, n2, n4, n6, n8, n1, n3, n5, n7]
normals = [nbt, nbt, nbt, nbt, nt, nt, nt, nt, nl, nl, nl, nl, nr, nr, nr, nr, nf, nf, nf, nf, nbk, nbk, nbk, nbk]

for i in range(8): # for each vertex
    for j in range(3): # for each direction
        k = 3*i+j

        eqs = []
        for n in range(24):
            eqs.append(np.dot(BDDF1,normals[n]).subs({x:nodes[n][0],y:nodes[n][1],z:nodes[n][2]}))

        eqs[k] -= 1 # the k^th functions should be a 1, rest are 0
        sol = solve(eqs)
        ux = BDDF1[0].subs(sol)
        uy = BDDF1[1].subs(sol)
        uz = BDDF1[2].subs(sol)

        if libceed:
            def _f(fcn):
                fcn = fcn.replace("x**2","x*x")
                fcn = fcn.replace("y**2","y*y")
                fcn = fcn.replace("z**2","z*z")
                fcn = fcn.replace("x","x[0]")
                fcn = fcn.replace("y","x[1]")
                fcn = fcn.replace("z","x[2]")
                if "/8" in fcn: fcn = "(%s)*0.125;" % (fcn.replace("/8",""))
                if "/16" in fcn: fcn = "(%s)*0.0625;" % (fcn.replace("/16",""))
                return fcn
            print("Bx[%2d] = " % (k) + _f("%s" % (ux)),";")
            print("By[%2d] = " % (k) + _f("%s" % (uy)),";")
            print("Bz[%2d] = " % (k) + _f("%s" % (uz)),";")


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

def RK2(func, X0,tmin,tmax,h,cons):
    N=int((tmax-tmin)/h)
    t = np.linspace(tmin,tmax,N)
    X  = np.zeros([N, len(X0)])
    X[0] = X0
    for i in range(N-1):
        k1 =h* func(t[i],X[i],cons)
        k2 = h*func( t[i] + h,X[i] +  k1,cons)
        X[i+1] = X[i] +  (k1 +k2 )/2
    return X,t

def f2(t,X,cons):
    k,m,b=cons
    x,z=X
    dx_dt=z
    dx2_dt2=-b*z/m-k*x/m
    return np.array([dx_dt, dx2_dt2])

def f3(t,X,cons):
    g,L=cons
    th,z=X
    dth_dt=z
    dth2_dt2=-g*th/L
    return np.array([dth_dt, dth2_dt2])

def f4(t,X,cons):
    k,m,g,L=cons
    w0=np.sqrt(g/L)
    xa,z1,xb,z2=X
    xb=np.radians(xb)
    xa=np.radians(xa)
    e1=z1
    e2=-(w0**2)*xa-(k/m)*(xa-xb)
    e3=z2
    e4=-(w0**2)*xb+(k/m)*(xa-xb)
    return np.array([e1,e2,e3,e4])

#3a
IC=[2,0];tmin=0;m=0.5;k=4;h=0.01;b=0#SI units
tp=2*np.pi*np.sqrt(m/k)
tmax=5*tp
cons=(k,m,b)
S=RK2(f2,IC,tmin,tmax,h,cons)
x,v= S[0].T
t=S[1]
t1=t/tp
fig, axs = plt.subplots(2)
fig.suptitle('SIMPLE HARMONIC OSCILLATOR',c="r")
axs[0].plot(t1,x,marker="*",c="orange")
axs[0].set(xlabel="time/time period ",ylabel="Dispacement")
axs[0].grid()
axs[1].plot(t1,v,marker="*",c="cyan")
axs[1].set(xlabel=" Time/Time Period ",ylabel="Velocity")
axs[1].grid()
plt.show()
print("----------------------------- SIMPLE HARMONIC OSCILLATOR---------------------------------")
d={"No. of time periods":t1,"Displacement":x,"Velocity":v}
print(pd.DataFrame(d))

#3b https://scipy-lectures.org/intro/scipy/auto_examples/plot_odeint_damped_spring_mass.html (ref)
IC=[1,0];tmin=0;h=0.1;m=0.5;k=4 #SI units
tmax=100
dis=[];vel=[];time=[]
ba=[0.2,np.sqrt(4*m*k),3]
for b in ba:
     cons=(m,k,b)
     S=RK2(f2,IC,tmin,tmax,h,cons)
     x,v= S[0].T
     t=S[1]
     dis.append(x)
     vel.append(v)
fig, ax = plt.subplots(3,2)
fig.suptitle('DAMPED HARMONIC OSCILLATOR',c="r")
ax[0,0].plot(t,dis[0],marker="*",c="r")
ax[0,0].set(xlabel="time ",ylabel="Dispacement",title="Underdamped")
ax[0,0].grid()
ax[0,1].plot(t,vel[0],marker="*",c="green")
ax[0,1].set(xlabel="time ",ylabel="Velocity",title="Underdamped")
ax[0,1].grid()
ax[1,0].plot(t,dis[1],marker="*",c="purple")
ax[1,0].set(xlabel="time ",ylabel="Dispacement",title="Critically Damped")
ax[1,0].grid()
ax[1,1].plot(t,vel[1],marker="*",c="darkblue")
ax[1,1].set(xlabel="time ",ylabel="Velocity",title="Critically Damped")
ax[1,1].grid()
ax[2,0].plot(t,dis[2],marker="*",c="brown")
ax[2,0].set(xlabel="time ",ylabel="Dispacement",title="Overdamped")
ax[2,0].grid()
ax[2,1].plot(t,vel[2],marker="*",c="violet")
ax[2,1].set(xlabel="time ",ylabel="Velocity",title="Overdamped")
ax[2,1].grid()
plt.show()
print("---------------------- DAMPED HARMONIC OSCILLATOR  (UNDERDAMPED)-----------------------")
d={"Time":t,"Displacement":dis[0],"Velocity":vel[0]}
print(pd.DataFrame(d))
print("---------------------- DAMPED HARMONIC OSCILLATOR  (CRITICALLY DAMPED)-----------------------")
d={"Time":t,"Displacement":dis[1],"Velocity":vel[1]}
print(pd.DataFrame(d))
print("---------------------- DAMPED HARMONIC OSCILLATOR  (OVERDAMPED)-----------------------")
d={"Time":t,"Displacement":dis[2],"Velocity":vel[2]}
print(pd.DataFrame(d))

#3c
IC=[1,0];tmin=0;g=9.8;L=2  #SI units
cons1=(g,L)
tp=2*np.pi*np.sqrt(L/g)
tmax=10*tp
h=tp/50
S=RK2(f3,IC,tmin,tmax,h,cons1)
x,v= S[0].T
t=S[1]
t1=t/tp
fig, axs = plt.subplots(2)
fig.suptitle('SIMPLE PENDULUM',c="r")
axs[0].plot(t1,x,marker="*",c="r")
axs[0].set(xlabel="Time/Time Period ",ylabel="Angular Dispacement")
axs[0].grid()
axs[1].plot(t1,v,marker="*",c="green")
axs[1].set(xlabel=" Time/Time Period ",ylabel="Angular Velocity")
axs[1].grid()
plt.show()
print("------------------------------ SIMPLE PENDULUM------------------------------")
d={"No. of time periods":t,"Angular Displacement":x,"Angular Velocity":v}
print(pd.DataFrame(d))

#3d
IC=[10,0,-10,0];tmax=80;tmin=0;h=0.01;m1=10;k1=90;g1=9.8;l1=10  #SI units
cons1=(k1,m1,g1,l1)
S=RK2(f4,IC,tmin,tmax,h,cons1)
x1,v1,x2,v2= S[0].T
t3=S[1]
fig, axs = plt.subplots(2, 2)
fig.suptitle('COUPLED SYSTEM',c="r")
axs[0, 0].plot(t3,x1,c="r")
axs[0,0].set(xlabel=" Time",ylabel="Displacement",title="Mass A")
axs[0,0].grid()
axs[0, 1].plot(t3,v1,c="green")
axs[0,1].set(xlabel=" Time",ylabel="Velocity",title="Mass A")
axs[0,1].grid()
axs[1, 0].plot(t3,x2,c="magenta")
axs[1,0].set(xlabel=" Time",ylabel="Displacement",title="Mass B")
axs[1,0].grid()
axs[1, 1].plot(t3,v2,c="green")
axs[1,1].set(xlabel=" Time",ylabel="Velocity",title="Mass B")
axs[1,1].grid()
plt.show()
print("---------------------- COUPLED PENDULUM-----------------------")
print("-------------------------------- MASS A ---------------------------------")
d={"Time":t3,"Displacement":x1,"Velocity":v1}
print(pd.DataFrame(d))
print("-------------------------------- MASS B ---------------------------------")
d={"Time":t3,"Displacement":x2,"Velocity":v2}
print(pd.DataFrame(d))


import numpy as np
from matplotlib import pyplot
from scipy.integrate import solve_ivp as ode45
from scipy.interpolate import CubicSpline

def seirmodel(t, y, gamma, sigma, eta, Rstar):
    n = 10**7
    dy = np.zeros(5)
    #beta(t) et e(t) donc 5eqn, page 2 equations
    dy[0] = (-y[4]*y[0]*y[2])/n                        #s
    dy[1] = (y[4]*y[0]*y[2])/n - (sigma*y[1])          #e
    dy[2] = (sigma*y[1]) - (gamma*y[2])                #x
    dy[3] = gamma*y[2]                                 #r
    dy[4] = eta*((gamma*Rstar(t))-y[4])                #beta
    return dy

"""
on utilise les interpolations cubiques car si n > 10, les intervalles
[0, 1] et [ n-1, n] seraient problématiques dû à la création d'un  polynôme de degré n-1
"""
def q5(t, R, gamma = 0.06, eta = 0.1, sigma = 0.2, n = 10**7, y0 = [10**7-100, 0, 100, 0, 4*0.06], t_span = [0,400]):
    try:
        if( len(t) != len(R)):
            raise ValueError('Les données ne sont pas valides !')
        if not t or not R:
            raise ValueError('Aucune donnée !')
        R_func = CubicSpline(t,R, bc_type='clamped')
        solution = ode45(lambda t_ode, y : seirmodel(t_ode, y, gamma, sigma, eta, R_func), t_span, y0)
        #Dessiner graphique
        pyplot.figure()
        pyplot.plot(solution.t, solution.y[0,:], label="Personnes suceptibles")
        pyplot.plot(solution.t, solution.y[1,:], label="Personnes exposées")
        pyplot.plot(solution.t, solution.y[2,:], label="Personnes infectées")
        pyplot.plot(solution.t, solution.y[3,:], label="Personnes guéries")
        pyplot.plot(t, R, 'ro', label="Données")
        pyplot.legend(loc="best")
        pyplot.title("Scénario")
        pyplot.xlabel("Temps(t)")
        pyplot.ylabel("Nombre de personnes")
        pyplot.savefig("pic_q5.png")
        pyplot.figure()

        x = np.linspace(0,400, 4000)
        pyplot.plot(x, R_func(x), label="Nombre de reproduction de la maladie (R)")
        pyplot.plot(t, R, 'ro', label="Données")

        pyplot.legend(loc='best')
        pyplot.title("Évolution du nombre de repoducction de la maladie selon un scénario donné")
        pyplot.xlabel('Temps')
        pyplot.savefig('pic_r_q5.png')
        return
    except ValueError as e:
        print('Erreur dans les données : ', e)
        return
    except Exception as e:
        print('Quelque chose s\'est mal passé : ', e)
        return

def test_q5():
    t = [ 0, 30, 60, 100, 145, 190, 200, 240, 300, 400]
    R = [ 2.5, 4.5, 2.8,2,  2.5, 3, 2, 2.5, 3, 1]
    q5(t, R)

test_q5()


import math
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, precision_score, recall_score

class Random2DGaussian(object):

    min_x = 0
    max_x = 10
    min_y = 0
    max_y = 10

    def __init__(self):
        self.mean = [(self.max_x-self.min_x) * np.random.random_sample() + self.min_x,
                        (self.max_y-self.min_y) * np.random.random_sample() + self.min_y]

        eigval_x = (np.random.random_sample() * (self.max_x - self.min_x) / 5)**2
        eigval_y = (np.random.random_sample() * (self.max_y - self.min_y) / 5)**2
        D = [[eigval_x, 0], [0, eigval_y]]

        angle = 2 * math.pi * np.random.random_sample()
        R = [[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]]

        self.cov = np.dot(np.dot(R, D), np.transpose(R))

    def get_sample(self, n):
        return np.random.multivariate_normal(self.mean, self.cov, size=n)

def binlogreg_train(X, Y_, param_niter=1000, param_delta=0.001):
    w, b = np.random.randn(2, 1), 0
    for i in range(param_niter):
        scores = np.dot(X, w) + b

        probs = np.exp(scores) / (1 + np.exp(scores))
        loss = -np.sum(np.log(probs[Y_==1]) + np.log(1 - probs[Y_==0]))

        if i % 100 == 0:
            print "Iteration {}: loss: {}".format(i, loss)

        dl_dscores = probs - Y_

        grad_w = np.dot(np.transpose(X), dl_dscores)
        grad_b = np.sum(dl_dscores)

        w += -param_delta * grad_w
        b += -param_delta * grad_b

    return w, b

def binlogreg_classify(X, w, b):
    s = np.dot(X, w) + b
    probs = np.exp(s) / (1 + np.exp(s))

    return probs

def sample_gauss_2d(C, N):

    X = []
    Y = []

    for i in range(C):
        gauss = Random2DGaussian()
        x = gauss.get_sample(N)
        X.extend(x)

        y=[[i] for j in range(N)]
        Y.extend(y)

    return np.array(X), np.array(Y)

def eval_perf_binary(Y, Y_):
    acc = accuracy_score(Y, Y_)
    prec = precision_score(Y, Y_, average='binary')
    rec = recall_score(Y, Y_, average='binary')
    return acc, prec, rec

def decision_func(X):
    scores = X[:,0] + X[:,1] - 5
    return scores>0.5

def graph_data(X, Y, Y_):
    correct_idx = (Y==Y_).T.flatten()
    wrong_idx = (Y!=Y_).T.flatten()

    plt.scatter(X[correct_idx, 0], X[correct_idx, 1], c=Y_[correct_idx], marker='o', label='correct')
    plt.scatter(X[wrong_idx, 0], X[wrong_idx, 1], c=Y_[wrong_idx], marker='s', label='wrong')
    plt.legend(loc='upper left')
    plt.show()

if __name__ == '__main__':
    X, Y_ = sample_gauss_2d(2, 100)
    w, b = binlogreg_train(X, Y_)
    probs = binlogreg_classify(X, w, b)
    Y = probs>0.5

    print eval_perf_binary(Y, Y_)

    graph_data(X, Y, Y_)


from .util import computeD1D2

from scipy.stats import norm
import numpy as np

def blackScholesCall(current: float, volatility: float, ttm: float,
                     strike: float, rf: float) -> float:
    """Function to compute the Black-Scholes-Merton price of a European Call
    Option, parameterized by the current underlying asset price, volatility,
    time to expiration, strike price, and risk-free rate.

    Arguments:
        current {float} -- Current price of the underlying asset.
        volatility {float} -- Volatility of the underlying asset price.
        ttm {float} -- Time to expiration (in years).
        strike {float} -- Strike price of the option contract.
        rf {float} -- Risk-free rate (annual).

    Returns:
        float -- Price of a European Call Option contract.
    """

    d1, d2 = computeD1D2(current, volatility, ttm, strike, rf)

    call = (current * norm.cdf(d1)) \
        - (strike * np.exp(-1 * rf * ttm) * norm.cdf(d2))

    return call


import numpy as np
import sympy

def createDFTMatrix(n):
    w = sympy.cos(-(2*sympy.pi)/n) + 1j*sympy.sin(-(2*sympy.pi/n))
    dftMatrix = [[None for _ in range(n)] for _ in range(n)]

    cnt = 0
    for i in range(n):
        for j in range(n):
            dftMatrix[i][j] = w**(cnt*j)
        cnt += 1

    return np.asarray(dftMatrix, dtype=np.complex)

if __name__ == "__main__":
    n = int(input("Antall elementer i vektoren: "))
    print("Elementene i vektoren:")
    x = []
    for _ in range(n):
        x.append(int(input()))

    x = np.asarray(x, dtype=np.complex)
    dftMatrix = createDFTMatrix(n)

    dft = np.asarray(1/sympy.sqrt(n)*dftMatrix.dot(x), dtype=np.complex)

    print(f"DFT(x) = {dft}")


import numpy as np
import pandas as pd

def sound_speed(T, gamma=1.4, Rg=287):
    a = np.sqrt(T * gamma * Rg)
    return a

def gas_ideal_density(T, P, Rg=287):
    rho = P / (Rg * T)
    return rho

def total_entalphy(h, U):
    ho = h + U**2 / 2
    return ho

def total_temperature(Ma, T, gamma):
    To = (1 + (gamma - 1) * Ma**2 / 2) * T
    return To

def total_pressure(Ma, P, gamma):
    Po = ((1 + (gamma - 1) * Ma**2 / 2)**(gamma / (gamma - 1))) * P
    return Po

def total_density(Ma, rho, gamma):
    rho_o = ((1 + (gamma - 1) * Ma ** 2 / 2) ** (1 / (gamma - 1))) * rho
    return rho_o

def downstream_prop(ps_rel, **kwargs):  # ps_rel stands for post shock relations
    gas_dict = kwargs.copy()
    for k in gas_dict.keys():
        if k == 'Pressure': gas_dict[k].append(gas_dict[k][0] * ps_rel.p21)
        if k == 'Temperature': gas_dict[k].append(gas_dict[k][0] * ps_rel.T21)
        if k == 'Density': gas_dict[k].append(gas_dict[k][0] * ps_rel.rho21)
        if k == 'Total_Pressure': gas_dict[k].append(gas_dict[k][0] * ps_rel.Po21)
        if k == 'Speed': gas_dict[k].append(gas_dict[k][0] * ps_rel.U21)
        gas_df = pd.DataFrame.from_dict(gas_dict, orient='index')
    return gas_df


# -*- coding:utf-8 -*-
"""
Project   : numpy
File Name : 5_19_to_5_20
Author    : Focus
Date      : 8/21/2021 3:13 PM
Keywords  :
Abstract  :
Param     :
Usage     : py 5_19_to_5_20
Reference :
"""
import sys
import numpy as np
from matplotlib.pyplot import show
from matplotlib.pyplot import plot

t = np.linspace(-np.pi, np.pi, 201)
k = np.arange(1, float(sys.argv[1]))
f = np.zeros_like(t)

"""
Generate sawtooth wave.
"""
for i in range(len(t)):
    f[i] = np.sum(np.sin(2 * np.pi * k * t[i]) / k) * (-2 / np.pi)

plot(t, f, lw=1.0)

"""
Generate triangle wave.
"""
for i in range(len(t)):
    f[i] = np.sum(np.sin(2 * np.pi * k * t[i]) / k)
    f[i] = np.abs(f[i] * (-2 / np.pi))

plot(t, f, lw=2.0)
show()


import numpy as np

def horizontal_avg(field, yidx=slice(None), xidx=slice(None)):
    """
    Do a horizontal average

    Parameters
    ----------
    field : array,  2D or 3D
         the field that you want to average

    xidx, yidx : slice, optional
         ranging over the interior points, excluding the halo
         default is all points

    Returns
    -------
    res : float or 1D array
       the horizontal average

    """
    if len(field.shape) == 2:
        res = np.mean(field[yidx, xidx])
    elif len(field.shape) == 3:
        res = np.mean(np.mean(field[:, yidx, xidx], axis=2), axis=1)
    else:
        raise ValueError("field should be a 2D or 3D array")
    return res

def azimuthal_avg(field, x, y, dx, dy):

    # todo: make the function accepts 3D or 4D array for field
    assert len(field.shape) == 2
    assert len(x.shape) == 2
    assert len(y.shape) == 2

    d = np.sqrt(x**2+y**2)

    xmax = np.max(x)
    ymax = np.max(y)
    rmax = min([xmax, ymax])

    dr = min(dx, dy)
    nr = int(rmax/dr+0.5)

    r = np.arange(nr)*dr
    # mid ring radial distance
    rm = 0.5*(r[1:]+r[:-1])

    res = np.zeros(nr-1)
    # average over rings
    for k in range(nr-1):
        res[k] = np.mean(field[(d >= r[k]) & (d < r[k+1])])

    return rm, res

def test_azimuthal_avg():
    nx, ny = 201, 201
    x = np.linspace(-1,1,nx)
    y = np.linspace(-1,1,ny)

    xx, yy = np.meshgrid(x, y)

    r = np.sqrt(xx**2+yy**2)

    f = lambda x: np.sin(x*5)

    phi = f(r)
    dx = x[1]-x[0]
    dy = y[1]-y[0]
    rm, phiavg = azimuthal_avg(phi, xx, yy, dx, dy)

    plt.clf()
    plt.plot(rm, phiavg, '+', label='computed')
    plt.plot(rm, f(rm), label='exact')
    plt.legend()
    plt.xlabel('r')

if __name__ == '__main__':

    import matplotlib.pyplot as plt

    plt.ion()

    test_azimuthal_avg()


# Import libraries
import matplotlib.pyplot as plt
import numpy as np

# Set custom color codes
color_codes = ['#1f77b4','#ff7f0e','#2ca02c',
               '#d62728','#9467bd','#8c564b','#e377c2',
               '#7f7f7f','#bcbd22','#17becf']

# Createing vectors X and Y
x = np.linspace(-3, 3, 100)

#define functions to plot
def y1(x):
    return x

def y2(x):
    return x**2

def y3_1(x):
    return np.where(x <= 1.5, x**2, np.nan)

def y3_2(x):
    return np.where(x > 1.5, x, np.nan)

# Define the figure options
fig, (ax1, ax2, ax3) = plt.subplots(1, 3)
fig = plt.figure(figsize = (12, 7))

# Add a title
plt.title('Some functions that tends to a limit')

# first plot
# Create the plot
ax1.plot(x, y1(x), alpha = 0.8, label ='y = x',
         color = color_codes[1],
         linewidth = 2,
         markersize = 5)

# Setting plot limits
ax1.axis(xmin=-2,xmax=2)

# Add X and y Label
ax1.set_xlabel('x')
ax1.set_ylabel('y')

# Add a grid
ax1.grid(alpha =.6, linestyle ='--')

# Add hvlines
ax1.axhline(y=1.5, color="gray", linestyle ='--')
ax1.axvline(x=1.5, color="gray", linestyle ='--')

# Add a Legend
ax1.legend()

# second plot
# Create the plot
ax2.plot(x, y2(x), alpha = 0.8, label ='y = x**2',
         color = color_codes[1],
         linewidth = 2,
         markersize = 5)

# Add and arrow
# ax2.arrow(1, y2(1), 1.1, y2(1.1), head_width=0.2, head_length=0.15, color = color_codes[1])

# Setting plot limits
ax2.axis(xmin=-3,xmax=3)

# Add X and y Label
ax2.set_xlabel('x')
ax2.set_ylabel('y')

# Add a grid
ax2.grid(alpha =.6, linestyle ='--')

# Add hvlines
ax2.axhline(y=1.5**2, color="gray", linestyle ='--')
ax2.axvline(x=1.5, color="gray", linestyle ='--')

# Add a Legend
ax2.legend()

# third plot
# Create the plot
ax3.plot(x, y3_1(x), alpha = 0.8, label ='y = x**2 if x < 2 \n x if x >= 2',
         color = color_codes[1],
         linewidth = 2,
         markersize = 5)

ax3.plot(x, y3_2(x), alpha = 0.8,
         color = color_codes[1],
         linewidth = 2,
         markersize = 5)

# Add and arrow
# ax2.arrow(1, y2(1), 1.1, y2(1.1), head_width=0.2, head_length=0.15, color = color_codes[1])

# Setting plot limits
ax3.axis(xmin=-3,xmax=3)

# Add X and y Label
ax3.set_xlabel('x')
ax3.set_ylabel('y')

# Add a grid
ax3.grid(alpha =.6, linestyle ='--')

# Add hvlines
ax3.axhline(y=1.5**2, color="gray", linestyle ='--')
ax3.axhline(y=1.5, color="gray", linestyle ='--')
ax3.axvline(x=1.5, color="gray", linestyle ='--')

# Add a Legend
ax3.legend()

# Show the plot
plt.show()


# -----------------------------------------------------
# -*- coding: utf-8 -*-
# @Time    : 8/23/2018 3:07 PM
# @Author  : sunyonghai
# @Software: ZJ_AI
# -----------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

def intersection_rect(a, b):
    x1 = max(a[0], b[0])
    y1 = max(a[1], b[1])

    x2 = min(a[2], b[2])
    y2 = min(a[3], b[3])
    return [x1,y1,x2,y2]

def intersection(a, b):
    i_rect = intersection_rect(a,b)
    i_area = area(i_rect)
    return i_area

def area(a):
    w = a[2] - a[0]
    h = a[3] - a[1]
    # print('area:', w * h)
    return  w * h

def union(a, b):
    area_a = area(a)
    area_b = area(b)
    area_i = intersection(a, b)
    return   area_a + area_b - area_i

def calc_iou(a, b):
    try:
        u = union(a, b)
        i = intersection(a, b)
        if abs(u)  < 0.000000001:
            return 1

        iou =  i / u
    except Exception as ex:
        print("{}，{},{}".format(i, u, ex))
    return iou

def test_calc_iou():
    # curr_iou = calc_iou([10,10,100,100], [50, 50, 200, 200])
    curr_iou = calc_iou([10,10,100,100], [10,10,100,100])
    return curr_iou

if __name__ == '__main__':
   print(test_calc_iou())


import numpy as np

def softmax(predictions):
    '''
    Computes probabilities from scores

    Arguments:
      predictions, np array, shape is either (N) or (batch_size, N) -
        classifier output

    Returns:
      probs, np array of the same shape as predictions -
        probability for every class, 0..1
    '''
    # TODO implement softmax
    x = predictions.copy()
    if len(predictions.shape) == 1:
        x -= np.max(x)
        return np.exp(x) / np.sum(np.exp(x))
    else:
        x -= np.max(x, axis=1, keepdims=True)
        return np.exp(x) / np.sum(np.exp(x), axis=1, keepdims=True)

def cross_entropy_loss(probs, target_index):
    '''
    Computes cross-entropy loss

    Arguments:
      probs, np array, shape is either (N) or (batch_size, N) -
        probabilities for every class
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)

    Returns:
      loss: single value
    '''
    # TODO implement cross-entropy
    if len(probs.shape) == 1:
        return -np.log(probs[target_index])
    else:
        n_samples = probs.shape[0]
        return np.mean(-np.log(probs[np.arange(n_samples), target_index]))

def l2_regularization(W, reg_strength):
    '''
    Computes L2 regularization loss on weights and its gradient

    Arguments:
      W, np array - weights
      reg_strength - float value

    Returns:
      loss, single value - l2 regularization loss
      gradient, np.array same shape as W - gradient of weight by l2 loss
    '''
    # TODO: Copy from previous assignment
    loss = reg_strength * np.sum(W ** 2)
    grad = reg_strength * 2 * W

    return loss, grad

def softmax_with_cross_entropy(predictions, target_index):
    '''
    Computes softmax and cross-entropy loss for model predictions,
    including the gradient

    Arguments:
      predictions, np array, shape is either (N) or (batch_size, N) -
        classifier output
      target_index: np array of int, shape is (1) or (batch_size) -
        index of the true class for given sample(s)

    Returns:
      loss, single value - cross-entropy loss
      dprediction, np array same shape as predictions - gradient of predictions by loss value
    '''
    # TODO copy from the previous assignment
    probes = softmax(predictions)
    loss = cross_entropy_loss(probes, target_index)
    dprediction = probes.copy()

    if len(predictions.shape) == 1:
        dprediction[target_index] -= 1
    else:
        n_samples = probes.shape[0]
        dprediction[np.arange(n_samples), target_index] -= 1
        dprediction /= n_samples

    return loss, dprediction

class Param:
    '''
    Trainable parameter of the model
    Captures both parameter value and the gradient
    '''
    def __init__(self, value):
        self.value = value
        self.grad = np.zeros_like(value)

class ReLULayer:
    def __init__(self):
        self.mask_ = None

    def forward(self, X):
        # TODO copy from the previous assignment
        self.mask_ = X > 0
        return X * self.mask_

    def backward(self, d_out):
        # TODO copy from the previous assignment
        d_result = self.mask_ * d_out
        return d_result

    def params(self):
        return {}

class FullyConnectedLayer:
    def __init__(self, n_input, n_output):
        self.W = Param(0.001 * np.random.randn(n_input, n_output))
        self.B = Param(0.001 * np.random.randn(1, n_output))
        self.X = None

    def forward(self, X):
        # TODO copy from the previous assignment
        self.X = np.copy(X)
        return np.dot(X, self.W.value) + self.B.value

    def backward(self, d_out):
        # TODO copy from the previous assignment

        self.W.grad = np.dot(self.X.T, d_out)
        self.B.grad = d_out.sum(axis=0)[None, :]
        d_input = np.dot(d_out, self.W.value.T)

        return d_input

    def params(self):
        return {'W': self.W, 'B': self.B}

class ConvolutionalLayer:
    def __init__(self, in_channels, out_channels,
                 filter_size, padding, verbose=False):
        '''
        Initializes the layer

        Arguments:
        in_channels, int - number of input channels
        out_channels, int - number of output channels
        filter_size, int - size of the conv filter
        padding, int - number of 'pixels' to pad on each side
        '''

        self.filter_size = filter_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.W = Param(
            np.random.randn(filter_size, filter_size,
                            in_channels, out_channels)
        )
        self.B = Param(np.zeros(out_channels))
        self.padding = padding
        self.verbose = verbose
        self.X = None

    def forward(self, X):
        batch_size, height, width, channels = X.shape

        out_height = height - self.filter_size + 1 + self.padding * 2
        out_width = width - self.filter_size + 1 + self.padding * 2

        # TODO: Implement forward pass
        # Hint: setup variables that hold the result
        # and one x/y location at a time in the loop below
        out = np.zeros((batch_size, out_height, out_width, self.out_channels))
        if self.padding:
            X_padding = np.zeros((batch_size, height + self.padding * 2,
                                  width + self.padding * 2, channels))
            X_padding[:, self.padding:height + self.padding, self.padding:width + self.padding, :] = X
            self.X = X_padding.copy()
        else:
            self.X = X.copy()

        # It's ok to use loops for going over width and height
        # but try to avoid having any other loops
        for y in range(out_height):
            for x in range(out_width):
                # TODO: Implement forward pass for specific location
                current_X = self.X[:, y:y + self.filter_size, x:x + self.filter_size, :].reshape((batch_size, -1))
                current_W = self.W.value.reshape((-1, self.out_channels))
                if self.verbose:
                    print('X', current_X.shape)
                    print('W', current_W.shape)
                out[:, y, x, :] = np.dot(current_X, current_W) + self.B.value

        return out

    def backward(self, d_out):
        # Hint: Forward pass was reduced to matrix multiply
        # You already know how to backprop through that
        # when you implemented FullyConnectedLayer
        # Just do it the same number of times and accumulate gradients

        batch_size, height, width, channels = self.X.shape
        _, out_height, out_width, _ = d_out.shape

        # TODO: Implement backward pass
        # Same as forward, setup variables of the right shape that
        # aggregate input gradient and fill them for every location
        # of the output
        d_input = np.zeros_like(self.X)

        # Try to avoid having any other loops here too
        for y in range(out_height):
            for x in range(out_width):
                # TODO: Implement backward pass for specific location
                # Aggregate gradients for both the input and
                # the parameters (W and B)
                current_X = self.X[:, y:y + self.filter_size, x:x + self.filter_size, :].reshape((batch_size, -1))
                current_W = self.W.value.reshape((-1, self.out_channels))
                if self.verbose:
                    print('X', current_X.shape)
                    print('W', current_W.shape)

                self.W.grad += np.dot(current_X.T, d_out[:, y, x, :]).reshape(self.W.value.shape)
                self.B.grad += d_out[:, y, x, :].sum(axis=0)
                d_input[:, y:y + self.filter_size, x: x + self.filter_size, :] += \
                    np.dot(d_out[:, y, x, :], current_W.T).reshape((batch_size, self.filter_size,
                                                                    self.filter_size, channels))
        if self.padding:
            d_input = d_input[:, self.padding:height-self.padding, self.padding:width-self.padding, :]

        return d_input

    def params(self):
        return {'W': self.W, 'B': self.B}

class MaxPoolingLayer:
    def __init__(self, pool_size, stride):
        '''
        Initializes the max pool

        Arguments:
        pool_size, int - area to pool
        stride, int - step size between pooling windows
        '''
        self.pool_size = pool_size
        self.stride = stride
        self.X = None

    def forward(self, X):
        batch_size, height, width, channels = X.shape
        # TODO: Implement maxpool forward pass
        # Hint: Similarly to Conv layer, loop on
        # output x/y dimension
        out_height = (height - self.pool_size) // self.stride + 1
        out_width = (width - self.pool_size) // self.stride + 1

        out = np.zeros((batch_size, out_height, out_width, channels))
        self.X = X.copy()

        for y in range(out_height):
            for x in range(out_width):
                current_X = self.X[:,
                                   y * self.stride:y * self.stride + self.pool_size,
                                   x * self.stride:x * self.stride + self.pool_size,
                                   :]
                out[:, y, x, :] = np.max(current_X, axis=(1, 2))
        return out

    def backward(self, d_out):
        # TODO: Implement maxpool backward pass
        batch_size, height, width, channels = self.X.shape
        _, out_height, out_width, _ = d_out.shape

        d_input = np.zeros_like(self.X)

        for y in range(out_height):
            for x in range(out_width):
                current_X = self.X[:,
                                   y * self.stride:y * self.stride + self.pool_size,
                                   x * self.stride:x * self.stride + self.pool_size,
                                   :]
                d_input_part = np.equal(np.ones_like(current_X) *
                                        current_X.max(axis=(1, 2)).reshape(batch_size, 1, 1, channels),
                                        current_X) * d_out[:, y, x, :].reshape(batch_size, 1, 1, channels)
                d_input[:,
                        y * self.stride:y * self.stride + self.pool_size,
                        x * self.stride:x * self.stride + self.pool_size,
                        :] += d_input_part
        return d_input

    def params(self):
        return {}

class Flattener:
    def __init__(self):
        self.X_shape = None

    def forward(self, X):
        batch_size, height, width, channels = X.shape

        # TODO: Implement forward pass
        # Layer should return array with dimensions
        # [batch_size, hight*width*channels]
        self.X_shape = X.shape
        return X.reshape(batch_size, height * width * channels)

    def backward(self, d_out):
        # TODO: Implement backward pass
        return d_out.reshape(self.X_shape)

    def params(self):
        # No params!
        return {}


# -*- coding: utf-8 -*-

import numpy as np, pandas as pd, matplotlib.pyplot as plt, seaborn as sns, arviz as az, scipy
from cmdstanpy import CmdStanModel

#%% data

dfhogg = pd.DataFrame(dict(
	x       = [ 201, 244,  47,  287,  203,  58,  210,  202,  198,  158, 165,  201,  157, 131, 166,  160, 186, 125, 218,   146],
	y       = [ 592, 401, 583,  402,  495, 173,  479,  504,  510,  416, 393,  442,  317, 311, 400,  337, 423, 334, 533,   344],
	sigma_x = [   9,   4,  11,    7,    5,   9,    4,    4,   11,    7,   5,    5,    5,   6,   6,    5,   9,   8,   6,    5.],
	sigma_y = [  61,  25,  38,   15,   21,  15,   27,   14,   30,   16,  14,   25,   52,  16,  34,   31,  42,  26,  16,    22],
	rho_xy  = [-.84, .31, .64, -.27, -.33, .67, -.02, -.05, -.84, -.69,  .3, -.46, -.03,  .5, .73, -.52,  .9,  .4, -.78, -.56]
), dtype = "float32")
dfhogg.plot.scatter(x = "x", y = "y", xerr = "sigma_x", yerr = "sigma_y", c = "b")

X, Y, sigmaY = dfhogg['x'].values, dfhogg['y'].values, dfhogg['sigma_y'].values
model_data_dict = {"N": len(dfhogg), "X": X, "Y": Y, "sigmaY": sigmaY}
model_data_stan = """
	data {
		int<lower=0> N;
		vector[N] X;
		vector[N] Y;
		vector<lower=0>[N] sigmaY;
	}
"""
model_trans_params_stan = """
	transformed parameters {
		vector[N] Yhat = b0 + b1 * X;
	}
"""
var_name = ["b0", "b1"]

#%% models

# Simple Linear Model with no Outlier Correction
modelfile_ols = "mdl_ols.stan"
with open(modelfile_ols, "w") as file: file.write(model_data_stan + """
	parameters { // discrete parameters impossible
		real b0;
		real b1;
	}
""" + model_trans_params_stan + """
	model {
		Y ~ normal(Yhat, sigmaY);
	}
""")

sm_ols = CmdStanModel(stan_file = modelfile_ols)
fit_ols = sm_ols.sample(
	data = model_data_dict, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit_ols.draws().shape # iterations, chains, parameters
fit_ols.summary().loc[var_name] # pandas DataFrame
print(fit_ols.diagnose())

posterior_ols = {k: fit_ols.stan_variable(k) for k in var_name}

az_trace_ols = az.from_cmdstanpy(fit_ols)
az.summary(az_trace_ols).loc[var_name] # pandas DataFrame
az.plot_trace(az_trace_ols, var_names = var_name)

gd = sns.jointplot(
	x = posterior_ols["b0"], y = posterior_ols["b1"],
	marginal_kws = {"kde": True, "kde_kws": {"cut": 1}},
)
gd.plot_joint(sns.kdeplot, zorder = 2, n_levels = 10, cmap = "gray_r")
gd.fig.suptitle("Posterior joint distribution (OLS)", y = 1.02)

# Simple Linear Model with Robust Student-T Likelihood: outliers to have a smaller influence in the likelihood estimation
modelfile_studentt = "mdl_studentt.stan"
with open(modelfile_studentt, "w") as file: file.write(model_data_stan + """
	parameters { // discrete parameters impossible
		real b0;
		real b1;
		real<lower=0> df;
	}
""" + model_trans_params_stan + """
	model {
		df ~ inv_gamma(1, 1);
		Y ~ student_t(df, Yhat, sigmaY);
	}
""")
var_name_studentt = var_name + ["df"]

sm_studentt = CmdStanModel(stan_file = modelfile_studentt)
fit_studentt = sm_studentt.sample(
	data = model_data_dict, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit_studentt.draws().shape # iterations, chains, parameters
fit_studentt.summary().loc[var_name_studentt] # pandas DataFrame
print(fit_studentt.diagnose())

posterior_studentt = {k: fit_studentt.stan_variable(k) for k in var_name_studentt}

az_trace_studentt = az.from_cmdstanpy(fit_studentt)
az.summary(az_trace_studentt).loc[var_name_studentt] # pandas DataFrame
az.plot_trace(az_trace_studentt, var_names = var_name_studentt)

# Linear Model with Custom Likelihood to Distinguish Outliers: Hogg Method
# idea: mixture model whereby datapoints can be: normal linear model vs outlier (for convenience also be linear)
modelfile_hogg = "mdl_hogg.stan"
with open(modelfile_hogg, "w") as file: file.write(model_data_stan + """
	parameters { // discrete parameters impossible
		real b0;
		real b1;
		real Y_outlier; // mean for all outliers
		real<lower=0> sigmaY_outlier; // additional variance for outliers
		simplex[2] cluster_prob; // mixture ratio
	}
""" + model_trans_params_stan + """
	model {
		b0 ~ normal(0, 5); // weakly informative Normal priors (L2 ridge reg) for inliers
		b1 ~ normal(0, 5); // likewise

		Y_outlier ~ normal(0, 10);
		sigmaY_outlier ~ normal(0, 10); // half-normal because of above constraint

		for (n in 1:N) { // custom mixture model: cluster 1 = inlier, 2 = outlier
			real cluster1 = log(cluster_prob[1]) + normal_lpdf(Y[n] | Yhat[n], sigmaY[n]);
			real cluster2 = log(cluster_prob[2]) + normal_lpdf(Y[n] | Y_outlier, sigmaY[n] + sigmaY_outlier);
			target += log_sum_exp(cluster1, cluster2);
		}
	}
""")
var_name_hogg_array = var_name + ["Y_outlier", "sigmaY_outlier", "cluster_prob[1]", "cluster_prob[2]"]
var_name_hogg_combi = var_name + ["Y_outlier", "sigmaY_outlier", "cluster_prob"]

sm_hogg = CmdStanModel(stan_file = modelfile_hogg)
fit_hogg = sm_hogg.sample(
	data = model_data_dict, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit_hogg.draws().shape # iterations, chains, parameters
fit_hogg.summary().loc[var_name_hogg_array] # pandas DataFrame
print(fit_hogg.diagnose())

posterior_hogg = {k: fit_hogg.stan_variable(k) for k in var_name_hogg_combi}

az_trace_hogg = az.from_cmdstanpy(fit_hogg)
az.summary(az_trace_hogg) # pandas DataFrame
az.plot_trace(az_trace_hogg, var_names = var_name_hogg_combi)

#%% some plots

Xrange = np.array([X.min() - 1, X.max()])
dfhogg.plot.scatter(x = "x", y = "y", xerr = "sigma_x", yerr = "sigma_y", c = "b", figsize = (8, 6))
for x, y, z in zip(
	[posterior_ols, posterior_studentt, posterior_hogg],
	["r", "g", "m"],
	["OLS", "Student-T", "Hogg method"]
):
	b0, b1 = x["b0"].mean(), x["b1"].mean()
	plt.plot(Xrange, b0 + b1 * Xrange, c = y, label = f"{z}: $ y = {b0:.1f} + {b1:.1f}x $")
plt.legend(loc = "lower right")

for var in var_name:
	a, b, c = posterior_ols[var], posterior_studentt[var], posterior_hogg[var]
	data = pd.DataFrame(data = {
		"y": np.concatenate((a, b, c)),
		"model": np.repeat([
			f"OLS: {a.mean():.1f}", f"Student-T: {b.mean():.1f}", f"Hogg method: {c.mean():.1f}"
		], 40000) # nb of data pts: 4chains × 50000 iter ÷ 5 thin
	})
	g = sns.displot(data, x = "y", hue = "model", bins = "sqrt", kde = True, palette = ["r", "g", "m"])
	g.fig.suptitle(f"Posterior distribution: {var}")
	g.fig.set_figwidth(8)
	g.fig.set_figheight(6)

#%% declare outliers

# Compute the un-normalized log probabilities for each cluster
cluster_0_log_prob = scipy.stats.norm.logpdf(
	np.expand_dims(Y, axis = 1),
	loc = posterior_hogg["b0"] + posterior_hogg["b1"] * np.expand_dims(X, axis = 1),
	scale = np.expand_dims(sigmaY, axis = 1)
) + np.log(posterior_hogg["cluster_prob"][:, 0])
cluster_1_log_prob = scipy.stats.norm.logpdf(
	np.expand_dims(Y, axis = 1),
	loc = posterior_hogg["Y_outlier"],
	scale = np.expand_dims(sigmaY, axis = 1) + posterior_hogg["sigmaY_outlier"]
) + np.log(posterior_hogg["cluster_prob"][:, 1])

# Bayes rule to compute the assignment probability: P(cluster = 1 | data) ∝ P(data | cluster = 1) P(cluster = 1)
log_p_assign_1 = cluster_1_log_prob - np.logaddexp(cluster_0_log_prob, cluster_1_log_prob)

# Average across the MCMC chain
log_p_assign_1bis = scipy.special.logsumexp(log_p_assign_1, axis=-1) - np.log(log_p_assign_1.shape[-1])

p_assign_1 = np.exp(log_p_assign_1bis)
dfhogg["is_outlier"] = [f"{100*i:.2f} %" for i in p_assign_1]
dfhogg["classed_as_outlier"] = p_assign_1 >= .95

#%% remove outliers + full model with all variances

no_outliers = dfhogg[~dfhogg.classed_as_outlier]

mdl_full_data = dict(
	N = len(no_outliers),
	X = no_outliers['x'].values,
	Y = no_outliers['y'].values,
	sigmaX = no_outliers['sigma_x'].values,
	sigmaY = no_outliers['sigma_y'].values,
	rhoXY = no_outliers['rho_xy'].values
)

modelfile_full = "mdl_full.stan"
with open(modelfile_full, "w") as file: file.write("""
	data {
		int<lower=0> N;
		vector[N] X;
		vector[N] Y;
		vector<lower=0>[N] sigmaX;
		vector<lower=0>[N] sigmaY;
		vector<lower=-1,upper=1>[N] rhoXY;
	}

	transformed data {
		real angle90 = pi()/2; // a cste
		vector[2] Z[N]; // data pt in vector form
		matrix[2,2] S[N]; // each data point’s covariance matrix
		for (i in 1:N) {
			Z[i] = [X[i], Y[i]]';
			real covXY = rhoXY[i]*sigmaX[i]*sigmaY[i];
			S[i] = [[sigmaX[i]^2, covXY], [covXY, sigmaY[i]^2]];
		}
	}

	parameters { // discrete parameters impossible
		real<lower=-angle90,upper=angle90> theta; // angle of the fitted line
		real b; // intercept
	}

	transformed parameters {
		vector[2] v = [-sin(theta), cos(theta)]'; // unit vector orthogonal to the line
		vector[N] lp; // log prob
		for (i in 1:N) {
			real delta = dot_product(v, Z[i]) - b*v[2]; // orthogonal displacement of each data point from the line
			real sigma2 = quad_form(S[i], v); // orthogonal variance of projection of each data point to the line
			lp[i] = delta^2/2/sigma2; // sum(lp) is faster than +=
		}
	}

	model {
		theta ~ uniform(-angle90, angle90);
		target += -sum(lp); // ATTENTION sign
	}

	generated quantities {
		real m = tan(theta); // slope
	}
""")
var_name_full = ["m", "b", "theta"]

sm_full = CmdStanModel(stan_file = modelfile_full)
fit_full = sm_full.sample(
	data = mdl_full_data, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit_full.draws().shape # iterations, chains, parameters
fit_full.summary().loc[var_name_full] # pandas DataFrame
print(fit_full.diagnose())

posterior_full = {k: fit_full.stan_variable(k) for k in var_name_full}

az_trace_full = az.from_cmdstanpy(fit_full)
az.summary(az_trace_full) # pandas DataFrame
az.plot_trace(az_trace_full, var_names = var_name_full)

b0, b1 = posterior_full["b"].mean(), posterior_full["m"].mean() # name changed
no_outliers.plot.scatter(x = "x", y = "y", xerr = "sigma_x", yerr = "sigma_y", c = "b")
plt.plot(Xrange, b0 + b1 * Xrange, label = f"$ y = {b0:.1f} + {b1:.1f}x $")
plt.legend(loc = "lower right")

#%% full model with intrinsic scatter

modelfile_full_intrinsic = "mdl_full_intrinsic.stan"
with open(modelfile_full_intrinsic, "w") as file: file.write("""
	data {
		int<lower=0> N;
		vector[N] X;
		vector[N] Y;
		vector<lower=0>[N] sigmaX;
		vector<lower=0>[N] sigmaY;
		vector<lower=-1,upper=1>[N] rhoXY;
	}

	transformed data {
		real angle90 = pi()/2; // a cste
		vector[2] Z[N]; // data pt in vector form
		matrix[2,2] S[N]; // each data point’s covariance matrix
		for (i in 1:N) {
			Z[i] = [X[i], Y[i]]';
			real covXY = rhoXY[i]*sigmaX[i]*sigmaY[i];
			S[i] = [[sigmaX[i]^2, covXY], [covXY, sigmaY[i]^2]];
		}
	}

	parameters { // discrete parameters impossible
		real<lower=-angle90,upper=angle90> theta; // angle of the fitted line
		real b; // intercept
		real<lower=0> V; // intrinsic Gaussian variance orthogonal to the line
	}

	transformed parameters {
		vector[2] v = [-sin(theta), cos(theta)]'; // unit vector orthogonal to the line
		vector[N] lp; // log prob
		for (i in 1:N) {
			real delta = dot_product(v, Z[i]) - b*v[2]; // orthogonal displacement of each data point from the line
			real sigma2 = quad_form(S[i], v); // orthogonal variance of projection of each data point to the line
			real tmp = sigma2 + V; // intermediary result
			lp[i] = .5*(log(tmp) + delta^2/tmp); // sum(lp) is faster than +=
		}
	}

	model {
		theta ~ uniform(-angle90, angle90);
		target += -sum(lp); // ATTENTION sign
	}

	generated quantities {
		real m = tan(theta); // slope
		real move_up = sqrt(V) / v[2];
	}
""")
var_name_full_intrinsic = var_name_full + ["V", "move_up"]

sm_full_intrinsic = CmdStanModel(stan_file = modelfile_full_intrinsic)
fit_full_intrinsic = sm_full_intrinsic.sample(
	data = mdl_full_data, show_progress = True, chains = 4,
	iter_sampling = 50000, iter_warmup = 10000, thin = 5
)

fit_full_intrinsic.draws().shape # iterations, chains, parameters
fit_full_intrinsic.summary().loc[var_name_full_intrinsic] # pandas DataFrame
print(fit_full_intrinsic.diagnose())

posterior_full_intrinsic = {k: fit_full_intrinsic.stan_variable(k) for k in var_name_full_intrinsic}

az_trace_full_intrinsic = az.from_cmdstanpy(fit_full_intrinsic)
az.summary(az_trace_full_intrinsic) # pandas DataFrame
az.plot_trace(az_trace_full_intrinsic, var_names = var_name_full_intrinsic)

b0, b1 = posterior_full_intrinsic["b"].mean(), posterior_full_intrinsic["m"].mean()
move_up = posterior_full_intrinsic["move_up"].mean()
no_outliers.plot.scatter(x = "x", y = "y", xerr = "sigma_x", yerr = "sigma_y", c = "b")
plt.plot(Xrange, b0 + b1 * Xrange, label = f"$ y = {b0:.1f} + {b1:.1f}x $")
plt.plot(Xrange, b0 + b1 * Xrange + move_up, linestyle = "--")
plt.plot(Xrange, b0 + b1 * Xrange - move_up, linestyle = "--")
plt.legend(loc = "lower right")


import numpy as np
import matplotlib.pyplot as plt

# Plot 2D points
def displaypoints2d(points):
    plt.figure(0)
    plt.plot(points[0, :], points[1, :], '.b')
    plt.xlabel('Screen X')
    plt.ylabel('Screen Y')

# Plot 3D points
def displaypoints3d(points):
    fig = plt.figure(1)
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[0, :], points[1, :], points[2, :], 'b')
    ax.set_xlabel("World X")
    ax.set_ylabel("World Y")
    ax.set_zlabel("World Z")

def cart2hom(points):
    """ Transforms from cartesian to homogeneous coordinates.

    Args:
      points: a np array of points in cartesian coordinates

    Returns:
      points_hom: a np array of points in homogeneous coordinates
    """

    #
    # You code here
    #
    ones = np.ones((1, np.shape(points)[1]))
    hom = np.vstack((points, ones))
    return hom

def hom2cart(points):
    """ Transforms from homogeneous to cartesian coordinates.

    Args:
      points: a np array of points in homogenous coordinates

    Returns:
      points_hom: a np array of points in cartesian coordinates
    """

    #
    # You code here
    #
    cart = points[0:-1, :] / points[-1, :]
    return cart

def gettranslation(v):
    """ Returns translation matrix T in homogeneous coordinates for translation by v.

    Args:
      v: 3d translation vector

    Returns:
      T: translation matrix in homogeneous coordinates
    """

    #
    # You code here
    #
    temp = np.reshape(v, (3, 1))
    hom = np.vstack((temp, 0))
    return hom

def getxrotation(d):
    """ Returns rotation matrix Rx in homogeneous coordinates for a rotation of d degrees around the x axis.

    Args:
      d: degrees of the rotation

    Returns:
      Rx: rotation matrix
    """

    #
    # You code here
    #
    return np.array([[1, 0, 0],
                     [0, np.cos(d), -np.sin(d)],
                     [0, np.sin(d), np.cos(d)]])

def getyrotation(d):
    """ Returns rotation matrix Ry in homogeneous coordinates for a rotation of d degrees around the y axis.

    Args:
      d: degrees of the rotation

    Returns:
      Ry: rotation matrix
    """

    #
    # You code here
    #
    return np.array([[np.cos(d), 0, np.sin(d)],
                     [0, 1, 0],
                     [-np.sin(d), 0, np.cos(d)]])

def getzrotation(d):
    """ Returns rotation matrix Rz in homogeneous coordinates for a rotation of d degrees around the z axis.

    Args:
      d: degrees of the rotation

    Returns:
      Rz: rotation matrix
    """

    #
    # You code here
    #
    return np.array([[np.cos(d), -np.sin(d), 0],
                     [np.sin(d), np.cos(d), 0],
                     [0, 0, 1]])

def getcentralprojection(principal, focal):
    """ Returns the (3 x 4) matrix L that projects homogeneous camera coordinates on homogeneous
    image coordinates depending on the principal point and focal length.

    Args:
      principal: the principal point, 2d vector
      focal: focal length

    Returns:
      L: central projection matrix
    """

    #
    # You code here
    #
    return np.array([[focal, 0, principal[0], 0],
                     [0, focal, principal[1], 0],
                     [0, 0, 1, 0]])

def getfullprojection(T, Rx, Ry, Rz, L):
    """ Returns full projection matrix P and full extrinsic transformation matrix M.

    Args:
      T: translation matrix
      Rx: rotation matrix for rotation around the x-axis
      Ry: rotation matrix for rotation around the y-axis
      Rz: rotation matrix for rotation around the z-axis
      L: central projection matrix

    Returns:
      P: projection matrix
      M: matrix that summarizes extrinsic transformations
    """

    #
    # You code here
    #
    R = Rz @ (Rx @ Ry)
    t = -R @ T[0:-1]
    temp = np.hstack((R, t))
    M = np.vstack((temp, [0, 0, 0, 1]))
    P = L @ M
    return P, M

def projectpoints(P, X):
    """ Apply full projection matrix P to 3D points X in cartesian coordinates.

    Args:
      P: projection matrix
      X: 3d points in cartesian coordinates

    Returns:
      x: 2d points in cartesian coordinates
    """

    #
    # You code here
    #
    temp_x = cart2hom(X)
    temp = P @ temp_x
    return hom2cart(temp)

def loadpoints():
    """ Load 2D points from obj2d.npy.

    Returns:
      x: np array of points loaded from obj2d.npy
    """

    #
    # You code here
    #
    points = np.load('data/obj2d.npy')

    return points

def loadz():
    """ Load z-coordinates from zs.npy.

    Returns:
      z: np array containing the z-coordinates
    """

    #
    # You code here
    #
    return np.load('data/zs.npy')

def invertprojection(L, P2d, z):
    """
    Invert just the projection L of cartesian image coordinates P2d with z-coordinates z.

    Args:
      L: central projection matrix
      P2d: 2d image coordinates of the projected points
      z: z-components of the homogeneous image coordinates

    Returns:
      P3d: 3d cartesian camera coordinates of the points
    """

    #
    # You code here
    #
    size_points = np.shape(P2d)
    ones = np.ones((1, size_points[1]))
    hom_points = np.vstack((P2d, ones))
    # check
    z_points = hom_points * z  # check
    # p_xy=L[0:2,2] #slicing construct no array
    # ??? why get [8, -10]
    p_xy = L[np.arange(2), 2]  # Integer array indexing

    p_xy_reshape = np.reshape(np.append(p_xy, 0), (3, 1))
    z_p_xy = p_xy_reshape * z
    inv_principal = z_p_xy + z_points
    L_truncted = L[np.arange(3), 0:3]
    P3d = np.linalg.inv(L_truncted) @ inv_principal

    return P3d

def inverttransformation(M, P3d):
    """ Invert just the model transformation in homogeneous coordinates
    for the 3D points P3d in cartesian coordinates.

    Args:
      M: matrix summarizing the extrinsic transformations
      P3d: 3d points in cartesian coordinates

    Returns:
      X: 3d points after the extrinsic transformations have been reverted
    """

    #
    # You code here
    #
    ones = np.ones((1, np.shape(P3d)[1]))
    p4d_C = np.vstack((P3d, ones))
    p4d_W = np.linalg.inv(M) @ p4d_C

    # p3d_W=p4d_W[0:-1,:]
    return p4d_W

def p3multiplecoice():
    '''
    Change the order of the transformations (translation and rotation).
    Check if they are commutative. Make a comment in your code.
    Return 0, 1 or 2:
    0: The transformations do not commute.
    1: Only rotations commute with each other.
    2: All transformations commute.
    '''

    return 0


r"""
Solve Poisson equation on (0, 2pi)x(0, 2pi) with periodic bcs

    \nabla^2 u = f, u(2pi, y) = u(0, y), u(x, 2pi) = u(x, 0)

Use Fourier basis and find u in VxV such that

    (v, div(grad(u))) = (v, f)    for all v in VxV

where V is the Fourier basis span{exp(1jkx)}_{k=-N/2}^{N/2-1} and
VxV is a tensorproductspace.

"""
import os
from sympy import symbols, cos, sin, lambdify
import numpy as np
from shenfun import inner, grad, TestFunction, TrialFunction, Array, FunctionSpace, \
    TensorProductSpace, Function, dx, comm, la

# Use sympy to compute a rhs, given an analytical solution
x, y = symbols("x,y", real=True)
ue = cos(4*x) + sin(8*y)
fe = ue.diff(x, 2) + ue.diff(y, 2)

# Size of discretization
N = (64, 64)

K0 = FunctionSpace(N[0], family='F', dtype='D', domain=(-2*np.pi, 2*np.pi))
K1 = FunctionSpace(N[1], family='F', dtype='d', domain=(-2*np.pi, 2*np.pi))
T = TensorProductSpace(comm, (K0, K1), axes=(0, 1))
u = TrialFunction(T)
v = TestFunction(T)

# Get f on quad points
fj = Array(T, buffer=fe)

# Compute right hand side
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)

# Solve Poisson equation
u_hat = Function(T)
#A = inner(v, div(grad(u)))
A = inner(grad(v), grad(u))
sol = la.SolverDiagonal(A)
u_hat = sol(-f_hat, u_hat, constraints=((0, 0),))

uq = Array(T)
uq = T.backward(u_hat, uq, fast_transform=True)

uj = Array(T, buffer=ue)
assert np.allclose(uj, uq)

# Test eval at point
point = np.array([[0.1, 0.5], [0.5, 0.6]])
p = T.eval(point, u_hat)
ul = lambdify((x, y), ue)
assert np.allclose(p, ul(*point))
p2 = u_hat.eval(point)
assert np.allclose(p2, ul(*point))
print(np.sqrt(dx((uj-uq)**2)))

if 'pytest' not in os.environ:
    import matplotlib.pyplot as plt
    plt.figure()
    X = T.local_mesh(True)
    plt.contourf(X[0], X[1], uq)
    plt.colorbar()

    plt.figure()
    plt.contourf(X[0], X[1], uj)
    plt.colorbar()

    plt.figure()
    plt.contourf(X[0], X[1], uq-uj)
    plt.colorbar()
    plt.title('Error')
    plt.show()


import numpy as np

##################### Points #####################

def compute_distance(p1,p2=np.array([0,0])):
    """
        Euclidean distance between p1 and p2
    """
    p1 = np.array(p1)
    p2 = np.array(p2)
    return np.sqrt(np.sum((p1-p2)**2))

def random_displ(max_pix=5,min_pix=0,shape=None):
    """
        Get random displacements between min_pix and max_pix
    """
    disp = np.random.randint(min_pix,max_pix,shape)
    disp *= np.random.choice([-1,1],shape)
    return disp

##################### Lines #####################

def line_polar(p1,p2):
    """
        Compute distance coordinates and angle for the line connecting p1 and p2
    """
    dy = (p2[1]-p1[1])
    dx = (p2[0]-p1[0])
    alpha = np.arctan2(dy,dx)
    l = np.arange(compute_distance(p1,p2)+1)
    return l,alpha

def line_polar_to_cart(l,alpha,point=(0,0)):
    """
        Line points from polar to cartesian coordinates
    """
    return point[0]+l*np.cos(alpha),point[1]+l*np.sin(alpha)

def find_intersect(alpha1,p1,alpha2,p2,ret_xy=True):
    """
        Compute intersection between two lines in polar coordinates. Returns either x,y of the intersection or
        its distance from p1 along line1
    """
    if alpha1 == alpha2:
        if ret_xy:
            return np.inf,np.inf
        else:
            return np.inf

    p1 = np.array(p1)
    p2 = np.array(p2)

    p21 = p2-p1
    l_inter = (np.tan(alpha2)*p21[0]-p21[1])/(np.tan(alpha2)*np.cos(alpha1)-np.sin(alpha1))
    if ret_xy:
        return l_inter*np.cos(alpha1)+p1[0],l_inter*np.sin(alpha1)+p1[1]
    else:
        return l_inter

##################### Other #####################

def PolyArea(x,y):
    """
        Compute area of polygon defined by x,y points
    """
    return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))