{%- macro render_class_model(class) %}
class {{ class.name }}(BaseModel):
    """
    {{ class.description }}
    """

    def __init__(self):
        super().__init__()

        # declare constants
        {%- for const in class.c %}
            {%- set comp = class.components[const] %}
        {{ comp.name }} = {{ comp.start_value.data[0] }}
        {%- endfor %}

        # declare parameters
        {{ render_symbols(class, class.p) | indent(4) }}
        {{ render_vector("self.p", class, class.p) | indent(12) }}

        # declare parameter start values
        {{ render_vector_start_dict("self.p0", class, class.p) | indent(4) }}

        # declare discrete states
        {{ render_symbols(class, class.z) | indent(4) }}
        {{ render_vector("self.z", class, class.z) | indent(12) }}

        # declare discrete states start values
        {{ render_vector_start_dict("self.z0", class, class.z) | indent(4) }}

        # declare inputs
        {{ render_symbols(class, class.u) | indent(4) }}
        {{ render_vector("self.u", class, class.u) | indent(12) }}

        # declare outputs
        {{ render_zero_symbols(class, class.y) | indent(4) }}
        {{ render_vector("self.y", class, class.y) | indent(12) }}

        # declare internal variables
        {{ render_symbols(class.w) | indent(4) }}

        # declare continuous states
        {{ render_symbols(class, class.x) | indent(4) }}
        {{ render_vector("self.x", class, class.x) | indent(12) }}

        # declare continuous states start values
        {{ render_vector_start_dict("self.x0", class, class.x) | indent(4) }}

        # alg
        {{ render_statement_list(class.alg) | indent(4) }}

        # ode
        {{ render_ordermap_expression("self.ode", class.ode, class.x) | indent(8) }}

{%- endmacro %}

{%- macro render_class_function(class) %}
class __{{ class.name }}(BaseModel):
    """
    {{ class.description }}
    """

    def __init__(self):
        super().__init__()

        # declare constants
        {%- for var in class.c %}
        {{ key }} = {{ render_expression(expr=val.start) }}
        {%- endfor %}

        # declare inputs
        {{ render_symbols(class.u) | indent(4) }}
        {{ render_vector("self.u", class.u) | indent(4) }}

        # declare internal variables
        {{ render_symbols(class.w) | indent(4) }}

        # declare outputs
        {{ render_zero_symbols(class.y) | indent(4) }}

        # alg
        {{ render_statement_list(class.alg) | indent(4) }}

        # build function
        {{ render_vector("self.y", class.y) | indent(4) }}
        self.F = ca.Function("{{ class.name }}", [self.u], [self.y], ['u'], ['y'])

    def __call__(self, *args):
        return self.F(ca.vertcat(*args))


{{ class.name }} = __{{ class .name }}()
{%- endmacro %}

{%- macro render_class(class) %}
    {% if class.class_type == "Model" %}
        {{ render_class_model(class) }}
    {% elif class.class_type == "Function" %}
        {{ render_class_function(class) }}
    {% else %}
        {{ warn("unhandled class type: " + class.class_type) }}
    {% endif %}
{%- endmacro %}

{%- macro render_symbols(class, var_list) -%}
    {%- for var in var_list -%}
    {{ var }} = ca.SX.sym('{{ var }}'
        {%- set comp = class.components[var] -%}
        {%- if comp.array_subscripts | length > 0 -%}{{ ", " }}
            {%- for sub in comp.array_subscripts -%}
                {{ render_subscript(sub) -}}{% if not loop.last %},{% endif -%}
            {%- endfor -%}
        {%- endif -%})
    {% endfor -%}
{%- endmacro -%}

{%- macro render_zero_symbols(class, var_list) -%}
    {%- for var in var_list -%}
        {%- set comp = class.components[var] -%}
        {{- var }} = ca.SX.zeros(
        {%- if comp.array_subscripts | length > 0 -%}
            {%- for sub in comp.array_subscripts -%}
            {{ render_subscript(sub) }}{% if not loop.last %},{% endif -%}
            {%- endfor -%}
        {%- endif -%})
    {% endfor -%}
{%- endmacro -%}

{%- macro render_vector(target, class, var_list) -%}
    {{ target }} = ca.vertcat(
    {%- for var in var_list %}{{"\n"}}
        {{- var }}{% if not loop.last %},{% endif %}
    {%- endfor -%})
{%- endmacro -%}

{%- macro render_vector_start_dict(target, class, var_list) -%}
    {{ target }} = {
    {%- for var in var_list %}
        {%- set comp = class.components[var]%}
        '{{ var }}': {{ comp.start_value.data[0] | indent(8) -}}
        {% if not loop.last %},{% endif %}
    {%- endfor %}
    }
{%- endmacro -%}

{%- macro render_ordermap_expression(target, map, var_list) -%}
    {{ target }} = ca.vertcat(
    {% for var in var_list -%}
        {%- set expr = map[var] -%}
        {{ render_expression(expr=expr) }}{% if not loop.last %},{% endif %}
    {% endfor -%})
{%- endmacro -%}

{%- macro render_vector_expression(target, expr_list) -%}
    {{ target }} = ca.vertcat(
    {% for expr in expr_list -%}
        {{ render_expression(expr=expr) }}{% if not loop.last %},{% endif %}
    {% endfor -%})
{%- endmacro -%}

{%- macro render_binary(op, expr) -%}
    {{ render_expression(expr=expr.lhs) }}{{ op }}{{ render_expression(expr=expr.rhs) -}}
{%- endmacro -%}

{%- macro render_unary(op, expr) -%}
    {{ op }}{{ render_expression(expr=expr.rhs) -}}
{%- endmacro -%}

{%- macro render_subscript(sub) -%}
    {%- for key, value in sub | items -%}
        {%- if key == "Expression" -%}
            {{ render_expression(expr=sub.Expression) | int }}
        {%- elif key == "Colon" -%}
            :
        {%- endif -%}
    {%- endfor -%}
{%- endmacro -%}

{%- macro render_component_reference(comp) -%}
    {{ comp.name }}
    {%- if comp.array_subscripts | length > 0 %}[
        {%- for sub in comp.array_subscripts -%}
            {#- handles index from 1 to 0 from Modelica to python-#}
            {{ render_subscript(sub) | int -1 }}
        {%- endfor -%}]
    {%- endif -%}
{%- endmacro -%}

{%- macro render_expression(expr) -%}
    {%- for key, value in expr | items -%}
        {%- if key == "UnsignedReal" -%}
            {{ value }}
        {%- elif key == "UnsignedInteger" -%}
            {{ value }}
        {%- elif key == "Ref" -%}
            {{ render_component_reference(comp=value.comp) }}
        {%- elif key == "Add" -%}
            {{ render_binary(op=" + ", expr=value) -}}
        {%- elif key == "Sub" -%}
            {{ render_binary(op=" - ", expr=value) -}}
        {%- elif key == "Mul" -%}
            {{ render_binary(op="*", expr=value) -}}
        {%- elif key == "Div" -%}
            {{ render_binary(op="/", expr=value) -}}
        {%- elif key == "Exp" -%}
            {{- render_binary(op="^", expr=value) -}}
        {%- elif key == "Negative" -%}
            {{ render_unary(op="-", expr=value) -}}
        {%- elif key == "Parenthesis" -%}
            ( {{- render_expression(expr=value.rhs) -}} )
        {%- elif key == "Not" %}
            {{- render_unary(op="!", expr=value) -}}
        {%- elif key == "Or" -%}
            ca.logic_or(
                {{- render_expression(expr=value.lhs) -}},
                {{- render_expression(expr=value.rhs) -}})
        {%- elif key == "And" -%}
            ca.logic_and(
                {{- render_expression(expr=value.lhs) -}},
                {{- render_expression(expr=value.rhs) -}})
        {%- elif key == "Equal" -%}
                {{- render_expression(expr=value.lhs) -}} ==
                {{- render_expression(expr=value.rhs) -}}
        {%- elif key == "ArrayArguments" -%}
            ca.vertcat(
            {%- for arg in value.args -%}
                {{- "\n       " }} {{ render_expression(expr=arg) }} {%- if not loop.last -%},{% endif %}
            {%- endfor -%})
        {%- elif key == "FunctionCall" -%}
            {{ value.comp.name }}(
                {%- for arg in value.args -%}
                    {{ render_expression(expr=arg) }} {% if not loop.last -%}, {% endif %}
                {%- endfor -%})
        {%- else -%}
            {{ warn("unknown expression: " + key + value | pprint) }}
        {%- endif -%}
    {%- endfor -%}
{%- endmacro -%}

{%- macro render_statement_list(stmt_list) -%}
    {%- for stmt in stmt_list -%}
        {{ render_statement(stmt=stmt) }}
    {% endfor -%}
{%- endmacro -%}

{%- macro render_statement(stmt) -%}
    {%- for key, value in stmt | items -%}
        {%- if key == "Assignment" -%}
            {{ render_component_reference(comp=value.comp) }} = {{ render_expression(expr=value.rhs) }}
        {%- else -%}
            {{ warn("unknown statement: " + key) }}
        {%- endif %}
    {%- endfor -%}
{%- endmacro -%}
# rumoca pkg version : {{ def.rumoca_version }}
# rumoca git version : {{ def.rumoca_git_hash }}
# template md5       : {{ def.template_md5 }}
# model md5          : {{ def.model_md5 }}

import casadi as ca
import numpy as np

class BaseModel:

    def __init__(self):
        pass

    def __repr__(self):
        return repr(self.__dict__)

    def simulate(self, t=None, u=None):
        if t is None:
            t = np.arange(0, 1, 0.01)
        if u is None:
            u = 0

        p0 = np.array([self.p0[k] for k in self.p0.keys()])
        z0 = np.array([self.z0[k] for k in self.z0.keys()])
        x0 = np.array([self.x0[k] for k in self.x0.keys()])

        F = ca.integrator(
            'F', 'idas',
            {'x': self.x, 'z': self.z, 'p': self.p, 'u': self.u, 'ode': self.ode, 'alg': self.alg},
            t[0], t)

        res = F(x0=x0, z0=z0, p=p0, u=u)
        return {
            't': t,
            'x': res['xf'].T
        }
    
    def linearize(self):
        A = ca.jacobian(self.ode, self.x)
        B = ca.jacobian(self.ode, self.u)
        C = ca.jacobian(self.y, self.x)
        D = ca.jacobian(self.y, self.u)
        return (A, B, C, D)

def cat(axis, *args):
    return ca.vertcat(*args)

{%- for key, val in def.classes | items %}
{{ render_class(class=val) }}
{%- endfor %}
