import numpy as np

def normalize_features(X, mean_vector=None, std_vector=None):
    """
        :param X            : two-dimensional array of our dataset, with shape NxD. N is the number of rows (each row is a sample) and D is the number of columns
        :param mean_vector  : optionally given, precomputed vector of mean values.
        :param std_vector   : optionally given, precomputed vector of standard deviation values.

        If the function is called passing *only X* as input argument, then the mean and std values will be computed from X, and then X will be normalized using these values.
        If the function is called passing *all three* input arguments, X array will be normalized using the precomputed mean and std values that were passed as input arguments.

        Normalization is performed by subtracting the mean and dividing by the standard deviation: x_normalized = (x - x_mean)/(x_std+ε)
    """

    if (mean_vector is None) and (std_vector is None):
        return_mean_std = True
        # get the mean for every column
        mean_vector = np.mean(X, axis=0)
        # get the std for every column
        std_vector = np.std(X, axis=0)

        # insert extra singleton dimension, to obtain 1xD shape
        mean_vector = np.expand_dims(mean_vector, axis=0)
        std_vector = np.expand_dims(std_vector, axis=0)
    else:
        return_mean_std = False

    # get the number of samples in the dataset
    N = X.shape[0]

    # repeat N times across first dimension, to obtain NxD shape
    repeated_mean = np.repeat(mean_vector, N, axis=0)
    repeated_std = np.repeat(std_vector, N, axis=0)

    #if np.sum(1.0*(repeated_std==0))>0:
    if np.any(repeated_std==0):
        print('Adding epsilon to avoid division by zero during normalization')
        repeated_std += np.finfo(float).eps
    # subtract column mean and divide each element by column standard deviation
    X_normalized = (X - repeated_mean) / repeated_std
    print('Dataset normalization complete.')

    if return_mean_std==True:
        return X_normalized, mean_vector, std_vector
    else:
        return X_normalized


# coding: utf-8

# # Test pyIAST for match with competitive Langmuir model

# In the case that the pure-component isotherms $N_{i,pure}(P)$ follow the Langmuir model with the same saturation loading $M$:
#
# $N_{i,pure} = M \frac{K_iP}{1+K_iP},$
#
# The mixed gas adsorption isotherm follows the competitive Langmuir isotherm:
#
# $N_i = M \frac{K_i p_i}{1 + \sum_j K_jp_j},$
#
# where $p_i$ is the partial pressure of component $i$. Here, we generate synthetic pure-component adsorption isotherm data and confirm that pyIAST agrees with the competitive Langmuir isotherm for 3 components.

# In[1]:

from __future__ import absolute_import
import numpy as np
import pyiast
import pandas as pd
import matplotlib.pyplot as plt
from six.moves import range
plt.style.use('fivethirtyeight')

colors = ['b', 'g', 'r']  # for representing each component
component_names = {0: 'A', 1: 'B', 2: 'C'}

# ## Generate synthetic pure-component isotherm data, fit Langmuir models to them.

# Model parameters ($M$, $\{K_i\}$)

# In[2]:

M = 1.0
langmuirKs = [2.0, 10.0, 20.0]  # K_i

# Generate data according to Langmuir model, store in list of Pandas DataFrames

# In[3]:

pressure = np.logspace(-3, np.log10(10), 20)
dfs = [
    pd.DataFrame({
        'P':
        pressure,
        'L':
        M * langmuirKs[i] * pressure / (1.0 + langmuirKs[i] * pressure)
    }) for i in range(3)
]

# Use pyIAST to fit Lanmguir models to the data, then plot fits

# In[4]:

isotherms = [
    pyiast.ModelIsotherm(
        dfs[i], pressure_key='P', loading_key='L', model='Langmuir')
    for i in range(3)
]

for i in range(len(isotherms)):
    isotherms[i].print_params()
    pyiast.plot_isotherm(isotherms[i])

# Plot synthetic data all in one plot for paper

# In[5]:

p_plot = np.logspace(-3, np.log10(11))  # for plotting

fig = plt.figure(facecolor='w')
for i in range(len(isotherms)):
    plt.scatter(dfs[i]['P'], dfs[i]['L'], color=colors[i], s=50, label=None)
    plt.plot(
        p_plot,
        M * langmuirKs[i] * p_plot / (1.0 + langmuirKs[i] * p_plot),
        color=colors[i],
        linewidth=2,
        label=r'$N_%s(P) = \frac{%d P}{1+%dP}$' %
        (component_names[i], langmuirKs[i], langmuirKs[i]))

plt.xlim([-.05 * 10, 1.05 * 10])
plt.ylim([-.05 * M, M * 1.05])

plt.xlabel('Pressure (bar)')
plt.ylabel('Gas uptake (mmol/g)')

plt.legend(loc='lower right')

plt.tight_layout()
plt.savefig(
    'pure_component_Langmuir.png',
    format='png',
    dpi=300,
    facecolor=fig.get_facecolor())
plt.show()

# ## Compare pyIAST predicted component loadings to that of competitive Langmuir

# Let us consider a tertiary mixture of components 0, 1, and 2 above at a total pressure of `total_pressure` bar.

# In[6]:

total_pressure = 1.0

# We will explore gas phase composition space ($\{y_i\}$) by generating random compositions and checking that they are within the triangle. We do not want to get too close to a pure phase boundary becuase of numerical instability, so we keep a distance `dx` away from pure phases. We will perform `num_tests` tests.

# In[15]:

dx = 0.0001
num_tests = 100

# Generate the compositions and store in list `compositions`

# In[16]:

compositions = []
test_no = 0
while test_no < num_tests:
    # generate random compoisitions
    y1 = np.random.uniform(dx, 1.0 - dx)
    y2 = np.random.uniform(dx, 1.0 - dx)
    y3 = 1.0 - y2 - y1
    # check that composition is within the triangle
    if y3 < dx:
        continue
    # viable composition
    compositions.append([y1, y2, y3])
    # keep generating until we have num_tests
    test_no += 1

# Next, we assert that pyIAST gives the same result as the competitive Langmuir isotherm for each of these compositions.

# Function to compute loading according to competitive Langmuir

# In[17]:

def competitive_langmuir_loading(partial_pressures, i):
    """
    Calculate loading of component i according to competitive Langmuir
    """
    return M * langmuirKs[i] * partial_pressures[i] / (
        1.0 + np.dot(langmuirKs, partial_pressures))

# Function to compute loading according to pyIAST

# In[10]:

def iast_loading(partial_pressures, i):
    """
    Calculate loading of component i according to IAST

    partial_pressures: Array, partial pressures of each component
    i: component in the mixture
    """
    component_loadings = pyiast.iast(partial_pressures, isotherms)
    return component_loadings[i]

# Loop over compositions, assert pyIAST agrees with competitive Langmuir for each component. If this runs, then there is agreement!

# In[14]:

for i in range(num_tests):
    partial_pressure = np.array(compositions[i]) * total_pressure
    # for each component...
    for c in range(len(langmuirKs)):
        np.testing.assert_almost_equal(
            competitive_langmuir_loading(partial_pressure, c),
            iast_loading(partial_pressure, c),
            decimal=4)

# ### This is using a custom library to plot the phase diagrams for the paper.

# Use ternary to plot phase diagram
# https://github.com/marcharper/python-ternary

# In[19]:

import ternary
scale = 10  # resolution in triangle
axis_colors = {'l': colors[1], 'r': colors[0], 'b': colors[2]}
cmaps = ["Blues", "Greens", "Reds"]
iast_or_lang = 'lang'

for c in range(3):
    if iast_or_lang == 'lang':
        f = lambda p: competitive_langmuir_loading(p, c)
    else:
        f = lambda p: iast_loading(p, c)

    # loop over component
    fig, ax = plt.subplots(facecolor='w')
    ax.axis("off")
    figure, tax = ternary.figure(ax=ax, scale=scale)

    tax.heatmapf(
        f,
        boundary=False,
        style="hexagonal",
        cmap=plt.cm.get_cmap(cmaps[c]),
        vmax=M,
        vmin=0.0,
        cbarlabel="%s uptake (mmol/g)" % component_names[c])

    tax.boundary(linewidth=2.0, color_dict=axis_colors)

    tax.left_axis_label("$p_1$ (bar)", color=axis_colors['l'], offset=0.16)
    tax.right_axis_label("$p_0$ (bar)", color=axis_colors['r'], offset=0.16)
    tax.bottom_axis_label("$p_2$ (bar)", color=axis_colors['b'], offset=-0.06)

    tax.gridlines(
        color="blue",
        multiple=1,
        linewidth=2,
        horizontal_kwargs={'color': axis_colors['b']},
        left_kwargs={'color': axis_colors['l']},
        right_kwargs={'color': axis_colors['r']},
        alpha=0.7)  # Every 5th gridline, can be a float

    tax.ticks(
        axis='rlb',
        linewidth=1,
        locations=np.arange(scale + 1),
        clockwise=True,
        color_dict=axis_colors,
        ticks=["%.1f" % (1.0 - 1.0 * i / scale) for i in range(scale + 1)],
        offset=0.03)
    tax.clear_matplotlib_ticks()

    tax._redraw_labels()
    # if iast_or_lang == 'iast':
    #     tax.set_title("IAST uptake, component %d" % c, y=1.08, fontsize=14)
    # if iast_or_lang == 'lang':
    #     tax.set_title("Competitive Langmuir uptake, component %d" % c, y=1.08, fontsize=14)
    plt.tight_layout()
    if iast_or_lang == 'iast':
        plt.savefig(
            "Tertiary_diagram_IAST_component_%d.png" % c,
            format='png',
            dpi=300,
            facecolor=fig.get_facecolor())
    if iast_or_lang == 'lang':
        plt.savefig(
            "Tertiary_diagram_Langmuir_component_%d.png" % c,
            format='png',
            dpi=300,
            facecolor=fig.get_facecolor())
    tax.show()

# In[ ]:

manual_tests/python_scripts/Test IAST for Langmuir case.py


#Exploring fancy indexing
import numpy as np
rand = np.random.RandomState(42)

x = rand.randint(100, size=10)
print(x)

ind = [3,7,4]
print(x[ind])

#shape of the result reflects the indexing array
ind = np.array([[3,7],[4,5]])
print("x:", x)
print(x[ind])

#this works in multiple dimesnions
X = np.arange(12).reshape((3,4))
print("X:",X)

row = np.array([0,1,2])
col = np.array([2,1,3])
print("X[row,col]:", X[row,col])
#prints (0,2), (1,1), (2,3)

print(X[row[:, np.newaxis], col])
print(row[:, np.newaxis] * col)

#combined indexing

print("X:",X)
#combining fancy and simple indicies
print(X[2, [2,0,1]])
print(X[1:, [2,0,1]])
print("we can combine fancy indices with masking as well")
mask = np.array([1,0,1,0], dtype=bool)
print(X[row[:, np.newaxis], mask])

print("Example: selecting random points")
mean = [0,0]
cov = [[1,2],[2,5]]
X= rand.multivariate_normal(mean, cov, 100)
print(X.shape) #(100,2)

import matplotlib
matplotlib.use("TKagg")
import matplotlib.pyplot as plt
import seaborn; seaborn.set() # for plot styling

plt.scatter(X[:, 0],X[:,1])
plt.show()

#we can use fancy indixing to select 20 random points

indicies = np.random.choice(X.shape[0], 20, replace=False)
print(indicies)
selection = X[indicies] #fancy indexing
print(selection.shape) #(20,2)

plt.scatter(selection[:,0], selection[:,1], facecolor='red', s=200)
plt.scatter(X[:,0],X[:,1], alpha=0.3)
plt.show()

print("#Modifying values with fancy indexing")

x = np.arange(10)
i = np.array([2,1,8,4])
x[i] = 99
print(x)
x[i] -=10
print(x)

x = np.zeros(10)
x[[0,0]] = [4,6]
print(x)
i = [2,3,3,4,4,4]
x[i] +=1
print(x)

x = np.zeros(10)
np.add.at(x, i, 10)
print(x)

print("Example: Binnng Data")
np.random.seed(42)
x=np.random.randn(100)

#compute a histogram by hand
bins = np.linspace(-5,5,20)
counts = np.zeros_like(bins)

#find the appropriate bin for each x
i = np.searchsorted(bins,x)

#add 1 to each of these bins
np.add.at(counts, i, 1)

#plot the results
plt.plot(bins, counts, linestyle='steps')
plt.show()

#or you can use plt.hist()

plt.hist(x, bins, histtype='step')
plt.show()


# -*- coding: utf-8 -*-
"""
Regula Falsi

Controllo su:
    - nIter
    - Residuo
    - Intervallo con tolleranza + eps * max(a,b)
"""
import numpy as np
import math

def sign(x) : return math.copysign(1, x)

def regula_falsi(fname, a, b, tol, nMax):
    # A e B devono avere segni discordi
    if sign(fname(a)) == sign(fname(b)):
        print("Algoritmo non applicabile")
        return [], 0, []
    else:
        xks = []
        it = 0
        eps = np.spacing(1)
        fk = fname(a)
        while it < nMax and abs(b-a) >= tol + eps * max(abs(a), abs(b)) and abs(fk) >= tol:
            # Calcolo il nuovo punto Xk
            xk = a - fname(a) * ((b - a)/(fname(b) - fname(a)))
            xks.append(xk)
            fk = fname(xk)
            if sign(xks[it]) == sign(a):
                a = xks[it]
            elif sign(xks[it]) == sign(b):
                b = xks[it]
            elif fname(xks[it]) == 0:
                break;
            it += 1
        return xk, it+1, xks


import os
import pandas as pd
import math
import numpy as np
import scipy.optimize as opt
import matplotlib.pyplot as plt

path = os.getcwd() + '/data/ex2data2.txt'
data = np.loadtxt(path, delimiter=",")
temp = np.ones(((data.shape)[0],1), dtype=np.float64)
data = np.insert(data, [0], temp, axis=1)

def addParams(data):
	new_params = np.zeros((data.shape[0],25), dtype=np.float64)
	count = 0
	for i in range(2, 7):
		for j in range(i + 1):
			x1 = data[:,1]
			x2 = data[:,2]
			new_params[:, count] = np.power(x1, i - j) * np.power(x2, j)
			count += 1
	data = np.insert(data, [3], new_params, axis=1)
	return data

def g(z):
	#sigmoid function
	return 1.0 / (1 + math.pow(math.e, -z))

g = np.vectorize(g, otypes=[np.float64])
#now, when g is applied to a vector, it is applied to each cell

def J_Reg(theta, x, y, lamda):
	#using the vector notation
	ones = np.ones((y.shape[0],), dtype=np.float64)
	h = g(x @ theta)
	tmp = y.transpose() @ np.log(h) + (ones - y).transpose() @ np.log(ones - h)
	reg_term = lamda / (2 * x.shape[0]) * (theta[1:].transpose() @ theta[1:])
	j = -(tmp) / x.shape[0] + reg_term
	return j

def gradientReg(theta, x, y, lamda):
	tmp = x.transpose() @ (g(x @ theta) - y)
	reg_term = (lamda / x.shape[0]) * theta
	reg_term[0] = 0
	grad = tmp + reg_term
	return grad

data = addParams(data)
x = data[:, :-1]
y = data[:, -1]
theta = np.zeros((data.shape[1] - 1,), dtype=np.float64)
lamda = 1
#print(J_Reg(theta, x, y, lamda))
result = opt.fmin_tnc(func=J_Reg, x0=theta, fprime=gradientReg, args=(x, y, lamda))
theta = result[0]

delta = 0.025
xrange = np.arange(-5.0, 20.0, delta)
yrange = np.arange(-5.0, 20.0, delta)
X1, X2 = np.meshgrid(xrange,yrange)

# F is one side of the equation, G is the other
F = 0
count = 0
for i in range(7):
	for j in range(i + 1):
		F += theta[count] * np.power(X1, i - j) * np.power(X2, j)
		count += 1
'''
F = theta[0] + theta[1]*X1 + theta[2]*X2 + theta[3]*X1*X1 + theta[4]*X1*X2 + theta[5]*X2*X2 + \
theta[6]*X1*X1*X1 + theta[7]*X1*X1*X2 + theta[8]*X1*X2*X2 + theta[9]*X2*X2*X2 + \
theta[10]*X1*X1*X1*X1 + theta[11]*X1*X1*X1*X2 + theta[12]*X1*X1*X2*X2 + theta[13]*X1*X2*X2*X2 + theta[14]*X2*X2*X2*X2 + \
theta[15]*X1*X1*X1*X1*X1 + theta[16]*X1*X1*X1*X1*X2 + theta[17]*X1*X1*X1*X2*X2 + theta[18]*X1*X1*X2*X2*X2 + theta[19]*X1*X2*X2*X2*X2 + theta[20]*X2*X2*X2*X2*X2 + \
theta[21]*X1*X1*X1*X1*X1*X1 + theta[22]*X1*X1*X1*X1*X1*X2 + theta[23]*X1*X1*X1*X1*X2*X2 + theta[24]*X1*X1*X1*X2*X2*X2 + theta[25]*X1*X1*X2*X2*X2*X2 + theta[26]*X1*X2*X2*X2*X2*X2 + theta[27]*X2*X2*X2*X2*X2*X2
'''

plt.contour(X1, X2, (F), [0], colors=["k"])
#x1 on x-axis and x2 on y-axis
plt.plot([row[1] for row in data if row[-1] == 0], [row[2] for row in data if row[-1] == 0], "yo")
plt.plot([row[1] for row in data if row[-1] == 1], [row[2] for row in data if row[-1] == 1], "k+")
plt.ylabel('admissions')
plt.show()

'''
positive = data[data['Accepted'].isin([1])]
negative = data[data['Accepted'].isin([0])]

fig, ax = plt.subplots(figsize=(12,8))
ax.scatter(positive['Test 1'], positive['Test 2'], s=50, c='b', marker='o', label='Accepted')
ax.scatter(negative['Test 1'], negative['Test 2'], s=50, c='r', marker='x', label='Rejected')
ax.legend()
ax.set_xlabel('Test 1 Score')
ax.set_ylabel('Test 2 Score')
plt.show()
'''


# Gradient Descent application for Optimization Problem
# using Rosenbrock Function

# we ask help for mathematics calculation
import numpy as np

# First case: we set a=0 and b=100
def DerrivRosenbrock0 ( point ):
    dx = 2*point[0] - 400*point[0]*(point[1] - (point[0]**2))
    dy = 200*(point[1] - (point[0]**2))
    return dx, dy

# Second case: we set a=1 and b=100
def DerrivRosenbrock1 ( point ):
    dx = (-2*(1 - point[0]) - 400*(point[1] - (point[0]**2)**2))
    dy = 200*(point[1] - (point[0]**2))
    return dx, dy

def main():
    # set the learning rate first
    # this learning set is ad hoc corresponds to objective function
    lrate = 0.002
    # initialize a point
    a = np.array([-.5, .2])
    # set number of epochs (can be changed as you want)
    epoch = 100000
    # we have to record all tuples of the points and its function
    ai = []
    for i in range(epoch):
        # objective function for first case (uncomment code below)
        # f = (a[0]**2) + (100*((a[1] - a[0]**2)**2))
        # objective function for second case
        f = ((1 - a[0])**2) + (100*((a[1] - a[0]**2)**2))
        # append the point and its obj function to ai as 1D list
        ai.append([a,f])
        # compute its derrivative on point a
        # Derrivative for first case (uncomment code below)
        # fi = np.array(DerrivRosenbrock0(a))
        # Derrivative for second case
        fi = np.array(DerrivRosenbrock1(a))
        # set the new point based on its
        a = a - np.dot(lrate,fi)

    # convert ai into a numpy array
    ai = np.array(ai)
    # print(ai.shape) 'uncomment this to check ai shape and delete this sentences'
    # print the last 10 of ai
    print(ai[-10:-1])
    # the minimum value of the function is just the last element of ai
    print(f'the minimum is: {ai[-1, 1]} at point: {ai[-1,0]}')

if __name__ == '__main__':
    main()


"""
SYS-611: Buffon's Needle Experiment Example with Antithetic Variables.

This example performs a Monte Carlo simulation of Buffon's Needle Experiment
to estimate the probability of a needle of certain length crossing lines
on a floor with certain spacing. This probability is proportional to the
mathematical constant pi.

@author: Paul T. Grogan <pgrogan@stevens.edu>
"""

# import the python3 behavior for importing, division, and printing in python2
from __future__ import absolute_import, division, print_function

# import the matplotlib pyplot package and refer to it as `plt`
# see http://matplotlib.org/api/pyplot_api.html for documentation
import matplotlib.pyplot as plt

# import the scipy stats package and refer to it as `stats`
# see http://docs.scipy.org/doc/scipy/reference/stats.html for documentation
import scipy.stats as stats

# import the numpy package and refer to it as `np`
# see http://docs.scipy.org/doc/numpy/reference/ for documentation
import numpy as np

# define the line width and needle length for buffon's experiment
line_width = 3.0
needle_length = 2.5

# define a process generator for the event if a needle crosses a line
def drop_needle():
    r_1 = np.random.rand()
    r_2 = np.random.rand()
    # generate distance between needle centroid and nearest line from uniform
    # distribution between 0 and line_width/2
    d_1 = r_1*line_width/2
    d_2 = (1-r_1)*line_width/2
    # generate acute angle between needle and line from uniform distribution
    # between 0 and pi/2 radians
    theta_1 = r_2*np.pi/2
    theta_2 = (1-r_2)*np.pi/2

    # for each antithetic variable, record 1 if d < needle_length/2*sin(theta)
    # otherwise record 0
    x_1 = 1 if d_1 < needle_length/2*np.sin(theta_1) else 0
    x_2 = 1 if d_2 < needle_length/2*np.sin(theta_2) else 0
    # return the average of the two antithetic variables
    return (x_1+x_2)/2.

# set the random number generator seed to 0
np.random.seed(0)

# generate 850 samples
samples = [drop_needle() for i in range(850)]

# compute the lower and upper-bounds using a 95% confidence interval
confidence_level = 0.05
z_crit = stats.norm.ppf(1-confidence_level/2)

print('P(X) = {:.3f} +/- {:.3f} (95% CI)'.format(
        np.average(samples),
        z_crit*stats.sem(samples)
    ))

# compute the exact solution, as solved by calculus
solution = 2*needle_length/(line_width*np.pi)

# compute running statistics for mean and confidence interval
mean_estimate = np.array([np.average(samples[0:i]) for i in range(len(samples))])
confidence_int = z_crit*np.array([stats.sem(samples[0:i]) for i in range(len(samples))])

# create a plot to show the mean estimate with 95% confidence interval bounds
plt.figure()
plt.plot(range(len(samples)), mean_estimate,
         'b', label='Mean Estimate')
plt.plot(range(len(samples)), mean_estimate-confidence_int,
         'g', label='95% CI Lower Bound')
plt.plot(range(len(samples)), mean_estimate+confidence_int,
         'r', label='95% CI Upper Bound')
plt.plot([0, len(samples)], [solution, solution],
          '-k', label='Analytical Solution')
plt.xlabel('Sample')
plt.ylabel('Estimate of $P(x)$')
plt.legend(loc='best')

#%%

# transform the mean estimate to estimate pi using the solution form
pi_estimate = 2*needle_length/(line_width*mean_estimate)
pi_lower_bound = 2*needle_length/(line_width*(mean_estimate+confidence_int))
pi_upper_bound = 2*needle_length/(line_width*(mean_estimate-confidence_int))

print('pi = {:.3f} +/- {:.3f} (95% CI)'.format(
        pi_estimate[-1],
        pi_upper_bound[-1] - pi_estimate[-1]
    ))

# create a plot to show the pi estimate with 95% confidence interval bounds
plt.figure()
plt.plot(range(len(samples)), pi_estimate,
         'b', label='Mean Estimate')
plt.plot(range(len(samples)), pi_lower_bound,
         'g', label='95% CI Lower Bound')
plt.plot(range(len(samples)), pi_upper_bound,
         'r', label='95% CI Upper Bound')
plt.plot([0, len(samples)], [np.pi, np.pi],
         '-k', label='Analytical Solution')
plt.xlabel('Sample')
plt.ylabel('Estimate of $\pi$')
plt.legend(loc='best')


# -*- coding: utf-8 -*-
"""This module provides functions to convert between different attitude
representations, as well as computing
"""

# Imports
import logging
import copy
import math
import numpy as np

# Brahe Imports
from   brahe.utils import logger
import brahe.constants as _const
from   brahe.epoch import Epoch

#####################
# Rotation Matrices #
#####################

def Rx(angle:float, use_degrees:bool=False) -> np.ndarray:
    """Rotation matrix, for a rotation about the x-axis.

    Args:
        angle (float): Counter-clockwise angle of rotation as viewed
            looking back along the postive direction of the rotation axis.
        use_degrees (bool): Handle input and output in degrees. Default: ``False``

    Returns:
        r (np.ndarray): Rotation matrix

    References:

        1. O. Montenbruck, and E. Gill, *Satellite Orbits: Models, Methods and Applications*, 2012, p.27.
    """

    if use_degrees:
        angle *= math.pi/180.0

    c = math.cos(angle)
    s = math.sin(angle)

    return np.array([[1.0,  0.0,  0.0],
                      [0.0,   +c,   +s],
                      [0.0,   -s,   +c]])

def Ry(angle:float, use_degrees:bool=False) -> np.ndarray:
    """Rotation matrix, for a rotation about the y-axis.

    Args:
        angle (float): Counter-clockwise angle of rotation as viewed
            looking back along the postive direction of the rotation axis.
        use_degrees (bool): Handle input and output in degrees. Default: ``False``

    Returns:
        r (np.ndarray): Rotation matrix

    References:

        1. O. Montenbruck, and E. Gill, *Satellite Orbits: Models, Methods and Applications*, 2012, p.27.
    """

    if use_degrees:
        angle *= math.pi/180.0

    c = math.cos(angle)
    s = math.sin(angle)

    return np.array([[ +c,  0.0,   -s],
                      [0.0, +1.0,  0.0],
                      [ +s,  0.0,   +c]])

def Rz(angle:float, use_degrees:bool=False) -> np.ndarray:
    """Rotation matrix, for a rotation about the z-axis.

    Args:
        angle (float): Counter-clockwise angle of rotation as viewed
            looking back along the postive direction of the rotation axis.
        use_degrees (bool): Handle input and output in degrees. Default: ``False``

    Returns:
        r (np.ndarray): Rotation matrix

    References:

        1. O. Montenbruck, and E. Gill, *Satellite Orbits: Models, Methods and Applications*, 2012, p.27.
    """

    if use_degrees:
        angle *= math.pi/180.0

    c = math.cos(angle)
    s = math.sin(angle)

    return np.array([[ +c,   +s,  0.0],
                      [ -s,   +c,  0.0],
                      [0.0,  0.0,  1.0]])


# -*- coding: utf-8 -*-
"""
A function to compute the cost.
"""

import numpy as np
from functions import sigmoid

def compute_loss(y, tx, w, fn="mse"):
    """compute the loss given the function 'fn'"""

    assert len(y) == len(tx), "y and tx must have the same number of rows"

    # cost given a binary prediction (0 or 1)
    def with_pred(pred):
        """compute the cost by negative log likelihood"""
        return - np.asscalar((y.T @ np.log(pred)) + ((1 - y).T @ np.log(1 - pred)))

    if fn.lower() == "mse":
        N = len(y)
        e = y - (tx @ w)
        return (e.T @ e) / (2 * N)

    if fn.lower() == "rmse":
        return np.sqrt(2 * compute_loss(y, tx, w, fn="mse"))

    if fn.lower() == "mae":
        N = len(y)
        e = y - tx @ w
        return np.mean(np.abs(e))

    if fn.lower() == "sig":
        pred = sigmoid(tx @ w)
        return with_pred(pred)

    raise NameError("Not such fn exists \"" + fn + "\"")

def calculate_loss(e, fn="mse"):
    """Calculate the mse for vector e."""

    if fn.lower() == "mse":
        return (e.T @ e) / (2 * len(e))

    if fn.lower() == "rmse":
        return np.sqrt(2 * calculate_loss(e, fn="mse"))

    if fn.lower() == "mae":
        return np.mean(np.abs(e))

    raise NameError("Not such fn exists \"" + fn + "\"")


import numpy as np

def is_quadratic(num):
    return np.sqrt(float(num)).is_integer()

def matrix_argmax(mat):
    return np.unravel_index(mat.argmax(), mat.shape)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def inv_sigmoid(x, envelop=np.Inf):
    assert (envelop > 0)
    return np.maximum(np.minimum(np.log(x / (1 - x)), envelop), -envelop)

def linear_transform_01(A):
    max_A = np.max(A).astype(np.float)
    min_A = np.min(A).astype(np.float)
    return (A - min_A) / (max_A - min(A))

def Jsoftmax(w, t=1.0, axis=0):
    w = np.array(w)  # n, c, h, w
    ndim = len(w.shape)
    w_sh = np.array(w.shape)
    w_sh[axis] = 1
    reshapedim = np.array([1] * ndim, dtype=np.int)
    reshapedim[axis] = w.shape[axis]

    w_max = np.amax(w, axis=axis, keepdims=True)
    w -= np.tile(w_max, reshapedim)  # to avoid numerical issues (exploding exp)

    e = np.exp(np.array(w) / t)
    Z = np.sum(e, axis=axis, keepdims=True)

    out = e / np.tile(Z, reshapedim)
    return out

def proj_lp(v, xi, p):
    if np.linalg.norm(v.reshape(-1), p) > xi:
        if p == np.inf:
            v_projected = np.sign(v) * np.minimum(np.abs(v), xi)
        elif p == 2:
            v_projected = v / np.linalg.norm(v.reshape(-1)) * xi
        else:
            raise
        return v_projected
    else:
        return v

def compute_percentiles(x, thres=0.25):
    x = np.array(x)
    p = np.sort(x.reshape(-1))
    N = float(x.size)
    low_n, upp_n = int(thres * N), int((1 - thres) * N)
    return p[low_n], p[upp_n]


__author__ = "Haim Adrian"

import cv2
import numpy as np

def log(consumer, *message):
    text = ' '.join(message)
    if consumer is not None:
        consumer(text)

def do_progress(consumer, current_progress, step):
    if consumer is not None:
        current_progress += step
        consumer(current_progress)

    return current_progress

# I have created this method to compare matrices between py3.5 and py3.7 in order to find a bug I had in py3.5,
# where I used a uint8 matrix rather than float64 and had overflows occurred in py3.5
def dump_matrix_to_disk(mat, file_name):
    with open(file_name, 'wb') as f:
        f.write(str.encode(str(mat.shape)))
        f.write(str.encode('\n'))
        for line in mat:
            np.savetxt(f, line, fmt='%.2f')

def find_harris_corners(image, k, window_size, console_consumer=None):
    """
    Harris Corner Detector algorithm implementation.
    Harris detector responses are the CRF(x, y) values from the equation: CRF = det(M) − k*(tr(M)**2) (Corner Response Function)
    where det(M) = Lambda1*Lambda2, which is the value of the determinant, and tr(M) = Lambda1+Lambda2, which is the trace
    of the determinant. (Lambda1, Lambda2 = eigen values).
    When the CRF value of the pixel point is greater than the given threshold, it is determined that the target point is a corner point.
    :param image: Input single-channel 8-bit or floating-point image.
    :param k: Harris detector free parameter in the equation.
    :param window_size: Neighborhood size.
    :param console_consumer: Used for logging purposes. (Passing step updates to a listener)
    :return: Image to store the Harris detector responses. It has the same size as source image. dst(x,y) = detM(x,y) − k*(trM(x,y)**2)
    Corners in the image can be found as the local maxima of this response map.
    """
    if not isinstance(image, np.ndarray):
        log(console_consumer, "find_harris_corners: Not a tensor. Was: Image=", image.__class__)
        return None

    if image.ndim not in [2, 3]:
        log(console_consumer, 'find_harris_corners: Illegal image dimension. Image N-dimension can be 2 or 3 only. Was:', image.ndim)
        return None

    if image.ndim == 3:
        twoD = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        twoD = image

    # Find all eigen values and vectors for each pixel in the image.
    # Eigen values are used in order to find the determinant and trace, which are needed for Harris Corner equation
    eigen_vals_and_vectors = cv2.cornerEigenValsAndVecs(twoD, window_size, 5)

    # 0 is for first eigen value, and 1 is for second eigen value
    # Use batch operation rather than going in a very slow loop over them and calculate the response for each pixel
    # det(M) = lambda1 * lambda2
    # trace(M) = lambda1 + lambda2
    detM = eigen_vals_and_vectors[:, :, 0] * eigen_vals_and_vectors[:, :, 1]
    traceM = eigen_vals_and_vectors[:, :, 0] + eigen_vals_and_vectors[:, :, 1]
    harris_responses = detM - k * (traceM ** 2)

    return harris_responses

# I leave this code here as originally I have started from it until found cornerEigenValsAndVecs
# Behind the scenes of cv2.cornerEigenValsAndVecs:

# ! offset = int(window_size / 2)
# Pad the image so we can calculate Harris score for the whole image
# Use Copy Padding so we will not accidentally invent corners with a constant padding of 0's
# From now on, go ahead as float64. Otherwise this won't work in py3.5 (It does work in 3.7 however)
# ! padded = np.float64(cv2.copyMakeBorder(twoD, offset, offset, offset, offset, cv2.BORDER_REPLICATE))

# For 2D array, so the result of gradient is two arrays ordered by axis
# ! dy, dx = np.gradient(padded)

# Compute products of derivatives at every pixel
# ! Ixx = dx ** 2
# ! Ixy = dx * dy
# ! Iyy = dy ** 2

# Calculate step size in the area we have to progress in a progressbar
# ! progress_steps_left = max_progress - progress
# ! progress_step = progress_steps_left / image.shape[0]
# ! for x in range(offset, padded.shape[0] - offset):
# ! for y in range(offset, padded.shape[1] - offset):
# The window according to the Harris Equation
# ! window_Ixx = Ixx[x - offset: x + offset + 1, y - offset: y + offset + 1]
# ! window_Ixy = Ixy[x - offset: x + offset + 1, y - offset: y + offset + 1]
# ! window_Iyy = Iyy[x - offset: x + offset + 1, y - offset: y + offset + 1]

# Sum the squares in the sliding window.
# Here is the M = SUM_for-each-x-y(window(x, y)[[Ixx, Ixy], [Ixy, Iyy]]
# ! sum_xx = window_Ixx.sum()
# ! sum_xy = window_Ixy.sum()
# ! sum_yy = window_Iyy.sum()

# Calculate determinant and trace for the sum of squares matrix (M)
# ! detM = (sum_xx * sum_yy) - (sum_xy * sum_xy)  # Multiply main diagonal, minus multiplication of secondary diagonal
# ! traceM = sum_xx + sum_yy  # Main diagonal

# Calculate corner_response_function for Harris Corner equation
# ! corner_response_function = detM - k * (traceM ** 2)

# ! harris_responses[x - offset, y - offset] = corner_response_function

# ! progress = do_progress(progress_consumer, progress, progress_step)

# ! return harris_responses


#! /usr/bin/env python3

import sys
import csv
import argparse
import numpy as np

ROUND_OFF: int = 4
theta: float = 0.0
r: float = 0.0
z: float = 0.0
y: float = 0.0
x: float = 0.0

class DataConverter:
    def __init__(self):

        try:
            self.options = self.read_options(sys.argv[1:])
        except:
            print("Input data file required...")
            exit(1)

        if self.options.input == None:
            print("[+] Please enter the input/origianl data file\n")
            print(self.parser.print_help())
            exit(1)

        self.INPUT_FILE_PATH: str = self.options.input
        self.OUTPUT_FILE_PATH: str =  self.INPUT_FILE_PATH[:-3] + "Rec.csv" # "Rec_LidarScanCSVData.csv"

        try:
            self.cylindrical_file = csv.reader(open(self.INPUT_FILE_PATH, "r"))
        except FileNotFoundError as e:
            print(e)
            exit(1)

        self.rectangular_file = csv.writer(open(self.OUTPUT_FILE_PATH, "w"))
        self.preport("Completed reading contents of file: " + self.INPUT_FILE_PATH)
        self.preport("Created data output save file: " + self.OUTPUT_FILE_PATH)

        self.cylindrical_data: list = list(self.cylindrical_file)
        self.rectangular_data: list = self.cylindrical_data

        self.num_rows: int = len(self.cylindrical_data)
        self.preport("Data processing has began...\n[+] Please wait..........")

        self.data_conversion()

        # Save the new file with elements in rectangular coordinates
        self.rectangular_file.writerows(self.rectangular_data)
        self.preport("Output file '" + self.OUTPUT_FILE_PATH + "' Saved.")
        self.preport("Data conversion completed!!!")

    # declare function to define command-line arguments
    def read_options(self, args=sys.argv[1:]):
        """Returns dictionary of arguments values"""
        self.parser = argparse.ArgumentParser(description="The parsing commands lists.")
        self.parser.add_argument("input", help="Enter the original/input data file")
        return self.parser.parse_args(args)

    def deg2rad(self, pembe: float) -> float:
        return (np.pi * pembe) / 180.0

    def preport(self, mesg: str):
        print("[+] " + mesg + " \r")

    def dreport(self, mesg: str):
        print("[-] " + mesg + "\r")

    def data_conversion(self):
        first_data_row_index = 0

        if type(self.rectangular_data[0][0]) == str:
            first_data_row_index = 1
            self.rectangular_data[0][0] = "X"
            self.rectangular_data[0][1] = "Y"
            self.rectangular_data[0][2] = "Z"

        for row in range(first_data_row_index, self.num_rows):
            theta = float(self.cylindrical_data[row][0])
            r = float(self.cylindrical_data[row][1])
            z = float(self.cylindrical_data[row][2])

            theta = self.deg2rad(theta)

            x = round(r * np.cos(theta), 4)
            y = round(r * np.sin(theta), 4)
            z = round(z, 4)

            self.rectangular_data[row][0] = x
            self.rectangular_data[row][1] = y
            self.rectangular_data[row][2] = z

DataConverter()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Date    : 2017-10-14 20:54:56
# @Author  : jingray (lionel_jing@163.com)
# @Link    : http://www.jianshu.com/u/01fb0364467d
# @Version : $Id$

import random as rd

# Coin flips
def simulate_coin_flips(num_trials):
    heads = 0
    tails = 0
    p_heads = 0.5
    for i in range(num_trials):
        random_number = rd.random()
        if random_number < p_heads:
            heads = heads + 1
        else:
            tails += 1
    percent_heads = heads / num_trials
    return percent_heads

percentage = simulate_coin_flips(2000000) # calling the function
print(percentage)

# Simulating Probabilities

def simulate_dice_rolls(N):
    roll_counts = [0,0,0,0,0,0]
    for i in range(N):
        roll = rd.choice([1,2,3,4,5,6])
        index = roll - 1
        roll_counts[index] = roll_counts[index] + 1
    return roll_counts

def show_roll_data(roll_counts):
    number_of_sides_on_die = len(roll_counts)
    for i in range(number_of_sides_on_die):
        number_of_rolls = roll_counts[i]
        number_on_die = i+1
        print(number_on_die, "came up", number_of_rolls, "times")

roll_data = simulate_dice_rolls(60000)
show_roll_data(roll_data)

from matplotlib import pyplot as plt

def visualize_one_die(roll_data):
    roll_outcomes = [1,2,3,4,5,6]
    fig, ax = plt.subplots()
    ax.bar(roll_outcomes, roll_data)
    ax.set_xlabel("Value on Die")
    ax.set_ylabel("# rolls")
    ax.set_title("Simulated Counts of Rolls")
    plt.show()

roll_data = simulate_dice_rolls(500)
visualize_one_die(roll_data)

# Calculating Probabilities for Bayes' rule

# The complement P(~A) function takes in the probability of an event P(A).
def complement(p_A):
    complement = 1 - p_A
    return complement

# Test code - do not change
import solution

test_value = 0.4265

# This line calls your function and stores the output
comp = complement(test_value)
correct_comp = solution.complement(test_value)

# Assertion, comparison test
try:
    assert(comp == correct_comp)
    print('That\'s right!')
except:
    print ('Your code returned that the complement is: ' +str(comp)
           + ', which does not match the correct value: ' +str(correct_comp))

# Test 2
comp2 = complement(solution.test_value)
correct_comp2 = solution.complement(solution.test_value)

# Assertion, comparison test
try:
    assert(comp2 == correct_comp2)
    print('That\'s right!')
except:
    print ('For test 2, your code returned that the complement is: ' +str(comp2)
           + ', which does not match the correct value.')

## Complete this joint function
def joint(p_A, p_B):

    ## calculates the joint probability of
    ## any variables p_A, p_B, WHEN THOSE PROBABILITIES
    ## ARE INDEPENDENT (this code wouldn't work
    ## for probabilities that depend on each other).
    joint = p_A * p_B

    return joint

# Test values
test_A = 0.15
test_B = 0.42

# This line calls your function and stores the output
j = joint(test_A, test_B)
correct_j = solution.joint(test_A, test_B)

# Assertion, comparison test
try:
    assert(j == correct_j)
    print('That\'s right!')
except:
    print ('Your code returned that the joint probability is: ' +str(j)
           + ', which does not match the correct value.')

# The inputs to total probability are given descriptive names

def total_probability(p_disease, p_pos_given_disease, p_pos_given_no_disease):

    ## calculates the total probability of a positive test result
    ## You may use other variable in this function as well as long
    ## as total is correct
    total = p_disease * p_pos_given_disease + (1 - p_disease) * p_pos_given_no_disease

    return total

# Test code - do not change
import solution

p_disease = 0.1
p_pos_given_disease= 0.9
p_pos_given_no_disease= 0.2

# This line calls your function and stores the output
total = total_probability(p_disease, p_pos_given_disease, p_pos_given_no_disease)
correct_total = solution.total_probability(p_disease, p_pos_given_disease, p_pos_given_no_disease)

# Assertion, comparison test
try:
    assert(total == correct_total)
    print('That\'s right!')
except:
    print ('Your code returned that the total probability is: ' +str(total)
           + ', which does not match the correct value.')

# Given three input probabilities, complete this function
# so that it returns the posterior probability

def bayes(p_A, p_B_given_A, p_notB_given_notA):

    ## TODO: Calculate the posterior probability
    ## and change this value
    p_notA = 1-p_A
    p_B_given_notA = 1-p_notB_given_notA
    p_B = p_A*p_B_given_A + p_notA*p_B_given_notA
    posterior = p_A*p_B_given_A/p_B

    return posterior

# Test code - do not change
import solution

test_p_A = 0.4
test_p_B_given_A = 0.7
test_p_notB_given_notA = 0.9

# This line calls your function and stores the output
posterior = bayes(test_p_A, test_p_B_given_A, test_p_notB_given_notA)
correct_posterior = solution.bayes(test_p_A, test_p_B_given_A, test_p_notB_given_notA)

# Assertion, comparison test
try:

    assert(abs(posterior - correct_posterior) < 0.0001)
    print('That\'s right!')
except:
    print ('Your code returned that the posterior is: ' +str(posterior)
           + ', which does not match the correct value.')

# 2D Array
import numpy as np

# A simple robot world can be defined by a 2D array
# Here is a 5x6 world
world = np.array([ [0, 0, 0, 1, 0],
                   [0, 0, 0, 1, 0],
                   [0, 1, 1, 0, 0],
                   [0, 0, 0, 0, 1],
                   [1, 0, 0, 1, 0],
                   [1, 0, 0, 0, 0] ])

# Visualize the world
print(world)
# Print out some information about the world

print('The shape of this array is: ' + str(world.shape))
print('Notice that the x and y dimensions are in the opposite order than usual!')
print('It\'s height is: ' + str(world.shape[1]) +
      ', and it\'s length/width is: ' + str(world.shape[0]))

      # Access a location and read its value
value = world[3][0]
print('\n')
print('Value at index [3, 0] = ' +str(value))

# Read the first three items in the 3rd row
row = 2
column_index = 0
value_left = world[row, column_index]
value_middle = world[row, column_index + 1]
value_right = world[row, column_index + 2]

print('\nThe first three values in row 2 : ' +str(value_left)+', '
                                              +str(value_middle) +', '
                                              +str(value_right) )

# Compare the first two values and print the result
compare = world[0][0] == world[0][1]
print('\nDo the first two values match? ' + str(compare))

# Define a function to plant a tree
# and change the value of the array in that location
def plant_tree(y, x):
    # check that the indices are in the boundaries of the array dimensions
    if(y < world.shape[0] and x < world.shape[1]):
        world[y,x] = 1
        print('\n' + str(world)) # prints a newline and the current world

# Call the function at the location x = 3, and y = 2
# You can call this multiple times in a row
plant_tree(0, 2)

# This function uses nested for loops and knowledge
# about the shape of the array to print out each item with it's index
def iterate2D(world):
    # y-dimension (rows)
    for i in range(0, world.shape[0]):
        # x-dimension (columns)
        for j in range(0, world.shape[1]):
            print('Index ['+str(i)+']['+str(j)+'] = ' +str(world[i][j]))

# Call the iterate function
print('\n')
iterate2D(world)

# This function is similar to our iterate2D function,
# But looks for the first tree in the array and prints its location [x][y]
def first_tree(world):
    # iterates through all indices starting at the top-left [0][0]
    for i in range(0, world.shape[0]):
        for j in range(0, world.shape[1]):
            # check if a tree is found
            if(world[i][j] == 1):
                # if so, print the index and leave the loop with a return statement
                print('First tree found at location: ['+str(i)+']['+str(j)+']')
                return

# Call the first_tree function
print('\n')
first_tree(world)

import numpy as np

# A 5x4 robot world of characters 'o' and 'b'
world = np.array([ ['o', 'b', 'o', 'o', 'b'],
                   ['o', 'o', 'b', 'o', 'o'],
                   ['b', 'o', 'o', 'b', 'o'],
                   ['b', 'o', 'o', 'o', 'o'] ])

# Sensor measurement
measurement = ['b', 'o']

# This function takes in the world and the sensor measurement.
# Complete this function so that it returns the indices of the
# likely robot locations, based on matching the measurement
# with the color patterns in the world

def find_match(world, measurement):
    # Empty possible_locations list
    possible_locations = []

    ## TODO: Iterate through the world
    ## Look at two adjacent indices at a time - the square the robot is
    ## on top of and the square to its right
    ## (Making sure not to go past the bounds of the world)
    for i in range(0,world.shape[0]):
        for j in range(0,world.shape[1]):
            if j<world.shape[1]-1 and [world[i][j],world[i][j+1]] == measurement:
                possible_locations.append([i,j])
    ## TODO: If two adjacent colors in the world match
    ## the two colors in the sensor measurement
    ## Add those indices to the possible_locations list
    ## Append them in the format [row_index, column_index], i.e. [0, 0]

    return possible_locations

# This line runs the function and stores the output - do not delete
locations = find_match(world, measurement)

# In numpy, this is easy to print COLUMN 0?

import numpy as np

np_grid = np.array([
    [0, 1, 5],
    [1, 2, 6],
    [2, 3, 7],
    [3, 4, 8]
])

# The ':' usually means "*all values*
print(np_grid[:,0])

# What if you wanted to change the size of the array?

# For example, we can turn the 2D grid from above into a 1D array
# Here, the -1 means fit all values into the length of this array
np_1D = np.resize(np_grid, (1, -1))

print(np_1D)

# We can also create a 2D array of zeros or ones
# which is useful for car world creation and analysis

# Create a 5x4 array
zero_grid = np.zeros((5, 4))

print(zero_grid)


# %%
import numpy as np

# %%
vector = np.array([1, 2, 3, 4])

# %%
vector

# %%
vector.shape

# %%
mat = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

# %%
mat

# %%
mat.shape

# %%
mat.sum()

# %%
mat.sum(axis=0)

# %%
mat[0]

# %%
mat[1]

# %%
mat.sum(axis=1)

# %%
mat

# %%
mat.mean()

# %%
mat.mean(axis=1)

# %%
mat.mean(axis=0)

# %%
tensor = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])

# %%
tensor

# %%
tensor.shape

# %%
list(range(10))

# %%
np.array(range(10))

# %%
range(1.0, 10.0)

# %%
np.arange(10)

# %%
np.arange(10.0)

# %%
np.arange(10).dtype

# %%
np.arange(10.0).dtype

# %%
np.arange(2, 10)

# %%
np.arange(2, 10, 3)

# %%
np.linspace(0.1, 1.0, 10)

# %%
np.arange(0.1, 1.1, 0.1)

# %%
np.linspace(0.0, 1.0, 10)

# %%
np.zeros(3)

# %%
np.zeros(3).shape

# %%
np.zeros((3,))

# %%
np.zeros((3, 3))

# %%
np.zeros((2, 3, 4))

# %%
np.ones(3)

# %%
np.ones((1, 2, 3))

# %%
np.ones(1, 2, 3)

# %%
id = np.eye(4)

# %%
id

# %%
np.eye(3, 4)

# %%
np.identity(3)

# %%


import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import coint
import matplotlib.pyplot as plt
# Set a seed value to make the experience reproducible
np.random.seed(123)

# Generate Symbol1 daily returns
Symbol1_returns = np.random.normal(0, 1, 100)
# Create a series for Symbol1 prices
Symbol1_prices = pd.Series(np.cumsum(Symbol1_returns),\
                           name='Symbol1') + 10
Symbol1_prices.plot(figsize=(15,7))
plt.show()

# Create a series for Symbol2 prices
# We are going to mimikate the Symbol1 behavior
noise = np.random.normal(0, 1, 100)
Symbol2_prices = Symbol1_prices + 10 + noise
Symbol2_prices.name = 'Symbol2'

plt.title("Symbol 1 and Symbol 2 prices")
Symbol1_prices.plot()
Symbol2_prices.plot()
plt.legend()
plt.show()

def zscore(series):
    return (series - series.mean()) / np.std(series)

score, pvalue, _ = coint(Symbol1_prices, Symbol2_prices)
print(pvalue)
ratios = Symbol1_prices / Symbol2_prices
plt.title("Ration between Symbol 1 and Symbol 2 price")

ratios.plot()
plt.show()

#plt.axhline(ratios.mean())
#plt.legend([' Ratio'])

zscore(ratios).plot()
plt.title("Z-score evolution")
plt.axhline(zscore(ratios).mean(),color="black")
plt.axhline(1.0, color="red")
plt.axhline(-1.0, color="green")
plt.show()

ratios.plot()
buy = ratios.copy()
sell = ratios.copy()
buy[zscore(ratios)>-1] = 0
sell[zscore(ratios)<1] = 0
buy.plot(color="g", linestyle="None", marker="^")
sell.plot(color="r", linestyle="None", marker="v")
x1,x2,y1,y2 = plt.axis()
plt.axis((x1,x2,ratios.min(),ratios.max()))
plt.legend(["Ratio", "Buy Signal", "Sell Signal"])
plt.show()

symbol1_buy=Symbol1_prices.copy()
symbol1_sell=Symbol1_prices.copy()
symbol2_buy=Symbol2_prices.copy()
symbol2_sell=Symbol2_prices.copy()

Symbol1_prices.plot()
symbol1_buy[zscore(ratios)>-1] = 0
symbol1_sell[zscore(ratios)<1] = 0
symbol1_buy.plot(color="g", linestyle="None", marker="^")
symbol1_sell.plot(color="r", linestyle="None", marker="v")

Symbol2_prices.plot()
symbol2_buy[zscore(ratios)<1] = 0
symbol2_sell[zscore(ratios)>-1] = 0
symbol2_buy.plot(color="g", linestyle="None", marker="^")
symbol2_sell.plot(color="r", linestyle="None", marker="v")

x1,x2,y1,y2 = plt.axis()
plt.axis((x1,x2,Symbol1_prices.min(),Symbol2_prices.max()))
plt.legend(["Symbol1", "Buy Signal", "Sell Signal","Symbol2"])
plt.show()


# Value iteration, Policy iteration
import numpy as np

def policy_eval(pi: np.ndarray, v: np.ndarray, num_states: int, T: np.ndarray, R: np.ndarray, gamma,
               threshold: float = 1e-5) -> np.ndarray:
    """
    Given certain policy probabilities, return the associated values of all states.
    :param pi: Current policy function as an array, where (action = pi[state]).
    :param v: Current value function as an array, where (value = v[state]).
    :param num_states: Number of states in the environment.
    :param T: Transition probabilities, as a 2D array.
    :param R: Rewards for certain state transitions, as a 2D array.
    :param gamma: Discount factor.
    :param threshold:
    :return: The new value of all states.
    """
    v_new = np.copy(v)
    v += 100

    while np.max(np.abs(v_new - v)) > threshold:
        v = np.copy(v_new)
        for s in range(num_states):
            # Deterministic policy
            a = pi[s]
            # Use in-place updating for policy evaluation
            v_new[s] = np.sum(T[s, a, :] * (R[s, a, :] + gamma * v))

    return v_new

def policy_iter(pi: np.ndarray, v: np.ndarray, num_states: int, T: np.ndarray, R: np.ndarray, gamma=1.0) -> tuple[
    np.ndarray, np.ndarray]:
    """
    Perform one iteration of the policy iteration algorithm
    :param pi: Current policy function as an array, where (action = pi[state]).
    :param v: Current value function as an array, where (value = v[state]).
    :param num_states: Number of states in the environment.
    :param T: Transition probabilities, as a 2D array.
    :param R: Rewards for certain state transitions, as a 2D array.
    :param gamma: Discount factor.
    :return: Updated pi and v.
    """
    # Evaluate the policy
    v = policy_eval(pi=pi, v=v, num_states=num_states, T=T, R=R, gamma=gamma)

    # Improve the policy
    pi_new = np.copy(pi)
    pi += 100
    while np.any(pi_new != pi):
        pi = np.copy(pi_new)
        for s in range(num_states):
            pi_new[s] = np.argmax(np.sum(T[s, :, :] * (R[s, :, :] + gamma * v[np.newaxis, :]), axis=1))

    return pi, v

def value_iter(v: np.ndarray, num_states: int, T: np.ndarray, R: np.ndarray, gamma=1.0) -> tuple[np.ndarray, np.ndarray]:
    """
    Perform one iteration of the value iteration algorithm
    :param v: Current value function as an array, where (value = v[state].
    :param num_states: Number of states in the environment.
    :param T: Transition probabilities, as a 2D array.
    :param R: Rewards for certain state transitions, as a 2D array.
    :param gamma: Discount factor.
    :return: Updated pi and v.
    """
    delta = 0
    v_new = np.copy(v)
    for s in range(num_states):
        v = np.copy(v_new)
        v_new[s] = np.max(np.sum(T[s, :, :] * (R[s, :, :] + gamma * v[np.newaxis, :]), axis=1))
        delta = max([delta, np.abs(v[s] - v_new[s])])

    v = np.copy(v_new)

    # Assuming a deterministic policy, so pi in Z^num_states instead of pi in [0, 1]^(num_states * num_actions)
    pi = np.zeros(num_states, dtype=np.int32)
    for s in range(num_states):
        pi[s] = np.argmax(np.sum(T[s, :, :] * (R[s, :, :] + gamma * v[np.newaxis, :]), axis=1))

    return pi, v