'''
#########################################################
This file containts classes to handle data processing

Author: Julian Kates-Harbeck, jkatesharbeck@g.harvard.edu

This work was supported by the DOE CSGF program.
#########################################################
'''

from __future__ import print_function
import itertools

import numpy as np
from scipy.interpolate import UnivariateSpline
import sys

#interpolate in a way that doesn't use future information. 
#It simply finds the latest time point in the original array
#that is less than or equal than the time point in question
#and interpolates to there.
def time_sensitive_interp(x,t,t_new):
    indices = np.maximum(0,np.searchsorted(t,t_new,side='right')-1)
    return x[indices]

def resample_signal(t,sig,tmin,tmax,dt,precision_str='float32'):
    order = np.argsort(t)
    t = t[order]
    sig = sig[order,:]
    sig_width = sig.shape[1]
    tt = np.arange(tmin,tmax,dt,dtype=precision_str)
    sig_interp = np.zeros((len(tt),sig_width),dtype=precision_str)
    for i in range(sig_width):
        sig_interp[:,i] = time_sensitive_interp(sig[:,i],t,tt) #make sure to not use future information
        # f = UnivariateSpline(t,sig[:,i],s=0,k=1,ext=0)
        # sig_interp[:,i] = f(tt)
    if(np.any(np.isnan(sig_interp))):
        print("signal contains nan")
    if(np.any(t[1:] - t[:-1] <= 0)):
        print("non increasing")
        idx = np.where(t[1:] - t[:-1] <= 0)[0][0]
        print(t[idx-10:idx+10])

    return tt,sig_interp

def cut_signal(t,sig,tmin,tmax):
    mask = np.logical_and(t >= tmin,  t <= tmax)
    return t[mask],sig[mask,:]

def cut_and_resample_signal(t,sig,tmin,tmax,dt,precision_str):
    t,sig = cut_signal(t,sig,tmin,tmax)
    return resample_signal(t,sig,tmin,tmax,dt,precision_str)

def get_individual_shot_file(prepath,shot_num,ext='.txt'):
    return prepath + str(shot_num) + ext 

def append_to_filename(path,to_append):
    ending_idx = path.rfind('.')
    new_path = path[:ending_idx] + to_append + path[ending_idx:]
    return new_path

def train_test_split(x,frac,do_shuffle=False):
    if not isinstance(x,np.ndarray):
        return train_test_split_robust(x,frac,do_shuffle)
    mask = np.array(range(len(x))) < frac*len(x)
    if do_shuffle:
        np.random.shuffle(mask)
    return x[mask],x[~mask]

def train_test_split_robust(x,frac,do_shuffle=False):
    mask = np.array(range(len(x))) < frac*len(x)
    if do_shuffle:
        np.random.shuffle(mask)
    train = []
    test = []
    for (i,_x) in enumerate(x):
        if mask[i]:
            train.append(_x)
        else:
            test.append(_x)
    return train,test

def train_test_split_all(x,frac,do_shuffle=True):
    groups = []
    length = len(x[0])
    mask = np.array(range(length)) < frac*length
    if do_shuffle:
        np.random.shuffle(mask)
    for item in x:
        groups.append((item[mask],item[~mask]))
    return groups


def concatenate_sublists(superlist):
    return list(itertools.chain.from_iterable(superlist))

def get_signal_slices(signals_superlist):
    indices_superlist = []
    signals_so_far = 0
    for sublist in signals_superlist:
        indices_sublist = signals_so_far + np.array(range(len(sublist)))
        signals_so_far += len(sublist)
        indices_superlist.append(indices_sublist)
    return indices_superlist