# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/api/09_cli.ipynb.

# %% auto 0
__all__ = ['mytz', 'main_acto', 'main_esri']

# %% ../nbs/api/09_cli.ipynb 4
import sys
import torch
import argparse
import circadian
import numpy as np
import pandas as pd
from torch import jit
from pathlib import Path
from pytz import timezone
from datetime import datetime
from .metrics import *
import matplotlib.pyplot as plt
from .plots import Actogram
from .models import Hannay19
from .utils import phase_ic_guess
from .utils import simple_norm_stepshr_sleep_classifier
from .sleep import cluster_sleep_periods_scipy, sleep_midpoint
from .readers import load_json, load_csv

mytz = timezone('EST')

# %% ../nbs/api/09_cli.ipynb 6
def main_acto():
    parser = argparse.ArgumentParser(description="""Make an actogram""")

    parser.add_argument('-j', '--json',
                        required=False,
                        action='store',
                        type=str,
                        help="Specify a json file with the data"
                        )
    parser.add_argument('-c', '--csv',
                        required=False,
                        action='store',
                        type=str,
                        help="Specify a directory with csv data in it"
                        )
    parser.add_argument("-sm", "--sleepmodel",
                        required=False,
                        action='store',
                        type=str,
                        help="Machine learning pytorch model for steps and heartrate to wake prediction"
                        )
    parser.add_argument('-d', '--dlmo',
                        required=False,
                        action='store_true',
                        default=False,
                        help="Integrate the model and plot dlmo times"
                        )
    parser.add_argument('-cbt', '--cbt',
                        required=False,
                        action='store_true',
                        default=False,
                        help="Integrate the model and plot core body temperature times"
                        )
    parser.add_argument('--sleep',
                        required=False,
                        action='store_true',
                        default=False,
                        help="Add sleep midpoints"
                        )
    parser.add_argument('-t', '--threshold',
                        required=False,
                        action='store',
                        type=float,
                        help="Threshold for displaying as light on the actogram",
                        default=1.0
                        )
    parser.add_argument('-m', '--multiplier',
                        required=False,
                        action='store',
                        type=float,
                        help="Steps to light multiplier",
                        default=1.0
                        )
    parser.add_argument('-p', '--period',
                        required=False,
                        action='store',
                        default=23.84,
                        type=float,
                        help="Set the SPM period"
                        )
    parser.add_argument('-r', '--raw',
                        required=False,
                        action='store_true',
                        default=False,
                        help="Plot the raw steps and heartrate data"
                        )
    parser.add_argument('--scatter',
                        required=False,
                        action='store_true',
                        default=False,
                        help="Plot the raw steps and heartrate data"
                        )
    parser.add_argument('-t1', '--t1',
                        required=False,
                        type=float,
                        action='store',
                        default=None,
                        help="Trim data before this time, in days"
                        )
    parser.add_argument('-t2', '--t2',
                        required=False,
                        type=float,
                        action='store',
                        default=None,
                        help="Trim data after this time, in days"
                        )
    parser.add_argument('-s', '--sigma',
                        required=False,
                        type=float,
                        action='store',
                        default=0.5,
                        help="Smooth the light(steps) data"
                        )


    args = parser.parse_args()

    def generated_sleep(ml_model_path, steps, hr):
        ml_model = jit.load(ml_model_path)
        data = torch.vstack((torch.tensor(steps), torch.tensor(hr)))
        data = simple_norm_stepshr_sleep_classifier(data).unsqueeze(0).float()
        wake_predicted = torch.sigmoid(ml_model(data)).squeeze(
            0).squeeze(0).detach().numpy()
        len_out = wake_predicted.shape[0]
        wake_predicted = np.hstack((wake_predicted, len(steps) - len_out))
        print(wake_predicted)
        wake_predicted = np.where(wake_predicted > 0.50, 1.0, 0.0)
        return wake_predicted
    
    if args.json:
        awObj = read_standard_json(args.json)

    if args.csv:
        awObj = read_standard_csv(args.csv)

    if args.csv or args.json:
        if args.t1 or args.t2:
            t1 = args.t1 or 0
            t2 = args.t2 or np.floor(np.squeeze(awObj.time_total)[-1]/24.0)+1
            awObj.trim_by_hour(t1*24.0, t2*24.0, inplace=True)
        hr = awObj.heartrate
        ts = awObj.time_total
        steps = awObj.steps

        if args.sleepmodel:
            pred_sleep = generated_sleep(
                ml_model_path=args.sleepmodel, 
                steps=steps, 
                hr=hr)
            print(pred_sleep)
            awObj._dataframe['wake'][0:len(pred_sleep)] = pred_sleep

    if args.raw:
        awObj.plot()

    ic = np.array([0.70, phase_ic_guess(ts[0]), 0.0])

    if args.scatter:
        awObj.scatter_hr_steps()


    plt.figure()
    ax = plt.gca()
    acto = Actogram(np.hstack(ts),
                    np.hstack(steps),
                    ax=ax,
                    threshold=args.threshold,
                    opacity=1.0,
                    sigma=[args.sigma, args.sigma])

    if awObj.wake is not None and args.sleep:
        print("Adding sleep layer")
        acto = Actogram(np.hstack(ts),
                        np.hstack(awObj.wake),
                        ax=ax,
                        threshold=0.5,
                        smooth=False,
                        opacity=0.50,
                        color='green')

    if args.dlmo:
        ic = np.array([0.70, phase_ic_guess(ts[0]), 0.0])
        hmodel = Hannay19({'tau': args.period})
        traj = hmodel(ts, args.multiplier*steps, ic)
        dlmos = hmodel.dlmos(traj)
        acto.plot_phasemarker(
            dlmos, error=np.ones(len(dlmos)), color='blue')
        print(f"DLMO mean: {np.mean(np.fmod(dlmos, 24.0))}")
        print("Last 14 days of DLMOs")
        print(np.fmod(dlmos[-14:], 24.0))

    if args.cbt:
        ic = np.array([0.70, phase_ic_guess(ts[0]), 0.0])
        hmodel = Hannay19({'tau': args.period})
        traj = hmodel(ts, args.multiplier*steps, ic)
        cbts = hmodel.cbt(traj)
        acto.plot_phasemarker(cbts, error=np.ones(len(cbts)), color='red')
        print(f"CBT mean: {np.mean(np.fmod(cbts, 24.0))}")
        print("Last 14 days of CBTs")
        print(np.fmod(cbts[-14:], 24.0))


    # if args.sleep:
    #     ts_flat = np.hstack(ts)
    #     steps_flat = np.hstack(steps)
    #     if awObj.wake is None:
    #         spm2 = Hannay19({'tau' : args.period})
    #         sol= spm2.integrate_model(ts_flat, args.multiplier*steps_flat, phase_ic_guess(ts_flat[0]))
    #         wake_score = np.diff(sol[3,:], prepend=0) > 0
    #     else:
    #         print("Using the provided wake column")
    #         wake_score=np.hstack(awObj.wake)

    #     sleep_clusters = cluster_sleep_periods_scipy(wake_score,
    #                                            90.0,
    #                                            makeplot=False)
    #     sleep_mid, duration = sleep_midpoint(ts_flat, sleep_clusters)
    #     acto.plot_phasemarker(sleep_mid, error=(
    #         duration/2.0), color="green", alpha=1.0)
    #     print(f"The median sleep duration is {np.median(duration)}")
    #     print(f"The average bedtime is {np.mean(np.fmod(sleep_mid-0.5*duration, 24.0))}")
    #     print(f"The average waketime is {np.mean(np.fmod(sleep_mid,24.0)+0.5*duration)}")

    plt.show()

# %% ../nbs/api/09_cli.ipynb 7
def main_esri():
    mytz = timezone('EST')
    parser = argparse.ArgumentParser(
        description="""Compute the Entrainment Signal Regularity Index for a data set""")

    parser.add_argument("-a", "--actogram",
                        required=False,
                        action='store_true',
                        default=False
                        )

    parser.add_argument('-j', '--json',
                        required=False,
                        action='store',
                        type=str,
                        help="Specify a json file with the data"
                        )
    parser.add_argument('-c', '--csv',
                        required=False,
                        action='store',
                        type=str,
                        help="Specify a directory with csv data in it"
                        )

    parser.add_argument('-s', '--sigma',
                        required=False,
                        type=float,
                        action='store',
                        default=0.5,
                        help="Smooth the light(steps) data"
                        )

    parser.add_argument('-t1', '--t1',
                        required=False,
                        type=float,
                        action='store',
                        default=None,
                        help="Trim data before this time, in days"
                        )
    parser.add_argument('-t2', '--t2',
                        required=False,
                        type=float,
                        action='store',
                        default=None,
                        help="Trim data after this time, in days"
                        )

    parser.add_argument("-sm", "--sleepmodel",
                        required=False,
                        action='store',
                        type=str,
                        help="Machine learning pytorch model for steps and heartrate to wake prediction"
                        )

    parser.add_argument('-t', '--threshold',
                        required=False,
                        action='store',
                        type=float,
                        help="Threshold for displaying as light on the actogram",
                        default=1.0
                        )

    parser.add_argument('-o', '--outfile',
                        required=False,
                        action='store',
                        type=str,
                        help="Path to the outfile",
                        default="esri_results.csv"
                        )
    parser.add_argument('-np', '--noplot',
                        required=False,
                        action='store_true',
                        help="Use this to suppress the plot",
                        default=False
                        )


    args = parser.parse_args()

    def generated_sleep(ml_model_path, steps, hr):
        ml_model = jit.load(ml_model_path)
        data = torch.vstack((torch.tensor(steps), torch.tensor(hr)))
        data = simple_norm_stepshr_sleep_classifier(data).unsqueeze(0).float()
        wake_predicted = torch.sigmoid(ml_model(data)).squeeze(
            0).squeeze(0).detach().numpy()
        len_out = wake_predicted.shape[0]
        wake_predicted = np.hstack((wake_predicted, len(steps) - len_out))
        wake_predicted = np.where(wake_predicted > 0.50, 1.0, 0.0)
        return wake_predicted


    if args.json:
        awObj = load_json(args.json)

    if args.csv:
        awObj = load_csv(args.csv)

    if args.csv or args.json:
        if args.t1 or args.t2:
            t1 = args.t1 or 0
            t2 = args.t2 or np.floor(np.squeeze(awObj.time_total)[-1]/24.0)+1
            awObj.trim_data(t1*24.0, t2*24.0)
        hr = awObj.heartrate
        ts = awObj.time_total
        steps = awObj.steps

        if args.sleepmodel:
            awObj.wake = generated_sleep(
                ml_model_path=args.sleepmodel, steps=steps, hr=hr)


    def moving_average(x, w):
        return np.convolve(x, np.ones(w), 'valid') / w


    smooth_length = 100
    num_days_compactness = 4.5

    ts_compact, timestamps, compactness = esri_trajectory(
        awObj, gamma=0.0, multiplier=1.0, num_days=num_days_compactness)
    compactness_smooth = moving_average(compactness, smooth_length)
    plt.figure()
    ax1 = plt.gca()
    ax1.plot(np.array(compactness)*24.0, np.array(ts_compact) /
            24.0 + num_days_compactness,  alpha=0.50)
    ax1.plot(compactness_smooth*24.0,
            np.array(ts_compact[smooth_length-1:])/24.0 + num_days_compactness, color='red')
    ax1.set_title(f"ESRI")
    ax1.set_ylabel("Days")
    ax1.set_xlabel("ESRI Score")
    ax1.set_ylim((0, max(np.array(ts_compact) / 24.0)))
    ax1.invert_yaxis()



    print(f"Writing results to {args.outfile}")
    df = pd.DataFrame({'timestamp': timestamps, 'esri': compactness})
    df.to_csv(args.outfile, index=False)

    if not args.noplot:
        acto = Actogram(np.hstack(ts),
                        np.hstack(steps),
                        ax=ax1,
                        threshold=args.threshold,
                        opacity=1.0,
                        sigma=[args.sigma, args.sigma])

        plt.show()
