#!/usr/bin/env python
import cPickle
import numpy as np
import matplotlib.pyplot as plt
import canonical_system_id as csi
from dtk import bicycle, control

allRiders = ['Charlie', 'Jason', 'Luke']
# create the comparison models
canon = csi.load_benchmark_canon(allRiders)
H = csi.lateral_force_contribution(allRiders)

with open('idMatrices.p') as f:
    idMatrices = cPickle.load(f)

def plot(canon, H, riders, environments, idMats):
    filename = ''
    for rider in riders:
        filename += '-' + rider
    for env in environments:
        filename += '-' + env.replace(' ', '')

    filename = 'canonical-id-plots/' + filename[1:]

    print filename

    v0 = 0.
    vf = 10.
    num = 100

    mM, mC1, mK0, mK2, mH = csi.mean_canon(riders, canon, H)
             'Balance With Disturbance',
             'Track Straight Line With Disturbance']

runs = csi.select_runs(riders, maneuvers, environments)

# This gives the proportion of the lateral force which should be added to the
# steer torque and roll torque equations in the canonical equations.
H = csi.lateral_force_contribution(riders)

# try to load in all of the runs for the given run numbers.
trials, errors = csi.load_trials(runs, H)

rollParams = ['Mpd', 'C1pd', 'K0pd']
steerParams = ['Mdd', 'C1dp', 'C1dd',
               'K0dd', 'K2dd', 'HdF']
canon = csi.load_benchmark_canon(riders)

idMatrices = {}
covarMatrices = {}

print('Computing the estimate for all runs.')
runs = list(set(runs).difference(errors))
means = csi.mean_canon(riders, canon, H)
idMat, rollCovar, steerCovar = csi.enforce_symmetry(runs, trials, rollParams,
        steerParams, *means)
idMatrices['All'] = idMat
covarMatrices['All'] = (rollCovar, steerCovar)
print('Done.')

for rider in riders:
    idMatrices[rider] = {}