# import the modules
import sys
import GPy
import csv
import numpy as np
import cPickle as pickle
import scipy.stats as stats
import sklearn.metrics as metrics
from matplotlib import pyplot as plt

%matplotlib notebook

"""
Explanation: Experiment 7: TRo Journal

In this experiment, the generalization of cloth models to unseen T-shirts of the mannequin is verified. The evaluation is performed using RMSE, NRMSE, Pearson correlation as the parameters. In this notebook, the test inference is performed for pre-trained cloth models.
End of explanation
"""


# function to compute reconstruction error
def reconstructionError(model, valData, testData, mKey, kKey, optimizeFlag=False):    
    nSamplesVal = valData[mKey].shape[0]
    nSamplesTest = testData[mKey].shape[0]
    
    nDimIn = valData[kKey].shape[1]
    nDimOut = valData[mKey].shape[1]
    
    qDim = model.X.mean.shape[1]
    
    # computing reconstruction error for test1, test2 with variances
    predictVal = np.zeros((nSamplesVal,nDimOut))
    predictTest = np.zeros((nSamplesTest,nDimOut))

    for n in range(nSamplesVal):
        yIn = valData[kKey][n,:]
        yTrueOut = valData[mKey][n,:]
    
        [xPredict, infX] = model.Y0.infer_newX(yIn[None,:], optimize=False)
        yOut = model.predict(xPredict.mean, Yindex=1)    
        sys.stdout.write('.')
        
        predictVal[n,:] = yOut[0]
    sys.stdout.write('\n')
        
    for n in range(nSamplesTest):
        yIn = testData[kKey][n,:]
        yTrueOut = testData[mKey][n,:]
    
        [xPredict, infX] = model.Y0.infer_newX(yIn[None,:], optimize=optimizeFlag)
        yOut = model.predict(xPredict.mean, Yindex=1)    
        sys.stdout.write('.')
        
        predictTest[n,:] = yOut[0]
    sys.stdout.write('\n')
    results = {}
    valResults = {}
    testResults = {}
    
    valResults['pred'] = predictVal
    testResults['pred'] = predictTest
    
    valErrors = np.sqrt(metrics.mean_squared_error(valData[mKey],predictVal,multioutput='raw_values'))
    testErrors = np.sqrt(metrics.mean_squared_error(testData[mKey],predictTest,multioutput='raw_values'))

    valNormErrors = np.divide(np.sqrt(metrics.mean_squared_error(valData[mKey],predictVal,multioutput='raw_values')), 
                              valData[mKey].max(axis=0) - valData[mKey].min(axis=0))
    testNormErrors = np.divide(np.sqrt(metrics.mean_squared_error(testData[mKey],predictTest,multioutput='raw_values')), 
                               testData[mKey].max(axis=0) - testData[mKey].min(axis=0))

    valCorr = np.zeros((1,nDimOut))
    testCorr = np.zeros((1,nDimOut))
    for d in range(dims[1]):
        valCorr[0,d],_ = stats.pearsonr(valData[mKey][:,d],predictVal[:,d])
        testCorr[0,d],_ = stats.pearsonr(testData[mKey][:,d],predictTest[:,d])

    valResults['rmse'] = valErrors
    testResults['rmse'] = testErrors
    
    valResults['nrmse'] = valNormErrors
    testResults['nrmse'] = testNormErrors
    
    valResults['corr'] = valCorr
    testResults['corr'] = testCorr
        
    results['train'] = valResults
    results['test'] = testResults
    return results

"""
Explanation: Plotting and Analysis Functions
End of explanation
"""


# load all the files and create Data
nShr = 4
nPos = 6
names = []
for sInd in range(nShr):
    for pInd in range(nPos):
        names.append('K1S%dP%dT1' % (sInd+1,pInd+1))

        
# directory to save the results
dName = '../Results/Exp7'
if not os.path.exists(dName):
    os.makedirs(dName)
    
Data = pickle.load(open('../Data/Data.p','rb'))

# loop over the kinect keys
kinectExt = 'C'
kinectDim = 7500
kinectKey = 'Cloud'

mocapDim = 8
mocapExt = 'T'
mocapKey = 'TopCoord'
    
nTrials = 6
keys = [kinectKey,mocapKey]
expName = '%s%s' % (kinectExt,mocapExt)

ValInd = [[6,7,14,15,22,23],[12,13,20,21,4,5],[18,19,2,3,10,11],[0,1,8,9,16,17]]
TestInd = [[0,1,2,3,4,5],[6,7,8,9,10,11],[12,13,14,15,16,17],[18,19,20,21,22,23]]

for sInd in nShr:
    results = {}
    valData = {}
    testData = {}

    # load the trained MRD model
    mrdModel = pickle.load(open('../Models/Exp7/%s%d.p' % (expName,sInd+1),'rb'))
        
    for pInd,valInd,testInd in zip(range(nTrials),ValInd[sInd],TestInd[sInd]):
        print 'Cycle:%d,%d' % (sInd+1,pInd+1)
        print valInd, testInd
        
        for key in keys:
            valData[key] = Data[names[valInd]][key]
            testData[key] = Data[names[testInd]][key]
            
        # apply inference to test and val data
        results = reconstructionError(mrdModel,valData,testData,mocapKey,kinectKey,optimizeFlag=True)
        
        # save results to file
        pickle.dump(results,open('../Results/Exp7/Res%d%d.p' % (expName,sInd+1,pInd+1),'wb'))

# loop over the kinect keys
kinectExt = 'C'
kinectDim = 7500
kinectKey = 'Cloud'

mocapDim = 8
mocapExt = 'T'
mocapKey = 'TopCoord'
    
nTrials = 6
keys = [kinectKey,mocapKey]

ValInd = [[6,7,14,15,22,23],[12,13,20,21,4,5],[18,19,2,3,10,11],[0,1,8,9,16,17]]
TestInd = [[0,1,2,3,4,5],[6,7,8,9,10,11],[12,13,14,15,16,17],[18,19,20,21,22,23]]

for sInd in range(nShr):
    valData = {}
    testData = {}

    for pInd,valInd,testInd in zip(range(nTrials),ValInd[sInd],TestInd[sInd]):
        for key in keys:
            valData[key] = Data[names[valInd]][key]
            testData[key] = Data[names[testInd]][key]
            
        results = pickle.load(open('../Results/Exp7/Res%d%d.p' % (sInd+1,pInd+1),'rb'))
        
        newResults = {'train':{}, 'test':{}}
        newResults['train']['corr'] = results['val']['corr']
        newResults['train']['pred'] = results['val']['pred']
        newResults['train']['nrmse'] = results['val']['err']
        newResults['train']['rmse'] = np.sqrt(metrics.mean_squared_error(valData['TopCoord'],results['val']['pred'],multioutput='raw_values'))
        
        newResults['test']['corr'] = results['test']['corr']
        newResults['test']['pred'] = results['test']['pred']
        newResults['test']['nrmse'] = results['test']['err']
        newResults['test']['rmse'] = np.sqrt(metrics.mean_squared_error(testData['TopCoord'],results['test']['pred'],multioutput='raw_values'))
        pickle.dump(newResults, open('../Results/Exp7/Res%d%d.p' % (sInd+1,pInd+1), 'wb'))

"""
Explanation: Data Loading
End of explanation
"""


