base.py

import numpy as np
import copy
import casadi as cas
import utility
from collections import OrderedDict
from unit_system import time_scale_table
from collections import UserDict, Iterable
from typing import Union, Optional, Any
from data_type import TrajectoryList, TrajectoryDict, ValueDict, StrDict, SymDict, StrList, ValueList, BoolList, BoolDict
from data_type import VariableCodeType, ValueComponentType

def add_var(name, index_list=None, m=1, val=None, lbx=None, ubx=None):
    """Add a variable to the model
               This will also collect all the variables into a dictionary
            """

    # Define the variable
    var = Variable(name, index_list, m, val, lbx, ubx)
    var_sym = var.symbol

    return var_sym


class ModelBase:
    keywords = ['vars', 'odes', 'algs', 'inits', 'constraints', 'submodels']

    def __init__(self, name):
        self.name = name

    @staticmethod
    def _combine_eq(*args):
        """ Combine all the equations into one single list """
        eq_list = []
        for eq in args:
            if isinstance(eq, list) or isinstance(eq, tuple):
                eq_list += eq
            else:
                eq_list.append(eq)
        return eq_list
    
    def _validate_new_property_names(self, name):
        """All the added variables, equations and submodels will be checked by this function"""
        if name in self.__dict__ or name in self.keywords:
            raise KeyError(f'Property "{name}" has already existed in {self.name} model. Please use another name.')


class Model(ModelBase):

    def __init__(self, name, time_unit='s', time_unit_default='s'):
        super().__init__(name)

        # Specify time unit of the differential equations
        self.time_unit_default = time_unit_default
        self._time_unit = time_unit
        self.time_scale = time_scale_table[time_unit] / time_scale_table[time_unit_default]

        # Lists of all variables
        self.vars = OrderedDict()

        # List of all equations
        self.odes = []
        self.algs = []
        self.inits = []  # Equations for initialization

        # List of all submodels
        self.submodels = []

        # CasADi symbols after concatenation
        self.vars_code = {'x': [], 'z': [], 'p': [], 'fixed': [], 'free': []}
        self.vars_symbol = {'x': None, 'z': None, 'p': None, 'fixed': None, 'free': None}
        self.vars_symbol_out = {'x': None, 'z': None, 'p': None}
        self.vars_val = {'x': None, 'z': None, 'p': None, 'fixed': None, 'free': None}
        self.vars_lb = {'x': None, 'z': None, 'p': None, 'fixed': None, 'free': None}
        self.vars_ub = {'x': None, 'z': None, 'p': None, 'fixed': None, 'free': None}
        self.vars_name = {'x': None, 'z': None, 'p': None, 'fixed': None, 'free': None}

        self.ode_symbol = cas.SX()
        self.alg_symbol = cas.SX()
        self.init_symbol = cas.SX()

        self.symbol_output = False

        self.flag_override_build = True
        self.dae_prob = None

    def setup(self, **kwargs) -> dict[str, Union[cas.MX, cas.SX]]:

        self.build(**kwargs)
        self.assembly_dae_components()

        return self.dae_prob

    def build(self, **kwargs):
        """Add model variables and equations"""
        self.flag_override_build = False
        pass

    def assembly_dae_components(self):
        """Classify and assembly all the dae components"""
        # Get variable symbols and values
        xzp = self.classify_and_vectorize_xzp()
        x, z, p = xzp['x'], xzp['z'], xzp['p']

        # Get equations
        self.ode_symbol = cas.vcat(self.odes)
        self.alg_symbol = cas.vcat(self.algs)

        # ode problem
        self.dae_prob = {'x': x, 'z': z, 'p': p, 'ode': self.ode_symbol, 'alg': self.alg_symbol}

    def add_var(self,
                var_type: Union[str, list[str], dict[Any, list[str]]],
                name: str,
                index_list: Optional[list[str]] = None,
                m=1,
                val=1.,
                lbx=-np.inf,
                ubx=np.inf,
                fixed: Optional[Union[bool, list[bool], dict[Any, list[bool]]]] = None) -> cas.SX:
        """Add a variable to the model
           This will also collect all the variables into a dictionary
           domain is for var_type=xz only, which should be a list, tuple or range
        """
        # Avoid repeated variable name
        self._validate_new_property_names(name)
        self.__validate_var_name(name)

        # Define the variable
        if index_list is None:
            var = VariableList(name, var_type, index_list, m, val, lbx, ubx,
                               model=self, fixed=fixed)
        else:
            var = VariableDict(name, var_type, index_list, m, val, lbx, ubx,
                               model=self, fixed=fixed)
        var_sym = var.symbol

        self.vars[var.global_name] = var

        setattr(self, f'{name}_obj', var)

        return var_sym

    def add_submodel(self, submodel):
        self._validate_new_property_names(submodel.name)
        self.submodels.append(submodel)
        setattr(self, submodel.name, submodel)

    def __validate_var_name(self, name):
        """Check the name of the variable to avoid to crash name with an exiting variable"""
        variable_full_name = f'{self.name}.{name}'
        if variable_full_name in self.vars:
            raise KeyError(f'Current name "{name}" has already existed in variable list. Please use another name.')

    def classify_and_vectorize_xzp(self):
        self._classify_xzp_by_name()
        var_symbols = self.collect_var_symbol(['x', 'z', 'p'])
        return var_symbols

    def classify_and_vectorize_fixed(self):
        self._classify_fixed_free_by_name()
        var_symbols = self.collect_var_symbol(['fixed', 'free'])
        return var_symbols

    def _classify_xzp_by_name(self):
        """Classify xzp variables by putting their names into different list, and generate variable code"""
        vars_code = {'x': [], 'z': [], 'p': []}

        for global_name, var in self.vars.items():
            try:
                if isinstance(var.var_type, StrList) or isinstance(var.var_type, list):
                    for i, vi in enumerate(var.var_type):
                        vars_code[vi].append((global_name, i))

                elif isinstance(var.var_type, StrDict) or isinstance(var.var_type, dict):
                    for key, vtypes in var.var_type.items():
                        for i, vi in enumerate(vtypes):
                            vars_code[vi].append((global_name, key, i))
                else:
                    raise ValueError(f"var_type can only be list or dict but you're having {type(var.var_type)}.")
            except:
                print(f'Error is in variable', var.global_name)
                raise ValueError('Cannot register the variable!')

        for key, val in vars_code.items():
            self.vars_code[key] = val

    def _classify_fixed_free_by_name(self):
        """Classify fixed and free variables by putting names into different list, and generate variable code"""
        vars_code = {'fixed': [], 'free': []}

        for global_name, var in self.vars.items():
            try:
                if isinstance(var.var_type, StrList) or isinstance(var.var_type, list):

                    for i, assign_or_not in enumerate(var.fixed):
                        if assign_or_not:
                            vars_code['fixed'].append((global_name, i))
                        else:
                            vars_code['free'].append((global_name, i))

                elif isinstance(var.var_type, StrDict) or isinstance(var.var_type, dict):
                    for key, assign_or_not_list in var.fixed.items():
                        for i, assign_or_not in enumerate(assign_or_not_list):
                            if assign_or_not:
                                vars_code['fixed'].append((global_name, key, i))
                            else:
                                vars_code['free'].append((global_name, key, i))
                else:
                    raise ValueError(f"var_type can only be list or dict but you're having {type(var.var_type)}.")
            except:
                print(f'Error is in variable', var.global_name)
                raise ValueError('Cannot register the variable!')

        for key, val in vars_code.items():
            self.vars_code[key] = val

    def collect_var_symbol(self, var_types: Iterable = ('x', 'z', 'p')):
        self.vars_symbol = self._collect_var_prop('symbol', var_types)
        return self.vars_symbol

    def collect_var_val(self, var_types: Iterable = ('x', 'z', 'p')):
        self.vars_val = self._collect_var_prop('val', var_types)
        return self.vars_val

    def collect_var_lb(self, var_types: Iterable = ('x', 'z', 'p')):
        self.vars_lb = self._collect_var_prop('lb', var_types)
        return self.vars_lb

    def collect_var_ub(self, var_types: Iterable = ('x', 'z', 'p')):
        self.vars_ub = self._collect_var_prop('ub', var_types)
        return self.vars_ub

    def collect_var_symbol_out(self, var_types: Iterable = ('x', 'z', 'p')):
        self.vars_symbol_out = self._collect_var_prop('symbol_out', var_types)

        return self.vars_symbol_out

    def _collect_var_prop(self, prop: str, var_types: Iterable[str]):
        """Concatenate a property of all variables
           prop can be 'val', 'lb', 'ub', 'symbol'
        """
        prop_all_variables = {key: [] for key in var_types}
        for var_type in var_types:
            vars_code = self.vars_code[var_type]
            for code in vars_code:
                try:
                    prop_one_variable = getattr(self.vars[code[0]], prop)[code[1:]]
                except:
                    raise ValueError('Error when get {}'.format(prop), code)

                if isinstance(prop_one_variable, list):
                    prop_all_variables[var_type] += prop_one_variable
                else:
                    prop_all_variables[var_type].append(prop_one_variable)

        prop_all_variables_cat = {key: cas.vcat(val) for key, val in prop_all_variables.items()}
        return prop_all_variables_cat

    def collect_var_name(self, var_type='p'):
        """ Collect the full names of all the variables belonging to given var_type
            var_type can be p, x or z"""
        var_names = []

        for vtype, vars_code in self.vars_code.items():
            if vtype == var_type:
                for i, code in enumerate(vars_code):
                    name = [s if j == 0 else '[{}]'.format(s) for j, s in enumerate(code)]
                    var_name = ''.join(name)
                    var_names.append(var_name)

        self.var_names = var_names

        return var_names

    def combine_ode(self, *args):
        """Combine ODE equations into one single list"""
        eq_list = self._combine_eq(*args)
        eq_list = [self.time_scale * eq for eq in eq_list]
        return eq_list

    def combine_alg(self, *args):
        """Combine algebraic equations into one single list"""
        return self._combine_eq(*args)

    def combine_init_eq(self, *args):
        inits = self._combine_eq(*args)

        if len(self.algs) == 0:
            raise Warning("It seems algebraic equations haven't been collected yet!")
        inits += self.algs

        init_eq_for_estimation = self.add_init_eq_for_estimation()
        init_eq_for_estimation = self._combine_eq(*init_eq_for_estimation)
        inits += init_eq_for_estimation

        return inits

    def add_init_eq_for_estimation(self) -> list:
        return []

    @staticmethod
    def _combine_eq(*args):
        """ Combine all the equations into one single list """
        eq_list = []
        for eq in args:
            if isinstance(eq, list) or isinstance(eq, tuple):
                eq_list += eq
            else:
                eq_list.append(eq)
        return eq_list

    def load_var_values(self, sol_val, var_type):
        m, n = sol_val.shape
        if n == 1:
            sol_val = sol_val.T

        for i, code in enumerate(self.vars_code[var_type]):
            self.vars[code[0]].val[code[1:]] = sol_val[-1, i]

    def load_var_trajectory(self, sol_val: Union[np.ndarray, cas.MX, cas.DM],
                            var_type: str, batch_i=0):
        """
            sol_val, are trajectory values for all the x or z.
            It is 2d with row representing time and column for different indexed variables

            var_type, can be x or z
            n_t is the
        """
        n_t = sol_val.shape[0]  # the number of time points

        if list(self.vars.values())[0].trajectory_batch == batch_i:
            # More time points come, extend trajectory
            for var in self.vars.values():
                var.extend_val_trajectory(n_t, batch_i)

        # Set trajectory values
        for i, code in enumerate(self.vars_code[var_type]):
            self.set_var_component_trajectory(code, sol_val[:, i])

    def add_vars_from_submodels(self):
        """Add variables in all the submodel"""
        vars_submodel_new = OrderedDict()
        for subm in self.submodels:
            # Get variable dictionary
            vars_submodel = subm.vars

            # Rename all the variables in submodel
            new_names = []
            for key, var in vars_submodel.items():
                var.global_name = f'{self.name}.{key}'
                new_names.append(var.global_name)

            # Generate a new variable dictionary containing variables in a submodel
            vars_submodel_new.update(OrderedDict(zip(new_names, vars_submodel.values())))

        # Merge the two dictionaries
        self.vars.update(vars_submodel_new)

    def add_eqs_from_submodels(self):
        self.__add_eqs_from_submodels('odes')
        self.__add_eqs_from_submodels('algs')
        self.__add_eqs_from_submodels('inits')

    def __add_eqs_from_submodels(self, eq_type):
        eqs = getattr(self, eq_type)
        for submodel in self.submodels:
            eqs += getattr(submodel, eq_type)
        setattr(self, eq_type, eqs)

    def get_var_symbol(self, var_name: str) -> Union[cas.SX, SymDict]:
        return self.vars[f'{self.name}.{var_name}'].symbol

    def get_var_value(self, var_name: str) -> Union[ValueList, ValueDict]:
        return self.vars[f'{self.name}.{var_name}'].val

    def get_var_fixed(self, var_name: str) -> Union[ValueList, ValueDict]:
        return self.vars[f'{self.name}.{var_name}'].fixed

    def add_init_eq(self, *args):
        self.inits = self.combine_init_eq()

    def get_var_component_trajectory(self, var_code: VariableCodeType) -> Union[np.ndarray, cas.MX]:
        return self.vars[var_code[0]].val_trajectory.getitems(key=var_code[1:])

    def set_var_component_trajectory(self, var_code: VariableCodeType, val: Union[np.ndarray, cas.MX]):
        n_t = val.shape[0]
        self.vars[var_code[0]].val_trajectory.setitems(range(-n_t, 0), var_code[1:], val)

    def get_var_component_val(self, var_code: VariableCodeType) -> ValueComponentType:
        return self._get_var_component_prop(var_code, 'val')

    def set_var_component_val(self, var_code: VariableCodeType, val: ValueComponentType):
        self._set_var_component_prop(var_code, val, 'val')

    def get_var_component_sym(self, var_code: VariableCodeType) -> ValueComponentType:
        return self._get_var_component_prop(var_code, 'symbol_out')

    def set_var_component_sym(self, var_code: VariableCodeType, val: ValueComponentType):
        try:
            self._set_var_component_prop(var_code, val, 'symbol_out')
        except:
            print(var_code)

    def _get_var_component_prop(self, var_code: VariableCodeType, prop: str) -> ValueComponentType:
        # Note: type(prop) must be within ['val', 'lb', 'ub', 'symbol']
        return getattr(self.vars[var_code[0]], prop)[var_code[1:]]

    def _set_var_component_prop(self, var_code: VariableCodeType, val: ValueComponentType, prop: str):
        # Note: type(prop) must be within ['val', 'lb', 'ub', 'symbol']
        getattr(self.vars[var_code[0]], prop)[var_code[1:]] = val

    @staticmethod
    def get_alg_fun(x, z, algs):
        alg_fun = cas.Function('alg_fun', [x, z], [algs])
        return alg_fun

    @staticmethod
    def get_ode_rhs_fun(x, z, odes):
        ode_rhs_fun = cas.Function('ode_rhs_fun', [x, z], [odes])
        return ode_rhs_fun

    @property
    def time_unit(self):
        return self._time_unit

    @time_unit.setter
    def time_unit(self, time_unit):
        self._time_unit = time_unit
        self.time_scale = time_scale_table[time_unit] / time_scale_table[self.time_unit_default]


class Variable:
    def __init__(self, name, index_list=None, m=1, lb=-np.inf, ub=np.inf, model=None):

        """var_type can be a string, or a dictionary of string, or a dictionary of lists of string,
                but self.var_type must be a dictionary of lists of string
        """
        self.name = name
        self.global_name = f'{model.name}.{name}' if model is not None else name
        self.index_list = index_list
        self.m = m      # The length of a single casadi symbol
        self.model = model

        # Define symbol and get number of index
        if index_list is None:
            n_index = 1
        else:
            n_index = len(index_list)
        self.length = m * n_index
        self.size = (m, n_index)

        # Properties to be defined in the specific child class
        self._lb = None
        self._ub = None
        self._val = None
        self._var_type = None
        self._fixed = None

        self.trajectory_batch = 0
        self.val_trajectory = None

        self.symbol = None
        self.symbol_out = None

    def __call__(self):
        return self.symbol

    @property
    def val(self):
        return self._val

    @val.setter
    def val(self, value):
        self._val.resetitems(value)

    @property
    def lb(self):
        return self._lb

    @lb.setter
    def lb(self, lb):
        self._lb.resetitems(lb)

    @property
    def ub(self):
        return self._ub

    @ub.setter
    def ub(self, ub):
        self._ub.resetitems(ub)

    @property
    def var_type(self):
        return self._var_type

    @var_type.setter
    def var_type(self, var_type):
        self._var_type.resetitems(var_type)

    @property
    def fixed(self):
        return self._fixed

    @fixed.setter
    def fixed(self, fixed):
        self._fixed.resetitems(fixed)

    def extend_val_trajectory(self, n_t, batch):
        """Initialize or augment val_trajectory to store trajectory values"""

        if batch == 0:
            self.initialize_trajectory()
        self.val_trajectory.extend_trajectory(n_t)
        self.trajectory_batch += 1

    def initialize_trajectory(self):
        """Initialize val_trajectory when trajectory_batch==0"""
        raise NotImplementedError


class VariableList(Variable):
    def __init__(self, name,
                 var_type,
                 index_list=None,
                 m=1,
                 val=None,
                 lb=-np.inf,
                 ub=np.inf,
                 model=None,
                 fixed: Optional[Union[list[bool], tuple[bool], bool]] = None,
                 default=1.):

        """var_type can be a string, or a dictionary of string, or a dictionary of lists of string,
                but self.var_type must be a dictionary of lists of string
        """
        super().__init__(name, index_list, m, lb, ub, model)

        self.default_value = default
        self._val = ValueList(m, val, default)
        self._lb = ValueList(m, lb, -np.inf)
        self._ub = ValueList(m, ub, np.inf)
        self._var_type = StrList(m, var_type)

        self.symbol = cas.SX.sym(name, m)
        self.symbol_out = ValueList(m, val, default)

        if fixed is None:
            fixed = [True if v_type != 'z' else False for v_type in self._var_type]
        self._fixed = BoolList(m, fixed)

    def initialize_trajectory(self):
        self.val_trajectory = TrajectoryList(0, self.m)


class VariableDict(Variable):
    def __init__(self,
                 name,
                 var_type,
                 index_list=None,
                 m=1,
                 val=None,
                 lb=-np.inf,
                 ub=np.inf,
                 model=None,
                 fixed: Optional[Union[bool, dict[Any, bool], tuple[bool, ...], list[bool]]] = None,
                 default=1.):

        """var_type can be a string, or a dictionary of string, or a dictionary of lists of string,
                but self.var_type must be a dictionary of lists of string
        """
        super().__init__(name, index_list, m, lb, ub, model)

        self._var_type = StrDict(index_list, m, var_type)

        sym_dict = utility.SXDict(name, index_list, m)
        self.symbol = SymDict(index_list, m, sym_dict)

        self.default_value = default
        self._val = ValueDict(index_list, m, val, default)

        self.symbol_out = ValueDict(index_list, m, val, default)

        if fixed is None:
            self._fixed = copy.deepcopy(self._var_type)
            for key, vals in self.var_type.items():
                for i, vi in enumerate(vals):
                    if vi != 'z':
                        self._fixed[(key, i)] = True
                    else:
                        self._fixed[(key, i)] = False
        else:
            self._fixed = BoolDict(index_list, m, fixed)

    def initialize_trajectory(self):
        self.val_trajectory = TrajectoryDict(0, self.index_list, self.m)


if __name__ == '__main__':
    myval = ValueDict({'a': [1], 'b': [2]}, 1)
    a0_ori = myval[('a', 0)]
    myval[('a', 0)] = 5
    assert a0_ori != myval['a', 0]

    myval2 = ValueList(3, [1, 2, 3])
    print(myval)
    print(myval2)

    traj = TrajectoryList(3, 1)
    print(traj[1, 0])
    traj[1, 0] = 10
    print(traj[1, 0])

    traj.setitems(range(3), 0, np.array([100, 10, 200]))
    print(traj[2, 0])

    print(traj.getitems(list(range(3)), 0))

    traj.extend_trajectory(5)
    traj.setitems(range(3, 8), 0, np.array([1, 2, 3, 4, 5]))
    print(traj.getitems(list(range(-5, 0))))

    traj.setitems(range(-5, 0), 0, np.array([10, 20, 30, 40, 50]))
    print(traj.getitems(list(range(-5, 0))))