# -*- coding: utf-8 -*-
"""
Anharmonic Oscillations with an RK2 Method

Created on Wed Oct 23 10:18:14 2019

@author: gafeiden
"""
import numpy as np
import matplotlib.pyplot as plt

def springForce(k, x, p):
    """ Calculate the spring force """
    return -k*x**(p - 1)

def rk2Solve(m, x0, v0, k=1.0, p=2, dt=1.0e-2):
    """ Solve the anharmonic oscillator using RK2 """
    v_list = [v0]
    x_list = [x0]

    t = 0.0  # initial time [s]
    t_list = [t]
    for i in range(int(10/dt)):
        t += dt

        v12 = v0 + dt*springForce(k, x0, p)/2  # v0 + h*k1/2
        x12 = x0 + dt*v0/2

        v0 += dt*springForce(k, x12, p)
        x0 += dt*v12

        v_list.append(v0)
        x_list.append(x0)
        t_list.append(t)

    t_list = np.array(t_list)
    x_list = np.array(x_list)
    v_list = np.array(v_list)

    solution = np.column_stack((t_list, x_list))
    solution = np.column_stack((solution, v_list))

    return solution

def eulerSolve(m, x0, v0, k=1.0, p=2, dt=1.0e-2):
    """ Solve the anharmonic oscillator using Euler's method """
    v_list = [v0]
    x_list = [x0]

    t = 0.0  # initial time [s]
    t_list = [t]
    for i in range(int(10/dt)):
        t  += dt
        v0 += dt/m * springForce(k, x0, p)
        x0 += v0*dt

        v_list.append(v0)
        x_list.append(x0)
        t_list.append(t)

    t_list = np.array(t_list)
    x_list = np.array(x_list)
    v_list = np.array(v_list)

    solution = np.column_stack((t_list, x_list))
    solution = np.column_stack((solution, v_list))

    return solution

# initial conditions of mass
m  = 1.0  # object's mass [kg]
v0 = 0.0  # object's initial x velocity [m/s]
x0 = 1.0  # object's initial position [m]

# spring constant
k = 10.0   #

dt = 1.0e-2  # time step [s]

p02_rk = rk2Solve(m, x0, v0, k=k, p=2,  dt=dt)
p06_rk = rk2Solve(m, x0, v0, k=k, p=6,  dt=dt)
p10_rk = rk2Solve(m, x0, v0, k=k, p=10, dt=dt)

p02_em = eulerSolve(m, x0, v0, k=k, p=2,  dt=dt)
p06_em = eulerSolve(m, x0, v0, k=k, p=6,  dt=dt)
p10_em = eulerSolve(m, x0, v0, k=k, p=10, dt=dt)

fig, ax = plt.subplots(1, 2)
# Position vs Time, RK2
ax[0].plot(p02_rk[:,1], p02_rk[:,0], '-', lw=2)
ax[0].plot(p06_rk[:,1], p06_rk[:,0], '-', lw=2)
ax[0].plot(p10_rk[:,1], p10_rk[:,0], '-', lw=2)
ax[0].set_title("RK2 :: Position vs Time")
ax[0].set_xlabel("Position [m]")
ax[0].set_ylabel("Time [s]")

# Velocity vs Time, RK2
ax[1].plot(p02_rk[:,2], p02_rk[:,0], '-', lw=2)
ax[1].plot(p06_rk[:,2], p06_rk[:,0], '-', lw=2)
ax[1].plot(p10_rk[:,2], p10_rk[:,0], '-', lw=2)
ax[1].set_title("RK2 :: Velocity vs Time")
ax[1].set_xlabel("Velocity [m/s]")
ax[1].set_ylabel("Time [s]")

fig.tight_layout()
plt.show()

fig, ax = plt.subplots(1, 2)

# Position vs Time, Euler
ax[0].plot(p02_em[:,1], p02_em[:,0], '-', lw=2)
ax[0].plot(p06_em[:,1], p06_em[:,0], '-', lw=2)
ax[0].plot(p10_em[:,1], p10_em[:,0], '-', lw=2)
ax[0].set_title("Euler :: Position vs Time")
ax[0].set_xlabel("Position [m]")
ax[0].set_ylabel("Time [s]")

# Velocity vs Time, Euler
ax[1].plot(p02_em[:,2], p02_em[:,0], '-', lw=2)
ax[1].plot(p06_em[:,2], p06_em[:,0], '-', lw=2)
ax[1].plot(p10_em[:,2], p10_em[:,0], '-', lw=2)
ax[1].set_title("Euler :: Velocity vs Time")
ax[1].set_xlabel("Velocity [m/s]")
ax[1].set_ylabel("Time [s]")

fig.tight_layout()

plt.show()


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

# Load the dataframe
df = pd.read_csv(path)

#Code starts here
# TASK - 1
p_a = len(df[df['fico'].astype(float)>700])/len(df)
print("p_a : ", p_a)

p_b = len(df[df['purpose']=='debt_consolidation'])/len(df)
print("p_b : ",p_b)

df1 = df[df['purpose'] == 'debt_consolidation']

p_a_b = len(df1[df1['fico'].astype(float)>700])/len(df1)
print("p_a_b : ", p_a_b)

result = p_a_b == p_a
print("Result : ",result)

# TASK - 2
prob_lp = len(df[df['paid.back.loan'] == 'Yes'])/len(df)
print(prob_lp)

prob_cs = len(df[df['credit.policy'] == 'Yes'])/len(df)
print(prob_cs)

new_df = df[df['paid.back.loan'] == 'Yes']

prob_pd_cs = len(new_df[new_df['credit.policy'] == 'Yes'])/len(new_df)
print(prob_pd_cs)

bayes = round((prob_pd_cs * prob_lp) / prob_cs, 4)
print(bayes)

# TASK - 3
df['purpose'].value_counts().plot(kind='bar')
plt.title("Distribution Graph of 'Purpose'")
plt.show()

df1 = df[df['paid.back.loan'] == 'No']
print(df1.shape)

df1['purpose'].value_counts().plot(kind='bar')
plt.title("Purpose Distribution for Paid.Back.Loan = No")
plt.show()

# TASK - 4
inst_median = df['installment'].median()
print("inst_median : ", inst_median)

inst_mean = df['installment'].mean()
print("inst_mean : ", inst_mean)

df['installment'].plot(kind='hist')
plt.title("Distribution of Installment")
plt.axvline(x=inst_mean, color='red')
plt.axvline(x=inst_median, color='blue')
plt.show()

df['log.annual.inc'].plot(kind='hist')
plt.title("Distribution of Log Annual Income")
plt.show()


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

'''
Write down the function that we need to solve for
x' = t^{-2}(tx-x^2)
'''

def f(t,x):
    return (t-3.2)*x + 8*t*exp((t-3.2)**2/2)*cos(4*t**2)

'''
another derivative of the function function
'''

def dfx(t,x):
    return t-3.2

# analytical solution

def sol(t,t0,x0):

    return exp((t-3.2)**2/2)*(sin(4*t**2))

# backward Euler
def beuler(t0, tn, n, x0):
    h = abs(tn - t0)/n
    t = linspace(0, tn, n+1)
    x = zeros(n+1)
    x[0] = x0

    for k in range(0,n):
        err = 1
        zold = x[k] + h*f(t[k], x[k])
        I = 0

        while err > 10**(-10) and I < 5:
            F = x[k] + h*f(t[k+1], zold) - zold
            dF = h*dfx(t[k+1], zold)-1
            znew = zold - F/dF
            err = abs(znew- zold)
            I += 1
        x[k+1] = znew
    return x

def RK4(t0,tn,n,x0):
    h = abs(tn-t0)/n
    t = linspace(t0, tn, n+1)
    x = zeros(n+1)
    x[0] = x0
    for i in range(0,n):
        K1 = f(t[i], x[i])
        K2 = f(t[i]+h/2, x[i]+K1*h/2)
        K3 = f(t[i] + h/2, x[i]+K2*h/2)
        K4 = f(t[i] + h, x[i]+K3*h)

        x[i+1] = x[i] + h*(K1+2*K2+ 2*K3+K4)/6
    return x

def AdBash3(t0,tn,n,x0):
    h = abs(tn-t0)/n
    t = linspace(t0, tn, n+1)
    x = zeros(n+1)

    x[0:3] = RK4(t0, t0+2*h,2,x0)
    K1 = f(t[1], x[1])
    K2 = f(t[0], x[0])

    for i in range(2,n):
        K3= K2
        K2= K1
        K1 = f(t[i], x[i])

        # prediction_resul
        x[i+1] = x[i] + h*(23*K1 - 16*K2+5*K3)/12
    return x

def PreCorr3(t0,tn,n,x0):
    h = abs(tn-t0)/n
    t = linspace(t0, tn, n+1)
    x = zeros(n+1)

    x[0:3] = RK4(t0, t0+2*h,2,x0)
    K1 = f(t[1], x[1])
    K2 = f(t[0], x[0])

    for i in range(2,n):
        K3= K2
        K2= K1
        K1 = f(t[i], x[i])

        # prediction_resul
        x[i+1] = x[i] + h*(23*K1 - 16*K2+5*K3)/12

        K0 = f(t[i+1],x[i+1])

        # correctorical

        x[i+1] = x[i] + h*(9*K0 + 19*K1 - 5*K2 + K3)/24

    return x

# print

fg = 1
n= 300
t0 = 0
tn = 6
x0 = 0.75
t = linspace(t0, tn , n+1)
xe = AdBash3(t0,tn, n, x0)
xb = beuler(t0,tn,n,x0)
xpc = PreCorr3(t0, tn, n,x0)

#figure(fg)

plot(t,xe,'red',label ='Adam Bashforth 3')
plot(t, xb, 'black', label ='Backward Euler')
plot(t,xpc,'cyan',label='Predictor Corrector')

t = linspace(t0, tn, 401)
xsol = sol(t,t0,x0)

plot(t,xsol, color='green', label='Exact')

#title('n = %d' %n)
#axis([0,tn, -60,40])
legend(loc='lower left')
plt.savefig('Test4.png')


#Find Schröder Numbers

def getSchroderTriangle(n=5,b=True):
    from sympy import Float as mmmfloat
    from math import log10 as mloggg
    colk = [k for k in range(n+1)]
    a=[[ 1 if(j==0) else 0 for j in colk] for i in colk]
    a[1][1]=2
    for i in range(1,n+1):
        for j in range(1,i+1):
            a[i][j] = a[i][j-1]+a[i-1][j-1]+a[i-1][j]
            if(b):
                if((mloggg(a[i][j]))+1>15):
                    a[i][j] = mmmfloat(a[i][j])

    return a

#Schroder Sequence from triangular
def SchroderSequenceFT(n=5,returni=False):
    a = getSchroderTriangle(n,False)
    b = [[a[i][j] for j in range(len(a[i])) if(i==j) ][0] for i in range(len(a))]
    if(returni):
        return b
    else:
        print(b)

#Schroder Sequence from triangular (method 2)
def SchroderSequenceFT2(n=5,returni=False):
    colk = [k for k in range(n+1)]
    a=[[ 1 if(j==0) else 0 for j in colk] for i in colk]
    b = []
    a[1][1]=2
    for i in range(1,n+1):
        for j in range(i+1):
            if(j==i or j==i-2 or j==i-1):
                a[i][j] = a[i][j-1]+a[i-1][j-1]+a[i-1][j]
                if(i==j):
                    b.append(a[i][j])
            else:
                continue
    return b

#Main numbers from Schroder Sequence
def SchroderSequence(n=5,returni=False):
    a = [1,2]
    for i in range(2,n+2):
        b = 3*a[i-1]
        for k in range(1,i-1):
            b+=a[k]*a[i-k-1]
        a.append(b)
        if(not returni):
            print(b,end=", ")
    if(returni):
        return a

#Schroder Sequence by triangular
def SchroderSequenceBT(countt=5,returni=False):
    a = getSchroderTriangle(countt,False)
    b = [j for sub in a for j in sub if(j!=0)]
    if(returni):
        return b
    else:
        print(b)

def SchroderTriangle(countt=10):
    s = getSchroderTriangle(countt)
    r = range(len(s))
    from pandas import DataFrame as mmpd
    df=mmpd(s, columns=r, index=r)
    df=df.rename_axis(columns="n/k")
    return df

#getSchroderTriangle(5)
#SchroderTriangle(100)
#SchroderSequence(1000) #539 ms
#SchroderSequenceFT(1000) #363 ms
#SchroderSequenceFT2(1000) #158 ms
#SchroderSequenceBT(200)


#from scipy.interpolate import CubicSpline
#from scipy import linalg

#x_points = [1.0, 2.0, 3.0, 4.0, 5.0]
#y_points = [1.0, 2.0, 1.0, 1.5, 1.0]

#cs = CubicSpline(x_points, y_points, bc_type='natural')
#derivative = cs.derivative()
#print(derivative)
#print(linalg.solve(cs.c, derivative.c))

#!/usr/bin/env python3

import numpy as np
from scipy.interpolate import CubicSpline

# calculate 5 natural cubic spline polynomials for 6 points
# (x,y) = (0,12) (1,14) (2,22) (3,39) (4,58) (5,77)
x = np.array([1.0, 2.0, 3.0, 4.0, 5.0])
y = np.array([1.0, 2.0, 1.0, 1.5, 1.0])

# calculate natural cubic spline polynomials
cs = CubicSpline(x,y,bc_type='natural')

# show values of interpolation function at x=1.25
print('S(2) = ', cs(2))

## Aditional - find polynomial coefficients for different x regions

# if you want to print polynomial coefficients in form
# S0(0<=x<=1) = a0 + b0(x-x0) + c0(x-x0)^2 + d0(x-x0)^3
# S1(1< x<=2) = a1 + b1(x-x1) + c1(x-x1)^2 + d1(x-x1)^3
# ...
# S4(4< x<=5) = a4 + b4(x-x4) + c5(x-x4)^2 + d5(x-x4)^3
# x0 = 0; x1 = 1; x4 = 4; (start of x region interval)

# show values of a0, b0, c0, d0, a1, b1, c1, d1 ...
cs.c

# Polynomial coefficients for 0 <= x <= 1
a0 = cs.c.item(3,0)
b0 = cs.c.item(2,0)
c0 = cs.c.item(1,0)
d0 = cs.c.item(0,0)

# Polynomial coefficients for 1 < x <= 2
a1 = cs.c.item(3,1)
b1 = cs.c.item(2,1)
c1 = cs.c.item(1,1)
d1 = cs.c.item(0,1)

# 2
a2 = cs.c.item(3,2)
b2 = cs.c.item(2,2)
c2 = cs.c.item(1,2)
d2 = cs.c.item(0,2)

# ...

# Polynomial coefficients for 4 < x <= 5
#a4 = cs.c.item(3,4)
#b4 = cs.c.item(2,4)
#c4 = cs.c.item(1,4)
#d4 = cs.c.item(0,4)

# Print polynomial equations for different x regions
print('S0(0<=x<=1) = ', a0, ' + ', b0, '(x-0) + ', c0, '(x-0)^2  + ', d0, '(x-0)^3')
print('S1(1< x<=2) = ', a1, ' + ', b1, '(x-1) + ', c1, '(x-1)^2  + ', d1, '(x-1)^3')
print('...')
#print('S5(4< x<=5) = ', a4, ' + ', b4, '(x-4) + ', c4, '(x-4)^2  + ', d4, '(x-4)^3')

# So we can calculate S(2) by using equation S1(1< x<=2)
print('S(3) = ', a2 + b2*0.25 + c2*(0.25**2) + d2*(0.25**3))

# Cubic spline interpolation calculus example
    #  https://www.youtube.com/watch?v=gT7F3TWihvk


import numpy as np
import matplotlib.pyplot as plt

# This is a short python script to play around with JumpDiff
import jumpdiff as jd

# %% Let's first integrate a jump-diffusion process
# Define:
# integration time and time sampling
t_final = 10000
delta_t = 0.01

# A drift function
def a(x):
    return -0.5*x

# and a (constant) diffusion term
def b(x):
    return 0.75

# Now define a jump amplitude and rate
xi = 2.5
lamb = 1.75

# and let jdprocess integrate the stochastic path
X = jd.jd_process(t_final, delta_t, a=a, b=b, xi=xi, lamb=lamb)

# %% Plot the trajectory of the jump-diffusion process
fig, ax = plt.subplots(1,1,figsize=(6,3))

ax.plot(np.linspace(0,t_final,int(t_final/delta_t)), X, color = 'black')
ax.set_xlabel('t', fontsize=16)
ax.set_ylabel('x', fontsize=16)

# %% To retrieve the moments, use:
edges, moments = jd.moments(timeseries = X, bw = 0.35)
# and don't forget that the Kramers─Moyal coefficient need `moments/delta_t`

# %% Let us plot the first Kramers─Moyal coefficient 'moments[1,...]/delta_t'
fig, ax = plt.subplots(1,1,figsize=(6,3))

ax.plot(edges, moments[1,...]/delta_t, color = 'black', label = '1st Kramers─Moyal coefficient')
ax.plot(edges, a(edges), '--', color = 'black', label = 'Theoretical curve a = -0.5*x')
ax.set_xlim([-5,5]); ax.set_ylim([-5,5])

ax.set_xlabel('x', fontsize=16)
ax.set_ylabel('$D_1$(x)', fontsize=16)
ax.legend(fontsize=13)

# %% The second Kramers─Moyal coefficient 'moments[2,...]/delta_t'
fig, ax = plt.subplots(1,1,figsize=(6,3))

ax.plot(edges, moments[2,...]/delta_t, color = 'black', label = '2nd Kramers─Moyal coefficient')
ax.plot(edges, (b(0)**2 + xi*lamb)*np.ones_like(edges), '--', color = 'black', label = 'Theoretical curve $b^2+λξ$')
ax.set_xlim([-5,5]); ax.set_ylim([3,8])

ax.set_xlabel('x', fontsize=16)
ax.set_ylabel('$D_2$(x)', fontsize=16)
ax.legend(fontsize=13)

# %% And the fourth Kramers─Moyal coefficient 'moments[4,...]/delta_t'
fig, ax = plt.subplots(1,1,figsize=(6,3))

ax.plot(edges, moments[4,...]/delta_t, color = 'black', label = '4th Kramers─Moyal coefficient')
ax.plot(edges, (3*(xi**2)*lamb)*np.ones_like(edges), '--', color = 'black', label = 'Theoretical curve $3λξ^2$')
ax.set_xlim([-5,5]); ax.set_ylim([0,60])

ax.set_xlabel('x', fontsize=16)
ax.set_ylabel('$D_4$(x)', fontsize=16)
ax.legend(fontsize=13)

# %% Finally, we can use simply the 'jump_amplitude' and 'jump_rate' functions
# to recover the ξ and λ parameters

xi_est = jd.jump_amplitude(moments = moments, verbose = True)
print(xi_est)

lamb_est = jd.jump_rate(moments = moments, xi_est = xi, verbose = True)
print(lamb_est/delta_t)
# Don't forget that the jump rate λ needs to be divide by 'delta_t' to yield a
# comparible result.

# %% #########################################################################

# To understand the usage of the q_ratio function, let us generate to sample
# trajectories: one without jumps, denoted d_timeseries, and one with jumps
# denoted j_timeseries.

# integration time and time sampling
t_final = 10000
delta_t = 0.01

# Drift function
def a(x):
    return -0.5*x

# Diffusion function
def b(x):
    return 0.75

# generate 2 trajectories
d_timeseries = jd.jd_process(t_final, delta_t, a=a, b=b, xi=0, lamb=0)
j_timeseries = jd.jd_process(t_final, delta_t, a=a, b=b, xi=2.5, lamb=1.75)

# %% Subsequently we time a time scale to analyse, as
lag = np.logspace(0, 3, 25, dtype=int)

d_lag, d_Q = jd.q_ratio(lag, d_timeseries)
j_lag, j_Q = jd.q_ratio(lag, j_timeseries)

# %% we can then finally plot the results
fig, ax = plt.subplots(1,1,figsize=(6,3))

ax.loglog(d_lag, d_Q, '-', color = 'black', label='diffusion')
ax.loglog(j_lag, j_Q, 'o-', color = 'black', label='jump-diffusion')
# ax.set_xlim([-5,5]); ax.set_ylim([-5,5])

ax.set_xlabel('lag', fontsize=16)
ax.set_ylabel('$Q$-ratio', fontsize=16)
ax.legend(fontsize=13)

fig.tight_layout()
fig.savefig('q_ratio.png', dpi=300)

# %% If one wishes to check the formulae behind the corrections of the moments,
# simply choose the desired power:
power = 4

jd.m_formula(power = power)


import jax.numpy as jnp
from jax import random

class jaxKMeans():
    def __init__(self,n_clusters = 4, maxIter=100000, n_init = 10,init_idx = None):
        print("Initializing KMeans ....")
        self.key = random.PRNGKey(0)
        self.n_clusters = n_clusters
        self.maxIter = maxIter
        self.n_init = n_init
        self.init_idx = init_idx

    def initializeMeans(self,X):
        n = X.shape[0]
        means = X[:self.n_clusters]
        idx = random.randint(self.key, (self.n_clusters,),0,n)

        #for i in range(idx.shape[0]):
        #    means[i] = X[i]

        return means

    def euclideanDistance(self,X,means):
        temp = (X ** 2).sum(-1)[:,None] + (means ** 2).sum(-1)[None,:]
        temp = temp - 2 * X @ means.T
        return temp

    def classify(self,X,means):
        distances = self.euclideanDistance(X,means)
        clusters = jnp.argmin(distances,axis=1)
        return clusters

    def updateMeans(self,X,clusters,means):
        clusters = clusters.reshape(clusters.shape[0],1)
        n = X.shape[1]
        X = jnp.hstack((X,clusters))
        X = X[X[:,n].argsort()]

        spilited = jnp.split(X[:,:n], jnp.unique(X[:,n],
                                                 return_index=True)[1][1:])

        temp = [0 for j in range(len(spilited))] #jnp.zeros((len(spilited),n))
        for i in range(len(spilited)):
            temp[i] = jnp.mean(spilited[i],axis=0)

        temp = jnp.array(temp)
        newmean = (means + temp)/2

        return newmean

    def calculateMeans(self,X):
        means = self.initializeMeans(X)
        belongsTo = jnp.zeros((X.shape[0]))

        for iteration in range(self.maxIter):
            noChange = True
            clusters = self.classify(X,means)
            means = self.updateMeans(X,clusters,means)

            if(clusters != belongsTo).any():
                noChange = False

            belongsTo = clusters
            if noChange:
                break

        return means

    def calculateInertia(self,X,means):
        dis = self.euclideanDistance(X,means)
        return jnp.sum(dis.reshape(-1))

    def fit(self,X):
        bestMeans = self.calculateMeans(X)
        bestInertia = self.calculateInertia(X,bestMeans)

        for i in range(self.n_init-1):
            means = self.calculateMeans(X)
            currInertia = self.calculateInertia(X,means)
            if currInertia < bestInertia:
                bestMeans = means
                bestInertia = currInertia

        self.means = bestMeans
        return self

    def predict(self,X):
        return  self.classify(X,self.means)

    def fit_predict(self,X):
        return self.fit(X).predict(X)

    def getMeans(self):
        return self.means


"""
Dataset generation module.

author : Nishant Baheti <nishantbaheti.it19@gmail.com>
"""

import numpy as np
from typing import Tuple

def rotation_matrix_2d(theta: float) -> np.ndarray:
    """
    Create 2D data rotation matrix.

    Reference article
    ------------------

    https://en.wikipedia.org/wiki/Rotation_matrix

    Args:
        theta (float): angle for rotation.

    Returns:
        np.ndarray: rotation matrix.
    """

    mat = np.array([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ])
    return mat

def spiral_data(data_limit: int = 30, n_classes: int = 2,
                         n_samples_per_class=300) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate spiral data for classification problem.

    Args:
        data_limit (int, optional): range of data. Defaults to 30.
        n_classes (int, optional): number of classes for classification. Defaults to 2.
        n_samples_per_class (int, optional): number of samples per classes. Defaults to 300.

    Returns:
        Tuple[np.ndarray, np.ndarray]: X,y.
    """

    theta = np.pi * (2 / n_classes)
    rotation_mat = rotation_matrix_2d(theta=theta)

    features = [np.array([
        [
            (np.cos(i/8) * i) + np.random.randn(), (np.sin(i/8) * i) + np.random.randn()
        ] for i in np.linspace(2, data_limit, n_samples_per_class)
    ])]
    target = [np.ones(shape=(n_samples_per_class, 1), dtype=np.int32) * 0]

    for i in range(0, n_classes-1):
        features.append(features[-1] @ rotation_mat)
        target.append(np.ones(shape=(n_samples_per_class, 1), dtype=np.int32) * i+1)
    X = np.vstack(features)
    y = np.vstack(target)
    return X, y

def sine_wave_from_sample(
        n_samples: int, signal_freq: float, n_cycles: int = 10, amplitude: int = 1,
        amp_shift: int = 0, phase_shift: int = 0) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Sine wave generation with number of samples and signal frequency

    Reference:
        https://machinelearningexploration.readthedocs.io/en/latest/MathExploration/Fourier.html#Sine-wave

    Args:
        n_samples (int): number of samples.
        signal_freq (float): signal frequency.
        n_cycles (int, optional): number of cycles. Defaults to 10.
        amplitude (int, optional): signal amplitude. Defaults to 1.
        amp_shift (int, optional): amplitude shift. Defaults to 0.
        phase_shift (int, optional): phase shift. Defaults to 0.

    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: signal wave, time, freq.
    """
    sampling_freq = (n_samples * signal_freq) / n_cycles
    # n_samples = int((sampling_freq / signal_freq) * n_cycles)

    t_step = 1 / sampling_freq
    time = np.linspace(0, (n_samples-1) * t_step, n_samples)

    f_step = sampling_freq / n_samples
    freq = np.linspace(0, (n_samples-1) * f_step, n_samples)

    wave = amp_shift + amplitude * \
        np.sin((2 * np.pi * signal_freq * time) + phase_shift)

    return wave, time, freq

def sine_wave_from_timesteps(
        signal_freq: float, time_step: float, amplitude: int = 1, amp_shift: int = 0,
        phase_shift: int = 0) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Sine wave generation with time steps and signal frequency

    Args:
        signal_freq (float): singal frequency.
        time_step (float): time step.
        amplitude (int, optional): amplitude. Defaults to 1.
        amp_shift (int, optional): amplitude shift. Defaults to 0.
        phase_shift (int, optional): phase shift. Defaults to 0.

    Returns:
        Tuple[np.ndarray, np.ndarray, np.ndarray]: signal wave, time, freq.
    """
    time = np.arange(0, 1, time_step)
    n = len(time)
    freq = (1/(time_step * n)) * np.arange(n)

    wave = amp_shift + amplitude * \
        np.sin((2 * np.pi * signal_freq * time) + phase_shift)

    return wave, time, freq

if __name__ == "__main__":

    # import seaborn as sns
    # import matplotlib.pyplot as plt

    # X, y = generate_spiral_data(n_classes=1)

    # sns.scatterplot(x=X[..., 0], y=X[..., 1], hue=y[..., 0], palette='dark')
    # plt.show()

    pass


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.api as sm
import seaborn as sns
sns.set()

data = pd.read_csv("real_estate_price_size.csv")
data.head()

data.describe()

y = data["price"]
x1 = data["size"]

plt.scatter(x1, y)
plt.xlabel("size of real estate in sqm", fontsize=20)
plt.ylabel("price of real estate in USD", fontsize=20)
plt.show()

x = sm.add_constant(x1)
results = sm.OLS(y, x).fit()
results.summary()

plt.scatter(x1, y)
yhat = 223.1787*x1+1.019e+05
fig = plt.plot(x1, yhat, lw=3.5, c="green", label="Regression Line")
plt.xlabel("size of real estate in sqm", fontsize=20)
plt.ylabel("price of real estate in USD", fontsize=20)
plt.xlim(0)
plt.ylim(0)
plt.show()


#
# pulled from codearmo.com
# https://www.codearmo.com/python-tutorial/options-trading-greeks-black-scholes
#
import numpy as np
from scipy.stats import norm

N = norm.cdf
N_prime = norm.pdf

def bs_call_value(S, K, T, r, sigma):
    d1 = (np.log(S / K) + (r + sigma ** 2 / 2) * T) / (sigma * np.sqrt(T))
    d2 = d1 - sigma * np.sqrt(T)
    return S * N(d1) - K * np.exp(-r * T) * N(d2)

def bs_put_value(S, K, T, r, sigma):
    d1 = (np.log(S / K) + (r + sigma ** 2 / 2) * T) / (sigma * np.sqrt(T))
    d2 = d1 - sigma * np.sqrt(T)
    return K * np.exp(-r * T) * N(-d2) - S * N(-d1)

def bs_calldiv_value(S, K, T, r, q, sigma):
    d1 = (np.log(S / K) + (r - q + sigma ** 2 / 2) * T) / (sigma * np.sqrt(T))
    d2 = d1 - sigma * np.sqrt(T)
    return S * np.exp(-q * T) * N(d1) - K * np.exp(-r * T) * N(d2)

def bs_putdiv_value(S, K, T, r, q, sigma):
    d1 = (np.log(S / K) + (r - q + sigma ** 2 / 2) * T) / (sigma * np.sqrt(T))
    d2 = d1 - sigma * np.sqrt(T)
    return K * np.exp(-r * T) * N(-d2) - S * np.exp(-q * T) * N(-d1)

def d1(S, K, T, r, sigma):
    return (np.log(S / K) + (r + sigma ** 2 / 2) * T) / sigma * np.sqrt(T)

def d2(S, K, T, r, sigma):
    return d1(S, K, T, r, sigma) - sigma * np.sqrt(T)

def delta_call(S, K, T, r, sigma):
    return N(d1(S, K, T, r, sigma))

def delta_put(S, K, T, r, sigma):
    return -N(-d1(S, K, T, r, sigma))

def gamma(S, K, T, r, sigma):
    return N_prime(d1(S, K, T, r, sigma)) / (S * sigma * np.sqrt(T))

def vega(S, K, T, r, sigma):
    return S * np.sqrt(T) * N_prime(d1(S, K, T, r, sigma))

def theta_call(S, K, T, r, sigma):
    p1 = -S * N_prime(d1(S, K, T, r, sigma)) * sigma / (2 * np.sqrt(T))
    p2 = r * K * np.exp(-r * T) * N(d2(S, K, T, r, sigma))
    return p1 - p2

def theta_put(S, K, T, r, sigma):
    p1 = -S * N_prime(d1(S, K, T, r, sigma)) * sigma / (2 * np.sqrt(T))
    p2 = r * K * np.exp(-r * T) * N(-d2(S, K, T, r, sigma))
    return p1 + p2

def rho_call(S, K, T, r, sigma):
    return K * T * np.exp(-r * T) * N(d2(S, K, T, r, sigma))

def rho_put(S, K, T, r, sigma):
    return -K * T * np.exp(-r * T) * N(-d2(S, K, T, r, sigma))


import numpy as np

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


"""
Fixed income analytics for the yield curve

Implements various functions useful for yield curve analytics.
"""

import numpy as np

def bootstrap(cash_flows : np.ndarray, prices : np.ndarray) -> np.ndarray:
    """
    Calculates the spot yield curve from a series of bonds.

    Parameters
    ----------
    cash_flows : np.ndarray
        Nonsingular payoff matrix of a series of bonds.
    prices : np.ndarray
        Corresponding price of each bond.

    Returns
    -------
    np.ndarray
        The corresponding spot yield curve derived from no-arbitrage relationships.
    """

    assert cash_flows.shape[1] == prices.shape[0], "Check shapes of input matrices."

    discount_factors = np.linalg.inv(cash_flows) @ prices
    yields = np.power(discount_factors, np.reshape(-1 / np.arange(start=1, stop=discount_factors.shape[0]+1), newshape=discount_factors.shape)) - 1 # Formula: d = 1/(1 + y)^i
    return np.reshape(yields, newshape=(yields.shape[0])) # Flatten output

def regression():
    raise NotImplementedError("Regression not yet implemented.")

def spline():
    raise NotImplementedError("Spline not yet implemented.")

def nelson_siegel(T : np.ndarray, theta0 : float, theta1 : float, theta2 : float, lambda_ : float) -> np.ndarray:
    """
    Calculates the Nelson-Siegel yield curve.

    Parameters
    ----------
    T : np.ndarray
        List of times at which to calculate yield.
    theta0 : float
        A Nelson-Siegel Model parameter.
    theta1 : float
        A Nelson-Siegel Model parameter.
    theta2 : float
        A Nelson-Siegel Model parameter.
    lambda_ : float
        A Nelson-Siegel Model parameter.

    Returns
    -------
    np.ndarray
        The Nelson-Siegel yield curve.
    """

    return theta0 + (theta1 + theta2) * (1 - np.exp(-T / lambda_)) / (T / lambda_) - theta2 * np.exp(-T / lambda_)


# -*- coding: utf-8 -*-
"""
Created on Sun May 10 16:53:23 2020

@author: eltes
"""

'''
Exercício 01.
Dadas as seguintes séries temporais
    • s1 = 168, 398, 451, 337, 186, 232, 262, 349, 189, 204, 220, 220, 207, 239, 259, 258,
242, 331, 251, 323, 106, 1055, 170

    • s2 = 168, 400, 451, 300, 186, 200, 262, 349, 189, 204, 220, 220, 207, 239, 259, 258,
242, 331, 251, 180, 106, 1055, 200

Utilize funções NumPy para:
• calcular a distância euclidiana entre as séries
• calcular a série temporal com os valores médios entre s1 e s2
• calcular a série temporal com os valores máximos de cada instante entre s1 e s2
• calcular a série temporal com os valores mínimos de cada instante entre s1 e s2
'''

import numpy as np

s1 = (168, 398, 451, 337, 186, 232, 262, 349, 189, 204, 220, 220, 207, 239, 259, 258,
242, 331, 251, 323, 106, 1055, 170)

s2 = (168, 400, 451, 300, 186, 200, 262, 349, 189, 204, 220, 220, 207, 239, 259, 258,
242, 331, 251, 180, 106, 1055, 200)

s1_array = np.array((s1)) # convertendo a lista s1 para array
s2_array = np.array((s2)) # convertendo a lista s2 para array

# Cálculo da Distância Euclidiana
dst = np.linalg.norm(s1_array - s2_array)

# visualizando o resultado
print("s1: "+ str(s1) + "\n\n" + "s2" + str(s2))
print()
print("The Euclidean distance(s1,s2) is:",dst)

# calculo série temporal com os valores médios entre s1 e s2
s_mean= np.mean(np.array([s1_array,s2_array]),axis=0) # Ver referência 1 abaixo
print("\nMean values: " + str(s_mean))

# calculo série temporal com os valores máximos de cada instante entre s1 e s2
s_max = np.maximum(s1_array,s2_array)
print("\nMaximum values: " + str(s_max))

# calculo série temporal com os valores mínimos de cada instante entre s1 e s2
s_min = np.minimum(s1_array,s2_array)
print("\nMinimum values: " + str(s_min))

# Referência 1
# https://stackoverflow.com/questions/19886584/calculate-mean-of-arrays-in-numpy


# coding: utf-8

# # Time Optimal Velocity Profiles
#
# ***
#
# When the maze solver commands that the robot go forward, it can say that it must go forward one or more squares depending on what it knows about the maze. When we don't know what is after the square we pass through, we must be going slow enough to handle any scenario. In other words, there is some $V_f$ that we must reach by the end of our motion. We also begin motions at this speed, since between we arrived where we are we required that we reach $V_f$ to get there. Therefore, we start and end at $V_f$, and we want to cover some distance $d$ in the fast possible time. To do so, we accelerate at our fixed $a$ until we reach max speed, or until we need to start slowing down (whichever comes first). This gives us a trapezoid shaped velocity profile.

# ## Going Straight

# In[5]:

get_ipython().run_line_magic('load_ext', 'tikzmagic')

# In[6]:

get_ipython().run_cell_magic('tikz', '-s 400,400', '\\draw[->] (0,0) -- (10,0);\n\\draw[->] (0,0) -- (0,5);\n\n\\draw[line width=1] (0,0.5) -- (2.5,3);\n\\draw[line width=1] (2.5,3) -- (5.5,3);\n\\draw[line width=1] (5.5,3) -- (8,0.5);\n\\draw[dashed] (0,0.5) -- (10,0.5);\n\\draw[dashed] (0,3) -- (10,3);\n\\draw[dashed] (2.5,0) -- (2.5,5);\n\\draw[dashed] (5.5,0) -- (5.5,5);\n\\draw[dashed] (8,0) -- (8,5);\n\n\\draw (-0.5, 0.5) node {$V_{f}$};\n\\draw (-0.5, 3) node {$V_{max}$};\n\\draw (2.5, -0.5) node {$t_b$};\n\\draw (5.5, -0.5) node {$t_f-t_b$};\n\\draw (8, -0.5) node {$t_f$};')

# The time to accelerate from $V_f$ to $V_{max}$ is $t_b = \frac{V-V_f}{a}$. We can substitute this into newtons first equation of motion as follows.
#
# \begin{align}
# d &= Vt_b - \frac{1}{2}a{t_b}^2 \\
#   &= V\Big(\frac{V-V_f}{a}\Big) - \frac{1}{2}a\Big(\frac{V-V_f}{a}\Big)^2 \\
#   &= \Big(\frac{V^2-VV_f}{a}\Big) - \Big(\frac{a(V-V_f)^2}{2a^2}\Big) \\
#   &= \Big(\frac{2V^2-2VV_f}{2a}\Big) - \Big(\frac{V^2-2VV_f+{V_f}^2}{2a}\Big) \\
#   &= \frac{2V^2-2VV_f - V^2 + 2VV_f - {V_f}^2}{2a} \\
# d &= \frac{V^2-{V_f}^2}{2a} \\
# \end{align}
#
# For example, if you're at starting at $V_f=0.2\frac{m}{s}$, and you're ramping up to $V=0.5\frac{m}{s}$, and you're acceleration is fixed at the $a=2\frac{m}{s^2}$, the distance you'll need to do that is $d = \frac{0.5 - 0.2}{2*2} = 0.075m$

# ## Code that proves it

# In[7]:

# dependencies and global setup
import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(suppress=True, precision=3, linewidth=100)

LOG_LVL = 2

def debug(*args):
    if LOG_LVL <= 0:
        print(*args)

def info(*args):
    if LOG_LVL <= 1:
        print(*args)

def warning(*args):
    if LOG_LVL <= 2:
        print(*args)

def log(*args):
    if LOG_LVL < 100:
        print(*args)

# In[8]:

def profile(V0, Vf, Vmax, d, A, buffer=3e-4):
    v = V0
    x = 0
    a = A
    vs = [v]
    xs = [x]
    a_s = [a]

    dt = 0.001
    while x < d:
        x = x + v*dt + a*dt*dt/2.0
        v = v + a*dt
        ramp_d = (v*v+ - Vf*Vf) / (2.0*A)
        if (d-x) < ramp_d + buffer:
            a = -A
        elif v < Vmax:
            a = A
        else:
            a = 0

        if v > Vmax:
            v = Vmax
        elif v < Vf:
            v = Vf

        xs.append(x)
        vs.append(v)
        a_s.append(a)

    return xs, vs, a_s

def graph(title, idx):
    plt.figure()
    plt.title(title)
    # various max speeds
    Vs = [0.35, 0.5, 0.75, 1, 2]
    Vf = 0.02
    V0 = 0.2
    d = 0.35
    a = 2
    for V in Vs:
        results  = profile(V0, Vf, V, d, a)
        vs = results[1]
        if V == 2: # make V=2 dashed so we can see it over V=1
            plt.plot(results[idx], label='V={}'.format(V), linestyle='dashed')
        else:
            plt.plot(results[idx], label='V={}'.format(V))
        plt.legend(bbox_to_anchor=(1, 1), loc=2)

graph("position", 0)
graph("velocity", 1)
graph("acceleration", 2)
plt.show()

# ## General Form Trajectory Planning

# Let's start out with a generating trajectories that are not time optimal, but rely on specifying the final time $v_f$. For smartmouse, our state space is $[x, y, \theta]$, and a turn can be defined as starting at a point $[x_0, y_0, \theta_0]$ and going to $[x_f, y_f, \theta_0]$. Of course, we also want to specify the velocities at these point, $[\dot{x}_0, \dot{y}_0,\dot{\theta}_0]$ and $[\dot{x}_f, \dot{y}_f,\dot{\theta}_f]$. We have four constraints, so if we want to fit a smooth polynomial to those points we need a 4th order polynomial.
#
# $$q(t) = a_0 + a_1t + a_2t^2 + a_3t^3$$
# $$\dot{q}(t) = a_1 + 2a_2t + 3a_3t^2$$
#
# If we sub in our constraints, we get the following system of equations.
#
# \begin{align}
# q(0) &= a_0 \\
# \dot{q}(0) &= a_1 \\
# q(t_f) &= a_0 + a_1t_f + a_2{t_f}^2 + a_3{t_f}^3\\
# \dot{q}(t_f) &= a_1 + 2a_2t_f + 3a_3{t_f}^2\\
# \end{align}
#
# In matrix form that looks like:
# \begin{equation}
# \begin{bmatrix}
# 1 & 0 & 0 & 0 \\
# 0 & 1 & 0 & 0 \\
# 1 & t_f & t_f^2 & t_f^3 \\
# 0 & 1 & 2t_f & 3t_f^2 \\
# \end{bmatrix}
# \begin{bmatrix}
# a_0 \\
# a_1 \\
# a_2 \\
# a_3 \\
# \end{bmatrix} =
# \begin{bmatrix}
# q(0) \\
# \dot{q}(0) \\
# q(t_f) \\
# \dot{q}(t_f) \\
# \end{bmatrix}
# \end{equation}
#
# It can be shown that the matrix on the left is invertable, so long as $t_f-t_0 > 0$. So we can invert and solve this equation and get all the $a$ coefficients. We can then use this polynomial to generate the $q(t)$ and $\dot{q}(t)$ -- our trajectory.

# In[9]:

def simple_traj_solve(q_0, q_f, q_dot_0, q_dot_t_f, t_f):
    # Example: you are a point in space (one dimension) go from rest at the origin to at rest at (0.18, 0, 0) in 1 second
    q_0 = np.array([0])
    q_dot_0 = np.array([0])
    q_t_f = np.array([0.18])
    q_dot_t_f = np.array([0])

    b = np.array([q_0, q_dot_0, q_t_f, q_dot_t_f])
    a = np.array([[1,0,0,0],[0,1,0,0],[1, t_f, pow(t_f,2),pow(t_f,3)],[0,1,2*t_f,3*pow(t_f,2)]])
    log(a, b)
    coeff = np.linalg.solve(a, b)
    log(coeff)

    return coeff

simple_traj_info = (0, 0, 0.18, 0, 1)
simple_traj_coeff = simple_traj_solve(*simple_traj_info)

# Here you can see that the resulting coeffictions are $a_0=0$, $a_1=0$, $a_2=0.54$, $a_0=-0.36$. Intuitively, this says that we're going to have positive acceleration, but our acceleration is going to slow down over time. Let's graph it!

# In[10]:

def simple_traj_plot(coeff, t_f):
    dt = 0.01
    ts = np.array([[1, t, pow(t,2), pow(t,3)] for t in np.arange(0, t_f+dt,  dt)])
    qs = ts@coeff
    plt.plot(ts[:,1], qs, label="x")
    plt.xlabel("time (seconds)")
    plt.xlabel("X (meters)")
    plt.legend(bbox_to_anchor=(1,1), loc=2)
    plt.show()

simple_traj_plot(simple_traj_coeff, simple_traj_info[-1])

# **ooooooooooh so pretty**
#
# Let's try another example, now with our full state space of $[x, y, \theta]$.

# In[11]:

def no_dynamics():
    # In this example, we go from (0.18, 0.09, 0) to (0.27,0.18, -1.5707). Our starting and ending velocities are zero
    q_0 = np.array([0.09,0.09,0])
    q_dot_0 = np.array([0,0,0])
    q_f = np.array([0.27,0.18,-1.5707])
    q_dot_f = np.array([0,0,0])
    t_f = 1

    b = np.array([q_0, q_dot_0, q_f, q_dot_f])
    a = np.array([[1,0,0,0],[0,1,0,0],[1, t_f, pow(t_f,2),pow(t_f,3)],[0,1,2*t_f,3*pow(t_f,2)]])
    coeff = np.linalg.solve(a, b)
    log(coeff)

    dt = 0.1
    ts = np.array([[1, t, pow(t,2), pow(t,3)] for t in np.arange(0, t_f+dt,  dt)])
    qs = ts@coeff

    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')
    plt.gca().set_adjustable("box")
    plt.subplot(221)
    plt.plot(ts[:,1], qs[:,0])
    plt.xlabel("time (seconds)")
    plt.title("x")
    plt.subplot(222)
    plt.plot(ts[:,1], qs[:,1])
    plt.xlabel("time (seconds)")
    plt.title("y")
    plt.subplot(223)
    plt.plot(ts[:,1], qs[:,2])
    plt.xlabel("time (seconds)")
    plt.title(r"$\theta$")
    plt.subplot(224)
    plt.scatter(qs[:,0], qs[:,1])
    plt.axis('equal')
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.tight_layout()
    plt.show()

no_dynamics()

# Well, they are smooth, but these are not possible to execute! The robot cannot simply translate sideways.

# # Trajectory Planning With a Simple Dynamics Model
#
# ***

# In[12]:

get_ipython().run_cell_magic('tikz', '-s 100,100', '\n\\draw [rotate around={-45:(0,0)}] (-.5,-1) rectangle (0.5,1);\n\\filldraw (0,0) circle (0.125);\n\n\\draw [->] (0,0) -- (0,1.5);\n\\draw [->] (0,0) -- (1.5,0);\n\\draw [->] (0,0) -- (1.5,1.5);\n\\draw (1.2, -0.2) node {$x$};\n\\draw (-0.2, 1.2) node {$y$};\n\\draw (1, 1.2) node {$v$};')

#
# We need to change our constraints to the system of equations. Specifically, we need our dynamics model. For now, let's assume a simplified car model.
#
# $$ \dot{x} = v\cos(\theta) $$
# $$ \dot{y} = v\sin(\theta) $$
#
# This basically claims that for any instant in time the robot is moving a constant velocity along $\theta$. This isn't very accurate, but let's just start with that since the real dynamics of our robot are more complex.
#
# First we will bring in the constraints from before. We must satisfy specific initial and final positions in $[x, y, \theta]$. I've used new letters for cofficients to avoid confusion.
#
# \begin{align}
# x_0 &= c_0 + c_1(0) + c_2(0)^2 + c_3(0)^3 + c_3(0)^4 + c_3(0)^5 \\
# y_0 &= d_0 + d_1(0) + d_2(0)^2 + d_3(0)^3 + d_3(0)^4 + d_3(0)^5 \\
# x_{t_f} &= c_0 + c_1(t_f) + c_2(t_f)^2 + c_3(t_f)^3 + c_3(t_f)^4 + c_3(t_f)^5 \\
# y_{t_f} &= d_0 + d_1(t_f) + d_2(t_f)^2 + d_3(t_f)^3 + d_3(t_f)^4 + d_3(t_f)^5 \\
# \end{align}
#
# Notice here we have 12 unknowns, $c_0 \dots c_5$ and $d_0 \dots d_5$. So we're gonna need more equations for there to be a unique solution. Also notice we haven't defined any constraints related to our dynamics model. That would be a good place to get our other equations!
#
# First, we want to be able to specify initial velocity $v_0$ and final velocity $v_{t_f}$. It is easlier to just constrain $\dot{x}_0$,  $\dot{y}_0$,  $\dot{x}_{t_f}$,  $\dot{y}_{t_f}$. So if we want to specify that we start facing $\tfrac{\pi}{2}$ going 1m/s, we'd just specify $cos(\tfrac{\pi}{2})$ for $\dot{x}_0$ and $sin(\tfrac{\pi}{2})$ for $\dot{y}_0$.
#
# \begin{align}
# \dot{x}_0 &= c_1 \\
# \dot{y}_0 &= d_1 \\
# \dot{x}_{t_f} &= (0)c_0 + (1)c_1 + 2t_fc_2 + 3{t_f}^2c_3 + 4{t_f}^3c_4 + 5{t_f}^4c_5 \\
# \dot{y}_{t_f} &= (0)d_0 + (1)d_1 + 2t_fd_2 + 3{t_f}^2d_3 + 4{t_f}^3d_4 + 5{t_f}^4d_5
# \end{align}
#
# Let's also make sure x and y components obey trigonometry.
#
# \begin{align}
#   v\cos(\theta)\sin(\theta) + v\cos(\theta)\sin(\theta) &= v\sin(2\theta) \\
#   \dot{x}\sin(\theta) + \dot{y}\sin(\theta) &= v\sin(2\theta)
# \end{align}
#
# We can get two equations out of this by specifying initial and final velocities
#
# \begin{align}
# v_0\sin(2\theta_0) &= \dot{x}_0\sin(\theta_0) + \dot{y}_0\cos(\theta_0) \\
# v_{t_f}\sin(2\theta_{t_f}) &= \dot{x}_{t_f}\sin(\theta_{t_f}) + \dot{y}_{t_f}\cos(\theta_{t_f})
# \end{align}
#
# We should write out the full form though, to make things in terms of our coefficients.
#
# \begin{align}
# v(0)\sin(2\theta_0) &= \Big[c_1 + 2(0)c_2 + 3(0)^2c_3 + 4(0)^3c_4 + 5(0)^4c_5\Big]\sin(\theta_0) + \Big[d_1 + 2(0)d_2 + 3(0)^2d_3 + 4(0)^3d_4 + 5(0)^4d_5\Big]\cos(\theta_0) \\
# v(0)\sin(2\theta_0) &= \sin(\theta_0)c_1 + \cos(\theta_0)d_1
# \end{align}
#
# \begin{align}
# v(t_f)\sin(2\theta_{t_f}) &= \Big[c_1 + 2(t_f)c_2 + 3(t_f)^2c_3 + 4(t_f)^3c_4\ + 5(t_f)^4c_5\Big]\sin(\theta_{t_f}) + \Big[d_1 + 2(t_f)d_2 + 3(t_f)^2d_3 + 4(t_f)^3d_4 + 5(t_f)^4d_5\Big]\cos(\theta_{t_f}) \\
# v(t_f)\sin(2\theta_{t_f}) &= \sin(\theta_{t_f})c_1 + 2\sin(\theta_{t_f})t_fc_2 + 3\sin(\theta_{t_f}){t_f}^2c_3 + 4\sin(\theta_{t_f}){t_f}^3c_4  + 5\sin(\theta_{t_f}){t_f}^4c_5 + \cos(\theta_{t_f})d_1 + 2\cos(\theta_{t_f})t_fd_2 + 3\cos(\theta_{t_f}){t_f}^2d_3 + 4\cos(\theta_{t_f}){t_f}^3d_4 + 5\cos(\theta_{t_f}){t_f}^4d_5 \\
# \end{align}
#
# The last two equations constrains the robot from moving in any direction other than its heading. Of course it must relate $\dot{x}$ to $\dot{y}$. Still not totally sure how we got this equation so I'm just copying it from some slides$\dots$. However you can plug in some example values and check. For instance translating sideways violates this equation: set $\dot{x}=1$, $\dot{y}=0$, $v=1$, $\theta=\tfrac{\pi}{2}$.
#
# \begin{align}
# v\cos(\theta)\sin(\theta) - v\cos(\theta)\sin(\theta) &= 0 \\
# v\cos(\theta)\sin(\theta) - v\sin(\theta)\cos(\theta) &= 0 \\
# \dot{x}\sin(\theta) - \dot{y}\cos(\theta) &= 0
# \end{align}
#
# and again written out fully in terms of our coefficients
#
# \begin{align}
# \Big[c_1 + 2(0)c_2 + 3(0)^2c_3 + 4(0)^3c_4 + 5(0)^4c_5\Big]\sin(\theta_0) - \Big[d_1 + 2(0)d_2 + 3(0)^2d_3 + 4(0)^3d_4 + 5(0)^4d_5\Big]\cos(\theta_0) &= 0 \\
# \sin(\theta_0)c_1 - \cos(\theta_0)d_1 &= 0
# \end{align}
#
# \begin{align}
# \Big[c_1 + 2(t_f)c_2 + 3(t_f)^2c_3 + 4(t_f)^3c_4 + 5(t_f)^4c_5\Big]\sin(\theta_{t_f}) - \Big[d_1 + 2(t_f)d_2 + 3(t_f)^2d_3 + 4(t_f)^3d_4 + 5(t_f)^4d_5\Big]\cos(\theta_{t_f}) &= 0 \\
# \sin(\theta_{t_f})c_1 + 2\sin(\theta_{t_f})t_fc_2 + 3\sin(\theta_{t_f}){t_f}^2c_3 + 4\sin(\theta_{t_f}){t_f}^3c_4  + 5\sin(\theta_{t_f}){t_f}^4c_5 - \cos(\theta_{t_f})d_1 - 2\cos(\theta_{t_f})t_fd_2 - 3\cos(\theta_{t_f}){t_f}^2d_3 - 4\cos(\theta_{t_f}){t_f}^3d_4 - 5\cos(\theta_{t_f}){t_f}^4d_5 &= 0
# \end{align}
#
# Ok, that should work. Now let's write it out in matrix form. We use $c$ and $s$ to shorten $\sin$ and $\cos$.
#
# \setcounter{MaxMatrixCols}{20}
# \begin{equation}
# \begin{bmatrix}
# 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\
# 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0 \\
# 0 & s(\theta_0) & 0 & 0 & 0 & 0 & 0 & c(\theta_0) & 0 & 0 & 0 & 0\\
# 0 & s(\theta_0) & 0 & 0 & 0 & 0 & 0 & -c(\theta_0) & 0 & 0 & 0 & 0\\
# 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\
# 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0\\
# 1 & t & {t_f}^2 & {t_f}^3 & {t_f}^4 & {t_f}^5 & 0 & 0 & 0 & 0 & 0 & 0 \\
# 0 & 0 & 0 & 0 & 0 & 0 & 1 & t_f & {t_f}^2 & {t_f}^3 & {t_f}^4 & {t_f}^5 \\
# 0 & s(\theta_{t_f}) & 2s(\theta_{t_f})t_f & 3s(\theta_{t_f}){t_f}^2 & 4s(\theta_{t_f}){t_f}^3 & 5s(\theta_{t_f}){t_f}^4 & 0 & c(\theta_{t_f}) & 2c(\theta_{t_f}){t_f} & 3c(\theta_{t_f}){t_f}^2 & 4c(\theta_{t_f}){t_f}^3 & 5c(\theta_{t_f}){t_f}^4 \\
# 0 & s(\theta_{t_f}) & 2s(\theta_{t_f})t_f & 3s(\theta_{t_f}){t_f}^2 & 4s(\theta_{t_f}){t_f}^3 & 5s(\theta_{t_f}){t_f}^4 & 0 & -c(\theta_{t_f}) & -2c(\theta_{t_f}){t_f} & -3c(\theta_{t_f}){t_f}^2 & -4c(\theta_{t_f}){t_f}^3 & -5c(\theta_{t_f}){t_f}^4 \\
# 0 & 1 & 2t_f & 3{t_f}^2 & 4{t_f}^3 & 5{t_f}^4 & 0 & 0 & 0 & 0 & 0 & 0 \\
# 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1 & 2t_f & 3{t_f}^2 & 4{t_f}^3 & 5{t_f}^4
# \end{bmatrix}
# \begin{bmatrix}
# c_0 \\
# c_1 \\
# c_2 \\
# c_3 \\
# c_4 \\
# c_5 \\
# d_0 \\
# d_1 \\
# d_2 \\
# d_3 \\
# d_4 \\
# d_5
# \end{bmatrix} =
# \begin{bmatrix}
# x_0 \\
# y_0 \\
# 0 \\
# v_0s(2\theta_0) \\
# c(\theta_0)v_0 \\
# s(\theta_0)v_0 \\
# x_{t_f} \\
# y_{t_f} \\
# 0 \\
# v_{t_f}s(2\theta_{t_f}) \\
# c(\theta_{t_f})v_{t_f} \\
# s(\theta_{t_f})v_{t_f} \\
# \end{bmatrix}
# \end{equation}

# In[13]:

# Let's solve this in code like we did before
def plot_vars(traj_plan):
    dt = 0.001
    T = np.arange(0, traj_plan.get_t_f()+dt, dt)
    xts = np.array([[1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5), 0, 0, 0, 0, 0, 0] for t in T])
    xdts = np.array([[0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4), 0, 0, 0, 0, 0, 0] for t in T])
    yts = np.array([[0, 0, 0, 0, 0, 0, 1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5)] for t in T])
    ydts = np.array([[0, 0, 0, 0, 0, 0, 0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4)] for t in T])
    xs = xts@traj_plan.get_coeff()
    ys = yts@traj_plan.get_coeff()
    xds = xdts@traj_plan.get_coeff()
    yds = ydts@traj_plan.get_coeff()

    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')
    plt.rc('axes.formatter', useoffset=False)
    plt.figure(figsize=(10, 2.5))

    plt.subplot(141)
    plt.plot(T, xs, linewidth=3)
    plt.xlabel("time (seconds)")
    plt.title("X")

    plt.subplot(142)
    plt.plot(T, ys, linewidth=3, color='r')
    plt.xlabel("time (seconds)")
    plt.title("Y")

    plt.subplot(143)
    plt.plot(T, xds, linewidth=3, color='g')
    plt.xlabel("time (seconds)")
    plt.title("$\dot{x}$")
    plt.tight_layout()

    plt.subplot(144)
    plt.plot(T,yds, linewidth=3, color='y')
    plt.xlabel("time (seconds)")
    plt.title("$\dot{y}$")
    plt.tight_layout()
    plt.show()

def plot_traj(traj_plan):
    dt = 0.03
    T = np.arange(0, traj_plan.get_t_f()+dt, dt)
    xts = np.array([[1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5), 0, 0, 0, 0, 0, 0] for t in T])
    yts = np.array([[0, 0, 0, 0, 0, 0, 1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5)] for t in T])
    xs = xts@traj_plan.get_coeff()
    ys = yts@traj_plan.get_coeff()
    plot_traj_pts(xs, ys, T, traj_plan.waypoints)

def plot_traj_pts(xs, ys, T, waypoints):
    plt.figure(figsize=(5, 5))
    plt.title("Trajectory")
    plt.xlabel("X")
    plt.ylabel("Y")
    W = 2
    plt.xlim(0, W * 0.18)
    plt.ylim(0, W * 0.18)
    plt.xticks(np.arange(2*W+1)*0.09)
    plt.yticks(np.arange(2*W+1)*0.09)
    plt.grid(True)
    plt.gca().set_axisbelow(True)

    for t, pt in waypoints:
        arrow_dx = cos(pt.theta) * (pt.v) * 0.1
        arrow_dy = sin(pt.theta) * (pt.v) * 0.1
        plt.arrow(pt.x, pt.y, arrow_dx, arrow_dy, head_width=0.005, head_length=0.005, width=0.001, fc='k', ec='k')

    plt.scatter(xs, ys, marker='.', linewidth=0)

    plt.show()

# In[14]:

from math import sin, cos, pi
from collections import namedtuple

WayPoint = namedtuple('WayPoint', ['x', 'y', 'theta', 'v'])

class TrajPlan:
    def x_constraint(t):
        return [1, t, pow(t, 2), pow(t, 3), pow(t, 4), pow(t, 5), 0, 0, 0, 0, 0, 0]

    def y_constraint(t):
        return [0, 0, 0, 0, 0, 0, 1, t, pow(t, 2), pow(t, 3), pow(t, 4), pow(t, 5)]

    def non_holonomic_constraint(theta_t, t):
        s_t = sin(theta_t)
        c_t = cos(theta_t)
        t_2 = pow(t, 2)
        t_3 = pow(t, 3)
        t_4 = pow(t, 4)
        return [0, s_t, 2 * s_t * t, 3 * s_t * t_2, 4 * s_t * t_3, 5 * s_t * t_4, 0, c_t, 2 * c_t * t, 3 * c_t * t_2, 4 * c_t * t_3, 5 * c_t * t_4]

    def trig_constraint(theta_t, t):
        s_t = sin(theta_t)
        c_t = cos(theta_t)
        t_2 = pow(t, 2)
        t_3 = pow(t, 3)
        t_4 = pow(t, 4)
        return [0, s_t, 2 * s_t * t, 3 * s_t * t_2, 4 * s_t * t_3, 5 * s_t * t_4, 0, -c_t, -2 * c_t * t, -3 * c_t * t_2, -4 * c_t * t_3, -5 * c_t * t_4]

    def x_dot_constraint(t):
        return [0, 1, 2 * t, 3 * pow(t, 2), 4 * pow(t, 3), 5 * pow(t, 4), 0, 0, 0, 0, 0, 0]

    def y_dot_constraint(t):
        return [0, 0, 0, 0, 0, 0, 0, 1, 2 * t, 3 * pow(t, 2), 4 * pow(t, 3), 5 * pow(t, 4)]

    def solve(self, waypoints):
        # Setup the matrices to match the equation above
        A = []
        b = []

        for t, pt in waypoints:
            A += [TrajPlan.x_constraint(t),
                      TrajPlan.y_constraint(t),
                      TrajPlan.non_holonomic_constraint(pt.theta, t),
                      TrajPlan.trig_constraint(pt.theta, t),
                      TrajPlan.x_dot_constraint(t),
                      TrajPlan.y_dot_constraint(t)]
            b += [pt.x,
                      pt.y,
                      0,
                      pt.v*sin(2*pt.theta),
                      cos(pt.theta)*pt.v,
                      sin(pt.theta)*pt.v]

        A = np.array(A)
        b = np.array(b)
        rank = np.linalg.matrix_rank(A)

        if rank == A.shape[1]:
            if A.shape[0] == A.shape[1]:
                coeff = np.linalg.solve(A, b)
            else:
                warning("not square, using least squares.".format(A.shape))
                coeff, resid, rank, s = np.linalg.lstsq(A, b)
        else:
            warning("Ranks don't match! {} equations {} variables, using least squares".format(rank, A.shape[1]))
            coeff, resid, rank, s = np.linalg.lstsq(A, b)

        debug("rank {}".format(rank))
        debug("A: \n{}".format(A))
        debug("coeff: \n{}".format(coeff))
        error = np.sum(np.power(A@coeff - b, 2))
        if error > 1e-10:
            info("These two vectors should be equal! But there is error.")
            info("b is: \n{}".format(b))
            info("A@coeff is: \n{}".format(A@coeff))
        info("RMS Error of solution to equations")
        info(error)

        self.coeff = coeff
        self.waypoints = waypoints

    def get_coeff(self):
        return self.coeff

    def get_t_f(self):
        return self.waypoints[-1][0]

# ## Example Plots

# In[15]:

# forward 1 cell, start from rest, end at 40cm/s, do it in .5 seconds
LOG_LVL = 5
fwd_1 = TrajPlan()
fwd_1.solve([(0, WayPoint(0.09, 0.09, pi/2, 0)), (0.5, WayPoint(0.09, 0.27, pi/2, 0.6))])
plot_vars(fwd_1)
plot_traj(fwd_1)

# In[16]:

# continue by turning right 90 degrees
LOG_LVL = 1
turn_right = TrajPlan()
turn_right.solve([(0, WayPoint(0.09, 0.18, pi/2, 0.4)), (0.5, WayPoint(0.18, 0.27, 0, 0.4))])
plot_vars(turn_right)
plot_traj(turn_right)

# In[17]:

# 3 waypoints!
LOG_LVL = 1
turn_right = TrajPlan()
turn_right.solve([(0, WayPoint(0.09, 0.09, pi/2, 0.0)), (0.5, WayPoint(0.18, 0.18, 0, 0.35)), (1, WayPoint(0.27, 0.27, pi/2, 0))])
plot_vars(turn_right)
plot_traj(turn_right)

# **Note for this system of equations with 3 waypoints, there is no solution. However, the error of the solution found is very small.**
#
# Now let's find one that really sucks!

# In[18]:

# 4 waypoints!
LOG_LVL = 1
turn_right = TrajPlan()
turn_right.solve([(0, WayPoint(0.09, 0.0, pi/2, 0.1)),
                  (1, WayPoint(0.09, 0.18, pi/2, 0.1)),
                  (2, WayPoint(0.18, 0.27, 0, 0.1)),
                  (3, WayPoint(0.27, 0.27, 0, 0.1))])
plot_traj(turn_right)

# # Trajectory Following
#
# ***
#
# Now that we have a trajectory, we want to design a controller that will follow it as closely as possible. To do this, I'm just going to do a proportional controller. Later we will design an optimal controller. We want to make sure the robot is on the path, facing along the path, and going the right speed. When all of these are true the change in speed should be zero. Let's come up with an equation to relate current pose and velocity to the desired pose and velocity. Let our outputs be the linear velocity $v$ and the rotational velocity $w$.
#
# $$ w = \bar{w} + d*P_1 + (\bar{\theta} - \theta)P_2$$
# $$ v = \bar{v} + l*P_3$$
#
# where $v_d$ is desired velocity, $\theta_d$ is the desired angle, $d$ is signed distance to the planned trajectory (to the right of the plan is positive), $v_d$ and $w_d$ are the desired velocities of the robot, and $P_1$, $P_2$, and $P_3$ are constants. Essentially what we're saying with the first equation is that when you're far off the trajectory you need to turn harder to get back on to it, but you also need to be aligned with it. The second equation says if you're lagging behind your plan speed up, or slow down if you're overshooting.

# In[19]:

from math import atan2, sqrt

LOG_LVL = 5
def simulate(q_0, waypoints, P_1, P_2, P_3, A=3):
    traj = TrajPlan()
    traj.solve(waypoints)

    dt = 0.01
    x = q_0[0]
    y = q_0[1]
    theta = q_0[2]
    v = q_0[3]
    w = q_0[4]
    actual_v = q_0[3]
    actual_w = q_0[4]
    v_acc = A * dt
    TRACK_WIDTH_M = 0.0633
    w_acc = v_acc / (TRACK_WIDTH_M/2)
    T = np.arange(0, traj.get_t_f()+dt, dt)
    x_bar_list = []
    y_bar_list = []
    x_list = []
    y_list = []
    for t in T:
        x_bar = [1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5), 0, 0, 0, 0, 0, 0] @ traj.get_coeff()
        dx_bar = [0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4), 0, 0, 0, 0, 0, 0] @ traj.get_coeff()
        ddx_bar = [0, 0, 0, 0, 0, 0, 0, 0, 2, 6*t, 12*pow(t,2), 20*pow(t,3)] @ traj.get_coeff()
        y_bar = [0, 0, 0, 0, 0, 0, 1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5)] @ traj.get_coeff()
        dy_bar = [0, 0, 0, 0, 0, 0, 0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4)] @ traj.get_coeff()
        ddy_bar = [0, 0, 0, 0, 0, 0, 0, 0, 2, 6*t, 12*pow(t,2), 20*pow(t,3)] @ traj.get_coeff()
        theta_bar = atan2(dy_bar, dx_bar)
        v_bar = sqrt(dx_bar*dx_bar + dy_bar*dy_bar)
        w_bar = 1/v_bar * (ddy_bar*cos(theta_bar) - ddx_bar*sin(theta_bar));

        # simple Dubin's Car forward kinematics
        x += cos(theta) * actual_v * dt
        y += sin(theta) * actual_v * dt
        theta += actual_w * dt

        # control
        euclidian_error = np.sqrt(pow(x_bar - x, 2) + pow(y_bar - y, 2))
        transformed_x = (x - x_bar) * cos(-theta_bar) + (y - y_bar) * -sin(-theta_bar)
        transformed_y = (x - x_bar) * sin(-theta_bar) + (y - y_bar) * cos(-theta_bar)
        right_of_traj = transformed_y < 0
        signed_euclidian_error = euclidian_error if right_of_traj else -euclidian_error
        lag_error = -transformed_x
        w = w_bar + signed_euclidian_error * P_1 + (theta_bar - theta) * P_2
        v = v_bar + lag_error * P_3

        # simple acceleration model
        if v < actual_v:
            actual_v = max(v, actual_v - v_acc)
        elif v > actual_v:
            actual_v = min(v, actual_v + v_acc)
        if w < actual_w:
            actual_w = max(w, actual_w - w_acc)
        elif w > actual_w:
            actual_w = min(w, actual_w + w_acc)

        x_bar_list.append(x_bar)
        y_bar_list.append(y_bar)
        x_list.append(x)
        y_list.append(y)

    plt.figure(figsize=(5, 5))
    W = 3
    plt.scatter(x_bar_list, y_bar_list, marker='.', linewidth=0, c='black', label='desired traj')
    plt.scatter(x_list, y_list, marker='.', linewidth=0, c=T, label='robot traj')
    plt.xlim(0, W * 0.18)
    plt.ylim(0, W * 0.18)
    plt.xticks(np.arange(2*W+1)*0.09)
    plt.yticks(np.arange(2*W+1)*0.09)
    plt.grid(True)
    plt.gca().set_axisbelow(True)
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.title("Trajectory Tracking")
    plt.legend(bbox_to_anchor=(1,1), loc=2)

# In[20]:

test_P_1=300
test_P_2=50
test_P_3=10
robot_q_0 = (0.08, 0.18, pi/2, 0.3, 0)
traj = [(0, WayPoint(0.09, 0.18, pi/2, 0.5)), (0.5, WayPoint(0.18, 0.27, 0, 0.35)), (1, WayPoint(0.27, 0.36, pi/2, 0))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[21]:

robot_q_0 = (0.11, 0.18, pi/2, 0.2, 5)
traj = [(0, WayPoint(0.09, 0.18, pi/2, 0.2)), (1, WayPoint(0.18, 0.27, 0, 0.35))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[22]:

robot_q_0 = (0.0, 0.25, 0, 0.2, 0)
traj = [(0, WayPoint(0.0, 0.27, 0, 0.2)), (1.25, WayPoint(0.54, 0.27, 0, 0.2))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[23]:

robot_q_0 = (0.45, 0.05, pi+0.25, 0.3, 0)
traj = [(0, WayPoint(0.45, 0.09, pi, 0.4)), (0.75, WayPoint(0.27, 0.27, pi/2, 0.4))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[24]:

robot_q_0 = (0.0, 0.25, 0, 0.2, -5)
traj = [(0, WayPoint(0.0, 0.27, 0, 0.2)), (2, WayPoint(0.48, 0.36, pi/2, 0.2))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[25]:

robot_q_0 = (0.25, 0.28, -pi*4/7, 0.5, 0)
traj = [(0, WayPoint(0.27, 0.27, -pi/2, 0.8)), (0.35, WayPoint(0.45, 0.09, 0, 0.8))]
simulate(robot_q_0, traj, test_P_1, test_P_2, test_P_3, A=6)
plt.show()

# In[26]:

# no initial error
robot_q_0 = (0.11, 0.18, pi/2, 0.8, 0)
traj = [(0, WayPoint(0.09, 0.18, pi/2, 0.8)), (0.25, WayPoint(0.18, 0.27, 0, 0.6)), (.5, WayPoint(0.27, 0.36, pi/2, 0.4))]
simulate(robot_q_0, traj, 10, 1000, 6, A=5)
plt.show()

# **Note**: The code above has a bug if I use `-pi` instead of `pi` in `robot_q_0`

# # LQR - The Optimal Controller
#
# ***
#
# ## Overview of the Steps:
#
# ### 1. Write out the non-linear dynamics $\dot{\vec{x}} = f(\vec{x}, \vec{u})$
#
# Here we are interested in the full blown system dynamics of the actual smartmouse robot. The forward kinematics, which depend on the current state $x$, $y$, and $\theta$ and the velocity inputs of the wheels $v_l$, and $v_r$ are as follows. In the general case where the two wheels have different velocities, we have this:
#
# \begin{align}
#   R &= \frac{W(v_l+v_r)}{2(v_r-v_l)} && \text{radius of turn} \\
#   \theta &\leftarrow \theta + \dfrac{v_l}{R-\frac{W}{2}}\Delta t \\
#   x &\leftarrow x-R\Bigg(\sin{\Big(\frac{v_r-v_l}{W}\Delta t-\theta\Big)}+\sin{\theta}\Bigg) \\
#   y &\leftarrow y-R\Bigg(\cos{\Big(\frac{v_r-v_l}{W}\Delta t-\theta\Big)}-\cos{\theta}\Bigg)
# \end{align}
#
# And in the special case where we're going perfectly straight:
#
# \begin{align}
#   \theta &\leftarrow \theta \\
#   x &\leftarrow x + v\Delta t\cos(\theta) \\
#   y &\leftarrow y + v\Delta t\sin(\theta) \\
# \end{align}
#
# We can take these equations and write them in the form of $\dot{\vec{x}} = f(\vec{x},\vec{u})$. Confusingly, $\vec{x}$ here is the full state vector $[x, y, \theta]$. Most controls texts simply use $\vec{x}$, so I'm sticking with that. Also, we defined $u = [v_l, v_r]$
#
# \begin{align}
#   \dot{x} &= \begin{bmatrix}\dot{x}\\ \dot{y}\\ \dot{\theta}\end{bmatrix} \\
#           &= \begin{bmatrix}
#                 -R\Bigg(\sin{\Big(\frac{v_r-v_l}{W}\Delta t-\theta\Big)}+\sin{\theta}\Bigg) \\
#                 -R\Bigg(\cos{\Big(\frac{v_r-v_l}{W}\Delta t-\theta\Big)}-\cos{\theta}\Bigg) \\
#                 \frac{v_l}{R-\frac{W}{2}}\Delta t \\
#                 \end{bmatrix}
# \end{align}
#
# ### 2. Identify the points around which we linearize our system, $(\bar{u}, \bar{x})$
#
# Because we are tracking a trajectory, we want to linearize around the trajectory we are trying to track. That means $\bar{u}$ is the control input associated with the trajectory, which means solving for the feed forward inputs. Specifically, that means we need to compute the $v_l$ and $v_r$ that would follow the trajectory at the point $\bar{x}$. To do this we must pick velocities that make instantaneous turning radius $R$ equal the instantaneous radius of the trajcetory at $\bar{x}$, and make the linear velocity at $\bar{x}$ equal the instantaneous linear velocity of the robot center $v$. To do this, we go back to our basic kinematics equations which rely on the fact that all points on the robot (center, left wheel, right wheel) have the same rotational velocity $\omega$ around the ICC.
#
# \begin{align}
#   \omega = \frac{v}{R} &= \frac{v_l}{R-\frac{W}{2}} \\
#   \frac{v}{R}\bigg(R - \frac{W}{2}\bigg) &= v_l \\
#   \omega = \frac{v}{R} &= \frac{v_r}{R+\frac{W}{2}} \\
#   \frac{v}{R}\bigg(R + \frac{W}{2}\bigg) &= v_r \\
# \end{align}
#
# Using these equations we can solve for the velocities of the wheels, which together make up $\bar{u}$. We just need the $R$ and $v$. These should be derived from the equation of the trajectory we are tracking. These are well studied equations, for which [a proof can be found other places on the internet](http://mathworld.wolfram.com/Curvature.html).
#
# $$ R = \frac{\dot{x}\ddot{y}-\dot{y}\ddot{x}}{{\big({\dot{x}}^2 + {\dot{y}}^2\big)}^{\frac{3}{2}}} = \frac{(c_1+2c_2t+3c_3t^2+4c_4t^3)(2d_2+6d_3t+12d_4t^2) - (d_1+2d_2t+3d_3t^3+4d_4t^3)(2c_2+6c_3t+12c_4t^2)}{{\big({(c_1+2c_2t+3c_3t^2+4c_4t^3)}^2 + {(d_1+2d_2t+3d_3t^3+4d_4t^3)}^2\big)}^{\frac{3}{2}}} $$
#
# $$ v = \sqrt{{(\dot{x})}^2 + {(\dot{y})}^2}  = \sqrt{{(c_1+2c_2t+3c_3t^2+4c_4t^3)}^2 + {(d_1+2d_2t+3d_3t^3+4d_4t^3)}^2} $$
#
# We can plug in the coefficients of our polynomials and get values for $v$ and $R$. Then, we can plus these into the equations just above and get the feed forward wheel velocities.
#
# ### 3. Write the linearized dynamics around $\bar{x}$ as $\dot{\vec{x}} \approx A\delta_x + B\delta_u$, where $\delta_x = (\vec{x} - \bar{x})$ and $\delta_u = (\vec{u} - \bar{u})$
#
# To do this, we need the partial derivitive matrixes $A$ and $B$.
#
# $$ A = \begin{bmatrix}
#          \frac{\partial f_1}{\partial x}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_1}{\partial y}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_1}{\partial \theta}\big|_{\bar{x}\bar{u}} \\
#          \frac{\partial f_2}{\partial x}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_2}{\partial y}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_2}{\partial \theta}\big|_{\bar{x}\bar{u}} \\
#          \frac{\partial f_3}{\partial x}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_3}{\partial y}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_3}{\partial \theta}\big|_{\bar{x}\bar{u}} \\
#        \end{bmatrix}
#      = \begin{bmatrix}
#         0 & 0 & R\bigg(\cos\Big(\frac{\bar{v}_r-\bar{v}_l}{W}\Delta t - \bar{\theta}\Big) - \cos(\bar{\theta})\bigg) \\
#         0 & 0 & -R\bigg(\sin\Big(\frac{\bar{v}_r-\bar{v}_l}{W}\Delta t - \bar{\theta}\Big) + \sin(\bar{\theta})\bigg) \\
#         0 & 0 & 0 \\
#        \end{bmatrix}
# $$
#
# $$ B = \begin{bmatrix}
#          \frac{\partial f_1}{\partial v_l}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_1}{\partial v_r}\big|_{\bar{x}\bar{u}} \\
#          \frac{\partial f_2}{\partial v_l}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_2}{\partial v_r}\big|_{\bar{x}\bar{u}} \\
#          \frac{\partial f_3}{\partial v_l}\big|_{\bar{x}\bar{u}} &
#          \frac{\partial f_3}{\partial v_r}\big|_{\bar{x}\bar{u}} \\
#        \end{bmatrix}
#      = \begin{bmatrix}
#           R\cos\Big(\frac{(\bar{v}_r-\bar{v}_l)\Delta t}{W}-\bar{\theta}\Big)\frac{\Delta t}{W} &
#           -R\cos\Big(\frac{(\bar{v}_r-\bar{v}_l)\Delta t}{W}-\bar{\theta}\Big)\frac{\Delta t}{W} \\
#           -R\sin\Big(\frac{(\bar{v}_r-\bar{v}_l)\Delta t}{W}-\bar{\theta}\Big)\frac{\Delta t}{W} &
#           R\sin\Big(\frac{(\bar{v}_r-\bar{v}_l)\Delta t}{W}-\bar{\theta}\Big)\frac{\Delta t}{W} \\
#           \frac{\Delta t}{R-\frac{W}{2}} &
#           0 \\
#        \end{bmatrix}
# $$
#
# ### 4. Check if our system is controllable by looking at the rank of the controllability matrix $C = [B, AB, A^2B, \dots, A^{n-1}B]$
#
# We have three state variables so $n = 3$, which means $C = [B, AB, A^2B]$.
#
# $$
# AB = \begin{bmatrix}
#        R\bigg(\cos\Big(\frac{v_r-v_l}{W}\Delta t - \theta\Big) - \cos(\theta)\bigg)\frac{\Delta t}{R-\frac{W}{2}} & 0 \\
#        -R\bigg(\sin\Big(\frac{v_r-v_l}{W}\Delta t - \theta\Big) + \sin(\theta)\bigg)\frac{\Delta t}{R-\frac{W}{2}} & 0 \\
#        0 & 0 \\
#      \end{bmatrix}
# $$
#
#
# $$
# A^2B =
#      \begin{bmatrix}
#        0 & 0 & 0\\
#        0 & 0 & 0\\
#        0 & 0 & 0\\
#      \end{bmatrix}
#      * B
#      = \begin{bmatrix}
#        0 & 0 & 0\\
#        0 & 0 & 0\\
#        0 & 0 & 0\\
#      \end{bmatrix}
# $$
#
# $$ C = \begin{bmatrix}
#          \begin{bmatrix}
#            R\cos\Big(\frac{(\bar{v_r}-\bar{v_l})\Delta t}{W}-\theta\Big)\frac{\Delta t}{W} &
#            -R\cos\Big(\frac{(\bar{v_r}-\bar{v_l})\Delta t}{W}-\theta\Big)\frac{\Delta t}{W} \\
#            -R\sin\Big(\frac{(\bar{v_r}-\bar{v_l})\Delta t}{W}-\theta\Big)\frac{\Delta t}{W} &
#            R\sin\Big(\frac{(\bar{v_r}-\bar{v_l})\Delta t}{W}-\theta\Big)\frac{\Delta t}{W} \\
#            \frac{\Delta t}{R-\frac{W}{2}} &
#            0 \\
#          \end{bmatrix} &
#          \begin{bmatrix}
#            R\bigg(\cos\Big(\frac{v_r-v_l}{W}\Delta t - \theta\Big) - \cos(\theta)\bigg)\frac{\Delta t}{R-\frac{W}{2}} & 0 \\
#            -R\bigg(\sin\Big(\frac{v_r-v_l}{W}\Delta t - \theta\Big) + \sin(\theta)\bigg)\frac{\Delta t}{R-\frac{W}{2}} & 0 \\
#            0 & 0 \\
#          \end{bmatrix} &
#          \begin{bmatrix}
#            0 & 0 & 0\\
#            0 & 0 & 0\\
#            0 & 0 & 0\\
#          \end{bmatrix}
#        \end{bmatrix}
# $$
#
# What is the rank of this matrix? It seems to depend on specific values of $x, y, \theta$.
#
# ### 5. Pick cost parameters $Q$ and $R$
#
# These need to be tuned on the simulation or real system, but the identity matrices $I$ are good starting points.
#
# ### 6. Solve for $K$ given $LQR(A, B, Q, R)$
#
# We want to minimize the quadratice cost function $J$, which is defined as follows.
#
# $$ J = \sum_0^N(\vec{x}_t - \bar{x}_t)^TQ(\vec{x}_t - \bar{x}_t) + \sum_0^N(\vec{u}_t - \bar{u}_t)^TR(\vec{u}_t - \bar{u}_t) $$
#
# We can do this with least squares, or dynamics programming. DP is more efficient O(Nn^3). N is some finite horizon, and n is the number of state dimensions (3 for us).
#
# maybe we compute K once instead of at every time step. could be consistent within one motion primitive.
#
# ### 7. Apply our new controller of the form $\vec{u} = -K(\vec{x} - \bar{x}) + \bar{u}$
#

# In[27]:

from math import atan2
import scipy.linalg

# source: http://www.kostasalexis.com/lqr-control.html
def dlqr(A,B,Q,R):
    """Solve the discrete time lqr controller.

    p[k+1] = A p[k] + B u[k]

    cost = sum p[k].T*Q*p[k] + u[k].T*R*u[k]
    """
    #ref Bertsekas, p.151

    #first, try to solve the ricatti equation
    P = np.matrix(scipy.linalg.solve_discrete_are(A, B, Q, R))

    #compute the LQR gain
    K = np.matrix(scipy.linalg.inv(B.T*P*B+R)*(B.T*P*A))

    eigVals, eigVecs = scipy.linalg.eig(A-B*K)

    return K, P, eigVals

def follow_plan(q_0, waypoints, P_1, P_2, P_3):
    traj = TrajPlan()
    traj.solve(waypoints)

    dt = 0.01
    x = q_0[0]
    y = q_0[1]
    theta = q_0[2]
    vl = q_0[3]
    vr = q_0[3]
    actual_vl = vl
    actual_vr = vr
    v_acc = 2 * dt
    W = 0.0633
    T = np.arange(0, traj.get_t_f()+dt, dt)
    x_bar_list = []
    y_bar_list = []
    x_list = []
    y_list = []
    vl_list = []
    vr_list = []
    actual_vl_list = []
    actual_vr_list = []
    for t in T:
        x_bar = [1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5), 0, 0, 0, 0, 0, 0] @ traj.get_coeff()
        dx_bar = [0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4), 0, 0, 0, 0, 0, 0] @ traj.get_coeff()
        ddx_bar = [0, 0, 0, 0, 0, 0, 0, 0, 2, 6*t, 12*pow(t,2), 20*pow(t,3)] @ traj.get_coeff()
        y_bar = [0, 0, 0, 0, 0, 0, 1, t, pow(t,2), pow(t,3), pow(t,4), pow(t,5)] @ traj.get_coeff()
        dy_bar = [0, 0, 0, 0, 0, 0, 0, 1, 2*t, 3*pow(t,2), 4*pow(t,3), 5*pow(t,4)] @ traj.get_coeff()
        ddy_bar = [0, 0, 0, 0, 0, 0, 0, 0, 2, 6*t, 12*pow(t,2), 20*pow(t,3)] @ traj.get_coeff()
        theta_bar = atan2(dy_bar, dx_bar)

        # full forward kinematics
        if vr - vl < 1e-5:
            x = x + cos(theta) * vl * dt
            y = y + sin(theta) * vl * dt
        else:
            R = W*(vl + vr)/(2*(vr - vl))
            x = x - R * (sin((vr-vl)*dt/W - theta) + sin(theta))
            y = y - R * (cos((vr-vl)*dt/W - theta) - cos(theta))
            theta = theta + vl / (R - W/2) * dt

        # compute instanteneous Radius
        R_bar = (dx_bar*ddy_bar - dy_bar*ddy_bar)/pow((pow(dx_bar, 2) + pow(dy_bar, 2)), 3/2)

        # feed forward inputs
        v_bar = np.sqrt(dx_bar*dx_bar + dy_bar*dy_bar)
        vl_bar = v_bar/R_bar*(R_bar-W/2)
        vr_bar = v_bar/R_bar*(R_bar+W/2)

        A = np.array([[0, 0, R_bar*(cos((vr_bar - vl_bar)*dt/W - theta_bar) - cos(theta_bar))],
                      [0, 0, -R_bar*(sin((vr_bar - vl_bar)*dt/W - theta_bar) + sin(theta_bar))],
                      [0, 0, 0]])

        B = np.array([[R_bar*cos((vr_bar - vl_bar)*dt/W - theta_bar)*dt/W, -R_bar*cos((vr_bar - vl_bar)*dt/W - theta_bar)*dt/W],
                      [-R_bar*sin((vr_bar - vl_bar)*dt/W - theta_bar)*dt/W, R_bar*sin((vr_bar - vl_bar)*dt/W - theta_bar)*dt/W],
                      [dt/(R_bar-W/2), 0]]);

        Q= np.eye(3);
        R = np.eye(2);

        K, P, eigs = dlqr(A, B, Q, R)
        eigs = np.linalg.eig(A - B*K)
#         info("eigs", eigs[0])
#         debug("K", K)
        x_vec = np.array([[x],[y],[theta]])
        x_bar_vec = np.array([[x_bar],[y_bar],[theta_bar]])
        u = -K * (x_vec - x_bar_vec) + np.array([[vl_bar],[vr_bar]]);
        vl = u[0,0]
        vr = u[1,0]

        # simple acceleration model
        if vl < actual_vl:
            actual_vl = max(vl, actual_vl - v_acc)
        elif vl > actual_vl:
            actual_vl = min(vl, actual_vl + v_acc)
        if vr < actual_vr:
            actual_vr = max(vr, actual_vr - v_acc)
        elif vr > actual_vr:
            actual_vr = min(vr, actual_vr + v_acc)

        x_bar_list.append(x_bar)
        y_bar_list.append(y_bar)
        x_list.append(x)
        y_list.append(y)
        vr_list.append(vr)
        vl_list.append(vl)
        actual_vr_list.append(actual_vr)
        actual_vl_list.append(actual_vl)

    plt.figure(figsize=(5, 5))
    CELL_COUNT = 3
    plt.scatter(x_bar_list, y_bar_list, marker='.', linewidth=0, c='black', label='desired traj')
    plt.scatter(x_list, y_list, marker='.', linewidth=0, label='robot traj')
    plt.xlim(0, CELL_COUNT * 0.18)
    plt.ylim(0, CELL_COUNT * 0.18)
    plt.xticks(np.arange(CELL_COUNT+1)*0.18)
    plt.yticks(np.arange(CELL_COUNT+1)*0.18)
    plt.grid(True)
    plt.gca().set_axisbelow(True)
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.title("LQR Trajectory Tracking")
    plt.legend(bbox_to_anchor=(1,1), loc=2)

    plt.figure()
    plt.plot(vr_list, label="vr")
    plt.plot(vl_list, label="vl")
    plt.plot(actual_vr_list, label="actual vr")
    plt.plot(actual_vl_list, label="actual vl")
    plt.legend(bbox_to_anchor=(1,1), loc=2)

# In[28]:

LOG_LVL=1
robot_q_0 = (0.08, 0.18, pi/2, 0.3)
traj = [(0, WayPoint(0.09, 0.18, pi/2, 0.5)), (0.5, WayPoint(0.18, 0.27, 0, 0.35))]
follow_plan(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# In[29]:

LOG_LVL=1
robot_q_0 = (0.07, 0.18, pi/2, 0.2)
traj = [(0, WayPoint(0.09, 0.18, pi/2, 0.2)), (0.5, WayPoint(0.09, 0.36, pi/2, 0.2))]
follow_plan(robot_q_0, traj, test_P_1, test_P_2, test_P_3)
plt.show()

# # Tuning new PIDs

# In[30]:

import csv
reader = csv.reader(open("./pid_data.csv", 'r'))
setpoints = []
speeds = []
for row in reader:
    setpoints.append(float(row[0]))
    speeds.append(float(row[1]))

t = np.arange(0, len(setpoints)/100, 0.01)
plt.plot(t, setpoints, label="setpoint")
plt.plot(t, speeds, label="actual speed")
plt.xlabel("time (s)")
plt.ylabel("speed (m/s)")
plt.title("PID Performance")
plt.show()

# In[ ]:

docs/ipynb_gen_py/Time Optimal Smartmouse Controls.py


'''
Computer vision assignment 1 by Yoseob Kim
A1_corner_detection.py
Implementation of Harris corner detector.
GitHub Link: https://github.com/pseudowasabi/computer-vision-exercises/tree/master/CV_A1_
'''

import cv2
import numpy as np
import math
import time
import operator
import os
from filtering_by_yoseob import *

## ** initial settings to make result directory.
## Reference - https://www.geeksforgeeks.org/python-os-makedirs-method/?ref=lbp
try:
    os.makedirs('result', exist_ok=True)
except OSError as error:
    print("[NOTICE!] '/result' directory cannot be created.")
    print("Please CREATE the DIRECTORY MANUALLY to save created images.")

###
# 3-2. Corner response (define function - a, b, c and d)
###

def compute_corner_response(img):
    ## a) apply Sobel filters
    # Sobel filtering function is implemented in filtering_by_yoseob.py
    # sobel_img_x, sobel_img_y is derivatives along x and y direction respectively.
    sobel_img_x, sobel_img_y = my_sobel_filtering(img)

    # padded with 0 value to operate corner response
    sobel_img_x = image_padding_2d(sobel_img_x, 2, 1)   # (img, padd_width, type=1)
    sobel_img_y = image_padding_2d(sobel_img_y, 2, 1)   # type=1 means just padd with zeros.

    ## b) compute second moment matrix M
    uni_window = np.ones((5, 5))
    patch_img_x = np.zeros((5, 5))
    patch_img_y = np.zeros((5, 5))
    size0 = img.shape[0]
    size1 = img.shape[1]

    ## c) variables for computing corner responses
    R = np.zeros((size0, size1))
    _k = 0.04
    _max_val = 0

    elapsed_ = list(range(0, size0, size0 // 20))
    for x in range(size0):
        for y in range(size1):
            # i. subtract mean of each image patch
            _sum_x = 0.
            _sum_y = 0.
            for i in range(x, x+5):
                for j in range(y, y+5):
                    patch_img_x[i-x][j-y] = sobel_img_x[i][j]
                    _sum_x = operator.__add__(_sum_x, sobel_img_x[i][j])
                    patch_img_y[i-x][j-y] = sobel_img_y[i][j]
                    _sum_y = operator.__add__(_sum_y, sobel_img_y[i][j])
            _avg_x = operator.__truediv__(_sum_x, 25)
            _avg_y = operator.__truediv__(_sum_y, 25)

            sum_of_ix_ix = 0.
            sum_of_ix_iy = 0.
            sum_of_iy_iy = 0.
            for i in range(5):
                for j in range(5):
                    patch_img_x[i][j] = operator.__sub__(patch_img_x[i][j], _avg_x)
                    patch_img_y[i][j] = operator.__sub__(patch_img_y[i][j], _avg_y)

                    sum_of_ix_ix = operator.__add__(sum_of_ix_ix, operator.__mul__(patch_img_x[i][j], patch_img_x[i][j]))
                    sum_of_ix_iy = operator.__add__(sum_of_ix_iy, operator.__mul__(patch_img_x[i][j], patch_img_y[i][j]))
                    sum_of_iy_iy = operator.__add__(sum_of_iy_iy, operator.__mul__(patch_img_y[i][j], patch_img_y[i][j]))

            # ii. get second moment matrix
            # since we use uniform window, just calculated the summation of ix_ix, ix_iy, iy_iy respectively (above).
            M = np.array([[sum_of_ix_ix, sum_of_ix_iy], [sum_of_ix_iy, sum_of_iy_iy]])
            eigenvalues, _ = np.linalg.eig(M)
            #print(eigen_values)

            e1 = eigenvalues[0]
            e2 = eigenvalues[1]
            R[x][y] = e1 * e2 - _k * ((e1 + e2) ** 2)

            ## d) normalize responses
            # i. negative values to 0, otherwise normalize to range [0, 1]
            if R[x][y] < 0:
                R[x][y] = 0
            if R[x][y] > _max_val:
                _max_val = R[x][y]

        if x in elapsed_:
            print('.', end='')

    #print("_max_val:", _max_val)
    #normalizer = 1.
    #if _max_val != 0:
    #    normalizer = 1 / _max_val
    #normalizer = 1 / 255

    #normalizer = 1 / np.linalg.norm(R)
    for x in range(size0):
        for y in range(size1):
            R[x][y] = operator.__truediv__(R[x][y], _max_val)
    return R

def non_maximum_suppression_win(R, winSize):
    ## c) compute local maximas by NMS
    #* input argument R is already thresholded in corner_response_embedding func.
    x_bound = R.shape[0] - winSize + 1
    y_bound = R.shape[1] - winSize + 1

    # elapsed_ = list(range(0, R.shape[0], R.shape[1] // 20))
    for x in range(0, x_bound, winSize // 2):
        for y in range(0, y_bound, winSize // 2):
            local_maxima = R[x][y]
            lm_x = x
            lm_y = y
            for i in range(x, x + winSize):
                for j in range(y, y + winSize):
                    if R[i][j] > local_maxima:
                        local_maxima = R[i][j]
                        lm_x = i
                        lm_y = j
                    R[i][j] = 0
            # print(x, y, local_maxima, lm_x, lm_y)
            R[lm_x][lm_y] = local_maxima
        #if x in elapsed_:
        #    print('.', end='')
    return R

###
# 3-1. Gaussian filtering
###

img_lenna = cv2.imread('./lenna.png', cv2.IMREAD_GRAYSCALE)
img_shapes = cv2.imread('./shapes.png', cv2.IMREAD_GRAYSCALE)

print('Computer Vision A#1 // Yoseob Kim')
print('Part #3. Corner Detection\n')

print('3-1. apply gaussian filtering for shapes and lenna. (size=7, sigma=1.5)')

print('filtering... (for "shapes.png")', end='')
filtered_img_shapes = my_gaussian_filtering(img_shapes, 7, 1.5)
print(' ---> done.')

print('filtering... (for "lenna.png")', end='')
filtered_img_lenna = my_gaussian_filtering(img_lenna, 7, 1.5)
print(' ---> done.')
print()
print()

###
# 3-2. Corner response (execute requirements - ...)
###
print('3-2. compute corner responses for shapes and lenna respectively.')

# for "shapes.png"
print(' ** corner responses for "shapes.png" initiate.')
print(' [about 20 dots will be shown to be done]')

start_time = time.process_time()
R_shapes = compute_corner_response(filtered_img_shapes)
elapsed_time = time.process_time() - start_time
print(' ---> done. /elapsed time:', elapsed_time)

#print(np.min(R_shapes), np.max(R_shapes))
R_shapes = cv2.normalize(R_shapes, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_raw_shapes.png', R_shapes)
R_shapes = cv2.normalize(R_shapes, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
print(' * corner responses of "shapes.png" saved to ./result/ directory.')

print(' ## notice: press any key (on image window) to continue. !!do not close window!!\n')
cv2.imshow("corner responses of shapes", R_shapes)
cv2.waitKey(0)

# for "lenna.png"
print(' ** corner responses for "lenna.png" initiate.')
print(' [about 20 dots will be shown to be done]')

start_time = time.process_time()
R_lenna = compute_corner_response(filtered_img_lenna)
elapsed_time = time.process_time() - start_time
print(' ---> done. /elapsed time:', elapsed_time)

#print(np.min(R_lenna), np.max(R_lenna))
R_lenna = cv2.normalize(R_lenna, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_raw_lenna.png', R_lenna)
R_lenna = cv2.normalize(R_lenna, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
print(' * corner responses of "lenna.png" saved to ./result/ directory.')

print(' ## notice: press any key (on image window) to continue. !!do not close window!!\n')
cv2.imshow("corner responses of lenna", R_lenna)
cv2.waitKey(0)
print()

###
# 3-3. Thresholding, Non-maximum suppression (execute requirements)
###
print('3-3. thresholding and apply NMS to corner responses for each images.')

## a) change corner response to green
## b) show image and save

def corner_response_embedding(R, img):
    # convert gray scale to rgb channel
    normalized_img = cv2.normalize(img, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
    rgb_img = cv2.cvtColor(normalized_img, cv2.COLOR_GRAY2BGR)

    # thresholding (greater than 0.1)
    for i in range(R.shape[0]):
        for j in range(R.shape[1]):
            if R[i][j] > 0.1:
                cv2.circle(rgb_img, (j, i), 5, (0, 1.0, 0), 1)
            else:
                R[i][j] = 0

    return R, rgb_img

print('a, b) thresholding and corner response embedding to original images')

# for shapes.png
#print(np.min(img_shapes), np.max(img_shapes))
thresholded_R_shapes, rgb_img_shapes = corner_response_embedding(R_shapes, img_shapes)
#print(np.min(rgb_img_shapes), np.max(rgb_img_shapes))
rgb_img_shapes = cv2.normalize(rgb_img_shapes, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_bin_shapes.png', rgb_img_shapes)
print(' ** corner response embedding of "shapes.png" saved to ./result/ directory.')

print(' ## notice: press any key (on image window) to continue. !!do not close window!!\n')
cv2.imshow("corner response > 0.1 :: shapes.png", rgb_img_shapes)
cv2.waitKey(0)

# for lenna.png
#print(np.min(img_lenna), np.max(img_lenna))
thresholded_R_lenna, rgb_img_lenna = corner_response_embedding(R_lenna, img_lenna)
#print(np.min(rgb_img_lenna), np.max(rgb_img_lenna))
rgb_img_lenna = cv2.normalize(rgb_img_lenna, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_bin_lenna.png', rgb_img_lenna)
print(' ** corner response embedding of "lenna.png" saved to ./result/ directory.')

print(' ## notice: press any key (on image window) to continue. !!do not close window!!\n')
cv2.imshow("corner response > 0.1 :: lenna.png", rgb_img_lenna)
cv2.waitKey(0)
#print()

## c) nms
## d) show image and save

print("c, d) apply non-maximum suppression to corner responses")
_winSize = 11

# for "shapes.png"
print('...applying NMS to R (shapes.png)', end='')
start_time = time.process_time()
suppressed_R_shapes = non_maximum_suppression_win(thresholded_R_shapes, _winSize)
elapsed_time = time.process_time() - start_time
print(' ---> done. /elapsed time:', elapsed_time)

_, rgb_img_shapes2 = corner_response_embedding(suppressed_R_shapes, img_shapes)
#print(np.min(rgb_img_shapes2), np.max(rgb_img_shapes2))
rgb_img_shapes2 = cv2.normalize(rgb_img_shapes2, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_sup_shapes.png', rgb_img_shapes2)
print(' ** NMS applied R of "shapes.png" saved to ./result/ directory.')

print(' ## notice: press any key (on image window) to continue. !!do not close window!!\n')
cv2.imshow("non-maximum suppressed :: shapes.png", rgb_img_shapes2)
cv2.waitKey(0)

# for "lenna.png"
print('...applying NMS to R (lenna.png)', end='')
start_time = time.process_time()
suppressed_R_lenna = non_maximum_suppression_win(thresholded_R_lenna, _winSize)
elapsed_time = time.process_time() - start_time
print(' ---> done. /elapsed time:', elapsed_time)

_, rgb_img_lenna2 = corner_response_embedding(suppressed_R_lenna, img_lenna)
#print(np.min(rgb_img_lenna2), np.max(rgb_img_lenna2))
rgb_img_lenna2 = cv2.normalize(rgb_img_lenna2, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U)
cv2.imwrite('./result/part_3_corner_sup_lenna.png', rgb_img_lenna2)
print(' ** NMS applied R of "lenna.png" saved to ./result/ directory.')

print(' ## notice: P#3 done. press any key (on image window) to finish.\n')
cv2.imshow("non-maximum suppressed :: lenna.png", rgb_img_lenna2)
cv2.waitKey(0)

cv2.destroyAllWindows()


import numpy

class Polynomial(object):
    def __init__(self, coefficients):
        self.coeff = coefficients

    def __call__(self, x):
        """Evaluate the polynomial."""
        s = 0
        for i in range(len(self.coeff)):
            s += self.coeff[i]*x**i
        return s

    def __add__(self, other):
        """Return self + other as Polynomial object."""
        # Two cases:
        #
        # self:   X X X X X X X
        # other:  X X X
        #
        # or:
        #
        # self:   X X X X X
        # other:  X X X X X X X X

        # Start with the longest list and add in the other
        if len(self.coeff) > len(other.coeff):
            result_coeff = self.coeff[:]  # copy!
            for i in range(len(other.coeff)):
                result_coeff[i] += other.coeff[i]
        else:
            result_coeff = other.coeff[:] # copy!
            for i in range(len(self.coeff)):
                result_coeff[i] += self.coeff[i]
        return Polynomial(result_coeff)

    def __mul__(self, other):
        c = self.coeff
        d = other.coeff
        M = len(c) - 1
        N = len(d) - 1
        result_coeff = numpy.zeros(M+N+1)
        for i in range(0, M+1):
            for j in range(0, N+1):
                result_coeff[i+j] += c[i]*d[j]
        return Polynomial(result_coeff)

    def differentiate(self):
        """Differentiate this polynomial in-place."""
        for i in range(1, len(self.coeff)):
            self.coeff[i-1] = i*self.coeff[i]
        del self.coeff[-1]

    def derivative(self):
        """Copy this polynomial and return its derivative."""
        dpdx = Polynomial(self.coeff[:])  # make a copy
        dpdx.differentiate()
        return dpdx

    def __str__(self):
        s = ''
        for i in range(0, len(self.coeff)):
            if self.coeff[i] != 0:
                s += ' + %g*x^%d' % (self.coeff[i], i)
        # Fix layout
        s = s.replace('+ -', '- ')
        s = s.replace('x^0', '1')
        s = s.replace(' 1*', ' ')
        s = s.replace('x^1 ', 'x ')
        #s = s.replace('x^1', 'x') # will replace x^100 by x^00
        if s[0:3] == ' + ':  # remove initial +
            s = s[3:]
        if s[0:3] == ' - ':  # fix spaces for initial -
            s = '-' + s[3:]
        return s

    def simplestr(self):
        s = ''
        for i in range(0, len(self.coeff)):
            s += ' + %g*x^%d' % (self.coeff[i], i)
        return s

def test_Polynomial():
    p1 = Polynomial([1, -1])
    p2 = Polynomial([0, 1, 0, 0, -6, -1])
    p3 = p1 + p2
    p3_exact = Polynomial([1, 0, 0, 0, -6, -1])
    msg = 'p1 = %s, p2 = %s\np3=p1+p2 = %s\nbut wrong p3 = %s'%\
          (p1, p2, p3_exact, p3)
    assert p3.coeff == p3_exact.coeff, msg
    # Note __add__ applies lists only, here with integers, so
    # == for comparing lists is not subject to round-off errors

    p4 = p1*p2
    # p4.coeff becomes a numpy array, see __mul__
    p4_exact = Polynomial(numpy.array([0,  1, -1,  0, -6,  5,  1]))
    msg = 'p1 = %s, p2 = %s\np4=p1*p2 = %s\ngot wrong p4 = %s'%\
          (p1, p2, p4_exact, p4)
    assert numpy.allclose(p4.coeff, p4_exact.coeff, rtol=1E-14), msg

    p5 = p2.derivative()
    p5_exact = Polynomial([1, 0, 0, -24, -5])
    msg = 'p2 = %s\np5 = p2.derivative() = %s\ngot wrong p5 = %s'%\
          (p2, p5_exact, p5)
    assert p5.coeff == p5_exact.coeff, msg

    p6 = Polynomial([0, 1, 0, 0, -6, -1])  # p2
    p6.differentiate()
    p6_exact = p5_exact
    msg = 'p6 = %s\p6.differentiate() = %s\ngot wrong p6 = %s'%\
          (p2, p6_exact, p6)
    assert p6.coeff == p6_exact.coeff, msg

if __name__ == '__main__':
    import sys
    if len(sys.argv) >= 2 and sys.argv[1] == 'verify':
        test_Polynomial()


# coding: utf-8
import numpy

def calculateStep(a, b, n):
    return (b-a)/n # retorna o valor de h, onde a é o limite inferior, b o superior e n é a quantidade de intervalos

def trapezeRule(step, data):
    return (step/2) * (data[0] + 2 * sum(data[1:-1]) + data[-1]) #Calcula o método utilizando h/2 * (y0 + (y1 + y2 + y3 + ... + yn-1) + yn)

def simpsonByThree(step, data):
    even = [data[i] for i in range(2,len(data) - 1, 2)] #Separa os números que estão nos índices pares do vetor de valores
    odd = [data[i] for i in range(1,len(data) - 1, 2)]  #Separa os números que estão nos índices ímpares do vetor de valores

    return (step/3) * (data[0] + 4 * sum(odd) + 2 * sum(even) + data[-1]) #Calcula o método utilizando a fórmula h/3 * (y0 + 4 * (y1 + y3 + y5 + ... + y2n-1) + 2 * (y2 + y4 + y6 + ... + y2n) + yn)

option = 0
i = 0
j = 0
k = 0

while (k == 0):

    while (i == 0):
        i = int(input("\nEntre com 1 para calcular a partir de uma funcao f(x) e limites\nEntre com 2 para calcular a partir de uma tabela de dados\nSua opcao: "))
        if i == 1:
            a = float(input("\nEntre com o limite inferior: "))
            b = float(input("Entre com o limite superior: "))
            interval = float(input("Entre com o numero de intervalos desejados: "))
            func = input("Entre com a funcao desejada no seguinte formato - lambda x: f(x) \n-> ")
            funcFinal = eval(func)
            step = calculateStep(a, b, interval)
            values = [funcFinal(i) for i in numpy.arange(a, b + step, step)]
        elif i == 2:
            step = float(input("Insira o valor do passo: "))
            print("Insira os valores de f(x) separados por espaco: ")
            values = list(map(float,input().split()))
        else:
            i = 0

    while(j == 0):
        j = int(input("\nEntre com 1 para o metodo do trapezio\nEntre com 2 para 1/3 de Simpson\nSua opcao: "))
        if j == 1:
            option = 1
        elif j == 2:
            option = 2
        else:
            j = 0

    if option == 1:
        print("\n\nO resultado da integral definida pelo metodo dos trapezios para %d intervalos e f(x) ou tabela fornecidos sera\n-> %.6f" % ((len(values) -1 ), trapezeRule(step, values)))
    else:
        print("\n\nO resultado da integral definida pelo metodo de 1/3 de Simpson para %d intervalos e f(x) ou tabela fornecidos sera\n-> %.6f" % ((len(values) -1 ), simpsonByThree(step, values)))

    if k == 0:
        x = int(input("\n\nDeseja calcular uma nova integral?\n1 - Sim, calcular outra integral\n2 - Encerrar o programa\nInsira sua opcao: "))
        if x == 1:
            i = 0
            j = 0
            k = 0
        else:
            k = 1


"""
Function to return the theta (dimensionless temperature) profile at a certain
dimensionless point (r) in a 1D solid sphere, cylinder, or slab shape.

References:
1) Recktenwald 2006
2) Bergman, Lavine, Incropera, Dewitt 2011 from Ch. 5, pg.299-304
"""

# Modules and Other Required Functions
#------------------------------------------------------------------------------

import numpy as np
import scipy.special as sp
from funcRoots import roots

# First and Second Terms of the Theta Function
#------------------------------------------------------------------------------

def funcCn(root, b):
    """
    First term in the theta function
    root = root from the zeta, Bi equation
    b = shape factor where 2 sphere, 1 cylinder, 0 slab
    """
    if b == 2:
        Cn = 4*(np.sin(root)-root*np.cos(root)) / (2*root-np.sin(2*root))
    elif b == 1:
        Cn = (2/root) * (sp.j1(root) / (sp.j0(root)**2 + sp.j1(root)**2))
    elif b == 0:
        Cn = (4*np.sin(root)) / (2*root + np.sin(2*root))
    return Cn

def funcDn(r, root, b):
    """
    Second term in the theta function
    root = root from the zeta, Bi equation
    b = shape factor where 2 sphere, 1 cylinder, 0 slab
    """
    if b == 2:
        Dn = (1/(root*r)) * np.sin(root*r)
    elif b == 1:
        Dn = sp.j0(root * r)
    elif b == 0:
        Dn = np.cos(root * r)
    return Dn

# Theta Function
#------------------------------------------------------------------------------

def theta(r, b, z, Bi, Fo):
    """
    Dimensionless temperature for analytical solution of 1D transient heat
    conduction for a solid sphere, cylinder, or slab.
    r = dimensionaless length term to evaluate theta, (-)
    b = shape factor where 2 sphere or 1 cylinder or 0 slab, (-)
    z = range of zeta values to evaluate zeta, Bi equation for positive roots
    Bi = Biot number h*L/k, (-)
    Fo = Fourier number alpha*t/L^2, (-)
    """

    rts = roots(z, b, Bi)   # positive roots of the zeta, Bi equation
    n = len(rts)            # number of positive roots

    # initial dimensionless temperature at first root
    theta = funcCn(rts[0], b)*np.exp(-rts[0]**2 * Fo)*funcDn(r, rts[0], b)

    # summation of theta for the remaining roots
    for i in range(1, n):
        dTheta_o = funcCn(rts[i], b)*np.exp(-rts[i]**2 * Fo)*funcDn(r, rts[i], b)
        theta = theta + dTheta_o

    return theta    # theta temperature profile evaluated at r


import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pdb
import os
import skimage.io as io

from scipy.stats import pearsonr

# -------------------------- movies-metadata.csv -------------------------------
# Read csv
features = pd.read_csv("data/movies-metadata.csv", thousands=',')

# Drop first row
features = features.dropna(axis=0)

# Extract values of columns needed
revenue = features["Box_office"].values.astype(np.float)
imdb = features["imdbVotes"].values.astype(np.float)
imdbRating = features["imdbRating"].values.astype(np.float)

# Evaluation Pearson r correlation
correlation, _ = pearsonr(imdb, revenue)

# Figure 1: Shows correlation between columns
plt.figure(1, figsize = (20, 7))
plt.subplot(2, 2, 1)
plt.subplots_adjust(wspace = 0.2, hspace = 0.5)
plt.scatter(imdb, revenue, s = 1, marker = "o", facecolor = "none", edgecolor = "blue")
plt.title("Revenue vs imdb Votes", fontsize=16)
plt.xlabel("imdb Votes", fontsize=14)
plt.ylabel("Revenue", fontsize=14)
plt.annotate(f"Pearson-R = {correlation:.2f}", (np.min(imdb), 0.98*np.max(revenue)), fontsize=12) # plot the value on the graph

# Evaluation Pearson r correlation
correlation2, _ = pearsonr(imdbRating, revenue)

# Figure 2: Shows correlation between columns
plt.subplot(2, 2, 2)
plt.scatter(imdbRating, revenue, s = 1, marker = "o", facecolor = "none", edgecolor = "blue")
plt.title("Revenue vs imdb Rating", fontsize=16)
plt.xlabel("imdb Rating", fontsize=14)
plt.ylabel("Revenue", fontsize=14)
plt.annotate(f"Pearson-R = {correlation2:.2f}", (np.min(imdbRating), 0.98*np.max(revenue)), fontsize=12) # plot the value on the graph

# ------------------------------------------------------------------------------

# -------------------------- movies-metadata-cleaned.csv -----------------------

# Read csv
features = pd.read_csv("data/movies-metadata-cleaned.csv", thousands=',')

# Drop first row
features = features.dropna(axis=0)

# Extract values of columns needed
revenue = features["Box_office"].values.astype(np.float)
imdb = features["imdbVotes"].values.astype(np.float)
imdbRating = features["imdbRating"].values.astype(np.float)

# Evaluation Pearson r correlation
correlation, _ = pearsonr(imdb, revenue)

# Figure 3: Shows correlation between columns
plt.figure(2, figsize = (20, 7))
plt.subplot(2, 2, 1)
plt.subplots_adjust(wspace = 0.2, hspace = 0.5)
plt.scatter(imdb, revenue, s = 1, marker = "o", facecolor = "none", edgecolor = "blue")
plt.title("Revenue vs imdb Votes", fontsize=16)
plt.xlabel("imdb Votes", fontsize=14)
plt.ylabel("Revenue", fontsize=14)
plt.annotate(f"Pearson-R = {correlation:.2f}", (np.min(imdb), 0.98*np.max(revenue)), fontsize=12) # plot the value on the graph

# Evaluation Pearson r correlation
correlation2, _ = pearsonr(imdbRating, revenue)

# Figure 4: Shows correlation between columns
plt.subplot(2, 2, 2)
plt.scatter(imdbRating, revenue, s = 1, marker = "o", facecolor = "none", edgecolor = "blue")
plt.title("Revenue vs imdb Rating", fontsize=16)
plt.xlabel("imdb Rating", fontsize=14)
plt.ylabel("Revenue", fontsize=14)
plt.annotate(f"Pearson-R = {correlation2:.2f}", (np.min(imdbRating), 0.98*np.max(revenue)), fontsize=12) # plot the value on the graph

# ------------------------------------------------------------------------------

plt.show()


import matplotlib.pyplot as plt
import numpy as np
from scipy.fftpack import fft, fftshift
from scipy import signal

M = 32
N = 128
hN = N//2
hM = M//2
fftbuffer = np.zeros(N)
w = signal.blackman(M)

plt.figure(1, figsize=(9.5, 6))

plt.subplot(3,2,1)
plt.plot(np.arange(-hM, hM), w, 'b', lw=1.5)
plt.axis([-hM, hM-1, 0, 1.05])
plt.title('w1, M=32')

fftbuffer = np.zeros(N)
fftbuffer[:hM] = w[hM:]
fftbuffer[N-hM:] = w[:hM]
X = fft(fftbuffer)
mX = 20*np.log10(abs(fftshift(X)))
plt.subplot(3,2,3)
plt.plot(np.arange(-hN, hN), mX-max(mX), 'r', lw=1.5)
plt.axis([-hN//2,hN//2,-80,1])
plt.title('mW1')

pX = np.angle(fftshift(X))
plt.subplot(3,2,5)
plt.plot(np.arange(-hN, hN), pX, 'c', lw=1.5)
plt.axis([-hN,hN-1,-np.pi,np.pi])
plt.title('pW1')

M = 31
N = 128
hN = N//2
hM = (M+1)//2
fftbuffer = np.zeros(N)
w = signal.blackman(M)

plt.subplot(3,2,2)
plt.plot(np.arange(-hM, hM-1), w, 'b', lw=1.5)
plt.axis([-hM, hM, 0, 1.05])
plt.title('w2, M=31')

fftbuffer = np.zeros(N)
fftbuffer[:hM] = w[hM-1:]
fftbuffer[N-hM+1:] = w[:hM-1]
X = fft(fftbuffer)
mX = 20*np.log10(abs(fftshift(X)))
plt.subplot(3,2,4)
plt.plot(np.arange(-hN, hN), mX-max(mX), 'r', lw=1.5)
plt.axis([-hN/2,hN/2-1,-80,1])
plt.title('mW2')

pX = np.angle(fftshift(X))
plt.subplot(3,2,6)
plt.plot(np.arange(-hN, hN), pX, 'c', lw=1.5)
plt.axis([-hN,hN-1,-np.pi,np.pi])
plt.title('pW2')

plt.tight_layout()
plt.savefig('blackman-even-odd.png')
plt.show()


from sklearn.neighbors import KernelDensity
import numpy as np
import scipy as sp

def norm_cdf_int(mu, std, LB, UB):
    """ Return P(LB < X < UB) for X Normal(mu, std) """
    rv = sp.stats.norm(mu, std)
    return rv.cdf(UB) - rv.cdf(LB)

def norm_cdf_int_approx(mu, std, LB, UB):
    """
    Return P(LB < X < UB) for X Normal(mu, std) using approximation of Normal CDF

    Input: All inputs as 1-D arrays
    """
    l = normal_cdf_approx((LB - mu)/std)
    u = normal_cdf_approx((UB - mu)/std)
    return u - l

def normal_cdf_approx(x):
    """
    Approximation of standard normal CDF

    Input: x = array

    Polynomial approximation from Abramowitz and Stegun p. 932
    http://people.math.sfu.ca/~cbm/aands/frameindex.htm

    Absolute error < 7.5*10^-8
    """
    p = 0.2316419
    b = [0.319381530, -0.356563782, 1.781477937, -1.821255978, 1.330274429]

    xx = abs(x) # Approximation only works for x > 0, return 1 - p otherwise

    t = 1/(1 + p*xx)
    Z = (1/(np.sqrt(2*np.pi)))*np.exp(-(x*x)/2)
    pol = b[0]*t + b[1]*(t**2) + b[2]*(t**3) + b[3]*(t**4) + b[4]*(t**5)

    prob = 1 - Z*pol # For x > 0
    prob[x < 0] = 1 - prob[x < 0] # Change when x < 0

    return prob

def mode_from_samples(samples, bandwidth_fraction = 0.1):
    """
    Compute the mode for each set of samples in 'samples'

    Using kernel density estimation

    Input:
    samples -- m x n array with n samples for each of m univariate random variables
    bandwidth_fraction -- kde will use bandwidth = [range of dataseries (max - min)] * bandwidth_fraction
    """

    # Function to optimize for finding the mode
    # -- the kernel density estimator
    def optfun(x, *args):
        kde = args[0]
        return -kde.score_samples(x.reshape(-1, 1))

    mode = np.zeros(samples.shape[0])

    for i in range(samples.shape[0]):
        data = samples[i].T
        min_x, max_x = data.min(), data.max()
        bandwidth = bandwidth_fraction*(max_x - min_x)

        # Fit kde
        kde = KernelDensity(kernel='gaussian', bandwidth=bandwidth).fit(data)

        # Find argmax of density
        args = (kde, )
        bounds = [(min_x, max_x)]

        res = sp.optimize.differential_evolution(optfun, bounds = bounds, args = args)

        mode[i] = res.x[0]

    return mode

def trunc_norm_moments_approx_corrfree(mu, sigma, LB, UB, inf_num = 1E100):
    """
    Correlation free approximation of truncated moments of multivariate Gaussian

    If X~N(mu, sigma), compute expectation and variance of X | LB <= X <= UB

    Input:
    mu, LB, UB : 1D numpy arrays
    sigma : numpy matrix
    inf_num : inf values are replaced with this number in calculations

    Returns:
    tmu, tvar (expectation and variance of truncated variable)
    """

    s2 = np.diag(sigma)
    s = np.sqrt(s2)
    a = (LB - mu )/s
    b = (UB - mu )/s

    # Replace inf and -inf by numbers
    a[a == float('inf')] = inf_num
    a[a == float('-inf')] = -inf_num
    b[b == float('inf')] = inf_num
    b[b == float('-inf')] = -inf_num

    phi_a = sp.stats.norm.pdf(a)
    phi_b = sp.stats.norm.pdf(b)
    PHI_diff = normal_cdf_approx(b) - normal_cdf_approx(a)

    tmu = mu + s*(phi_a - phi_b)/PHI_diff
    tvar = s2*(1 + (a*phi_a - b*phi_b)/PHI_diff - ((phi_a - phi_b)/PHI_diff)**2)

    return tmu, tvar