# -*- coding: utf-8 -*-

"""

Created on Wed Dec 07 10:39:45 2016

@author: Fruit Flies

"""

from __future__ import division

import sys

#sys.path.append('C:\Users\Fruit Flies\Documents\Python Scripts\BayesChangePt')

from os import listdir

from os.path import isfile, join

from load_hdf5_data import load_hdf5

from my_wavelet_denoise import wavelet_denoise

import numpy as np

from scipy import signal, interpolate

from expresso_data_folders import dataFolders

#import bayesian_changepoint_detection.offline_changepoint_detection as offcd

#from functools import partial

import seaborn

from changepy import pelt

from changepy.costs import normal_mean

import matplotlib.pyplot as plt

from expresso_gui_params import analysisParams

#from bayesian_changepoint_detection import offline_changepoint_detectionas offcd

#from functools import partial

#---------------------------------------------------------------------------------------

folder_ind = 6 #0 = few big, 1 = few small, 2 = many small, 3 = no drinking, 4 = annotated files

datapath = dataFolders[folder_ind]

filenames = [f for f in listdir(datapath) if isfile(join(datapath, f))]

file_ind = 0

grpnum = 0

dsetnum = 3

wlevel = analysisParams['wlevel']

wtype = analysisParams['wtype']

medfilt_window = analysisParams['medfilt_window']

#var_scale = .5

#var_scale_testdata = 2*(1.3e-4)

mad_thresh = analysisParams['mad_thresh']

var_user = analysisParams['var_user']

#---------------------------------------------------------------------------------------

filename = join(datapath, filenames[file_ind])

print(filename)

dset, t = load_hdf5(filename,grpnum,dsetnum)

dset_check = (dset != -1)

if (np.sum(dset_check) == 0):

sys.exit("Bad dataset; try different group or dataset number")

frames = np.arange(0,dset.size)

dset = dset[dset_check]

frames = frames[np.squeeze(dset_check)]

new_frames = np.arange(0,np.max(frames)+1)

sp_raw = interpolate.InterpolatedUnivariateSpline(frames, dset)

dset = sp_raw(new_frames)

frames = new_frames

dset_denoised = wavelet_denoise(dset, wtype, wlevel)

#---------------------------------------------------------------------------------------

#sp_signal = interpolate.UnivariateSpline(frames, np.squeeze(dset_denoised))

#sp_der = sp_signal.derivative(n=1)

#dset_der = sp_der(frames)

dset_denoised_med = signal.medfilt(dset_denoised,medfilt_window)

sp_dset = interpolate.InterpolatedUnivariateSpline(frames, np.squeeze(dset_denoised_med))

sp_der = sp_dset.derivative(n=1)

dset_der = sp_der(frames)

#---------------------------------------------------------------------------------------

#Q, P, Pcp = offcd.offline_changepoint_detection(dset_der, partial(offcd.const_prior,

# l=(len(dset_der)+1)), offcd.gaussian_obs_log_likelihood, truncate=-50)

#---------------------------------------------------------------------------------------

dset_var = np.var(dset_der)

#Q, P, Pcp = offcd.offline_changepoint_detection(dset_var, \

# partial(offcd.const_prior, l=(len(dset_var)+1)), \

# offcd.gaussian_obs_log_likelihood, truncate=-40)

changepts = pelt(normal_mean(dset_der,var_user),len(dset_der)) #var_scale*dset_var #var_scale_testdata*len(dset_der)

#changepts = pelt(normal_meanvar(dset_der),len(dset_der))

N = len(dset_der) - 1

if 0 not in changepts:

changepts.insert(0,0)

#if len(dset_der) not in changepts:

# changepts.append(len(dset_der))

if N not in changepts:

changepts.append(N)

#print(changepts)

#---------------------------------------------------------------------------------------

piecewise_fits = np.empty(len(changepts)-1)

piecewise_fit_dist = np.empty_like(dset_der)

for i in range(0,len(changepts)-1):

ipt1 = changepts[i]

ipt2 = changepts[i+1] + 1

fit_temp = np.median(dset_der[ipt1:ipt2])

piecewise_fits[i] = fit_temp

piecewise_fit_dist[ipt1:ipt2] = fit_temp*np.ones_like(dset_der[ipt1:ipt2])

mean_pw_slope = np.mean(piecewise_fit_dist)

std_pw_slope = np.std(piecewise_fit_dist)

mad_slope = np.median(np.abs(np.median(dset_der)-dset_der))

piecewise_fits_dev = (piecewise_fits - np.median(dset_der)) / mad_slope

bout_ind = (piecewise_fits_dev < mad_thresh) #~z score of 1 #(mean_pw_slope - std_pw_slope)

bout_ind = bout_ind.astype(int)

bout_ind_diff = np.diff(bout_ind)

#plt.figure()

#plt.plot(bout_ind)

bouts_start_ind = np.where(bout_ind_diff == 1)[0] + 1

bouts_end_ind = np.where(bout_ind_diff == -1)[0] + 1

#print(bouts_start_ind)

#print(bouts_end_ind)

if len(bouts_start_ind) != len(bouts_end_ind):

minLength = np.min([len(bouts_start_ind), len(bouts_end_ind)])

bouts_start_ind = bouts_start_ind[0:minLength]

bouts_end_ind = bouts_end_ind[0:minLength]

#print(bouts_start_ind)

#print(bouts_end_ind)

changepts_array = np.asarray(changepts)

bouts_start = changepts_array[bouts_start_ind]

bouts_end = changepts_array[bouts_end_ind]

bouts = np.vstack((bouts_start, bouts_end))

print(bouts)

#print(changepts[bouts_start_ind])

#print(changepts[bouts_end_ind])

#bouts = changepts[bouts]

#---------------------------------------------------------------------------------------

plt.figure()

plt.plot(dset_denoised_med)

plt.plot(changepts[1:-1], dset_denoised_med[changepts[1:-1]], 'go')

f, (ax1, ax2) = plt.subplots(2, sharex=True, sharey=True)

ax1.plot(frames,dset)

for i in np.arange(bouts.shape[1]):

ax1.plot(frames[bouts[0,i]:bouts[1,i]], dset[bouts[0,i]:bouts[1,i]],'r-')

ax2.plot(frames, dset_denoised_med)

for i in np.arange(bouts.shape[1]):

ax2.plot(frames[bouts[0,i]:bouts[1,i]], dset_denoised_med[bouts[0,i]:bouts[1,i]],'r-')

#plt.plot(sp_signal(frames))

#plt.plot(changepts, dset_denoised[changepts], 'go')