#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# This file is part of casiopeia.
#
# Copyright 2014-2016 Adrian BÃ¼rger, Moritz Diehl
#
# casiopeia is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# casiopeia is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warrantime_points of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with casiopeia. If not, see <http://www.gnu.org/licenses/>.

# This example is an adapted version of the system identification example
# included in CasADi, for the original file see:
# https://github.com/casadi/casadi/blob/master/docs/examples/python/sysid.py

import pylab as pl

import casadi as ca
import casiopeia as cp

import os

# System setup

T = 5.0
N = 100

x = ca.MX.sym("x", 4)

u = ca.MX.sym("u", 1)
eps_u = ca.MX.sym("eps_u", 1)

p = ca.MX.sym("p", 3)

k_M = p[0]
c_M = p[1]
c_m = p[2]

M = 250.0
m = 50.0

k_M_true = 4.0
c_M_true = 4.0
c_m_true = 1.6

p_true = [k_M_true, c_M_true, c_m_true]

p_scale = [1e3, 1e4, 1e5]

f = ca.vertcat( \

        x[1], \
        (p_scale[0] * k_M / m) * (x[3] - x[1]) + (p_scale[1] * c_M / m) * (x[2] - x[0]) - (p_scale[2] * c_m / m) * (x[0] - (u + eps_u)), \
        x[3], \
        -(p_scale[0] * k_M / M) * (x[3] - x[1]) - (p_scale[1] * c_M / M) * (x[2] - x[0]) \

    )

phi = x

system = cp.system.System( \
    x = x, u = u, p = p, f = f, phi = phi, eps_u = eps_u)


time_points = pl.linspace(0, T, N+1)

u0 = 0.05
x0 = pl.zeros(x.shape)

udata = u0 * pl.sin(2 * pl.pi*time_points[:-1])

# udata = ca.DM([-0.0144505, -0.00970027, -0.0131704, -0.0162461, -0.0194673, -0.0226493, 0.00758111, -0.0377716, 0.00802565, -0.0376995, 0.00770359, -0.00442005, -0.000855387, 0.0022456, 0.00545773, 0.00859222, 0.0117409, 0.0148885, 0.0180565, 0.0212502, 0.0154754, 0.0304801, -0.00435828, 0.0413637, -0.00430136, 0.0415592, -0.00399078, 0.0419632, -0.00324539, 0.00905321, 0.0056727, 0.00271414, -0.000364645, -0.00338611, -0.00643686, -0.00950217, -0.00850269, -0.017266, -0.00206468, -0.025834, 0.0121024, -0.0336159, 0.0120527, -0.0338047, 0.0117481, -0.0342033, 0.0110075, -0.00128918, 0.00209328, 0.00505328, 0.00813338, 0.0111559, 0.0142075, 0.0172735, 0.017385, 0.0247696, -0.00750083, 0.038085, -0.0077066, 0.0380376, -0.00733884, 0.00481979, 0.0348082, -0.010399, 0.0355777, -0.00990138, 0.0360839, -0.00939524, 0.0365825, -0.00891039, 0.0370485, -0.00846729, 0.0374656, -0.00807851, 0.0378246, -0.00775023, 0.0381219, -0.00748364, 0.0383584, -0.00727628, 0.0385379, -0.00712312, 0.0386664, -0.0070176, 0.0387508, -0.00695237, 0.0387987, -0.0069199, 0.0388175, -0.00691297, 0.0388142, -0.00692494, 0.0387951, -0.00694997, 0.0387654, -0.00698312, 0.0387298, -0.00702037, 0.0386917, -0.00705856])

# Turn 1, p_2 = 2.2
# udata = ca.DM([-0.00829643, -0.00515627, -0.00894248, -0.012289, 0.0118338, -0.0250525, 0.0121442, -0.0250723, 0.0117754, 0.00237793, 0.00608787, 0.00934724, 0.0127382, 0.0160609, 0.0194053, 0.0227508, 0.026116, 0.00188827, 0.0387551, 0.00155363, 0.0387881, 0.00167901, 0.0389903, 0.00194365, 0.0393043, 0.00229499, 0.0396823, 0.00269017, 0.0400863, 0.00309572, 0.0404872, 0.00348681, 0.0408641, 0.00384619, 0.0412033, 0.00416307, 0.0414965, 0.00443194, 0.0417407, 0.00465139, 0.0419359, 0.00482309, 0.042085, 0.00495083, 0.0421926, 0.00503975, 0.0422643, 0.00509568, 0.042306, 0.00512464, 0.0423237, 0.00543714, 0.0147748, 0.0109789, 0.00762548, 0.00412762, 0.000691339, -0.00277157, -0.00623775, -0.00972388, -0.00675703, -0.0185972, -0.0106631, -0.0245226, 0.00368551, -0.0336932, 0.00326556, -0.0341991, 0.00241641, -0.00721143, -0.0037635, -0.000731318, 0.00243011, 0.00553794, 0.00867811, 0.0118334, 0.0096454, 0.0198257, -0.00760238, 0.0294271, -0.00785601, 0.0293103, -0.0078546, 0.0294124, -0.00766823, 0.0296678, -0.00735777, 0.0300206, -0.00668697, 0.00283895, 0.0269457, -0.00971403, 0.0276994, -0.00928816, 0.028094, -0.00892604, 0.0284231, -0.00862983, 0.0286869, -0.00839749])

# Turn 2, p2 = 1.7
#udata = ca.DM([-0.00846481, -0.00366597, -0.00754265, -0.0109401, -0.0144678, -0.0179162, -0.0213776, -0.0248305, -0.0282937, 0.00325144, -0.0095417, -0.00609477, -0.0031221, -5.54444e-05, 0.00292505, 0.00591331, 0.00889787, 0.0119029, 0.0149362, 0.00941118, 0.0238493, -0.0128281, 0.0355525, -0.0129672, 0.0355749, -0.0128076, 0.035849, -0.0124399, 0.0362909, -0.0119413, 0.0368306, -0.0113746, 0.0374119, -0.0107893, 0.0379919, -0.0102223, 0.0385395, -0.00969926, 0.039034, -0.00923651, 0.0394629, -0.00884259, 0.0398212, -0.00851997, 0.0401087, -0.00826659, 0.0403293, -0.00807717, 0.0404894, -0.00794444, 0.0405969, -0.007575, 0.00537229, 0.00174034, -0.00146114, -0.00480086, -0.00808428, -0.0113978, -0.0147215, -0.0180725, -0.0109479, -0.0281232, -0.0162854, -0.0330267, 0.0038699, -0.0448677, 0.00332037, -0.0455267, 0.00232515, -0.0109728, -0.00776897, -0.00496767, -0.0020501, 0.000817866, 0.00372071, 0.00664635, 0.00125526, 0.0152777, -0.00403337, 0.022473, -0.0183765, 0.0300578, -0.018412, 0.030176, -0.0181646, 0.0305297, -0.0177255, 0.031035, -0.016909, -0.00371346, -0.00704345, -0.00996377, 0.0223731, -0.0255345, 0.0232642, -0.0249706, 0.0237387, -0.0245792, 0.0240534, -0.0243344])

# udata weighted
#udata = ca.DM([-0.00520227, -0.00257014, -0.00633508, -0.00966998, -0.013155, -0.0165846, -0.0200461, -0.0235147, -0.0270079, -0.00511096, -0.0388559, -0.00495019, -0.039065, -0.0052395, -0.0394182, -0.00593371, -0.0144635, -0.0108782, -0.00772065, -0.00442712, -0.00119117, 0.00207486, 0.00535066, 0.00865524, 0.00139953, 0.0176823, 0.00475885, 0.0217964, -0.00569706, 0.0283258, -0.00557302, 0.0285491, -0.00497639, 0.00352673, -7.96868e-05, -0.00324342, -0.00653669, -0.00976441, -0.0130165, -0.0162732, -0.0195549, -0.0158583, -0.0279326, -0.015772, -0.0337181, -0.0058277, -0.0398193, -0.00589265, -0.0399889, -0.00614702, -0.0403113, -0.00681587, -0.015336, -0.0117419, -0.00857805, -0.0052794, -0.00203962, 0.00122869, 0.00450687, -0.0176328, 0.0159055, -0.0182095, 0.0157011, -0.018317, 0.0156807, -0.0182641, 0.0157959, -0.0180978, 0.016034, -0.0178583, 0.0162666, -0.0175786, 0.016556, -0.017285, 0.0168487, -0.0169973, 0.0171279, -0.0167294, 0.0173823, -0.0164904, 0.0176048, -0.0162853, 0.0177921, -0.016116, 0.0179437, -0.0159817, 0.0180612, -0.0158802, 0.0181476, -0.015808, 0.0182068, -.0157609, 0.0182429, -0.0157346, 0.0182604, -0.0157249, 0.0182634, -0.0157276, 0.0182558, -0.0157389]) 

# udata weighted less input noise
# udata = ca.DM([-0.0134171, -0.00842978, -0.0120243, -0.0151871, -0.0184921, -0.0217422, -0.0250236, -0.0283189, -0.0265121, -0.0368323, -0.00226643, -0.015246, -0.0117478, -0.00875782, -0.00567611, -0.00269435, 0.000287519, 0.0032575, 0.00624136, 0.00924666, 0.0122751, 0.00979179, 0.00380797, 0.0252326, 0.0056666, 0.0285892, 0.0134783, 0.0316668, -0.0055205, 0.00756258, 0.00416894, 0.00125839, -0.00174408, -0.00465683, -0.00757626, -0.0104915, -0.0134283, -0.0163937, -0.0194038, -0.00385087, -0.0308164, -0.00885794, -0.0331922, -0.0147686, -0.0362871, -0.0226967, -0.0397543, -0.00301911, -0.016135, -0.0127747, -0.00989409, -0.00691863, -0.0040304, -0.0011331, 0.00176243, 0.00468028, 0.00762952, 0.00787311, 0.00299108, 0.0112158, 0.0112296, 0.0134084, 0.0169374, 0.014243, 0.0280616, -0.00880904, 0.0397936, -0.00813525, 0.00506049, 0.00176208, -0.00108758, -0.00404475, -0.00693134, -0.00983916, -0.012756, -0.0157033, -0.0186856, -0.0108104, -0.0126291, -0.0301752, -0.0105492, -0.034156, -0.018307, -0.0370511, 0.000360684, -0.012709, -0.00929629, -0.00637321, -0.00335717, -0.000433004, 0.0024969, 0.00542114, 0.00836606, 0.0113369, 0.0143557, -0.0179136, 0.0300656, -0.018525, 0.0298629, -0.0185682])

# udata weighted tf = 5.0, sigma_u = 0.001
#udata = ca.DM([-0.0158205, -0.0306043, -0.0411917, -0.0259636, -0.013651, -0.00652367, 0.00412271, -0.011427, -0.0238322, -0.0316207, -0.0417728, -0.0262736, -0.0138369, -0.00611593, 0.00414417, 0.0193004, 0.00430615, -0.00735196, -0.0130593, -0.0241014, -0.008428, 0.00399964, 0.0120962, 0.0219796, 0.00649744, -0.00595596, -0.013656, -0.0239644, -0.0391078, -0.0241197, -0.0124612, -0.00676603, 0.00428699, 0.0184327, 0.0034787, -0.0082597, -0.0139619, -0.0251336, -0.00947455, 0.00294051, 0.0110091, 0.0208921, 0.0362513, 0.0212292, 0.00957465, 0.0038058, -0.00716848, -0.0213564, -0.0063961, 0.00534127, 0.0110565, 0.0222139, 0.0065559, -0.00586055, -0.0139289, -0.0238138, -0.0391726, -0.0241505, -0.0124961, -0.00672769, 0.00424692, 0.0184347, 0.00347452, -0.00826285, -0.0139781, -0.0251354, -0.00947695, 0.00293924, 0.0110077, 0.0208926, 0.0362515, 0.0212295, 0.00957511, 0.00380664, -0.00716786, -0.0213558, -0.00639561, 0.00534172, 0.0110572, 0.0222142, 0.00655521, -0.00586129, -0.0139283, -0.0238154, -0.039171, -0.0241471, -0.0124936, -0.00672427, 0.00424989, -0.0114193, -0.0238527, -0.0319485, -0.0418404, -0.0263568, -0.0139068, -0.00620566, 0.00409797, -0.0114114, -0.0238301, -0.00538917])

# udata weighted, tf = 5.0, sigma_u = 0.005
#udata = ca.DM([-0.00865264, -0.0233332, -0.0326838, -0.0499101, -0.0348144, -0.0232985, -0.0175077, -0.00878532, 0.0081495, -0.00702837, -0.01873, -0.0251372, -0.0354426, -0.0499991, -0.034981, -0.023254, -0.0174267, -0.00844182, 0.0085084, -0.00674725, -0.0183793, -0.0247268, -0.0350917, -0.0496195, -0.034606, -0.0228779, -0.0170612, -0.00601815, 0.00818965, -0.00676811, -0.0185001, -0.0219461, -0.03612, -0.0201172, -0.0397916, -0.0235488, -0.0137446, -0.0013322, -0.017535, 0.00185627, -0.0145145, -0.0272921, -0.036239, -0.0499404, -0.0347114, -0.0227824, -0.0162608, -0.00559902, 0.00897182, -0.00601006, -0.0177183, -0.021169, -0.0353102, -0.0193042, -0.00699581, 0.00175138, 0.0108202, 0.0265577, 0.0114605, -0.000190328, -0.00369578, -0.0177452, -0.00172504, 0.0106017, 0.0193907, 0.0274616, 0.0445343, 0.0292932, 0.0176914, 0.0114566, 0.00102187, -0.0134414, -0.0294636, -0.0149347, -0.0035984, 0.000680974, 0.0126283, 0.0260364, 0.0111385, -0.000642226, -0.0062548, -0.0175778, -0.031664, -0.0167297, -0.00499807, 0.000644059, 0.00973993, 0.0266885, 0.0114368, -0.000194072, -0.00653654, -0.0150836, -0.0320431, -0.0167653, -0.03694, -0.0209215, -0.0399135, -0.0225286, -0.0404631, -0.0224447])

# udata weighted, tf = 5.0, sigma_u = 0.005, p_L = [4.082, 3.973, 1.616] as initial params
#udata = ca.DM([-0.00750703, 0.00913669, 0.0220866, 0.0296195, 0.0460261, 0.0304416, 0.0499982, 0.033759, 0.0211894, 0.0128117, -0.00203445, -0.0169411, -0.00241727, 0.00882324, 0.0138102, 0.0258826, 0.0395631, 0.0247059, 0.0131711, 0.00709643, -0.00437765, -0.0186208, -0.00372299, 0.00777423, 0.0138853, 0.025267, 0.0395367, 0.0246305, 0.0131314, 0.00700164, -0.00436575, -0.0186428, -0.035017, -0.0205003, -0.00941461, -0.00464579, 0.00740074, 0.0209836, 0.00612024, -0.00542827, -0.00921961, -0.0237426, -0.00779888, -0.0276405, -0.0112354, 0.00134591, 0.0102192, -0.00547353, 0.0140357, -0.00255183, -0.0152869, -0.024784, -0.0386468, -0.02351, -0.0118429, -0.00501222, 0.00593656, 0.0205816, 0.00564462, -0.00583115, -0.00961946, -0.0240569, -0.00810298, 0.00394063, 0.0130636, 0.0223687, 0.0382465, 0.0231859, 0.0117751, 0.0054356, -0.00559484, -0.0200046, -0.0362104, -0.0217006, -0.0106011, -0.00582847, 0.00623502, 0.0198186, 0.00495697, -0.0065901, -0.0126904, -0.0241578, -0.0384144, -0.0235163, -0.0120208, -0.00590995, 0.00337668, 0.0204126, 0.00518239, -0.00620927, -0.0129587, -0.0217397, -0.0387916, -0.0235371, -0.0438555, -0.0276591, -0.0467201, -0.029174, -0.0472253, -0.0290896])

# udata weighted, tf = 5.0, sigma_u = 0.005, p_L = [4.14054, 3.96515, 1.64997] as initial params
#udata = ca.DM([-0.00478741, 0.0113049, 0.0239446, 0.0309362, 0.0475909, 0.0322277, 0.0499981, 0.033987, 0.0212807, 0.0127835, -0.00260174, -0.0174883, -0.00320861, 0.00796152, 0.0132514, 0.025497, 0.0392554, 0.0246066, 0.0131363, 0.00668194, -0.00492859, -0.0192749, -0.00458539, 0.00684527, 0.0133319, 0.0248526, 0.0392217, 0.024524, 0.0130912, 0.00658653, -0.00492167, -0.0192967, -0.0356644, -0.0213522, -0.0103167, -0.00513811, 0.00700795, 0.0207247, 0.00606411, -0.00541673, -0.00946968, -0.0242207, -0.00850737, -0.0281794, -0.0118957, 0.000675527, 0.00985312, -0.0056186, 0.0137282, -0.00273397, -0.0154534, -0.0228741, -0.0398174, -0.0245545, -0.0130772, -0.00526507, 0.00508129, 0.0201622, 0.00537374, -0.00603188, -0.0101182, -0.0247716, -0.00904385, 0.00290652, 0.012318, 0.0208662, 0.0380663, 0.0230788, 0.0117833, 0.00482458, -0.0060484, -0.0206898, -0.0367453, -0.0224467, -0.0113863, -0.00620238, 0.00597602, 0.019693, 0.0050357, -0.00644302, -0.0129266, -0.0245217, -0.038884, -0.0241934, -0.0127642, -0.0062763, 0.00316272, 0.0202251, 0.00521663, -0.00610962, -0.0131903, -0.022163, -0.0392395, -0.0242091, -0.0143356, -0.00531497, -0.0218453, -0.0030035, -0.0196556, -0.000967417])


simulation_true_parameters = cp.sim.Simulation( \
    system = system, pdata = p_true)

# simulation_true_parameters.run_system_simulation( \
#     x0 = x0, time_points = time_points, udata = udata)

# ydata = simulation_true_parameters.simulation_results.T

sigma_u = 0.005
sigma_y = pl.array([0.01, 0.01, 0.01, 0.01])

repetitions = 100

p_test =[]

for k in range(repetitions):

    udata_noise = udata + sigma_u * pl.randn(*udata.shape)

    simulation_true_parameters.run_system_simulation( \
        x0 = x0, time_points = time_points, udata = udata_noise)

    ydata = simulation_true_parameters.simulation_results.T

    ydata_noise = ydata + sigma_y * pl.randn(*ydata.shape)

    wv = (1.0 / sigma_y**2) * pl.ones(ydata.shape)
    weps_u = (1.0 / sigma_u**2) * pl.ones(udata.shape)

    pe_test = cp.pe.LSq(system = system, \
        time_points = time_points, \
        udata = udata, \
        pinit = [1.0, 1.0, 1.0], \
        # pinit = [4.0, 4.0, 2.2], \
        # pinit = [4.082, 3.973, 1.616], \
        ydata = ydata_noise, \
        xinit = ydata_noise, \
        wv = wv,
        weps_u = weps_u,
        discretization_method = "multiple_shooting")

    pe_test.run_parameter_estimation()

    p_test.append(pe_test.estimated_parameters)

p_mean = []
p_std = []

for j, e in enumerate(p_true):

    p_mean.append(pl.mean([k[j] for k in p_test]))
    p_std.append(pl.std([k[j] for k in p_test], ddof = 0))

pe_test.compute_covariance_matrix()

# Generate report

print("\np_mean         = " + str(ca.DM(p_mean)))
print("phat_last_exp  = " + str(ca.DM(pe_test.estimated_parameters)))

print("\np_sd           = " + str(ca.DM(p_std)))
print("sd_from_covmat = " + str(ca.diag(ca.sqrt(pe_test.covariance_matrix))))
print("beta           = " + str(pe_test.beta))

print("\ndelta_abs_sd   = " + str(ca.fabs(ca.DM(p_std) - \
    ca.diag(ca.sqrt(pe_test.covariance_matrix)))))
print("delta_rel_sd   = " + str(ca.fabs(ca.DM(p_std) - \
    ca.diag(ca.sqrt(pe_test.covariance_matrix))) / ca.DM(p_std)))


fname = os.path.basename(__file__)[:-3] + ".rst"

report = open(fname, "w")
report.write( \
'''Concept test: covariance matrix computation
===========================================

Simulate system. Then: add gaussian noise N~(0, sigma^2), estimate,
store estimated parameter, repeat.

.. code-block:: python

    y_randn = sim_true.simulation_results + sigma * \n
(np.random.randn(*sim_true.estimated_parameters.shape))

Afterwards, compute standard deviation of estimated parameters, 
and compare to single covariance matrix computation done in PECas.

''')

prob = "ODE, 4 states, 1 control + noise, 3 params, (quarter vehilce noise)"
report.write(prob)
report.write("\n" + "-" * len(prob) + "\n\n.. code-block:: python")

report.write( \
'''.. code-block:: python

    # ----------------------- casiopeia system definition ---------------------- #

    Starting system definition ...

    The system is a dynamic system defined by a set of explicit ODEs xdot
    which establish the system state x and by an output function phi which
    sets the system measurements:

    xdot = f(t, u, q, x, p, eps_e, eps_u),
    y = phi(t, u, q, x, p).

    Particularly, the system has:
    1 time-varying controls u
    0 time-constant controls q
    3 parameters p
    4 states x
    4 outputs phi

    where xdot is defined by 
    xdot[0] = x[1]
    xdot[1] = (((((1000*p[0])/50)*(x[3]-x[1]))+(((10000*p[1])/50)*(x[2]-x[0])))
        -(((100000*p[2])/50)*(x[0]-(u+eps_u))))
    xdot[2] = x[3]
    xdot[3] = ((-(((1000*p[0])/250)*(x[3]-x[1])))-(((10000*p[1])/250)*
        (x[2]-x[0])))

    and where phi is defined by 
    y[0] = x[0]
    y[1] = x[1]
    y[2] = x[2]
    y[3] = x[3]

''')

report.write("\n**Test results:**\n\n.. code-block:: python")

report.write("\n\n    repetitions    = " + str(repetitions))
# report.write("\n    sigma          = " + str(sigma))

report.write("\n\n    p_true         = " + str(ca.DM(p_true)))
report.write("\n\n    p_mean         = " + str(ca.DM(p_mean)))
report.write("\n    phat_last_exp  = " + \
    str(ca.DM(pe_test.estimated_parameters)))

report.write("\n\n    p_sd           = " + str(ca.DM(p_std)))
report.write("\n    sd_from_covmat = " \
    + str(ca.diag(ca.sqrt(pe_test.covariance_matrix))))
report.write("\n    beta           = " + str(pe_test.beta))

report.write("\n\n    delta_abs_sd   = " + str(ca.fabs(ca.DM(p_std) - \
    ca.diag(ca.sqrt(pe_test.covariance_matrix)))))
report.write("\n    delta_rel_sd   = " + str(ca.fabs(ca.DM(p_std) - \
    ca.diag(ca.sqrt(pe_test.covariance_matrix))) / ca.DM(p_std)) \
    + "\n")

report.close()

try:

    os.system("rst2pdf " + fname)

except:

    print("Generating PDF report failed, is rst2pdf installed correctly?")
