#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Aug  5 17:20:37 2018

@author: Bora
"""
# %% import libraries
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from numpy.linalg import inv
from Functions.kalman_filter import kalman_filter
from Functions.kalman_smoother import kalman_smoother
from Functions.plot_model import plot_model
from Functions.EM import EM
# %% load in and set up data
df = pd.read_csv('Data/subj1_series1_data.csv')     
y = df['O2'].values[0:1000]

# %% Initialize all parameters and set important values

#initialize model parameters
p = 5

#intialize Kalman filter parameters
A = np.identity(p)
Q = np.identity(p)*(10**-3)
e = 10**-2 #this parameter is tuned!

a_predicted = np.zeros((len(y), p, 1))
a_refined = np.zeros((len(y), p, 1))

P_predicted = np.zeros((len(y), p, p))
P_refined = np.zeros((len(y), p, p))

#initialize Kalman smoother parameters
J = np.zeros((len(y), p, p))

#initialize EM algorithm parameters
num_iter = 100

#get first points for modelling purposes
initial_points = y[:p]

T = np.arange(p, len(y)) #start p ahead

# %% Set the analysis that you want
test_filter = True
test_smoother = True
test_EM = True

# %% Run kalman filter
if test_filter:
    savepath = "Plots/kalmanfilter.jpg"
    a_predicted, a_refined, P_predicted, P_refined = kalman_filter(a_predicted, a_refined, P_predicted, P_refined, A, Q, e, T, y, p)
    plot_model(initial_points, a_refined, y, T, p, savepath)
    
# %% Run kalman smoother
if test_smoother:
    savepath = "Plots/kalmansmoother.jpg"
    T_reversed = np.flip(T, axis = 0)
    a_refined, P_refined, J = kalman_smoother(a_predicted, a_refined, P_predicted, P_refined, A, J, T_reversed)
    plot_model(initial_points, a_refined, y, T, p, savepath)

# %% Run EM algorithm
if test_EM:
    savepath = "Plots/EM.jpg"
    a_refined, P_refined = EM(y, p, num_iter)
    plot_model(initial_points, a_refined, y, T, p, savepath)