#################################################################################################################
# ewstools
# Description: Python package for computing, analysing and visualising 
# early warning signals (EWS) in time-series data
# Author: Thomas M Bury
# Web: https://www.thomasbury.net/
# Code repo: https://github.com/ThomasMBury/ewstools
# Documentation: https://ewstools.readthedocs.io/
#
# The MIT License (MIT)
#
# Copyright (c) 2019 Thomas Bury
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#################################################################################################################


#---------------------------------
# Import relevant packages
#--------------------------------

# For numeric computation and DataFrames
import numpy as np
import pandas as pd

def simulate_ricker(tmax=500, tburn=100,
                    r=0.75, k=10, h=0.75,
                    F=[0,2.7],
                    sigma=0.02,
                    x0=0.8):
    '''
    Run a numerical simulation of the Ricker model
    with a Holling Type II harvesting term and additive white noise.
    Allows for linearly increasing/decreasing harvesting rate.
    Default parameter configuration takes model through a Fold bifurcation.
    
    Model configuration is as in Bury et al. (2020) Roy. Soc. Interface
    https://royalsocietypublishing.org/doi/full/10.1098/rsif.2020.0482
    
    Parameters
    ----------
    tmax : int, optional
        Number of time steps. The default is 500.
    tburn : int, optional
        Number of time steps to use as a burn in period
        to remove transients. The default is 100.
    r : float or list, optional
        Intrinsic growth rate. Can be provided as a list containing the start
        and end value if linear change is desired. The default is 0.75.
    k : float, optional
        Population carrying capacity. The default is 10.
    h : float, optional
        Half-saturation constant of the harvesting expression. The default is 0.75.
    F : float or list, optional
        Maximum harvesting rate. Can be provided as a list containing the start
        and end value if linear change is desired. The default is 0.
    sigma : float, optional
        Noise amplitude. The default is 0.02.
    x0 : float, optional
        Initial condition. The default is 0.8.

    Returns
    -------
    pd.Series
        Trajectory indexed by time.  

    '''
    
    # Define the map
    def de_fun(x,r,k,f,h,xi):
        return x*np.exp(r*(1-x/k)+xi) - f*x**2/(x**2+h**2)
            
        
    # Initialise arrays for time and state values
    t = np.arange(tmax)
    x = np.zeros(tmax)
    
    # Get r values, that may be linearly changing
    if (type(r) == float) or (type(r)==int):
        rvals = np.ones(tmax)*r
    else:
        rvals = np.linspace(r[0], r[1], tmax)
    
    # Get F values, that may be linearly changing
    if (type(F) == float) or (type(F)==int):
        Fvals = np.ones(tmax)*F
    else:
        Fvals = np.linspace(F[0], F[1], tmax)
    
    # Array of noise values (normal random variables with variance sigma^2)
    dW_burn = np.random.normal(loc=0, scale=sigma, size = tburn) # burn-in period
    dW = np.random.normal(loc=0, scale=sigma, size = tmax) # monitored period
    
    # Run burn-in period starting from intiial condition x0
    for i in range(int(tburn)):
        x0 = de_fun(x0,rvals[0],k,Fvals[0],h,dW_burn[i])
    
    # State value post burn-in period. Set as starting value.
    x[0]=x0
    
    # Run simulation
    for i in range(tmax-1):
        x[i+1] = de_fun(x[i],rvals[i],k,Fvals[i],h,dW[i])
        # Make sure that state variable stays >= 0
        if x[i+1] < 0:
            x[i+1] = 0
            
    # Store in pd.Series indexed by time
    series = pd.Series(data=x, index=t)
    series.index.name='time'
    
    return series





def simulate_rosen_mac(tmax=500, dt=0.01, tburn=100,
                       r=4, k=1.7, h=0.15, e=0.5, m=2,
                       a=[12,16],
                       sigma_x=0.01, sigma_y=0.01,
                       x0=1, y0=0.4,
                       ):
    '''
    Run a numerical simulation of the Rosenzweig-MacArthur model
    with additive white noise.
    
    
    Allows for linearly increasing/decreasing attack rate (a).
    
    The model at the default parameter settings has a transcritical 
    bifurcation at a=5.60 and a Hopf bifurcation at a=15.69.    
    
    Model configuration is as in Geller et al. (2016), Theoretical Ecology
    https://link.springer.com/article/10.1007/s12080-016-0303-2.
    
    Parameters
    ----------
    tmax : int, optional
        Total simulation time. The default is 500.
    dt : float, optional
        Time increment for each iteration of the Euler Maruyama scheme.
        The default is 0.01.
    tburn : int, optional
        Total burn in time to remove transients. 
        The default is 100.
    r : float, optional
        Intrinsic growth rate of the resource (x)
        The default is 4.
    k : float, optional
        Carrying capacity
        The default is 1.7.
    h : float, optional
        Handling time
        The default is 0.15.
    e : float, optional
        Conversion factor
        The default is 0.5.
    m : float, optional
        Per captia consumer mortality rate.
        The default is 2.
    a : float or list, optional
        Attack rate of the consumer (y). 
        Can be provided as a list containing the start
        and end value if linear change is desired.
        The default is [12,16].
    sigma_x : float, optional
        Noise amplitude for the resource (x). The default is 0.01.
    sigma_y : float, optional
        Noise amplitude for the consumer (y). The default is 0.01.        
    x0 : float, optional
        Initial condition for the resource. The default is 1.
    y0 : float, optional
        Initial condition for the consumer. The default is 0.4.
        
        
    Returns
    -------
    pd.DataFrame
        Trajectories of resource (x) and consumer (y) indexed by time.

    '''
    
    
    
    # Define the map
    def de_fun_x(x,y,r,k,a,h):
        return r*x*(1-x/k) - (a*x*y)/(1+a*h*x)
    
    def de_fun_y(x,y,e,a,h,m):
        return e*a*x*y/(1+a*h*x) - m*y
    
    
    
    # Initialise arrays for time and state values
    t = np.arange(0,tmax,dt)
    x = np.zeros(len(t))
    y = np.zeros(len(t))
    
    
    # Get a values, that may be linearly changing
    if (type(a) == float) or (type(a)==int):
        avals = np.ones(len(t))*a
    else:
        avals = np.linspace(a[0], a[1], len(t))
    
    
    # Array of noise values (normal random variables with variance sigma^2)
    dW_x_burn = np.random.normal(loc=0, scale=sigma_x*np.sqrt(dt), size = int(tburn/dt))
    dW_x = np.random.normal(loc=0, scale=sigma_x*np.sqrt(dt), size = len(t))
    
    dW_y_burn = np.random.normal(loc=0, scale=sigma_y*np.sqrt(dt), size = int(tburn/dt))
    dW_y = np.random.normal(loc=0, scale=sigma_y*np.sqrt(dt), size = len(t))
    
    
    # Run burn-in period
    for i in range(int(tburn/dt)):
        x0 = x0 + de_fun_x(x0,y0,r,k,avals[0],h)*dt + dW_x_burn[i]
        y0 = y0 + de_fun_y(x0,y0,e,avals[0],h,m)*dt + dW_y_burn[i]
    
    # Initial condition post burn-in period
    x[0]=x0
    y[0]=y0
    
    
    # Run simulation
    for i in range(len(t)-1):
        x[i+1] = x[i] + de_fun_x(x[i],y[i],r,k,avals[i],h)*dt + dW_x[i]
        y[i+1] = y[i] + de_fun_y(x[i],y[i],e,avals[i],h,m)*dt + dW_y[i]
        # make sure that state variable remains >= 0 
        if x[i+1] < 0:
            x[i+1] = 0
        if y[i+1] < 0:
            y[i+1] = 0
    
    # Store data in a dataframe indexed by time
    df = pd.DataFrame({'time':t, 'x':x, 'y':y})
    df.set_index('time', inplace=True)
    
    return df




def simulate_may(tmax=500, 
                 dt=0.01, 
                 tburn=100,
                 r=1, k=1, s=0.1,
                 h=[0.15,0.27],
                 sigma=0.01,
                 x0=0.8,
                 ):
    '''
    Run a numerical simulation of May's harvesting model
    with additive white noise.
    
    
    Allows for linearly increasing/decreasing harvesting rate (h)
    
    The model at the default parameter settings has a fold 
    bifurcation at h=0.260 
    

    Parameters
    ----------
    tmax : int, optional
        Total simulation time. The default is 500.
    dt : float, optional
        Time increment for each iteration of the Euler Maruyama scheme.
        The default is 0.01.
    tburn : int, optional
        Total burn in time to remove transients. 
        The default is 100.
    r : float, optional
        Intrinsic growth rate
        The default is 1.
    k : float, optional
        Carrying capacity
        The default is 1.
    s : float, optional
        Half-saturation constant of harvesting function
        The default is 0.1.
    h : float or list, optional
        Harvesting rate.
        Can be provided as a list containing the start
        and end value if linear change is desired.
        The default is [0.15,0.27].
    sigma : float, optional
        Noise amplitude. The default is 0.01.    
    x0 : float, optional
        Initial condition. The default is 0.8.
        
        
    Returns
    -------
    pd.DataFrame
        Trajectories of state variable (x) indexed by time.

    '''
    
    
    
    # Define the map
    def de_fun(x,r,k,h,s):
        return r*x*(1-x/k)  - h*(x**2/(s**2 + x**2))
    
    # Initialise arrays for time and state values
    t = np.arange(0,tmax,dt)
    x = np.zeros(len(t))
    
    
    # Get a values, that may be linearly changing
    if (type(h) == float) or (type(h)==int):
        hvals = np.ones(len(t))*h
    else:
        hvals = np.linspace(h[0], h[1], len(t))
    
    
    # Create brownian increments (s.d. sqrt(dt))
    dW_burn = np.random.normal(loc=0, scale=sigma*np.sqrt(dt), size = int(tburn/dt))
    dW = np.random.normal(loc=0, scale=sigma*np.sqrt(dt), size = len(t))
    
    # Run burn-in period on x0
    for i in range(int(tburn/dt)):
        x0 = x0 + de_fun(x0,r,k,hvals[0],s)*dt + dW_burn[i]
 
    
    # Initial condition post burn-in period
    x[0]=x0
    
    # Run simulation
    for i in range(len(t)-1):
        x[i+1] = x[i] + de_fun(x[i],r,k, hvals[i], s)*dt + dW[i]
        # make sure that state variable remains >= 0
        if x[i+1] < 0:
            x[i+1] = 0
    
    # Store data in a dataframe indexed by time
    df = pd.DataFrame({'time':t, 'x':x})
    df.set_index('time', inplace=True)
    
    return df









