# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_metrics.ipynb.

# %% auto 0
__all__ = ['esri', 'esri_trajectory']

# %% ../nbs/04_metrics.ipynb 3
import numpy as np
from .models import Hannay19
from .utils import phase_ic_guess
from .readers import WearableData

# %% ../nbs/04_metrics.ipynb 4
def esri(awObj: WearableData, 
                gamma: float = 0.0,
                multiplier: float = 1.0,
                num_days: float = 4.5):

    spmodel = Hannay19({'K': 0.0, 'gamma': gamma})
    psi0 = phase_ic_guess(awObj.time_total[0])

    idx = awObj.time_total < awObj.time_total[0]+24*num_days
    sol = spmodel.integrate_model(
        awObj.time_total[idx], multiplier*awObj.steps[idx],
        np.array([0.10, psi0, 0.0])
    )
    return sol.ts, sol.states



# %% ../nbs/04_metrics.ipynb 5
def esri_trajectory(awObj: WearableData,
                           gamma: float = 0.0,
                           multiplier: float = 1.0,
                           num_days: float = 4.5,
                           ):

    spmodel = Hannay19({'K': 0.0, 'gamma': gamma})
    compactness_trajectory = []
    time_trajectory = []
    timestamps = []
    timeStart = awObj.time_total[0]
    timestampStart = awObj.timestamp[0]
    while timeStart < awObj.time_total[-1] - 24*num_days:
        try:
            psi0 = phase_ic_guess(timeStart)
            idxStart = (awObj.time_total > timeStart)
            tsFilter = awObj.time_total[idxStart]
            idx = tsFilter < np.array(tsFilter[0])+24*num_days
            stepsFilter = awObj.steps[idxStart]
            trajectory = spmodel.integrate_model(
                tsFilter[idx], multiplier*stepsFilter[idx],
                np.array([0.10, psi0, 0.0])
            )
            sol = trajectory.states
            if sol[0, -1] > 0.0:
                compactness_trajectory.append(sol[0, -1])
                time_trajectory.append(timeStart)
                timestamps.append(
                    3600.0*(timeStart - awObj.time_total[0]) + timestampStart)
        except:
            print("Error in trajectory")
        timeStart += 1.0
    return time_trajectory, timestamps, compactness_trajectory

