'''
Introduction to Uncertainty Quantification: Assignment 1 Question 3
Student ID: 1100121

'''


import numpy as np
from numpy import linalg as LA
import matplotlib.pyplot as plt

############################# QUESTION 3.3

Δt = 0.1 #s
N = 1200 #Number of time-steps (so N+1 = 1201 total time points)
g = 9.807 #N/kg
b = 1e-4 #drag parameter

    # X = [x1,x2,v1,v2] = [position, velocity]
    # F: R^4 -> R^4
    #     [ I | [Δt-(bΔt^2)/2]I ]
    # F = [---+-----------------] (in block form)
    #     [ 0 |    [1-bΔt]I     ]
    # so that X_{k+1} = F.X_k + affine part + noise

F = np.array([
         [1  , 0.  ,  Δt - b*Δt*Δt/2.   ,  0.               ],
         [0. , 1.  ,  0                 ,  Δt - b*Δt*Δt/2.  ],
         [0. , 0.  ,  1.-b*Δt           ,  0.               ],
         [0. , 0.  ,  0.                ,1.-b*Δt            ]
         ])

Gu = np.array([ 0,
               -g*Δt*Δt/2.,
                0,
               -g*Δt]) #control term is a constant vector (as timestep is constant)

wind = 5.*np.random.randn(2*(N+2))+10.
    # wind at position X_k is wind[k] ~ N(10,5^2), k = 0,...,N+1
    # make 2 times as much wind to ensure the projective reaches the ground.

def nextIter(X,n):
    global Gu,wind
    # X_{K} = F.X_{K-1} + Gu + wind_K
    # wind only affects velocity horizontally

    return F.dot(X) + Gu + np.array([0.,0., wind[n] - wind[n-1] ,0.])

def trajectory():
    global N,wind
    initial_condition = np.array([0,0,300. + wind[0],600.])
    #initialise array of points
    graph = np.array([[0., 0., 0., 0.] for _ in range(N+1)]) #<- MUST NOT initialise
    graph[0] = initial_condition                              #   as [0,0,0,0]s

    for n in range(N):
        graph[n+1] = nextIter(graph[n],n+1)
    return graph

graph = trajectory()
# plt.plot(graph[:,0], graph[:,1],'b')
# plt.show()

############################# QUESTION 3.4

#observe at timesteps 400≤t≤600
default_start = 400
default_num_obs = 201 #note |{400,...,600}| = 201
default_end = default_start + default_num_obs -1 #the last time point

def polar(x):
    # polar:R^4 -> R^2 takes us from cartesian (x1,x2,v1,v2) space to polar (φ,r) space [(m,m,m/s,m/s)->(radians, m)]
    # (forgetting velocity components)

    return np.array([np.arctan2( x[1],30000. - x[0]) ,
                     np.sqrt((30000. - x[0])**2 + x[1]**2) ])

def cartesian(y):
    # cartesian:R^2 -> R^2 is the inverse of polar:R^2 -> R^2

    return np.array([30000. + y[1]*np.cos(np.pi-y[0]) , y[1]*np.sin(np.pi-y[0]) ])


def dPolar(x):
    # give the matrix representing the linear map best approximating polar() at x

    r = np.sqrt((30000-x[0])**2 + x[1]**2)
    return np.array([
        [x[1]/r**2       ,  (30000-x[0])/r**2 , 0 , 0],
        [(x[0]-30000)/r  ,  x[1]/r            , 0 , 0]
        ])


#sanity check for dPolar():
'''
def dPolarError(x,h):
    # given two 2-vectors, returns the size of the error term polar(x+h) - polar(x) - dPolar(x).h ; this should be small for h small
    return LA.norm(polar(x[0:2]+h[0:2]) -  polar(x[0:2])  - dPolar(x[0:2]).dot([ h[0], h[1], 0, 0]))

def dPolarSanityCheck(x, K = 50):
    # print a rapidly decreasing sequence to calm my nerves
    for n in range(K-30,K):
        print(dPolarError(x,x/np.exp(n)))

dPolarSanityCheck( np.array([1,1]) )
'''

def radians(deg):
    #converts degrees to radians
    return deg * np.pi / 180.

def observations(graph,start = default_start, num_obs = default_num_obs):
    # this is the array of observation 2-vectors (in polar coordinates) with noise

    φ_error = radians(5.)*np.random.randn(num_obs)   #N(0,radians(5)^2) noise
    r_error = 500.*np.random.randn(num_obs)          #N(0,500^2) noise

    obs = np.array([[0.,0.] for i in range(num_obs)])
    for i in range(num_obs):
        obs[i] = polar([graph[start+i,0],graph[start+i,1]]) + np.array([φ_error[i],r_error[i]])

    return obs


y = observations(graph)

# convert back to cartesian coordinates for later use
xhat = np.array([cartesian(pt) for pt in y])

(obs_post_x , obs_post_y) = ([30000.], [0.]) #in meters

# observations sanity check:
'''
# points should be 'centered on true trajectory'
plt.plot(graph[:,0], graph[:,1],'b', obs_post_x,obs_post_y,'bo', xhat[:,0],xhat[:,1],'rx' )
plt.show()
'''
############################# QUESTION 3.5

# covariance matrix of noise in dynamics(wind) at time k is
Q = np.diag([0.,0.,25.,0.])

# covariance matrix of noise in observations at time k is
R = np.diag( [radians(5.)**2 , 500.**2 ] )
# but only its inverse appears in the formula:
#invR = np.diag( [1/radians(5.)**2 , 1/50.**2])

#initialise length N+1 (one for each time point) array of covariance matrices
#
#      [   C_{0|0}   ]
#      [   C_{1|1}   ]
#  C = [   C_{2|2}   ] where each C_{k|k} is a 4x4 matrix
#      [     ...     ]
#      [   C_{N|N}   ]
#      [ C_{N+1|N+1} ]
#
C = np.array(
     [
      [[0.]*4]*4    # <--  4x4 zero matrix
     ]*(N+1)        # <-- N+1 rows
    )


# initialise length N+1 array of 4-vectors, X[k] = X_{k|k}
X = np.array(
     [
      [0.]*4  # <-- 4-vectors
     ]*(N+1)  # <-- N+1 rows
    )


#Evolve the system from t = default_start to t = default_end via the Extended Kalman Filter equations

############# INITIALISATION

# initial uncertainty = huge; lets just pretend the the noises are independent
C[default_start] = np.eye(4)*100000.

# start off with (xhat[0,0], xhat[0,1], something, something) because why not?
X[default_start] = np.array([xhat[0,0], xhat[0,1], 500,500])


for t in range(default_start,default_end+1):

    ############ STEP 2: PREDICTION
    Xt = F.dot(X[t]) + Gu                       #X_{t+1|t} in notes
    Ct = F.dot(C[t]).dot(F.T) + Q               #C_{t+1|t} in notes

    ############ STEP 3: CORRECTION VIA NEWTON-RAPHSON
    Ht = dPolar(Xt)
    C[t+1] = LA.inv( LA.inv(Ct) + (Ht.T).dot(LA.inv(R)).dot(Ht) )
    X[t+1] = Xt - C[t+1].dot(Ht.T).dot(LA.inv(R)).dot(polar(Xt) - y[t-default_start])


#plot the predicted data along with the observations and the true values
plt.plot(
    graph[:,0], graph[:,1], 'b',
    obs_post_x, obs_post_y, 'bo',
    xhat[:,0], xhat[:,1], 'gx',
    X[default_start:default_end,0], X[default_start:default_end,1], 'r'
    )
plt.show()


C_x_norm = np.array([LA.norm(c[0:2,0:2]) for c in C[default_start:default_end] ])
x_error = np.array([LA.norm(
    np.array([graph[t,0] - X[t,0], graph[t,1] - X[t,1]])
    ) for t in range(default_start,default_end+1)])

plt.plot(np.log(x_error), 'b', np.log(C_x_norm),'ro')
plt.show()

#repeat the plots but for velocity now
plt.plot(
    graph[:,0], graph[:,2], 'b',
    graph[:,0], graph[:,3], 'r',
    X[default_start:default_end,0], X[default_start:default_end,2], 'bx',
    X[default_start:default_end,0], X[default_start:default_end,3], 'rx'
    )
plt.show()

C_v_norm = np.array([LA.norm(c[2:4,2:4]) for c in C[default_start:default_end] ])
v_error = np.array([LA.norm(
    np.array([graph[t,2] - X[t,2], graph[t,3] - X[t,3]])
    ) for t in range(default_start,default_end+1)])

plt.plot(np.log(v_error), 'b', np.log(C_v_norm),'ro')
plt.show()

C_norm = np.array([LA.norm(c) for c in C[default_start:default_end]])
X_error = np.array([LA.norm(X[t]-graph[t]) for t in range(default_start,default_end-1)])
'''
#total C and X norm plots

plt.plot(X_error, 'g', C_norm, 'r')
plt.show()

plt.plot(np.log(X_error), 'g', np.log(C_norm), 'r')
plt.show()
'''
############################# QUESTION 3.6
'''
#run out of observations; continue with prediction steps until you hit the earth

#first check if you have hit the earth already
earth_flag =  False
impact_timestep = -1
for t in range(default_start, default_end+1):
    if X[t,1] < 0:
        earth_flag = True
        impact_timestep = t
        break

if earth_flag is False:

    for t in range(default_end,2*N): #one is unlikely to avoid the earth over 2N steps
        X[t+1] = F.dot(X[t]) + Gu
        C[t+1] = F.dot(C[t]).dot(F.T) + Q

        if X[t+1,1] < 0:
            earth_flag = True
            impact_timestep = (t+1)
            break

if earth_flag is False:
    print('wow.')
else:
    plt.plot(
      graph[:,0], graph[:,1], 'b',
      obs_post_x, obs_post_y, 'bo',
      xhat[:,0], xhat[:,1], 'gx',
      X[default_start:,0], X[default_start:,1], 'r'
    )
plt.show()
'''