"""Under construction: Functions to predict sleep"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/06_sleep.ipynb.

# %% auto 0
__all__ = ['TwoProcessModel', 'sleep_midpoint', 'cluster_sleep_periods_scipy']

# %% ../nbs/api/06_sleep.ipynb 4
from math import *
import numpy as np
import pylab as plt
from .plots import *
from .utils import *
from .models import * 
from .lights import *
from scipy.optimize import minimize

# %% ../nbs/api/06_sleep.ipynb 6
class TwoProcessModel:
    "Implementation of the two-process model of sleep regulation."
    def __init__(self, steps_wake_threshold: float = 10.0):
        self.steps_wake_threshold = steps_wake_threshold
        self.awake = True

    @staticmethod
    def check_wake_status(awake: bool, # current wake status
                          h: float,  # homeostat value
                          phase: float #circadian phase
                          ) -> bool:

        c = np.cos(phase)
        H_minus = 0.17
        H_plus = 0.6
        homeostat_a = 0.10

        upper = (H_plus + homeostat_a * c)
        lower = (H_minus + homeostat_a * c)
        above_threshold = h > upper
        below_threshold = h <= lower

        if above_threshold:
            return False
        elif below_threshold:
            return True
        else:
            return awake

    def dhomeostat(self, 
                   homeostat: float, # homeostat value
                   steps: float, # steps value
                   phase: float, #circadian phase
                   ) -> float:

        h = homeostat[0]
        tau_s = 4.2  # hours
        tau_w = 18.2  # hours
        mu_s = 1
        
        step_awake = (steps > self.steps_wake_threshold) or self.awake
        dH = (mu_s - h) / tau_w if step_awake else -h / tau_s
        self.awake = TwoProcessModel.check_wake_status(
            self.awake, 
            h, 
            phase
            )
        return np.array([dH]) 
    
    # implement a RK4 solver
    def step_rk4(self,
                 state: np.ndarray,
                 steps: float,
                 phase: float,
                 dt=0.10):

        k1 = self.dhomeostat(state, steps, phase) 
        k2 = self.dhomeostat(state + k1 * dt / 2.0, steps, phase)
        k3 = self.dhomeostat(state + k2 * dt / 2.0, steps, phase)
        k4 = self.dhomeostat(state + k3 * dt, steps, phase)
        state = state + (dt / 6.0) * (k1 + 2.0*k2 + 2.0*k3 + k4)
        return state

    def __call__(self, 
                 ts: np.ndarray,
                 phase: np.ndarray,
                 steps: np.ndarray,
                 initial_value: np.array = np.array([0.50]) #  initial value for the homeostat
                 ) -> DynamicalTrajectory:
        sol = np.zeros((len(initial_value), len(ts)))
        current_state = initial_value
        sol[:,0] = current_state
        for idx in range(1,len(ts)):
            current_state = self.step_rk4(current_state, steps[idx], phase[idx])
            sol[:,idx] = current_state
        return(DynamicalTrajectory(ts, np.array(sol.T)))

# %% ../nbs/api/06_sleep.ipynb 8
def sleep_midpoint(timetotal: np.ndarray, 
                   Wake: np.ndarray, 
                   durations=True):
    """
        Given a wearable data frame with a Wake column which takes 
        the values 0, 1, missing this routine will create a sleep phase 
        column which is based on constant phase accumulation between sleep 
        midpoints. 

        The sleep midpoints are found using the criteria than they the median 
        time where 

    """
    sleep_start = []
    sleep_end = []
    awake = Wake[0] > 0.50

    if not awake:
        sleep_start.append(timetotal[1])

    for k in range(1, len(Wake)):
        if (Wake[k] > 0.50 and not awake):
            awake = True
            sleep_end.append(timetotal[k])

        if (Wake[k] <= 0.50 and awake):
            awake = False
            sleep_start.append(timetotal[k])

    if Wake[-1] <= 0.50:
        sleep_end.append(timetotal[-1])

    assert len(sleep_start) == len(sleep_end)
    sleep_midpoints = []
    sleep_durations = []
    for (s1, s2) in zip(sleep_start, sleep_end):
        sleep_midpoints += [(s2-s1)/2+s1]
        sleep_durations += [s2-s1]

    if durations:
        return np.array(sleep_midpoints), np.array(sleep_durations)
    else:
        return np.array(sleep_midpoints)

# %% ../nbs/api/06_sleep.ipynb 9
def cluster_sleep_periods_scipy(wake_data: np.ndarray, 
                                epsilon: float,
                                makeplot: bool = False,
                                max_sleep_clusters=None, 
                                min_sleep_clusters=None):
    """
        Given a binary vector wake_data which gives a prediction for the sleep/wake  
        status and a regularization penalty ε this function will create smoothed 
        sleep-wake periods. This helps as preprocessing to remove erroneous short sleep 
        periods (and wake) which may mess up calculations like the sleep midpoint for 
        the day

        cluster_sleep_periods(wake_data : np.ndarray, epsilon: float, makeplot: bool=False):
    """
    np.nan_to_num(wake_data, 0.50)

    def objective(w):
        return sum(w * (1 - wake_data)) + sum((1 - w)*wake_data) + epsilon*sum((w[1:]-w[0:-1])**2)

    max_clusters = max_sleep_clusters or len(wake_data)

    def constraint1(x):
        return max_clusters-sum((x[1:]-x[0:-1])**2)  # geq 0

    min_clusters = min_sleep_clusters or 0

    print(
        f"The max clusters are {max_clusters} and the min clusters are {min_clusters}")

    def constraint2(x):
        return sum((x[1:]-x[0:-1])**2)-min_clusters  # geq 0

    bnds = (0.0, 1.0)
    all_bnds = [bnds for b in range(len(wake_data))]

    constraint1d = {'type': 'ineq', 'fun': constraint1}
    constraint2d = {'type': 'ineq', 'fun': constraint2}
    all_cons = [constraint1d, constraint2d]

    x0 = wake_data
    sol = minimize(objective, x0, method='SLSQP', bounds=all_bnds)

    if makeplot:
        pl = plt.scatter(range(len(wake_data)), wake_data + 0.1 *
                         np.random.randn(len(wake_data)), label="", color="blue")
        plt.plot(range(len(wake_data)), convert_binary(
            sol.x), lw=2.0, label="", color="red")
        plt.show()

    print(
        f"The max clusters are {max_clusters} takes value {constraint1(sol.x)}>=0.0")
    print(
        f"The min clusters are {min_clusters} and takes the value {constraint2(sol.x)}>=0.0")
    return(convert_binary(sol.x))
