# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/00_models.ipynb.

# %% auto 0
__all__ = ['DynamicalTrajectory', 'CircadianModel', 'Forger99', 'Hannay19', 'Hannay19TP', 'Jewett99']

# %% ../nbs/api/00_models.ipynb 4
import warnings
import numpy as np
from abc import ABC
from typing import Tuple
from scipy.signal import find_peaks
from fastcore.basics import patch_to
from .lights import LightSchedule

# %% ../nbs/api/00_models.ipynb 7
def _time_input_checking(time):
    "Checks if time is a valid input"
    if not isinstance(time, np.ndarray):
        raise TypeError("time must be a numpy array")
    if time.ndim != 1:
        raise ValueError("time must be a 1D array")
    if not np.issubdtype(time.dtype, np.number):
        raise TypeError("time must be numeric")
    if not np.all(np.diff(time) > 0):
        raise ValueError("time must be monotonically increasing")
    return True

def _state_input_checking(state, time):
    "Checks if state is a valid input"
    if not isinstance(state, np.ndarray):
        raise TypeError("states must be a numpy array")
    if state.ndim < 1:
        raise ValueError("states must have at least 1 dimension")
    if state.ndim > 3:
        raise ValueError("states must have at most 3 dimensions")
    if state.shape[0] != len(time):
        raise ValueError("states must have the same length as time")
    if not np.issubdtype(state.dtype, np.number):
        raise TypeError("states must be numeric")
    return True

# %% ../nbs/api/00_models.ipynb 9
class DynamicalTrajectory:
    "A class to store solutions of differential equation models that contains both the time points and the states"
    def __init__(self, 
                 time: np.ndarray, # time points
                 states: np.ndarray # state at time points
                 ) -> None:
        # input checking
        _time_input_checking(time)
        _state_input_checking(states, time)
        
        self.time = time
        self.states = states
        self.num_states = states.shape[1]
        if states.ndim >= 3:
            self.batch_size = states.shape[2]
        else:
            self.batch_size = 1

# %% ../nbs/api/00_models.ipynb 10
@patch_to(DynamicalTrajectory)
def __call__(self, timepoint: float) -> np.ndarray: # state of the system
    "Return the state at time t, linearly interpolated"
    # timepoint input checking
    if not isinstance(timepoint, (int, float)):
        raise TypeError("timepoint must be int or float")
    if timepoint < self.time[0] or timepoint > self.time[-1]:
        raise ValueError("timepoint must be within the time range")
    
    if self.batch_size == 1:
        values = np.zeros(self.num_states)
        for idx in range(self.num_states):
            values[idx] = np.interp(timepoint, self.time, self.states[..., idx])
    else:
        values = np.zeros((self.num_states, self.batch_size))
        for idx in range(self.num_states):
            for batch_idx in range(self.batch_size):
                values[idx, batch_idx] = np.interp(timepoint, self.time, self.states[..., idx, batch_idx])
    
    return values

# %% ../nbs/api/00_models.ipynb 11
@patch_to(DynamicalTrajectory)
def __getitem__(self, time_idx: int) -> Tuple[float, np.ndarray]:
    "Return the time and state at index idx"
    # idx input checking
    if not isinstance(time_idx, int):
        raise TypeError("idx must be int")
    if time_idx < -1 or time_idx >= len(self.time):
        raise ValueError(f"idx must be within 0 and {len(self.time)-1}, got {time_idx}")
    
    return self.time[time_idx], self.states[time_idx, ...]

# %% ../nbs/api/00_models.ipynb 12
@patch_to(DynamicalTrajectory)
def __len__(self) -> int:
    return len(self.time)

# %% ../nbs/api/00_models.ipynb 13
@patch_to(DynamicalTrajectory)
def get_batch(self, batch_idx: int) -> 'DynamicalTrajectory':
    "Obtain the trajectory for a single batch"
    # batch_idx input checking
    if not isinstance(batch_idx, int):
        raise TypeError("batch_idx must be int")
    if batch_idx < -1 or batch_idx >= self.batch_size:
        raise ValueError(f"batch_idx must be within -1 and {self.batch_size-1}, got {batch_idx}")
    if self.states.ndim >= 3:
        return DynamicalTrajectory(self.time, self.states[:, :, batch_idx])
    else:
        # no batch dimension
        return DynamicalTrajectory(self.time, self.states)

# %% ../nbs/api/00_models.ipynb 14
@patch_to(DynamicalTrajectory)
# String method
def __str__(self) -> str:
    time_str = np.array2string(self.time, precision=2, separator=", ", threshold=20)
    states_str = np.array2string(self.states, precision=2, separator=", ", threshold=20)
    output = f"Time:\n{time_str}\n"
    output += f"States:\n{states_str}"
    return output

# %% ../nbs/api/00_models.ipynb 17
def _parameter_input_checking(parameters):
    "Checks if parameters is a valid input for a circadian model"
    if not isinstance(parameters, dict):
        raise TypeError("parameters must be a dictionary")
    if len(parameters) == 0:
        raise ValueError("parameters must not be empty")
    for key, value in parameters.items():
        if not isinstance(key, str):
            raise TypeError("keys of parameters must be strings")
        if not isinstance(value, (int, float)):
            raise TypeError("values of parameters must be numeric")
    return True


def _positive_int_checking(number, name):
    "Checks if number is a positive integer"
    if not isinstance(number, int):
        raise TypeError(f"{name} must be an integer")
    if number < 1:
        raise ValueError(f"{name} must be positive")
    return True


def _initial_condition_input_checking(initial_condition, num_states):
    "Checks if initial_condition is a valid input for a circadian model"
    if not isinstance(initial_condition, np.ndarray):
        raise TypeError("initial_condition must be a numpy array")
    if not np.issubdtype(initial_condition.dtype, np.number):
        raise TypeError("initial_condition must be numeric")
    if initial_condition.shape[0] != num_states:
        raise ValueError(f"initial_condition must have length {num_states}")
    return True


def _model_input_checking(input, num_inputs, time):
    "Checks if input has the correct shape and values for a circadian model"
    if not isinstance(input, np.ndarray):
        raise TypeError("input must be a numpy array")
    if not np.issubdtype(input.dtype, np.number):
        raise TypeError("input must be numeric")
    if input.shape[0] != len(time):
        raise ValueError(f"input's first dimension must have length {len(time)} based on the time array provided")
    if num_inputs > 1:
        if input.shape[1] != num_inputs:
            raise ValueError(f"input must have {num_inputs} columns")

    
def _light_input_checking(light):
    "Checks if light is a valid input for a circadian model"
    if not isinstance(light, np.ndarray):
        raise TypeError("light must be a numpy array")
    if light.ndim != 1:
        raise ValueError("light must be a 1D array")
    if not np.issubdtype(light.dtype, np.number):
        raise TypeError("light must be numeric")
    if not np.all(light >= 0):
        raise ValueError("light intensity must be nonnegative")
    return True


def _wake_input_checking(wake):
    "Checks if wake is a valid input for a circadian model"
    if not isinstance(wake, np.ndarray):
        raise TypeError("wake must be a numpy array")
    if wake.ndim != 1:
        raise ValueError("wake must be a 1D array")
    if not np.issubdtype(wake.dtype, np.number):
        raise TypeError("wake must be numeric")
    if not np.all(wake >= 0) and not np.all(wake <= 1):
        raise ValueError("wake must be between 0 and 1")
    return True

# %% ../nbs/api/00_models.ipynb 19
class CircadianModel(ABC):
    "Abstract base class for circadian models that defines the common interface for all implementations"
    def __init__(self, 
                 default_params: dict, # default parameters for the model
                 num_states: int, # number of independent variables in the model
                 num_inputs: int, # number of inputs to the model such as light or wake state
                 default_initial_condition: np.ndarray # default initial conditions for the model
                 ) -> None:
        "Creates a new instance of the model"
        _parameter_input_checking(default_params)
        _positive_int_checking(num_states, "num_states")
        _positive_int_checking(num_inputs, "num_inputs")
        _initial_condition_input_checking(default_initial_condition, num_states)
        self.__default_params = default_params
        self._parameters = default_params
        for param_name, param_val in default_params.items():
            setattr(self, param_name, param_val)
        self.__num_states = num_states
        self.__num_inputs = num_inputs
        self.__default_initial_condition = default_initial_condition
        self._trajectory = None
        self._initial_condition = default_initial_condition


    @property
    def parameters(self) -> dict:
        return self._parameters

    @parameters.setter
    def parameters(self, value):
        for param_name, param_val in value.items():
            setattr(self, param_name, param_val)
        self._parameters = value

    @property
    def _default_params(self) -> dict:
        return self.__default_params

    @_default_params.setter
    def _default_params(self, value):
        self._parameters = value
        self.__default_params = value

    @property
    def _num_states(self) -> int: # number of independent variables in the model
        return self.__num_states

    @_num_states.setter
    def _num_states(self, value):
        self.__num_states = value

    @property
    def _num_inputs(self) -> int: # number of inputs to the model such as light or wake state
        return self.__num_inputs
    
    @_num_inputs.setter
    def _num_inputs(self, value):
        self.__num_inputs = value

    @property
    def _default_initial_condition(self) -> np.ndarray:
        return self.__default_initial_condition
    
    @_default_initial_condition.setter
    def _default_initial_condition(self, value):
        self.__default_initial_condition = value

    @property
    def trajectory(self) -> DynamicalTrajectory:
        return self._trajectory
    
    @trajectory.setter
    def trajectory(self, value):
        self._trajectory = value

    @property
    def initial_condition(self) -> np.ndarray:
        return self._initial_condition
    
    @initial_condition.setter
    def initial_condition(self, value):
        self._initial_condition = value

# %% ../nbs/api/00_models.ipynb 20
@patch_to(CircadianModel)
def derv(self,
         state: np.ndarray, # dynamical state of the model
         input: np.ndarray, # inputs to the model such as light or wake state
         ) -> np.ndarray:
    "Right-hand-side of the differential equation model with state and light as inputs"
    return NotImplementedError("derv is not implemented for this model")

# %% ../nbs/api/00_models.ipynb 21
@patch_to(CircadianModel)
def step_rk4(self,
            state: np.ndarray, # dynamical state of the model
            input: np.ndarray, # inputs to the model such as light or wake state
            dt: float, # step size in hours 
            ) -> np.ndarray:
    "Integrate the state of the model for one timestep using a fourth-order Runge-Kutta algorithm. Assumes a constant light value for the time step"
    k1 = self.derv(state, input)
    k2 = self.derv(state + k1 * dt / 2.0, input)
    k3 = self.derv(state + k2 * dt / 2.0, input)
    k4 = self.derv(state + k3 * dt, input)
    state = state + (dt / 6.0) * (k1 + 2.0*k2 + 2.0*k3 + k4)
    return state

# %% ../nbs/api/00_models.ipynb 22
@patch_to(CircadianModel)
def integrate(self,
              time: np.ndarray, # time points for integration. Time difference between consecutive values determines step size of the solver
              initial_condition: np.ndarray=None, # initial state of the model
              input: np.ndarray=None, # model input (such as light or wake) for each time point 
              ) -> DynamicalTrajectory:
    "Solve the model for specific timepoints given initial conditions and model inputs"
    # input checking
    _time_input_checking(time)
    if input is None:
        raise ValueError("a model input must be provided via the input argument")
    else:
        _model_input_checking(input, self._num_inputs, time)
    if initial_condition is None:
        initial_condition = self._default_initial_condition
    else:
        _initial_condition_input_checking(initial_condition, self._num_states)
    
    self.initial_condition = initial_condition
    
    n = len(time)
    sol = np.zeros((n, *initial_condition.shape))
    sol[0,...] = initial_condition
    state = initial_condition

    for idx in range(1, n):
        dt = time[idx] - time[idx-1]
        input_value = input[idx,...]
        state = self.step_rk4(state, input_value, dt)
        sol[idx,...] = state
    
    self._trajectory = DynamicalTrajectory(time, sol)
    return self._trajectory

# %% ../nbs/api/00_models.ipynb 23
@patch_to(CircadianModel)
def __call__(self,
             time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
             initial_condition: np.ndarray=None, # initial state of the model
             input: np.ndarray=None, # model input (such as light or wake) for each time point
             ):
    "Wrapper to integrate"
    return self.integrate(time, initial_condition, input)

# %% ../nbs/api/00_models.ipynb 24
@patch_to(CircadianModel)
def get_parameters_array(self)-> np.array:
    "Returns the parameters for the model as a numpy array"
    parameter_array = np.zeros(len(self.parameters))
    for idx, value in enumerate(self.parameters.values()):
        parameter_array[idx] = value
    return parameter_array

# %% ../nbs/api/00_models.ipynb 25
@patch_to(CircadianModel)
def phase(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the phase for. If None, the phase is calculated for the current trajectory 
          time: float=None # timepoint to calculate the phase at. If None, the phase is calculated for the entire trajectory
          ) -> float:
    "Calculates the phase of the model at a given timepoint"
    raise NotImplementedError("phase is not implemented for this model")

# %% ../nbs/api/00_models.ipynb 26
@patch_to(CircadianModel)
def amplitude(self,
              trajectory: DynamicalTrajectory=None, # trajectory to calculate the amplitude for. If None, the amplitude is calculated for the current trajectory 
              time: float=None, # timepoint to calculate the amplitude at. If None, the amplitude is calculated for the entire trajectory
              ) -> float:
    "Calculates the amplitude of the model at a given timepoint"
    raise NotImplementedError("amplitude is not implemented for this model")

# %% ../nbs/api/00_models.ipynb 27
@patch_to(CircadianModel)
def cbt(self,
        trajectory: DynamicalTrajectory=None, # trajectory to calculate the cbt for. If None, the cbt is calculated for the current trajectory
        ) -> np.ndarray: # array of times when the cbt occurs
    "Finds the core body temperature minumum markers along a trajectory"
    raise NotImplementedError("cbt is not implemented for this model")

# %% ../nbs/api/00_models.ipynb 28
@patch_to(CircadianModel)
def dlmos(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the dlmos for. If None, the dlmos are calculated for the current trajectory
          ) -> np.ndarray: # array of times when dlmo occurs 
    "Finds the Dim Light Melatonin Onset (DLMO) markers along a trajectory"
    raise NotImplementedError("dlmo is not implemented for this model")

# %% ../nbs/api/00_models.ipynb 29
@patch_to(CircadianModel)
def equilibrate(self,
                time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
                input: np.ndarray, # model input (such as light or wake) for each time point
                num_loops: int=10 # number of times to loop the input
                ) -> np.ndarray: # final state of the model
    "Equilibrate the model by looping the given light_estimate. Assumes the schedule repeats periodically after it ends"
    # input checking
    _time_input_checking(time)
    _model_input_checking(input, self._num_inputs, time)
    _positive_int_checking(num_loops, "num_loops")
    
    initial_condition = self._default_initial_condition
    dlmos = []
    for _ in range(num_loops):
        sol = self.integrate(time, initial_condition, input).states
        dlmos.append(self.dlmos())
        initial_condition = sol[-1, ...]
    # check that the output is equilibrated
    is_equilibrated = np.all(np.isclose(dlmos[-1][-1], dlmos[-2][-1], atol=1e-3))
    if not is_equilibrated:
        warnings.warn("The model did not equilibrate. Try increasing the number of loops.")
    final_state = sol[-1, ...]
    return final_state

# %% ../nbs/api/00_models.ipynb 30
def _get_default_initial_condition(
        model: CircadianModel, # model to calculate the default initial condition for
        num_loops: int=10 # number of times to loop the regular schedule
        ) -> np.ndarray:
    "Calculates a default initial condition by simulating the given model on a 16 hour light, 8 hour darkness schedule until convergence"
    # input checking
    if not isinstance(model, CircadianModel):
        raise TypeError("model must be a CircadianModel")
    model._default_initial_condition = 0.5 * np.ones(model._num_states)
    schedule = LightSchedule.Regular(lights_on=8, lights_off=24)    
    time = np.arange(0.0, 24.0, 0.1)
    light_input = schedule(time)
    default_initial_condition = model.equilibrate(time, light_input, num_loops)
    return default_initial_condition

def _check_cbtmin_spacing(
        cbtmin_times: np.ndarray, # array of times when the cbtmin occurs
        min_spacing: float=6.0, # minimum spacing between cbtmin markers
        ) -> bool:
    "Checks if the spacing between cbtmin markers is valid"
    cbtmin_spacing = np.diff(cbtmin_times)
    if np.any(cbtmin_spacing < min_spacing):
        # raise a warning
        warnings.warn(f"The data contains cbtmin markers that are spaced by less than {min_spacing} hours. Removal of duplicate cbtmin markers is recommended.")

# %% ../nbs/api/00_models.ipynb 32
class Forger99(CircadianModel): 
    "Implementation of Forger's 1999 model from the article 'A simpler model of the human circadian pacemaker'"
    def __init__(self, params=None):
        default_params = {
            'taux': 24.2, 'mu': 0.23, 'G': 33.75, 
            'alpha_0': 0.05, 'beta': 0.0075, 'p': 0.50, 
            'I0': 9500.0, 'k': 0.55, 'cbt_to_dlmo': 7.0,
            }
        num_states = 3 # x, xc, n
        num_inputs = 1 # light
        default_initial_condition = np.array([-0.0843259, -1.09607546, 0.45584306]) # condition at midnight for 16L, 8D schedule
        super(Forger99, self).__init__(default_params, num_states, num_inputs, default_initial_condition)
        if params is not None:
            self.parameters = params

    def integrate(self,
                  time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
                  initial_condition: np.ndarray=None, # initial state of the model
                  input: np.ndarray=None, # model input (such as light or wake) for each time point
                  ) -> DynamicalTrajectory:
        "Solve the model for specific timepoints given initial conditions and model inputs"
        # input checking for Forger99
        if input is not None:
            _light_input_checking(input)
        return super().integrate(time, initial_condition, input)

    def __repr__(self) -> str:
        return self.__str__()
    
    def __str__(self) -> str:
        return "Forger99"

# %% ../nbs/api/00_models.ipynb 33
@patch_to(Forger99)
def derv(self, 
         state: np.ndarray, # dynamical state (x, xc, n)
         input: float # light intensity in lux
         ) -> np.ndarray: # derivative of the state
     "Right-hand-side of the differential equation model"
     x = state[0,...]
     xc = state[1,...]
     n = state[2,...]
     light = input

     alpha = self.alpha_0 * pow((light / self.I0), self.p)
     Bhat = self.G * (1.0 - n) * alpha * (1 - 0.4 * x) * (1 - 0.4 * xc)
     mu_term = self.mu * (xc - 4.0 / 3.0 * pow(xc, 3.0))
     taux_term = pow(24.0 / (0.99669 * self.taux), 2.0) + self.k * Bhat

     dydt = np.zeros_like(state)
     dydt[0,...] = np.pi / 12.0 * (xc + Bhat)
     dydt[1,...] = np.pi / 12.0 * (mu_term - x * taux_term)
     dydt[2,...] = 60.0 * (alpha * (1.0 - n) - self.beta * n)

     return dydt

# %% ../nbs/api/00_models.ipynb 34
@patch_to(Forger99)
def phase(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the phase. If None, the current trajectory is used
          time: float=None # a time point to calculate the phase at. If None, the phase is calculated for the entire trajectory
          ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        x = trajectory.states[:, 0]
        y = -1.0 * trajectory.states[:, 1]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = state[0] 
            y = -1.0 * state[1]
    return np.angle(x + complex(0,1) * y)

# %% ../nbs/api/00_models.ipynb 35
@patch_to(Forger99)
def amplitude(self,
              trajectory: DynamicalTrajectory=None, # trajectory to calculate the amplitude. If None, the current trajectory is used
              time: float=None # a time point to calculate the amplitude at. If None, the amplitude is calculated for the entire trajectory
              ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        x = trajectory.states[:, 0]
        y = -1.0 * trajectory.states[:, 1]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = state[0] 
            y = -1.0 * state[1]
    return np.sqrt(x**2 + y**2)

# %% ../nbs/api/00_models.ipynb 36
@patch_to(Forger99)
def cbt(self,
        trajectory: DynamicalTrajectory=None, # trajectory to calculate the cbt. If None, the current trajectory is used
        ) -> np.ndarray:
    "Finds the core body temperature minumum markers for the model along a trajectory as the minimum of x"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    inverted_x = -1*trajectory.states[:,0]
    cbt_min_idxs, _ = find_peaks(inverted_x)
    cbtmin_times = trajectory.time[cbt_min_idxs]
    _check_cbtmin_spacing(cbtmin_times)
    return cbtmin_times

# %% ../nbs/api/00_models.ipynb 37
@patch_to(Forger99)
def dlmos(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the dlmo. If None, the current trajectory is used 
          ) -> np.ndarray:
    "Finds the Dim Light Melatonin Onset (DLMO) markers for the model along a trajectory"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    return self.cbt(trajectory) - self.cbt_to_dlmo

# %% ../nbs/api/00_models.ipynb 40
class Hannay19(CircadianModel):
    "Implementation of Hannay's 2019 single population model from the article 'Macroscopic models for human circadian rhythms'"
    def __init__(self, params=None):
        default_params = {
            'tau': 23.84, 'K': 0.06358, 'gamma': 0.024, 
            'Beta1': -0.09318, 'A1': 0.3855, 'A2': 0.1977, 
            'BetaL1': -0.0026, 'BetaL2': -0.957756, 'sigma': 0.0400692, 
            'G': 33.75, 'alpha_0': 0.05, 'delta': 0.0075, 'p': 1.5, 'I0': 9325.0,
            'cbt_to_dlmo': 7.0}
        num_states = 3 # R, Psi, n
        num_inputs = 1 # light
        default_initial_condition = np.array([0.82041911, 1.71383697, 0.52318122]) # condition at midnight for 16L, 8D schedule
        super(Hannay19, self).__init__(default_params, num_states, num_inputs, default_initial_condition)
        if params is not None:
            self.parameters = params

    def integrate(self,
                time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
                initial_condition: np.ndarray=None, # initial state of the model
                input: np.ndarray=None, # model input (such as light or wake) for each time point
                ) -> DynamicalTrajectory:
        "Solve the model for specific timepoints given initial conditions and model inputs"
        # input checking for Hannay19
        if input is not None:
            _light_input_checking(input)
        return super().integrate(time, initial_condition, input)

    def __repr__(self) -> str:
        return self.__str__()
    
    def __str__(self) -> str:
        return "Hannay19"

# %% ../nbs/api/00_models.ipynb 41
@patch_to(Hannay19)
def derv(self,
         state: np.ndarray, # dynamical state (R, Psi, n)
         input: float # light intensity in lux
         ) -> np.ndarray: # derivative of the state
    "Right-hand-side of the differential equation model"
    R = state[0,...]
    Psi = state[1,...]
    n = state[2,...]
    light = input   

    alpha = self.alpha_0 * pow(light, self.p) / (pow(light, self.p) + self.I0)

    Bhat = self.G * (1.0 - n) * alpha
    A1_term_amp = self.A1 * 0.5 * Bhat * (1.0 - pow(R, 4.0)) * np.cos(Psi + self.BetaL1)
    A2_term_amp = self.A2 * 0.5 * Bhat * R * (1.0 - pow(R, 8.0)) * np.cos(2.0 * Psi + self.BetaL2)
    LightAmp = A1_term_amp + A2_term_amp
    A1_term_phase = self.A1 * Bhat * 0.5 * (pow(R, 3.0) + 1.0 / R) * np.sin(Psi + self.BetaL1)
    A2_term_phase = self.A2 * Bhat * 0.5 * (1.0 + pow(R, 8.0)) * np.sin(2.0 * Psi + self.BetaL2)
    LightPhase = self.sigma * Bhat - A1_term_phase - A2_term_phase

    dydt = np.zeros_like(state)
    dydt[0,...] = -1.0 * self.gamma * R + self.K * np.cos(self.Beta1) / 2.0 * R * (1.0 - pow(R, 4.0)) + LightAmp
    dydt[1,...] = 2*np.pi/self.tau + self.K / 2.0 * np.sin(self.Beta1) * (1 + pow(R, 4.0)) + LightPhase
    dydt[2,...] = 60.0 * (alpha * (1.0 - n) - self.delta * n)

    return dydt

# %% ../nbs/api/00_models.ipynb 42
@patch_to(Hannay19)
def phase(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the phase. If None, the current trajectory is used
          time: float=None # a time point to calculate the phase at. If None, the phase is calculated for the entire trajectory
          ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        x = np.cos(trajectory.states[:, 1])
        y = np.sin(trajectory.states[:, 1])
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = np.cos(state[1])
            y = np.sin(state[1])
    return np.angle(x + complex(0,1) * y)

# %% ../nbs/api/00_models.ipynb 43
@patch_to(Hannay19)
def amplitude(self,
              trajectory: DynamicalTrajectory=None, # trajectory to calculate the amplitude. If None, the current trajectory is used
              time: float=None # a time point to calculate the amplitude at. If None, the amplitude is calculated for the entire trajectory
              ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        amplitude = trajectory.states[:, 0]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            amplitude = state[0] 
    return amplitude

# %% ../nbs/api/00_models.ipynb 44
@patch_to(Hannay19)
def cbt(self,
        trajectory: DynamicalTrajectory=None # trajectory to calculate the cbt. If None, the current trajectory is used
        ) -> np.ndarray:
    "Finds the core body temperature minumum markers for the model along a trajectory as the times where the phase is pi"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    inverted_x = -np.cos(trajectory.states[:,1])
    cbt_min_idxs, _ = find_peaks(inverted_x)
    cbtmin_times = trajectory.time[cbt_min_idxs]
    _check_cbtmin_spacing(cbtmin_times)
    return cbtmin_times

# %% ../nbs/api/00_models.ipynb 45
@patch_to(Hannay19)
def dlmos(self,
          trajectory: DynamicalTrajectory=None # trajectory to calculate the dlmo. If None, the current trajectory is used
          ) -> np.ndarray:
    "Finds the Dim Light Melatonin Onset (DLMO) markers for the model along a trajectory"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    return self.cbt(trajectory) - self.cbt_to_dlmo

# %% ../nbs/api/00_models.ipynb 48
class Hannay19TP(CircadianModel):
    "Implementation of Hannay's 2019 two population model from the article 'Macroscopic models for human circadian rhythms'"
    def __init__(self, params=None):
        default_params = {
            'tauV': 24.25, 'tauD': 24.0, 'Kvv': 0.05, 
            'Kdd': 0.04, 'Kvd': 0.05, 'Kdv': 0.01, 
            'gamma': 0.024, 'A1': 0.440068, 'A2': 0.159136, 
            'BetaL': 0.06452, 'BetaL2': -1.38935, 'sigma': 0.0477375, 
            'G': 33.75, 'alpha_0': 0.05, 'delta': 0.0075, 'p': 1.5, 
            'I0': 9325.0, 'cbt_to_dlmo': 7.0,
            }
        num_states = 5 # Rv, Rd, Psiv, Psid, n
        num_inputs = 1 # light
        default_initial_condition = np.array([0.82423745, 0.82304996, 1.75233424, 1.863457, 0.52318122]) # condition at midnight for 16L, 8D schedule
        super(Hannay19TP, self).__init__(default_params, num_states, num_inputs, default_initial_condition)
        if params is not None:
            self.parameters = params

    def integrate(self,
                time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
                initial_condition: np.ndarray=None, # initial state of the model
                input: np.ndarray=None, # model input (such as light or wake) for each time point
                ) -> DynamicalTrajectory:
        "Solve the model for specific timepoints given initial conditions and model inputs"
        # input checking for Hannay19TP
        if input is not None:
            _light_input_checking(input)
        return super().integrate(time, initial_condition, input)

    def __repr__(self) -> str:
        return self.__str__()
    
    def __str__(self) -> str:
        return "Hannay19TP"

# %% ../nbs/api/00_models.ipynb 49
@patch_to(Hannay19TP)
def derv(self,
         state: np.ndarray, # dynamical state (Rv, Rd, Psiv, Psid, n)
         input: float # light intensity in lux
         ) -> np.ndarray: # derivative of the state
     "Right-hand-side of the differential equation model"
     Rv = state[0,...]
     Rd = state[1,...]
     Psiv = state[2,...]
     Psid = state[3,...]
     n = state[4,...]
     light = input

     alpha = self.alpha_0 * pow(light, self.p) / (pow(light, self.p) + self.I0)
     Bhat = self.G * (1.0 - n) * alpha

     A1_term_amp = self.A1 * 0.5 * Bhat * (1.0 - pow(Rv, 4.0)) * np.cos(Psiv + self.BetaL)
     A2_term_amp = self.A2 * 0.5 * Bhat * Rv * (1.0 - pow(Rv, 8.0)) * np.cos(2.0 * Psiv + self.BetaL2)
     LightAmp = A1_term_amp + A2_term_amp
     A1_term_phase = self.A1 * Bhat * 0.5 * (pow(Rv, 3.0) + 1.0 / Rv) * np.sin(Psiv + self.BetaL)
     A2_term_phase = self.A2 * Bhat * 0.5 * (1.0 + pow(Rv, 8.0)) * np.sin(2.0 * Psiv + self.BetaL2)
     LightPhase = self.sigma * Bhat - A1_term_phase - A2_term_phase

     dydt = np.zeros_like(state)
     dydt[0,...] = -self.gamma * Rv + self.Kvv / 2.0 * Rv * (1 - pow(Rv, 4.0)) + self.Kdv / 2.0 * Rd * (1 - pow(Rv, 4.0)) * np.cos(Psid - Psiv) + LightAmp
     dydt[1,...] = -self.gamma * Rd + self.Kdd / 2.0 * Rd * (1 - pow(Rd, 4.0)) + self.Kvd / 2.0 * Rv * (1.0 - pow(Rd, 4.0)) * np.cos(Psid - Psiv)
     dydt[2,...] = 2.0 * np.pi / self.tauV + self.Kdv / 2.0 * Rd * (pow(Rv, 3.0) + 1.0 / Rv) * np.sin(Psid - Psiv) + LightPhase
     dydt[3,...] = 2.0 * np.pi / self.tauD - self.Kvd / 2.0 * Rv * (pow(Rd, 3.0) + 1.0 / Rd) * np.sin(Psid - Psiv)
     dydt[4,...] = 60.0 * (alpha * (1.0 - n) - self.delta * n)

     return dydt

# %% ../nbs/api/00_models.ipynb 50
@patch_to(Hannay19TP)
def phase(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the phase. If None, the current trajectory is used
          time: float=None # a time point to calculate the phase at. If None, the phase is calculated for the entire trajectory
          ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        time = trajectory.time
        x = np.cos(trajectory.states[:, 2])
        y = np.sin(trajectory.states[:, 2])
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = np.cos(state[2])
            y = np.sin(state[2])
    return np.angle(x + complex(0,1) * y)

# %% ../nbs/api/00_models.ipynb 51
@patch_to(Hannay19TP)
def amplitude(self,
              trajectory: DynamicalTrajectory=None, # trajectory to calculate the amplitude. If None, the current trajectory is used
              time: float=None # a time point to calculate the amplitude at. If None, the amplitude is calculated for the entire trajectory
              ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        time = trajectory.time
        amplitude = trajectory.states[:, 0]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            amplitude = state[0] 
    return amplitude

# %% ../nbs/api/00_models.ipynb 52
@patch_to(Hannay19TP)
def cbt(self,
        trajectory: DynamicalTrajectory=None, # trajectory to calculate the cbt. If None, the current trajectory is used
        ) -> np.ndarray:
    "Finds the core body temperature minumum markers for the model along a trajectory as the times where the phase is pi"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    inverted_x = -np.cos(trajectory.states[:,2])
    cbt_min_idxs, _ = find_peaks(inverted_x)
    cbtmin_times = trajectory.time[cbt_min_idxs]
    _check_cbtmin_spacing(cbtmin_times)
    return cbtmin_times

# %% ../nbs/api/00_models.ipynb 53
@patch_to(Hannay19TP)
def dlmos(self,
          trajectory: DynamicalTrajectory=None # trajectory to calculate the dlmo. If None, the current trajectory is used
          ) -> np.ndarray:
    "Finds the Dim Light Melatonin Onset (DLMO) markers for the model along a trajectory"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    return self.cbt(trajectory) - self.cbt_to_dlmo

# %% ../nbs/api/00_models.ipynb 56
class Jewett99(CircadianModel):
    "Implementation of Jewett's 1999 model from the article 'Revised Limit Cycle Oscillator Model of Human Circadian Pacemaker'"
    def __init__(self, params=None):
        default_params = {
            'taux': 24.2, 'mu': 0.13, 'G': 19.875,
            'beta': 0.013, 'k': 0.55, 'q': 1.0/3,
            'I0': 9500, 'p': 0.6, 'alpha_0': 0.16,
            'phi_ref': 0.8, 'cbt_to_dlmo': 7.0}
        num_states = 3 # x, xc, n
        num_inputs = 1 # light
        default_initial_condition = np.array([-0.10097101, -1.21985662, 0.50529415]) # condition at midnight for 16L, 8D schedule
        super(Jewett99, self).__init__(default_params, num_states, num_inputs, default_initial_condition)
        if params is not None:
            self.parameters = params

    def integrate(self,
                  time: np.ndarray, # time points for integration. Time difference between each consecutive pair of values determines step size of the solver
                  initial_condition: np.ndarray=None, # initial state of the model
                  input: np.ndarray=None, # model input (such as light or wake) for each time point
                  ) -> DynamicalTrajectory:
        "Solve the model for specific timepoints given initial conditions and model inputs"
        # input checking for Jewett99
        if input is not None:
            _light_input_checking(input)
        return super().integrate(time, initial_condition, input)

    def __repr__(self) -> str:
        return self.__str__()
    
    def __str__(self) -> str:
        return "Jewett99"

# %% ../nbs/api/00_models.ipynb 57
@patch_to(Jewett99)
def derv(self,
         state: np.ndarray, # dynamical state (x, xc, n)
         input: float # light intensity in lux
         ) -> np.ndarray: # derivative of the state
    "Right-hand-side of the differential equation model"
    x = state[0,...]
    xc = state[1,...]
    n = state[2,...]
    light = input

    alpha = self.alpha_0 * (light / self.I0) ** self.p
    Bhat = self.G * alpha * (1 - n) * (1 - 0.4 * x) * (1 - 0.4 * xc)
    mu_term = self.mu * (1.0/3.0 * x + 4.0/3.0 * x**3 - 256.0/105.0 * x**7)
    taux_term = pow(24.0 / (0.99729 * self.taux), 2) + self.k * Bhat 

    dydt = np.zeros_like(state)
    dydt[0,...] = np.pi/12 * (xc + mu_term + Bhat)
    dydt[1,...] = np.pi/12 * (self.q * Bhat * xc - x * taux_term)
    dydt[2,...] = 60.0 * (alpha * (1 - n) - self.beta * n)
    
    return dydt

# %% ../nbs/api/00_models.ipynb 58
@patch_to(Jewett99)
def phase(self,
          trajectory: DynamicalTrajectory=None, # trajectory to calculate the phase. If None, the current trajectory is used
          time: float=None # a time point to calculate the phase at. If None, the phase is calculated for the entire trajectory
          ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        x = trajectory.states[:, 0]
        y = -1.0 * trajectory.states[:, 1]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = state[0] 
            y = -1.0 * state[1]
    return np.angle(x + complex(0,1) * y)

# %% ../nbs/api/00_models.ipynb 59
@patch_to(Jewett99)
def amplitude(self,
              trajectory: DynamicalTrajectory=None, # trajectory to calculate the amplitude. If None, the current trajectory is used
              time: float=None # a time point to calculate the amplitude at. If None, the amplitude is calculated for the entire trajectory
              ) -> float:
    if trajectory is None:
        trajectory = self.trajectory
    else: 
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    if time is None:
        x = trajectory.states[:, 0]
        y = -1.0 * trajectory.states[:, 1]
    else:
        if not isinstance(time, (float, int)):
            raise ValueError("time must be a float or an int")
        else:
            state = trajectory(time)
            x = state[0] 
            y = -1.0 * state[1]
    return np.sqrt(x**2 + y**2)

# %% ../nbs/api/00_models.ipynb 60
@patch_to(Jewett99)
def cbt(self,
        trajectory: DynamicalTrajectory=None, # trajectory to calculate the cbt. If None, the current trajectory is used
        ) -> np.ndarray:
    "Finds the core body temperature minumum markers for the model along a trajectory as the corrected minimum of x"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    inverted_x = -1*trajectory.states[:,0]
    cbt_min_idxs, _ = find_peaks(inverted_x)
    cbtmin_times = trajectory.time[cbt_min_idxs] + self.phi_ref
    _check_cbtmin_spacing(cbtmin_times)
    return cbtmin_times

# %% ../nbs/api/00_models.ipynb 61
@patch_to(Jewett99)
def dlmos(self,
          trajectory: DynamicalTrajectory=None # trajectory to calculate the dlmo. If None, the current trajectory is used
          ) -> np.ndarray:
    "Finds the Dim Light Melatonin Onset (DLMO) markers for the model along a trajectory"
    if trajectory is None:
        trajectory = self.trajectory
    else:
        if not isinstance(trajectory, DynamicalTrajectory):
            raise ValueError("trajectory must be a DynamicalTrajectory")
    return self.cbt(trajectory) - self.cbt_to_dlmo
