
import os
import numpy as np
from scipy.interpolate import interp1d

directory = os.path.dirname(__file__)

class DDE:
    """The class DDE uses the C++ software "oddtool" to solve delay differential equations.
    
    :param time_grid: An array which contains the time points for which the states
        of the system should be stored. If ``time_grid`` is ``None``, the solution states
        at the time points chosen by oddtool's adaptive stepper will be stored.
    :type time_grid: numpy.ndarray or None
    :param dimension: The dimension of the system.
    :type dimension: int
    :param rhs: Contains the right hand sides
        of the system (i.e. the dynamics) written in C++ code. The list
        must have the length ``dimension``. If ``dimension`` is ``1``, ``rhs`` can also be a string.
        For a system of two delay-coupled Mackey-Glass oscillators ``rhs`` would be
        ``['p[0] * xd[0][1] / (1 + pow(xd[0][1], p[1])) - p[2] * x[0]',
        'p[0] * xd[0][0] / (1 + pow(xd[0][0], p[1])) - p[2] * x[1]']``.
        The C++ variable names ``p``, ``x`` and ``xd`` are reservered for the parameter array,
        the system state array and the delayed state array. The values ``p[0]``, ``p[1]``, ...
        must be provides by the below argument rhs_parameters. The variables ``x[0]``, ``x[1]``, ...
        are the first, second, ... component of the current systems state.
        The variable ``xd[j][i]`` denotes the :math:`i`-th component of the system state at
        time :math:`t - \\tau _j`.
    :type rhs: list of strings or single string
    :param rhs_parameters: An array which contains the parameters for the right hand sides.
        For the example with two delay-coupled Mackey-Glass oscillators this array could
        look as follows: ``array([2.2, 10.0, 1.0])``.
    :type rhs_parameters: numpy.ndarray of numpy floats or numpy ints
    :param delays: An array which contains the constant delay times.
        If there is only one delay time, this parameter can also be a float.
        For the example with two coupled Mackey-Glass oscillators this parameter could
        look as foolows: ``array([20])``. Since there is only one delay time for this
        system, the parameter can also be a simple float or int: ``20``.
    :type delays: numpy.ndarray of numpy floats or numpy ints or single float or int
    :param history: Contains history terms written in C++ code. The history term are functions
        for the system states at and before the initial time point.
        The list must have the length ``dimension``. For the example with two
        coupled Mackey-Glass oscillators this list could look as follows:
        ``['abs(cos(t)) + 1', '1 + sin(t)']``.
    :type history: list of strings or single string
    :param additional_lines: If a string is provided, it will be inserted in the
        rhs function in oddtool's C++ main file. (See example 2 below.)
    :type additional_lines: None or string
    :param t_start: If ``time_grid`` is ``None``, ``t_start`` will be used as the first
        time of the solution. If ``time_grid`` is not ``None``, ``t_start`` does not matter.
    :type t_start: float, int or None
    :param t_end: If ``time_grid`` is ``None``, ``t_end`` will be used as the last
        time of the solution. If ``time_grid`` is not ``None``, ``t_end`` does not matter.
    :type t_end: float, int or None
    :param h_max: Maximal time step for the integration. It should be
        smaller than the distance of time points in the parameter ``time_grid``.
    :type h_max: float or int
    :param h_min: Minimal time step for the integration. It should be
        smaller than or equal to ``h_max``.
    :type h_min: float or int
    :type terminal_output: bool
    :param terminal_output: If ``True``, oddtool's terminal output will be printed.
        
    :attributes:
        * **solution** (*numpy.ndarray*) -- The solution values and time points are stored here
          after running the method ``solve()``. The time points are stored in the 0-th columns, and the
          i-th component of system states are stored in the i-th column of the array. 
    """
    def __init__(self, time_grid, dimension, rhs, rhs_parameters, delays,
                 history, additional_lines=None, t_start=None, t_end=None, h_max=1, h_min=0.00001,
                 terminal_output=False):
        self.__initiation_done = False
        self.time_grid = time_grid
        self.dimension = dimension
        self.rhs = rhs
        self.rhs_parameters = rhs_parameters
        self.delays = delays
        self.history = history
        self.additional_lines = additional_lines
        self.t_start = t_start
        self.t_end = t_end
        self.h_max = h_max
        self.h_min = h_min
        self.terminal_output = terminal_output        
        self.__initiation_done = True
    
    @property
    def time_grid(self):
        return self.__time_grid
        
    @time_grid.setter
    def time_grid(self, time_grid):
        if time_grid is not None and type(time_grid) != np.ndarray:
            raise TypeError('time_grid must be a numpy.ndarray or None')
        if time_grid is not None and time_grid.dtype != np.float64 and time_grid.dtype != np.int64:
            raise TypeError('The time_grid array must have the dtype numpy.float64 or numpy.int64.')
        self.__time_grid = time_grid
        
    @property
    def dimension(self):
        return self.__dimension
    
    @dimension.setter
    def dimension(self, dimension):
        if type(dimension) != int:
            raise TypeError('dimension must be an int.')
        if dimension < 1:
            raise ValueError('dimension must be positive.')
        if self.__initiation_done:
            raise AttributeError('dimension cannot be changed. ' + \
                                 'Create a new instance of DDE if ' + \
                                 'you want to integrate a system with a different dimension.')
        self.__dimension = dimension
        
    @property
    def rhs(self):
        return self.__rhs
    
    @rhs.setter
    def rhs(self, rhs):
        if self.dimension != 1 and type(rhs) != list:
            raise TypeError('If dimesion is not 1, rhs must be a list.')
        if type(rhs) != list and type(rhs) != str:
            raise TypeError('rhs must be a list or a string.')
        if type(rhs) == list:
            for entry in rhs:
                if type(entry) != str:
                    raise TypeError('The list rhs must contain strings.')
            if len(rhs) != self.dimension:
                raise ValueError('The list rhs must have the length dimension.')
        if self.__initiation_done:
            raise AttributeError('rhs cannot be changed. ' + \
                                 'Create a new instance of DDE if ' + \
                                 'you want to integrate a different system.')
        self.__rhs = rhs
        
    @property
    def rhs_parameters(self):
        return self.__rhs_parameters
    
    @rhs_parameters.setter
    def rhs_parameters(self, rhs_parameters):
        if type(rhs_parameters) != np.ndarray:
            raise TypeError('rhs_parameters must be a numpy.ndarray')
        if rhs_parameters.dtype != np.float64 and rhs_parameters.dtype != np.int64:
            raise TypeError('The rhs_parameters array must have the dtype numpy.float64 or numpy.int64.')
        self.__rhs_parameters = rhs_parameters
        
    @property
    def delays(self):
        return self.__delays
    
    @delays.setter
    def delays(self, delays):
        if type(delays) != np.ndarray and type(delays) != float and type(delays) != int:
            raise TypeError('delays must be a numpy.ndarray, a float or an int.')
        if type(delays) == np.ndarray and delays.dtype != np.float64 and delays.dtype != np.int64:
            raise TypeError('delays array must have the dtype numpy.float64 or numpy.int64.')
        self.__delays = delays
    
    @property
    def history(self):
        return self.__history
    
    @history.setter
    def history(self, history):
        if self.dimension != 1 and type(history) != list:
            raise TypeError('If dimesion is not 1, history must be a list.')
        if type(history) != list and type(history) != str:
            raise TypeError('history must be a list or a string.')
        if type(history) == list:
            for entry in history:
                if type(entry) != str:
                    raise TypeError('The list history must contain strings.')
            if len(history) != self.dimension:
                raise ValueError('The list history must have the length dimension.')
        self.__history = history
    
    @property
    def additional_lines(self):
        return self.__additional_lines
    
    @additional_lines.setter
    def additional_lines(self, additional_lines):
        if additional_lines is not None and type(additional_lines) != str:
            raise TypeError('additional_lines must be None or a string.')
        self.__additional_lines = additional_lines
    
    @property
    def t_start(self):
        return self.__t_start
    
    @t_start.setter
    def t_start(self, t_start):
        if t_start is not None and type(t_start) != float and type(t_start) != int:
            raise TypeError('t_start must be a float, an int or None.')
        if self.time_grid is None and t_start is None:
            raise TypeError('If time_grid is None, t_start must be a number.')
        self.__t_start = t_start
        
    @property
    def t_end(self):
        return self.__t_end
    
    @t_end.setter
    def t_end(self, t_end):
        if t_end is not None and type(t_end) != float and type(t_end) != int:
            raise TypeError('t_end must be a float, an int or None.')
        if self.time_grid is None and t_end is None:
            raise TypeError('If time_grid is None, t_end must be a number.')
        if((type(self.t_start) == int or type(self.t_start) == float) and
           (type(t_end) == int or type(t_end) == float) and
           (self.t_start >= t_end)):
            raise ValueError('t_end must be larger than t_start.')
        self.__t_end = t_end
    
    @property
    def h_max(self):
        return self.__h_max
    
    @h_max.setter
    def h_max(self, h_max):
        if type(h_max) != float and type(h_max) != int:
            raise TypeError('h_max must be a float or an int.')
        self.__h_max = h_max
    
    @property
    def h_min(self):
        return self.__h_min
    
    @h_min.setter
    def h_min(self, h_min):
        if type(h_min) != float and type(h_min) != int:
            raise TypeError('h_min must be a float or an int.')
        if h_min > self.h_max:
            raise ValueError('h_min must be smaller than or equal to h_max.')
        self.__h_min = h_min
        
    @property
    def terminal_output(self):
        return self.__terminal_output
    
    @terminal_output.setter
    def terminal_output(self, terminal_output):
        if type(terminal_output) != bool:
            raise TypeError('terminal_output must be bool.')
        self.__terminal_output = terminal_output


    # prepare "oddtool/ODD_parameters.txt" file
    def _prepare_parameters_file(self):
        """This method creates the string with the parameters for oddtool
        in the desired format and stores it to the files
        "oddtools/ODD_parameters.txt".
        
        :param a: A letter.
        :type a: str.
        :returns: None
        """
        param_string = "N = " + str(self.dimension) + "\n"
        param_string += "name = python_call\n"
        index = 0
        for rhs_param in self.rhs_parameters:
            param_string += "p" + str(index) + " = " + str(rhs_param) + "\n"
            index += 1
        if type(self.delays) != np.ndarray:
            param_string += "tau0 = " + str(self.delays) + "\n"
        else:
            index = 0
            for tau_j in self.delays:
                param_string += "tau" + str(index) + " = " + str(tau_j) + "\n"
                index += 1
        if self.time_grid is not None:
            param_string += "t_start = " + str(self.time_grid[0]) + "\n"
            param_string += "t_end = " + str(self.time_grid[-1] + 20*self.h_max) + "\n"
        else:
            param_string += "t_start = " + str(self.t_start) + "\n"
            param_string += "t_end = " + str(self.t_end) + "\n"
        param_string += "hmax = " + str(self.h_max) + "\n"
        param_string += "hmin = " + str(self.h_min) + "\n"
        
        with open(directory + "/oddtool/ODD_parameters.txt", "w") as f:
            f.write(param_string)
        return
    


    # prepare "oddtool/main.cpp" file
    def _prepare_cpp_main_file(self):
        """This method creates the code for the C++ main file
        and stores it to the file "oddtool/main.cpp".
        """
        if type(self.history) == str:
            history_code = 'result[0] = ' + self.history + ';\n'
        else:
            index = 0
            history_code = ''
            for line in self.history:
                history_code += 'result[' + str(index) + '] = ' + line + ';\n'
                index += 1
        if type(self.rhs) == str:
            if self.additional_lines is None:
                rhs_code = 'result[0] = ' + self.rhs + ';\n'
            else:
                rhs_code = self.additional_lines + '\n' + 'result[0] = ' + self.rhs + ';\n'
        else:
            if self.additional_lines is None:
                rhs_code = ''
            else:
                rhs_code = self.additional_lines + '\n'
            index = 0
            for line in self.rhs:
                rhs_code += 'result[' + str(index) + '] = ' + line + ';\n'
                index += 1
        code =  '// file main.cpp\n' + \
                '\n' + \
                '#include "ODDtool.h"\n' + \
                '\n' + \
                '\n' + \
                'using namespace std;\n' + \
                '\n' + \
                '// user-supplied history function\n' + \
                'void myHistory(double t, vector<double>& result) {\n' + \
                history_code + \
                '}\n' + \
                '\n' + \
                '// systems equations\n' + \
                'void rhs(double t, const vector<double>& x, ODD_delayed_values& xd, const vector<double>& p, vector<double>& result) {\n' + \
                rhs_code + \
                '}\n' + \
                '\n' + \
                'int main() { \n' + \
                '// create an integrator\n' + \
                'ODD_integrator<CKStepper> integrator(rhs, myHistory);\n' + \
                '\n' + \
                '// start integration over time interval specified in parameter file (ODD_parameters.txt) \n' + \
                'integrator.integrate(); \n' + \
                '}\n'        
        with open(directory + "/oddtool/main.cpp", "w") as f:
            f.write(code)
        return
        

    # compile the cpp source files
    def _compile_cpp_files(self):
        """This method makes a system call to compile the C++ source files
        in the "oddtool" directory.
        """
        if self.terminal_output:
            os_return_value = os.system("(cd " + directory + "/oddtool && make)")
        else:
            os_return_value = os.system("(cd " + directory + "/oddtool && make >/dev/null 2>&1)")
        if os_return_value != 0:
            raise OSError('Was not able to compile C++ file. Delete the following files ' +
                          '(if existing) and try again: the main.cpp, ' +
                          'the oddtool executable, the ODD_parameters.txt, and ' +
                          ' the main.o file in the oddtool directory, ' +
                          'all .o files in the oddtool/sources directory.')
        return


    # run the compiled program
    def _run_oddtool(self):
        """This method makes a system call to run the compiled oddtool file.
        An output files "oddtool/python_call_output/x.txt" will be created.
        """
        if self.terminal_output:
            os_return_value = os.system("(cd " + directory + "/oddtool && ./oddtool)")
        else:
            os_return_value = os.system("(cd " + directory + "/oddtool && ./oddtool >/dev/null 2>&1)")
        if os_return_value != 0:
            raise OSError('Was not able to run the oddtool executable.')
        return


    # read output file
    
    @property
    def solution(self):
        return self.__solution

    def solve(self):
        """This method integrates the DDE. The solution will be stored in the attribute ``DDE.solution``."""
        
        if self.time_grid is None:
            if self.t_start is None:
                raise ValueError('If time_grid is None, t_start must be a number.')
            if self.t_end is None:
                raise ValueError('If time_grid is None, t_end must be a number.')
            if self.t_start >= self.t_end:
                raise ValueError('t_start must be smaller than t_end.')
        
        if self.dimension == 1:
            if type(self.history) == list and len(self.history) != 1:
                raise ValueError('If dimension is 1, only 1 history function must be given.')
        else:
            if type(self.history) != list:
                raise ValueError('If dimension is not 1, history must be a list.')
            elif len(self.history) != self.dimension:
                raise ValueError('The history list must be of length dimension.')
        
        if self.h_max < self.h_min:
            raise ValueError('h_max must be larger than or equal to h_min.')
        
        self._prepare_parameters_file()
        self._prepare_cpp_main_file()
        self._compile_cpp_files()
        self._run_oddtool()
        
        results_list = []
        with open(directory + "/oddtool/python_call_output/x.txt", "r") as f:
            for line in f:
                results_list.append(line.strip().split(" "))
        results_array = np.array(results_list, dtype=np.float64)
        # 2-dimensional array
        # results_array[i, 0] is the i-th time point t_i used by oddtool
        # results_array[i, j] is the j-th component of the system state
        # at time t_i for 1 <= j <= dimension
        if self.time_grid is None:
            self.__solution = results_array
        else:
            sol = self.time_grid
            for j in range(self.dimension):
                interpolant = interp1d(results_array[:, 0], results_array[:, j + 1], kind='cubic')
                results_j = interpolant(self.time_grid)
                sol = np.vstack((sol, results_j))
            self.__solution = sol.T
        return
