'''
Motor Example in Hylaa-Continuous
'''

import numpy as np
from scipy.sparse import csr_matrix, csc_matrix

from hylaa.hybrid_automaton import LinearHybridAutomaton, bounds_list_to_init
from hylaa.engine import HylaaSettings
from hylaa.engine import HylaaEngine
from hylaa.settings import PlotSettings, SimulationSettings
from hylaa.star import Star

def define_ha(safe):
    '''make the hybrid automaton and return it'''

    ha = LinearHybridAutomaton()

    mode = ha.new_mode('mode')

    # dynamics:
     #y1 == x1
     #y2 == x5

     #x1' == x2
     #x2' == 8487.2*x3 - 1.0865*x2
     #x3' == -2592.1*x1 - 21.119*x2 - 698.91*x3 - 141399.0*x4
     #x4' == x1 - 1.0*u1
     #x5' == x6
     #x6' == 8487.2*x7 - 1.0865*x6
     #x7' == - 2592.1*x5 - 21.119*x6 - 698.91*x7 - 141399.0*x8
     #x8' == x5 - 1.0*u2

    a_matrix = csc_matrix(np.array( \
        [ \
         [0, 1, 0, 0, 0, 0, 0, 0], \
         [0, -1.0865, 8487.2, 0, 0, 0, 0, 0], \
         [-2592.1, -21.119, -698.91, -141399.0, 0, 0, 0, 0], \
         [1, 0, 0, 0, 0, 0, 0, 0], \
         [0, 0, 0, 0, 0, 1, 0, 0], \
         [0, 0, 0, 0, 0, -1.0865, 8487.2, 0], \
         [0, 0, 0, 0, -2592.1, -21.119, -698.91, -141399.0], \
         [0, 0, 0, 0, 1, 0, 0, 0], \
        ], dtype=float))


    col_ptr = [n for n in a_matrix.indptr]
    rows = [n for n in a_matrix.indices]
    data = [n for n in a_matrix.data]

    b_matrix = csc_matrix(np.array( \
    [ \
        [0, 0], \
        [0, 0], \
        [0, 0], \
        [-1.0, 0], \
        [0, 0], \
        [0, 0], \
        [0, 0], \
        [0, -1.0], \
    ], dtype=float))

    num_inputs = b_matrix.shape[1]

    for u in xrange(num_inputs):
        rows += [n for n in b_matrix[:, u].indices]
        data += [n for n in b_matrix[:, u].data]
        col_ptr.append(len(data))

    combined_mat = csc_matrix((data, rows, col_ptr), \
                              shape=(a_matrix.shape[0] + num_inputs, a_matrix.shape[1] + num_inputs))

    mode.set_dynamics(csr_matrix(combined_mat))

    error = ha.new_mode('error')
    dims = combined_mat.shape[0]

    # safe error condition:   x1 in [0.35, 0.4] and x5 in [0.45, 0.6]
    # unsafe error condition: x1 in [0.3, 0.4] and x5 in [0.4, 0.6]
    bounds_x1 = [0.35, 0.4] if safe else [0.3, 0.4]
    bounds_x5 = [0.45, 0.6] if safe else [0.4, 0.6]

    output_space = csr_matrix(([1., 1.], [0, 4], [0, 1, 2]), shape=(2, dims))
    mat = csr_matrix(([-1, 1, -1, 1], [0, 0, 1, 1], [0, 1, 2, 3, 4]), dtype=float, shape=(4, 2))
    rhs = np.array([-bounds_x1[0], bounds_x1[1], -bounds_x5[0], bounds_x5[1]], dtype=float)

    trans1 = ha.new_transition(mode, error)
    trans1.set_guard(output_space, mat, rhs)

    return ha

def make_init_star(ha, hylaa_settings):
    '''returns a Star'''

    n = ha.dims
    bounds_list = [] # bounds on each dimension

    for dim in xrange(n):
        if dim == 0: # x1
            lb = 0.002
            ub = 0.0025
        elif dim == 4: # x5
            lb = 0.001
            ub = 0.0015
        elif dim == 8: # u1
            lb = 0.16
            ub = 0.3
        elif dim == 9: # u2
            lb = 0.2
            ub = 0.4
        else:
            lb = ub = 0

        bounds_list.append((lb, ub))

    init_space, init_mat, init_mat_rhs = bounds_list_to_init(bounds_list)

    return Star(hylaa_settings, ha.modes['mode'], init_space, init_mat, init_mat_rhs)

def define_settings():
    'get the hylaa settings object'
    plot_settings = PlotSettings()
    plot_settings.plot_mode = PlotSettings.PLOT_NONE
    plot_settings.xdim_dir = 0
    plot_settings.ydim_dir = 4
    #plot_settings.max_shown_polys = None

    settings = HylaaSettings(step=0.001, max_time=20.0, plot_settings=plot_settings)
    settings.simulation.sim_mode = SimulationSettings.KRYLOV
    #settings.simulation.check_answer = True

    #settings.simulation.krylov_stdout = True
    settings.simulation.krylov_transpose = True

    return settings

def run_hylaa(safe):
    'Runs hylaa with the given settings, returning the HylaaResult object.'

    ha = define_ha(safe)
    settings = define_settings()
    init = make_init_star(ha, settings)

    engine = HylaaEngine(ha, settings)
    engine.run(init)

    return ha.dims, engine.result

if __name__ == '__main__':
    run_hylaa(False)
