'''
MNA1 Example in Hylaa-Continuous

Note: due to rounding differences between Tran's original version and the .mat file,
The time of the safety violation is different (15.585 vs 16.554)
'''

import numpy as np
from scipy.io import loadmat
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(is_safe):
    '''make the hybrid automaton and return it'''

    ha = LinearHybridAutomaton()

    mode = ha.new_mode('mode')
    dynamics = loadmat('MNA_1.mat')

    a_matrix = csc_matrix(dynamics['A'])
    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(dynamics['B'])
    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')

    output_space = csr_matrix(([1.], [0], [0, 1]), shape=(1, a_matrix.shape[0] + num_inputs))

    threshold = 0.5 if is_safe else 0.2

    mat = csr_matrix(([-1], [0], [0, 1]), dtype=float, shape=(1, 1))
    rhs = np.array([-threshold], dtype=float)
    trans = ha.new_transition(mode, error)
    trans.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 < 2:
            lb = 0.001
            ub = 0.0015
        elif dim < 578:
            lb = 0
            ub = 0
        elif dim >= n-9 and dim <= n-5: # u1-u5
            lb = 0.1
            ub = 0.1
        elif dim >= n-4 and dim <= n-1: # u6-u9
            lb = 0.2
            ub = 0.2
        else:
            raise RuntimeError('unknown dimension: {}'.format(dim))

        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 = None
    #plot_settings.ydim_dir = 88

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

    #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)
