#!/usr/bin/env python 
# -*- coding: utf-8 -*-

# Copyright (C) 2010 Modelon AB
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import numpy as N
cimport numpy as N
import time
import pylab as P

import itertools

from exception import *
from problem import Explicit_Problem, Delay_Explicit_Problem, Implicit_Problem, SingPerturbed_Problem

include "constants.pxi" #Includes the constants (textual include)

realtype = N.float 

cdef class ODE:
    """
    Base class for all our integrators.
    """
    
    def __init__(self, problem):
        """
        Defines general starting attributes for a simulation
        problem.
        """
        self.statistics = {} #Initialize the statistics dictionary
        self.options = {"report_continuously":False,"verbosity":NORMAL,"backward":False, "store_event_points":True, "time_limit":0, "clock_step":False}
        #self.internal_flags = {"state_events":False,"step_events":False,"time_events":False} #Flags for checking the problem (Does the problem have state events?)
        self.supports = {"state_events":False,"interpolated_output":False,"report_continuously":False,"sensitivity_calculations":False,"interpolated_sensitivity_output":False} #Flags for determining what the solver supports
        self.problem_info = {"dim":0,"dimRoot":0,"dimSens":0,"state_events":False,"step_events":False,"time_events":False
                             ,"jac_fcn":False, "sens_fcn":False, "jacv_fcn":False,"switches":False,"type":0,"jaclag_fcn":False,'prec_solve':False,'prec_setup':False}
        #Type of the problem
        #0 = Explicit
        #1 = Implicit
        
        #Data object for storing the event data
        self.event_data = []
        self._event_info = N.array([])
        
        if problem is None:
            raise ODE_Exception('The problem needs to be a subclass of a Problem.')
        
        #Check Problem for event functions
        if hasattr(problem, 'time_events'):
            self.problem_info["time_events"] = True
        
        if hasattr(problem, 'state_events'):
            self.problem_info["state_events"] = True
        
        if hasattr(problem, 'step_events'):
            self.problem_info["step_events"] = True
        
        if hasattr(problem, 'y0'):
            self.y0 = N.array(problem.y0,dtype=realtype) if len(N.array(problem.y0,dtype=realtype).shape)>0 else N.array([problem.y0],dtype=realtype)
            self.problem_info["dim"] = len(self.y0)
        else:
            raise ODE_Exception('y0 must be specified in the problem.')
            
        if hasattr(problem, 'neq'):  # relevant for overdetermined problems: neq=number of equations >= dim
            self.problem_info["neq"] = problem.neq
        else:
            self.problem_info["neq"] = self.problem_info["dim"]
        
        if hasattr(problem, "p0"):
            self.p0 = N.array(problem.p0,dtype=realtype) if len(N.array(problem.p0,dtype=realtype).shape)>0 else N.array([problem.p0],dtype=realtype)
            self.problem_info["dimSens"] = len(self.p0)
            self.p = self.p0.copy()
        
        if hasattr(problem, "sw0"):
            self.sw0 = N.array(problem.sw0,dtype=N.bool) if len(N.array(problem.sw0,dtype=N.bool).shape)>0 else N.array([problem.sw0],dtype=N.bool)
            self.problem_info["switches"] = True
            self.sw = self.sw0.tolist()
        
        if hasattr(problem, 't0'):
            self.t0 = float(problem.t0)
        else:
            self.t0 = 0.0
            
        if hasattr(problem, "jac"):
            self.problem_info["jac_fcn"] = True
        if hasattr(problem, "jacv"):
            self.problem_info["jacv_fcn"] = True
        if hasattr(problem, "jaclag"):
            self.problem_info["jaclag_fcn"] = True
        if hasattr(problem, "prec_solve"):
            self.problem_info["prec_solve"] = True
        if hasattr(problem, "prec_setup"):
            self.problem_info["prec_setup"] = True
            
        #Reset solution variables
        self._reset_solution_variables()
        
        #Specify storing of sensitivity to 0
        problem._sensitivity_result = 0
        
        #Initialize timer
        self.elapsed_step_time = -1.0
        self.clock_start = -1.0
        
    def __call__(self, double tfinal, int ncp=0, list cpts=None):
        return self.simulate(tfinal, ncp, cpts)
        
    cdef _reset_solution_variables(self):
        """
        Resets solution variables.
        """
        self.t_sol = []
        self.y_sol = []
        self.yd_sol = []
        self.p_sol = [[] for i in range(self.problem_info["dimSens"])]
        
        
    cpdef simulate(self, double tfinal, int ncp=0, object ncp_list=None):
        """
        Calls the integrator to perform the simulation over the given time-interval. 
        If a second call to simulate is performed, the simulation starts from the last
        given final time.
        
            Parameters::
            
                tfinal  
                        - Final time for the simulation
                
                        - Should be a float or integer greater than the initial time.
                        
                ncp     
                        - Default '0'. Number of communication points where the 
                          solution is returned. If '0', the integrator will return 
                          at its internal steps.
                          
                        - Should be an integer.
                        
                ncp_list
                        - Default None. A list of time points where the solution
                          should be returned. Note, requires that ncp == 0.
                          
                    Example:
                    
                        simulate(10.0, 100), 10.0 is the final time and 100 is the number
                                             communication points.
                 
        """
        t0 = self.t
        
        #Reset solution variables
        self._reset_solution_variables()
        
        #Error checking
        try:
            tfinal = float(tfinal)
        except ValueError:
            raise AssimuloException('Final time must be an integer or float.')
            
        if (self.t > tfinal) and not self.options["backward"]:
            raise AssimuloException('Final time must be greater than start time.')
        
        if not isinstance(ncp, int):
            raise AssimuloException('Number of communication points must be an integer')
        
        if ncp < 0:
            ncp = 0
            self.log_message('Number of communication points must be a positive integer, setting ncp = 0.',WARNING)
        
        #Check solver support against current problem
        if self.problem_info["state_events"] and self.supports["state_events"] is False:
            self.log_message("The current solver does not support state events (root functions). Disabling and continues.", WHISPER)
            self.problem_info["state_events"] = False
        
        if self.problem_info["step_events"] and self.supports["report_continuously"] is False:
            self.log_message("The current solver does not support step events (report continuously). Disabling step events and continues.", WHISPER)
            self.problem_info["step_events"] = False
        
        if self.supports["report_continuously"] is False and self.options["report_continuously"]:
            self.log_message("The current solver does not support to report continuously. Setting report_continuously to False and continues.", WHISPER)
            self.options["report_continuously"] = False
        
        if (ncp != 0 or ncp_list != None) and (self.options["report_continuously"] or self.problem_info["step_events"]) and self.supports["interpolated_output"] is False:
            self.log_message("The current solver does not support interpolated output. Setting ncp to 0 and ncp_list to None and continues.", WHISPER)
            ncp = 0
            ncp_list = None
            
        if (ncp != 0 or ncp_list != None) and self.problem_info["state_events"] and self.supports["report_continuously"] is False:
            self.log_message("The current solver does not support interpolated output together with state events. Setting ncp to 0 and ncp_list to None and continues.", WHISPER)
            ncp = 0
            ncp_list = None
        elif (ncp != 0 or ncp_list != None) and self.problem_info["step_events"] and self.supports["report_continuously"]:
            if not self.report_continuously:
                 self.log_message("The problem contains step events: report_continuously is set to True", WHISPER)
            self.report_continuously = True
            
        #Determine the output list
        if ncp != 0:
            output_list = N.linspace(t0,tfinal,ncp+1)[1:]
            output_index = 0
        elif ncp_list != None:
            output_list = N.array(ncp_list, dtype=realtype, ndmin=1)[N.array(ncp_list, dtype=realtype, ndmin=1)>t0][N.array(ncp_list, dtype=realtype, ndmin=1)<tfinal]
            if output_list[-1] < tfinal: #Add the last point if necessary!
                output_list = N.append(output_list, tfinal)
            output_index = 0
        else:
            output_list = None
            output_index = 0
        
        #Determine if we are using one step mode or normal mode
        if self.problem_info['step_events'] or self.options['report_continuously']:
            REPORT_CONTINUOUSLY = 1
        else:
            REPORT_CONTINUOUSLY = 0
        
        #Determine if the output should be interpolated or not
        if output_list == None:
            INTERPOLATE_OUTPUT = 0
        else:
            INTERPOLATE_OUTPUT = 1

        #Time and Step events
        TIME_EVENT = 1 if self.problem_info['time_events'] is True else 0

        #Simulation starting, call initialize
        self.problem.initialize(self)
        self.initialize()
        
        #Start of simulation, start the clock
        time_start = time.clock()
        
        #Start the simulation
        self._simulate(t0, tfinal, output_list, REPORT_CONTINUOUSLY, INTERPOLATE_OUTPUT, TIME_EVENT)
        
        #End of simulation, stop the clock
        time_stop = time.clock()
        
        #Simulation complete, call finalize
        self.finalize()
        self.problem.finalize(self)
        
        #Print the simulation statistics
        self.print_statistics(NORMAL)
        
        #Log elapsed time
        self.log_message('Simulation interval    : ' + str(t0) + ' - ' + str(self.t) + ' seconds.', NORMAL)
        self.log_message('Elapsed simulation time: ' + str(time_stop-time_start) + ' seconds.', NORMAL)
        
        #Return the results
        if isinstance(self.problem, Explicit_Problem) or isinstance(self.problem, Delay_Explicit_Problem) or isinstance(self.problem, SingPerturbed_Problem):
            return self.t_sol, N.array(self.y_sol)
        else:
            return self.t_sol, N.array(self.y_sol), N.array(self.yd_sol)
        
    def _simulate(self,t0, tfinal, output_list, REPORT_CONTINUOUSLY, INTERPOLATE_OUTPUT, TIME_EVENT):
         pass
         
    cpdef initialize(self):
        pass
    
    cpdef finalize(self):
        pass
    
    def _set_verbosity(self, verb):
        try:
            self.options["verbosity"] = int(verb)
        except:
            raise AssimuloException("Verbosity must be an integer.")
    
    def _get_verbosity(self):
        """
        This determines the level of the output. A smaller value
        means more output. The following values can be set:
            
            QUIET   = 50
            WHISPER = 40
            NORMAL  = 30
            LOUD    = 20
            SCREAM  = 10
        
            Parameters::
            
                verb  
                        - Default 30 (NORMAL)
                    
                        - Should be a integer.

        """
        return self.options["verbosity"]
    
    verbosity = property(_get_verbosity,_set_verbosity)
    
    def _set_time_limit(self, time_limit):
        if time_limit < 0:
            raise AssimuloException("The time limit must be positive or zero.")
        self.options["time_limit"] = time_limit
        
    def _get_time_limit(self):
        """
        This option can be used to limit the time of an integration. I.e
        to set an upper bound on the time allowed for the integration
        to be completed. The time limit is specified in seconds. For the
        limit to be checked, the option report_continuously must be True.
        
            Parameters::
            
                time_limit
                            - Default 0, i.e. NO limit.
        """
        return self.options["time_limit"]
    
    time_limit = property(_get_time_limit, _set_time_limit)
    
    def _set_report_continuously(self, report_continuously):
        self.options["report_continuously"] = bool(report_continuously)
    
    def _get_report_continuously(self):
        """
        This options specifies if the solver should report the solution 
        continuously after steps.
        
            Parameters::
            
                report_continuously
                  
                        - Default False
                    
                        - Should be a boolean.

        """
        return self.options["report_continuously"]
    
    report_continuously = property(_get_report_continuously,_set_report_continuously)
    
    def _set_store_event_points(self, store_event_points):
        self.options["store_event_points"] = bool(store_event_points)
    
    def _get_store_event_points(self):
        """
        This options specifies if the solver should save additional points
        at the events, :math:`t_e^-, t_e^+`.
        
            Parameters::
            
                store_event_points
                  
                        - Default True
                    
                        - Should be a Boolean.

        """
        return self.options["store_event_points"]
    
    store_event_points = property(_get_store_event_points,_set_store_event_points)
    
    def _set_clock_step(self, clock_step):
        self.options["clock_step"] = clock_step
    
    def _get_clock_step(self):
        """
        Specifies if the elapsed time of an integrator step should be
        timed or not. Not that this is only possible if running in 
        report continuously mode.
        """
        return self.options["clock_step"]
        
    clock_step = property(_get_clock_step, _set_clock_step)
    
    def _set_backward(self, backward):
        self.options["backward"] = bool(backward)
        
    def _get_backward(self):
        """
        Specifies if the simulation is done in reverse time. (NOTE:
        experimental!)
        
            Parameters::
            
                backward
                
                    - Default False
                    - Boolean
        """
        return self.options["backward"]
        
    backward = property(_get_backward,_set_backward)
    
    cpdef log_message(self, message,int level):
        if level >= self.options["verbosity"]:
            print(message)
            
    cpdef log_event(self,double time,object event_info, int level):
        if level >= self.options["verbosity"]:
            self.event_data.append([time,event_info])
            
    cpdef clear_logs(self):
        """
        Clears the currently stored log messages.
        """
        self.event_data = []
            
    cpdef get_options(self):
        """
        Returns the current solver options.
        """
        return self.options
        
    cpdef get_supports(self):
        """
        Returns the functionality which the solver supports.
        """
        return self.supports
        
    cpdef get_statistics(self):
        """
        Returns the run-time statistics (if any).
        """
        return self.statistics
        
    cpdef get_event_data(self):
        """
        Returns the event information (if any). If there exists information
        about events, it will be returned as a list of tuples where the
        first value is the time of the occured event and the second is the
        event information.
        """
        return self.event_data
        
    cpdef print_event_data(self):
        """
        Prints the event information (if any).
        """
        cdef i = 0
        for i in self.event_data:
            print 'Time, t = %e'%i[0]
            print '  Event info, ', i[1]
        print 'Number of events: ', len(self.event_data)
        
    def print_statistics(self, verbose=NORMAL):
        """
        This method should print the statistics of the solver.
        """
        pass
    
    cpdef get_elapsed_step_time(self):
        """
        Returns the elapsed time of a step. I.e. how long a step took.
        Note that this is only possible if running in report_continuously
        mode and should only be used to get a general idea of how long
        a step really took. 
        
        Returns::
        
            Elapsed time (note -1.0 indicates that it was not used)
        """
        return self.elapsed_step_time
        
    def _compact_atol(self):
        """
        Reduces atol to a scalar if it is an ndarray and  all entries are the same.
        Used for print solver options in a more compact way
        """
        if isinstance(self.atol,N.ndarray) and (self.atol==self.atol[0]).all():
                return self.atol[0]
        else:
                return self.atol
