%matplotlib inline

import json as json
import numpy as np
import matplotlib.pylab as plt
from statsmodels.tsa.api import ARMA, ARIMA
import statsmodels.api as sm

file_path = '/home/david/git/MIC-XRD-Polymer/visualization/Vis_data.json'

with open(file_path) as json_file:
    data_dicts = json.load(json_file)

keys = data_dicts[0].keys()
pt_key = keys[5]
strain_key = keys[4]
sample_key = keys[3]
stress_key = keys[2]

samples = []
strains = []
pc_points = []
stresses = []
for pt_dict in data_dicts:
    samples.append(pt_dict[sample_key])
    strains.append(pt_dict[strain_key])
    stresses.append(pt_dict[stress_key])
samples_numbers = sorted(set(samples))

sample_dict = {}
for sample in samples_numbers:    
    sub_data_dict = {}
    sub_strains = []
    sub_stress = []
    for strain, stress, samp, pts in zip(strains, stresses, samples, data_dicts):
        if samp == sample:
            if len(sub_data_dict) == 0:
                sub_data_dict['points'] = np.array(pts[pt_key])[..., None]
            else:
                sub_data_dict['points'] = np.concatenate((sub_data_dict['points'], np.array(pts[pt_key])[..., None]), axis=-1)
            sub_strains.append(strain)
            sub_stress.append(stress)
    sub_data_dict['strains'] = np.array(sub_strains)
    sub_data_dict['stresses'] = np.array(sub_stress)
    sample_dict[sample] = sub_data_dict

sample0 = sample_dict[samples_numbers[0]]
sample1 = sample_dict[samples_numbers[1]]
sample2 = sample_dict[samples_numbers[2]]

stresses0 = sample0['stresses']
strains0 = sample0['strains']
points0 = sample0['points'][0]
stresses1 = sample1['stresses']
strains1 = sample1['strains']
points1 = sample1['points'][0]
stresses2 = sample2['stresses']
strains2 = sample2['strains']
points2 = sample2['points'][0]

p = 4
exo0 = np.concatenate((np.zeros(p), strains0))
exo = np.concatenate((stresses0[:, None], strains0[:, None]))
model = ARIMA(endog=points0 - points0[0], order=(p, 0, 0), exog=exo0[p:])

results = model.fit()

print(results.summary())

strains_exo = np.concatenate((np.zeros(p), strains2))
stresses_exo = np.concatenate((np.zeros(p), stresses1))
print stresses1.shape



print points2.shape
print points2[0]
y_points2 = model.predict(params=results.params, exog=strains_exo, dynamic=True) + points2[0]

plt.figure()
plt.plot(y_points2, label='Predicted')
plt.plot(points2, label='Actual')
plt.legend(loc=2)


y_points0 = model.predict(params=results.params, exog=exo0) + points0[0]
plt.figure()
plt.plot(y_points0, label='Predicted')
plt.plot(points0, label='Actual')
plt.legend(loc=2)

l = 0
plt.plot(points0[:] - y_points0[l:])