from __future__ import  division

import json
import warnings


import numpy as np
import pandas as pd


from IOHMM import SemiSupervisedIOHMM
from IOHMM import OLS, CrossEntropyMNL


warnings.simplefilter("ignore")

"""
Explanation: This is the IOHMM model with the parameters learned in a semi-supervised way. By using some labeled data, we force the learning process in a certain direction. The unlabeled data will be estimated using EM algorithm iteratively. See notes in http://pages.cs.wisc.edu/~jerryzhu/pub/sslicml07.pdf
SemiSupervisedIOHMM
End of explanation
"""


speed = pd.read_csv('../data/speed.csv')
speed.head()

"""
Explanation: Load speed data
End of explanation
"""


states = {}
corr = np.array(speed['corr'])
for i in range(int(len(corr)/2)):
    state = np.zeros((4,))
    if corr[i] == 'cor':
        states[i] = np.array([0,1,0,0])
        speed.at[i, 'rt'] = 1
    else:
        states[i] = np.array([1,0,0,0])
        speed.at[i, 'rt'] = 0

print(speed.shape[0])
print(len(states))

for t in range(5):
    print(states[t])

"""
Explanation: Label some states
In our structure of the code, the states should be a dictionary, the key is the index in the sequence (e.g. 0, 5) and the value is a one-out-of-n code of array where the kth value is 1 if the hidden state is k. n is the number of states in total.
In the following example, we assume that the "corr" column gives the correct hidden states. Here we assume only the first half of the sequence is labeled.
To make sure that the semi supervised model works, we set the value of 'rt' in state 0 as 0 and we set the value of 'rt' in state 1 as 1, the other values are not changed. So after training, we should be able to see 4 states clearly
End of explanation
"""


# there should be 4 hidden states in this model
SHMM = SemiSupervisedIOHMM(num_states=4, max_EM_iter=200, EM_tol=1e-10)

# we set only one output 'rt' modeled by a linear regression model
SHMM.set_models(model_emissions = [OLS()], 
                model_transition=CrossEntropyMNL(solver='lbfgs'),
                model_initial=CrossEntropyMNL(solver='lbfgs'))

# we set no covariates associated with initial/transitiojn/emission models
SHMM.set_inputs(covariates_initial = [], covariates_transition = [], covariates_emissions = [[]])

# set the response of the emission model
SHMM.set_outputs([['rt']])

# set the data and ground truth states
SHMM.set_data([[speed, states]])

"""
Explanation: Set up a simple model manully
End of explanation
"""


SHMM.train()

"""
Explanation: Start training
End of explanation
"""


# the coefficients of the output model for each states
# clearly the enforcement worked since the coefficient of the first two states are 0, and 1
print(SHMM.model_emissions[0][0].coef)
print(SHMM.model_emissions[1][0].coef)
print(SHMM.model_emissions[2][0].coef)
print(SHMM.model_emissions[3][0].coef)

# the scale/dispersion of the output model of each states
# since we know the first two states perfectly, the dispersions are 0.
print(np.sqrt(SHMM.model_emissions[0][0].dispersion))
print(np.sqrt(SHMM.model_emissions[1][0].dispersion))
print(np.sqrt(SHMM.model_emissions[2][0].dispersion))
print(np.sqrt(SHMM.model_emissions[3][0].dispersion))

# the transition probability from each state
print(np.exp(SHMM.model_transition[0].predict_log_proba(np.array([[]]))))
print(np.exp(SHMM.model_transition[1].predict_log_proba(np.array([[]]))))
print(np.exp(SHMM.model_transition[2].predict_log_proba(np.array([[]]))))
print(np.exp(SHMM.model_transition[3].predict_log_proba(np.array([[]]))))

"""
Explanation: See the training results
End of explanation
"""


json_dict = SHMM.to_json('../models/SemiSupervisedIOHMM/')
json_dict

with open('../models/SemiSupervisedIOHMM/model.json', 'w') as outfile:
    json.dump(json_dict, outfile, indent=4, sort_keys=True)

"""
Explanation: Save the trained model
End of explanation
"""


SHMM_from_json = SemiSupervisedIOHMM.from_json(json_dict)

"""
Explanation: Load back the trained model
End of explanation
"""


# the coefficients of the output model for each states
# clearly the enforcement worked since the coefficient of the first two states are 0, and 1
print(SHMM_from_json.model_emissions[0][0].coef)
print(SHMM_from_json.model_emissions[1][0].coef)
print(SHMM_from_json.model_emissions[2][0].coef)
print(SHMM_from_json.model_emissions[3][0].coef)

"""
Explanation: See if the coefficients are any different
End of explanation
"""


with open('../models/SemiSupervisedIOHMM/config.json') as json_data:
    json_dict = json.load(json_data)

SHMM_from_config = SemiSupervisedIOHMM.from_config(json_dict)

"""
Explanation: Set up the model using a config file, instead of doing it manully
End of explanation
"""


SHMM_from_config.set_data([[speed, states]])
SHMM_from_config.train()

"""
Explanation: Set data and start training
End of explanation
"""


# the coefficients of the output model for each states
# clearly the enforcement worked since the coefficient of the first two states are 0, and 1
print(SHMM_from_config.model_emissions[0][0].coef)
print(SHMM_from_config.model_emissions[1][0].coef)
print(SHMM_from_config.model_emissions[2][0].coef)
print(SHMM_from_config.model_emissions[3][0].coef)

"""
Explanation: See if the training results are any different?
End of explanation
"""


