# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_utils.ipynb.

# %% auto 0
__all__ = ['NpEncoder', 'simple_norm_stepshr_sleep_classifier', 'phase_ic_guess', 'abs_hour_diff', 'cut_phases_12',
           'convert_binary', 'cal_days_diff', 'interpolateLinear', 'interpolateLinearExt', 'parse_dt', 'circular_mean',
           'phase_coherence', 'phase_coherence_clock', 'angle_difference', 'subtract_clock_times', 'circular_av_clock',
           'circular_scatter', 'times_to_angle', 'timezone_mapper', 'split_missing_data', 'split_drop_data']

# %% ../nbs/07_utils.ipynb 3
import numpy as np
import pandas as pd
import scipy as sp
from numba import jit
from scipy import interpolate
import torch 
import json
import pytz
import datetime
import copy
import scipy as sp
from scipy.integrate import solve_ivp
from math import *
import pylab as plt


# %% ../nbs/07_utils.ipynb 4
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        if isinstance(obj, datetime.datetime):
            return obj.isoformat()
        return super(NpEncoder, self).default(obj)

# %% ../nbs/07_utils.ipynb 5
pd.options.mode.chained_assignment = None  # default='warn'

def simple_norm_stepshr_sleep_classifier(t):
        t[0,:] = torch.tanh((t[0,:] - 10.0)/ 100.0)
        t[1, torch.nonzero(t[1,:])] = torch.tanh((t[1,torch.nonzero(t[1,:])] - 60.0) / 30.0)
        return t 

def phase_ic_guess(time_of_day: float):
    time_of_day = np.fmod(time_of_day, 24.0)

    # Wake at 8 am after 8 hours of sleep
    # State at 00:00
    psi = 1.65238233

    # Convert to radians, add to phase
    psi += time_of_day * np.pi / 12
    return psi

def abs_hour_diff(x, y):
    """
    function abs_hour_diff(x,y)

    Find the difference in hours between
    two clock times (wrapped)
    """
    a1 = min(x, y)
    a2 = max(x, y)
    s1 = a2-a1
    s2 = 24.0+a1-a2
    return(min(s1, s2))


def cut_phases_12(p):
    """
    Function to make the branch cut for the DLMO times be at 12 instead of 24.
    This is better because lots of DLMOs are near midnight, but many fewer are near
    noon.

        cut_phases_12(timept)
    """

    while (p < 0.0):
        p += 24.0

    p = np.fmod(p, 24.0)

    if p > 12:
        return p-24.0
    else:
        return p


def convert_binary(x, breakpoint: float = 0.50):
    x[x <= breakpoint] = 0.0
    x[x > breakpoint] = 1.0
    return x


def cal_days_diff(a, b):
    """Get the calander days between two time dates"""
    A = a.replace(hour=0, minute=0, second=0, microsecond=0)
    B = b.replace(hour=0, minute=0, second=0, microsecond=0)
    return (A - B).days


@jit(nopython=True)
def interpolateLinear(t, xvals, yvals):
    """Implement a faster method to get linear interprolations of the light functions"""

    if (t >= xvals[-1]):
        return (0.0)
    if (t <= xvals[0]):
        t += 24.0

    i = np.searchsorted(xvals, t) - 1
    ans = (yvals[i + 1] - yvals[i]) / \
          ((xvals[i + 1] - xvals[i]) * (t - xvals[i])) + yvals[i]
    return (ans)


@jit(nopython=True)
def interpolateLinearExt(t, xvals, yvals):
    """Implement a faster method to get linear interprolations of the light functions, exclude non-full days"""
    i = np.searchsorted(xvals, t) - 1
    ans = (yvals[i + 1] - yvals[i]) / \
          ((xvals[i + 1] - xvals[i]) * (t - xvals[i])) + yvals[i]
    return (ans)


def parse_dt(date, time):
    strDate = date + ' ' + time
    return pd.to_datetime(strDate, format='%m/%d/%Y %I:%M %p')


# %% ../nbs/07_utils.ipynb 7
def circular_mean(series):
    Z=complex(0,0)
    series=np.array(series)
    for i in range(len(series)):
        Z+=np.exp(series[i]*complex(0,1))

    Z=Z/float(len(series))

    ans=np.angle(Z)
    if (ans<0.0):
        ans+=2*sp.pi
    return(ans)

def phase_coherence(series):
    Z=complex(0,0)
    series=np.array(series)
    for i in range(len(series)):
        Z+=np.exp(series[i]*complex(0,1))

    Z=Z/float(len(series))

    ans=np.absolute(Z)
    return(ans)

def phase_coherence_clock(series):
    angles=np.pi/12.0*series
    return(phase_coherence(angles))

def angle_difference(c1, c2) -> float:
    """Find the angle between two angles given in radians
    angle_difference(c1, c2)
    c1-c2
    """

    return(np.angle(np.exp(complex(0,1)*(c1-c2))))


def subtract_clock_times(c1, c2):
    """Find the hour differences between two clock times new"""
    a1=sp.pi/12.0*c1
    a2=sp.pi/12.0*c2
    adiff=angle_difference(a1, a2)
    return(12.0/sp.pi*adiff)


def circular_av_clock(series):
    """Find the average time given a list of clock times"""
    angles=sp.pi/12.0*series
    ans_angle=circular_mean(angles)
    #back to clock time
    return(ans_angle*12.0/sp.pi)
    

def circular_scatter(ax, angles, clock_times=False, radius=1.0, color='blue'):
    """Adds a polar scatter plot of clock times to an axes with polar axis i.e.
        ax = plt.subplot(111, polar=True)
        Will also plot the circular mean angle and the phase coherence
    """
    
    angles=np.array(angles)
    radii=radius*np.ones(len(angles))


    if clock_times:
        angles=angles*sp.pi/12.0

    ax.scatter(angles, radii, color=color)
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_thetagrids(list(range(0,360,45)), list(range(0,24,3)))
    ax.set_rmax(1.2)
    ax.set_rticks([0.0,0.2,0.6,0.8,1.0])
    ax.annotate("", xytext=(0.0,0.0), xy=(circular_mean(angles),phase_coherence(angles)),arrowprops=dict(facecolor=color))
    
    
    


# %% ../nbs/07_utils.ipynb 9
class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(NpEncoder, self).default(obj)


def times_to_angle(time_vector: np.ndarray):
    """
        Take an array of times and return R, psi 
        giving the mean angle (psi) and amplitude (R)
    """
    rad_vector = np.fmod(time_vector, 24.0) * np.pi/12.0
    Z = np.sum(np.exp(rad_vector*1j))/len(rad_vector)
    return(np.abs(Z), np.angle(Z))


def timezone_mapper(dt_object: datetime, timezone: str = 'America/Detroit'):
    """
        Take in local time as datetime object and give back UTC with 
        day lights savings accounted for as a timestamp
    """

    local_timezone = pytz.timezone(timezone)
    return local_timezone.localize(dt_object).timestamp()


def split_missing_data(date_time, ts, y, hr=None, break_threshold=96.0):

    # Find idx at start and end of long periods of zeros

    idx_start = None
    idx_end = None
    in_region = False
    crop_regions = []

    for (k, t) in enumerate(ts):
        if y[k] <= 0.0 and not in_region:
            idx_start = k
            in_region = True
        if y[k] > 0.0 and in_region:
            idx_end = k-1
            in_region = False
            if ts[idx_end]-ts[idx_start] >= break_threshold:
                crop_regions += [idx_start, idx_end]
    ts_split = np.split(ts, crop_regions)
    y_split = np.split(y, crop_regions)

    if hr is not None:
        hr_split = np.split(hr, crop_regions)

    print(f"Splitting data into {len(y_split)} regions")

    if hr is not None:
        return np.split(date_time, crop_regions), ts_split, y_split, hr_split
    else:
        return np.split(date_time, crop_regions), ts_split, y_split


def split_drop_data(date_time, ts, steps, hr, wake, break_threshold=96.0, min_length: float = 30.0):
    """
        Used to split long JSON into contin data steaks of at 
        least X=30 days.

        Uses that missing data will be zeros for steps and hr  and 
        0.5 for the wake data. 

        min_length is in days
    """

    idx_start = None
    idx_end = None
    in_region = False
    crop_regions = []

    for (k, t) in enumerate(ts):
        if (steps[k] <= 0.0 or hr[k] <= 0 or wake[k] == 0.50) and not in_region:
            idx_start = k
            in_region = True
        if steps[k] > 0.0 and hr[k] > 0 and wake[k] != 0.50 and in_region:
            idx_end = k-1
            in_region = False
            if ts[idx_end]-ts[idx_start] >= break_threshold:
                crop_regions += [idx_start, idx_end]

    ts_split = np.split(ts, crop_regions)
    steps_split = np.split(steps, crop_regions)
    hr_split = np.split(hr, crop_regions)
    wake_split = np.split(wake, crop_regions)
    date_time = np.split(date_time, crop_regions)

    # Find idxs for regions which are longer than min_length

    idx_long = [k for (k, val) in enumerate(ts_split)
                if (val[-1]-val[0])/24.0 >= min_length]

    if len(idx_long) > 0:
        return ([date_time[i] for i in idx_long], [ts_split[idx] for idx in idx_long], [steps_split[i] for i in idx_long],
                [hr_split[i] for i in idx_long], [wake_split[i] for i in idx_long])
    else:
        return None
