import json
import numpy as np
from numpy.linalg import inv
import matplotlib.pyplot as plt


data = open("./position_list.json", encoding='utf-8')
position_list = json.load(data)

def get_lat(position):
    return float(position["latitude"])
    
def get_long(position):
    return float(position["longitude"])

def get_time(position):
    return int(position["time"])

def wrap_position(lat, long, time):
    position = {}
    position["latitude"] = lat
    position["longitude"] = long
    position["time"] = time
    return position


# Q
def covariance4d(interval):
    sigma = 0.0625
    part1 = sigma * ((interval ** 4) / 4.0)
    part2 = sigma * ((interval ** 3) / 2.0)
    part3 = sigma * (interval ** 2)
    
    cov_matrix = np.array([[part1, part2, 0, 0],
                           [part2, part3, 0, 0],
                           [0, 0, part1, part2],
                           [0, 0, part1, part2]])
    return cov_matrix


# Initial Estimation Covariance Matrix
P = covariance4d(0)


# Initial State Matrix
pre_lat = get_lat(position_list[0])
pre_long = get_long(position_list[0])
pre_time = get_time(position_list[0])

X = np.array([[pre_lat],
              [0],
              [pre_long],
              [0]])
n = 4
raw_position_list = [wrap_position(pre_lat, pre_long, pre_time)]
kalman_position_list = [wrap_position(pre_lat, pre_long, pre_time)]
for position in position_list[1:10]:
    
    t_interval = (get_time(position) - pre_time) / 1000.0
    A = np.array([[1, t_interval, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, t_interval],
                  [0, 0, 0, 1]])
    print("A State Matrix:\n", A)
    print("X State Matrix:\n", X)
    X = A.dot(X)
    print("X State Matrix:\n", X)
    P = A.dot(P).dot(A.T) + covariance4d(t_interval)
    print("P State Matrix:\n", P)
    
    R = np.diag([29, 29, 29, 29])
    S = P + R
    print("S State Matrix:\n", S)
    K = P.dot(inv(S))
    print("K State Matrix:\n", K)

    # Reshape the new data into the measurement space.
    curr_lat = get_lat(position)
    v_y = (curr_lat - pre_lat) / t_interval
    curr_long = get_long(position)
    v_x = (curr_long - pre_long) / t_interval
    Y = np.array([[curr_lat],[v_y],[curr_long],[v_x]])
    print("Y State Matrix:\n", Y)
    # Update the State Matrix
    # Combination of the predicted state, measured values, covariance matrix and Kalman Gain
    print("X State Matrix:\n", X)
    print("fix:\n", K.dot(Y - X))
    X = X + K.dot(Y - X)
    print("Kalman Filter State Matrix:\n", X)

    # Update Process Covariance Matrix
    P = (np.identity(len(K)) - K).dot(P)
    
    pre_lat = curr_lat
    pre_long = curr_long
    pre_time = get_time(position)
    
    raw_position_list.append(wrap_position(pre_lat, pre_long, pre_time))
    kalman_position_list.append(wrap_position(X[0][0], X[2][0], pre_time))
    

print("Raw Matrix:\n", raw_position_list)
print("Kalman Filter State Matrix:\n", kalman_position_list)