# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_data/00_core.ipynb.

# %% auto 0
__all__ = ['hdf_extensions', 'obj_in_lst', 'count_parameters', 'get_hdf_files', 'apply_df_tfms', 'CreateDict', 'ValidClmContains',
           'ValidClmIs', 'FilterClm', 'get_hdf_seq_len', 'df_get_hdf_seq_len', 'DfHDFGetSeqLen', 'DfResamplingFactor',
           'DfHDFCreateWindows', 'DfApplyFuncSplit', 'DfFilterQuery', 'DfDropClmExcept', 'calc_shift_offsets',
           'running_mean', 'downsample_mean', 'resample_interp', 'hdf_extract_sequence', 'Memoize', 'MemoizeMP',
           'HDF2Sequence', 'hdf2scalars', 'HDF2Scalars', 'ScalarSequenceElement', 'TensorSequences',
           'TensorSequencesInput', 'TensorSequencesOutput', 'toTensorSequencesInput', 'toTensorSequencesOutput',
           'TensorScalars', 'TensorScalarsInput', 'TensorScalarsOutput', 'plot_sequence', 'plot_seqs_single_figure',
           'plot_seqs_multi_figures', 'show_batch', 'show_results']

# %% ../../nbs/00_data/00_core.ipynb 3
from fastai.data.all import *
import h5py

# %% ../../nbs/00_data/00_core.ipynb 4
def obj_in_lst(lst,cls):
    '''retrieve first object of type cls from a list'''
    return next(o for o in lst if type(o) is cls)

# %% ../../nbs/00_data/00_core.ipynb 5
def count_parameters(model):
    '''retrieve number of trainable parameters of a model'''
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

# %% ../../nbs/00_data/00_core.ipynb 11
hdf_extensions = ['.hdf5']
def get_hdf_files(path,recurse=True, folders=None):
    "Get hdf5 files in `path` recursively, only in `folders`, if specified."
    return get_files(path, extensions=hdf_extensions, recurse=recurse, folders=folders)

# %% ../../nbs/00_data/00_core.ipynb 14
def apply_df_tfms(src,pd_tfms = None):
    '''Create Pandas Dataframe out of a list of items, with a list of df transforms applied'''
    if type(src) is pd.DataFrame:
        df = src
    else:
        df = pd.DataFrame(data=src.items,columns=['path'],dtype=str)
    if pd_tfms is not None:
        for t in pd_tfms:
            df = t(df)
    return df

# %% ../../nbs/00_data/00_core.ipynb 17
def CreateDict(pd_tfms = None):
    '''Create List of Dictionarys out of a list of items, with a list of df transforms applied'''
    def _inner(src):
        df = apply_df_tfms(src,pd_tfms)
#         df_dict_list = df.to_dict(orient='records') native to_dict is slower than self written approach
        df_values = df.values
        df_dict = {name:list(df_values[:,i]) for (i,name) in enumerate(df.columns)}
        df_dict_list = [{name: df_dict[name][i] for name in df_dict} for i in range(len(df))]
        return df_dict_list
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 19
def ValidClmContains(lst_valid):
    '''add validation column using a list of strings that are part of the validation frames'''
    def _inner(df):
        re_valid = '|'.join([re.escape(f) for f in lst_valid])
        df['valid'] = df.path.str.contains(re_valid)
        return df

    return _inner

# %% ../../nbs/00_data/00_core.ipynb 21
def ValidClmIs(lst_valid):
    '''adds validation column using a list of validation filenames'''
    def _inner(df):
        df['valid'] = df.path.isin([str(f) for f in lst_valid])
        return df

    return _inner

# %% ../../nbs/00_data/00_core.ipynb 23
def FilterClm(clm_name,func = lambda x:x):
    '''adds validation column using a list of validation filenames'''
    def _inner(df):
        return df[func(df[clm_name])]

    return _inner

# %% ../../nbs/00_data/00_core.ipynb 25
def get_hdf_seq_len(df,clm,ds=None):
    '''extract the sequence length of the dataset with the 'clm' name and 'f_path' path  '''
    with h5py.File(df['path'],'r') as f:
        ds = f if 'dataset' not in df else f[df['dataset']]
        f_len = max(ds[clm].shape)
    return f_len 

# %% ../../nbs/00_data/00_core.ipynb 26
def df_get_hdf_seq_len(df,clm,ds=None):
    '''extracts the sequence length of every file in advance to prepare repeated window extractions with 'DfHDFCreateWindows' '''
#     df['seq_len'] = ([get_hdf_seq_len(row.path,clm) for (idx, row) in df.iterrows()])
    df['seq_len'] = df.apply(lambda x: get_hdf_seq_len(x,clm),axis=1)
    return df

# %% ../../nbs/00_data/00_core.ipynb 27
def DfHDFGetSeqLen(clm):
    def _inner(df):
        return df_get_hdf_seq_len(df,clm)
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 30
import numbers


def DfResamplingFactor(src_fs,lst_targ_fs):
    if not isinstance(src_fs, numbers.Number) and type(src_fs) is not str: 
        raise ValueError('src_fs has to be a column name or a fixed number')
    
    def _inner(df):
        np_targ_fs = np.array(lst_targ_fs)
        pd.options.mode.chained_assignment = None #every row is a reference so we need to suppress the warning messages while copying

        #repeat entries for every target fs
        res_df = df.iloc[np.repeat(np.arange(len(df)),len(np_targ_fs))] 
        targ_fs = np.tile(np_targ_fs,len(df))
        res_df['targ_fs'] = targ_fs
        
        if isinstance(src_fs, numbers.Number):
            #src_fs is a fixed number
            res_df['resampling_factor'] = targ_fs/src_fs
        else:
            #src_fs is a column name of the df
            res_df['resampling_factor'] = targ_fs/res_df[src_fs]

        pd.options.mode.chained_assignment = 'warn'
        
        return res_df
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 32
def DfHDFCreateWindows(win_sz,stp_sz, clm, fixed_start = False, fixed_end = False):
    '''create windows of sequences, splits sequence into multiple items'''
    def _inner(df):
        if fixed_start and fixed_end: 
            raise Exception
        
        if 'seq_len' in df:
            np_f_len = df.seq_len.values
        else:
            np_f_len = np.array([get_hdf_seq_len(row,clm) for (idx, row) in df.iterrows()])
            
        if 'resampling_factor' in df: 
            np_f_len =(np_f_len*df.resampling_factor.values).astype(int)
            
        n_win = ((np_f_len-win_sz)//stp_sz)+1
        #cast array n_win to int and clip negative values to 0
        n_win = n_win.astype(int)
        n_win = np.clip(n_win,a_min=0,a_max=None) #remove negative values at instances where the winsize is smaller than the seq_len
        lst_idx = np.arange(len(np_f_len))
        
        pd.options.mode.chained_assignment = None #every row is a reference so we need to suppress the warning messages while copying
        
        res_df = df.iloc[np.repeat(lst_idx,n_win)]
#         res_df = df.loc[np.repeat(lst_idx,n_win)] #the loc variant as a little bit slower because it creates copies and returns wrong values with redundant indexes, but is more robust

        step_idx = np.concatenate([np.arange(x) for x in n_win])
    
        
        res_df['l_slc'] = step_idx*stp_sz if not fixed_start else None
        res_df['r_slc'] = step_idx*stp_sz + win_sz if not fixed_end else None
            
        pd.options.mode.chained_assignment = 'warn'
            
        return res_df
    
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 40
def DfApplyFuncSplit(split_func,func1,func2):
    '''apply two different functions on the dataframe, func1 on the first indices of split_func, func2 on the second indices.
        Split_func is a Training, Validation split function'''
    def _inner(df):
        (idxs1,idxs2) = split_func(df.path)
        df1= func1(df.iloc[idxs1])
        df2= func2(df.iloc[idxs2])
        return pd.concat((df1,df2))
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 42
def DfFilterQuery(query):
    def _inner(df):
        return df.query(query)
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 45
def DfDropClmExcept(clms = ['path','l_slc','r_slc','p_sample','resampling_factor']):
    '''drop unused dataframe columns as a last optional step to accelerate dictionary conversion'''
    def _inner(df):
        return df[[c for c in clms if c in df]]
    return _inner

# %% ../../nbs/00_data/00_core.ipynb 49
def calc_shift_offsets(clm_shift):
    clm_shift = np.array(clm_shift)
    l_offs = -min(clm_shift.min(),0)
    r_offs = -max(clm_shift.max(),0)
    l_shift = clm_shift+l_offs
    r_shift = clm_shift+r_offs
    dim_red = l_offs-r_offs
    return l_shift,r_shift,dim_red

# %% ../../nbs/00_data/00_core.ipynb 56
def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0,axis=0),axis=0) 
    return (cumsum[N:] - cumsum[:-N]) / float(N)

# %% ../../nbs/00_data/00_core.ipynb 57
def downsample_mean(x,N):
    trunc = -(x.shape[0] % N)
    trunc = trunc if trunc != 0 else None
    return x[:trunc,:].reshape((-1,N,x.shape[-1])).mean(axis=1)

# %% ../../nbs/00_data/00_core.ipynb 58
from scipy.signal import butter, lfilter, lfilter_zi


def resample_interp(x,resampling_factor,sequence_first=True, lowpass_cut=1.0, upsample_cubic_cut = None):
    '''signal resampling using linear or cubic interpolation
    
    x: signal to resample with shape: features x resampling_dimension or resampling_dimension x  features if sequence_first=True
    resampling_factor: Factor > 0 that scales the signal
    lowpass_cut: Upper boundary for resampling_factor that activates the lowpassfilter, low values exchange accuracy for performance, default is 0.7
    upsample_cubic_cut: Lower boundary for resampling_factor that activates cubic interpolation at high upsampling values. 
                        Improves signal dynamics in exchange of performance. None deactivates cubic interpolation
    '''
    
    if sequence_first:
        x = x.T
    
    fs_n = resampling_factor
    #if downsampling rate is too high, lowpass filter before interpolation
    if fs_n < lowpass_cut:
        b,a = butter(2, fs_n)
        zi = lfilter_zi(b,a)*x[:,:1] #initialize filter with steady state at first time step value
        x,_ = lfilter(b,a,x,axis=-1,zi=zi)

#         sos = butter(2, fs_n*1.2,output='sos')
# #         sos = signal.cheby2(2,20, fs_n,output='sos')
# #         import pdb;pdb.set_trace()
#         zi = np.swapaxes(signal.sosfilt_zi(sos)[...,None]*x[:,0],1,2)
#         x,_ = signal.sosfilt(sos, x,axis=-1,zi=zi)
        
    x_int = tensor(x)[None,...]
    targ_size = int(x.shape[-1]*fs_n)
    
#     if upsampling rate is too high, switch from linear to cubic interpolation
    if upsample_cubic_cut is None or fs_n <= upsample_cubic_cut:
        x = np.array(nn.functional.interpolate(x_int, size=targ_size, mode='linear',align_corners=False)[0])
    else:
        x = np(nn.functional.interpolate(x_int[...,None], size=[targ_size,1], mode='bicubic',align_corners=False)[0,...,0])
#     x = array(x_int)[0]
    
    if sequence_first:
        x = x.T
    
    return x

# %% ../../nbs/00_data/00_core.ipynb 60
from scipy.signal import resample


def hdf_extract_sequence(hdf_path,clms,dataset = None, l_slc = None, r_slc= None, resampling_factor=None, fs_idx =None,dt_idx =False,fast_resample=True):
    '''
    extracts a sequence with the shape [seq_len x num_features]
    
    hdf_path: file path of hdf file, may be a string or path type
    clms: list of dataset names of sequences in hdf file
    dataset: dataset root for clms. Useful for multiples sequences stored in one file.
    l_slc: left boundary for extraction of a window of the whole sequence
    r_slc: right boundary for extraction of a window of the whole sequence
    resampling_factor: scaling factor for the sequence length, uses 'resample_interp' for resampling
    fs_idx: clms list idx of fs entry in sequence. Will be scaled by resampling_factor after resampling
    dt_idx: clms list idx of dt entry in sequence. Will be scaled by resampling_factor after resampling
    fast_resample: if True, uses linear interpolation with anti-aliasing filter for faster resampling. Is less accurate than fft based resampling
    '''

    if resampling_factor is not None:
        seq_len = r_slc-l_slc if l_slc is not None and r_slc is not None else None #calculate seq_len for later slicing, necesary because of rounding errors in resampling
        if l_slc is not None: 
            l_slc= math.floor(l_slc/resampling_factor)
        if r_slc is not None: 
            r_slc= math.ceil(r_slc/resampling_factor)

    with h5py.File(hdf_path,'r') as f:
        ds = f if dataset is None else f[dataset]
        l_array = [(ds[n][l_slc:r_slc]) for n in clms]
        seq = np.stack(l_array,axis=-1)

    if resampling_factor is not None:
        if fast_resample:
            res_seq = resample_interp(seq,resampling_factor)
        else:
            res_seq = resample(seq,int(seq.shape[0]*resampling_factor),window=('kaiser', 14.0))

        if fs_idx is not None: 
            res_seq[:,fs_idx] = seq[0,fs_idx] * resampling_factor
        if dt_idx is not None: 
            res_seq[:,dt_idx] = seq[0,dt_idx] / resampling_factor

        seq = res_seq
        
        if seq_len is not None: 
            seq = seq[:seq_len] #cut the part of the sequence that is too long because of resampling rounding errors
        
    return seq

# %% ../../nbs/00_data/00_core.ipynb 61
class Memoize:
    def __init__(self, fn):
        self.fn = fn
        self.memo = {}

    def __call__(self, *args):
        if args not in self.memo:
            self.memo[args] = self.fn(*args)
        return self.memo[args]



# %% ../../nbs/00_data/00_core.ipynb 62
from multiprocessing import Lock, Manager, shared_memory


class MemoizeMP:
    
    def __init__(self, fn):
        self.fn = fn
        self.manager = Manager()
        self.results_dict = self.manager.dict()  # Stores metadata about computed results
        self.lock = Lock()  # Ensure atomic updates to the results_dict
        self.local_memo = {}  # Local cache for each process
        self.shared_memory_segments = []  # Track all shared memory segments
    def __call__(self, *args):
        if args in self.local_memo:
            return self.local_memo[args][0]
        with self.lock:
            if args in self.results_dict:
                result_info = self.results_dict[args]
                existing_shm = shared_memory.SharedMemory(name=result_info['name'])
                result = np.ndarray(result_info['shape'], dtype=result_info['dtype'], buffer=existing_shm.buf)
                self.local_memo[args] = (result, existing_shm)
                return result
        result = self.fn(*args)
        with self.lock:
            if args not in self.results_dict:
                result_shm = shared_memory.SharedMemory(create=True, size=result.nbytes)
                shm_array = np.ndarray(result.shape, dtype=result.dtype, buffer=result_shm.buf)
                shm_array[:] = result[:]
              
                self.results_dict[args] = {
                    'name': result_shm.name,
                    'shape': result.shape,
                    'dtype': result.dtype.str
                }
                self.local_memo[args] = (result, result_shm)
                # Track this shared memory segment for later cleanup
                self.shared_memory_segments.append(result_shm)
        return result
    def cleanup_shared_memory(self):
        """Explicitly cleanup all tracked shared memory segments."""
        for shm in self.shared_memory_segments:
            try:
                shm.close()
                shm.unlink()
            except FileNotFoundError:
                # The shared memory segment was already cleaned up
                pass
        # Clear the list after cleanup
        self.shared_memory_segments.clear()
    def __del__(self):
        self.cleanup_shared_memory()

# %% ../../nbs/00_data/00_core.ipynb 63
class HDF2Sequence(Transform):
    
    def __init__(self, clm_names,clm_shift=None,truncate_sz=None,to_cls=noop,cached=True, fs_idx =None,dt_idx =None,fast_resample=True):
        if clm_shift is not None:
            assert len(clm_shift)==len(clm_names) and all(isinstance(n, int) for n in clm_shift)
            self.l_shift,self.r_shift,_ = calc_shift_offsets(clm_shift)
        
        if not cached:
            self._exseq = self._hdf_extract_sequence 
        elif cached == 'local':
            self._exseq = Memoize(self._hdf_extract_sequence) 
        else :
            self._exseq = Memoize(self._hdf_extract_sequence) 
            #self._exseq = MemoizeMP(self._hdf_extract_sequence)

        self.cached = cached is not None
        store_attr('clm_names,clm_shift,truncate_sz,to_cls,fs_idx,dt_idx,fast_resample')
        
    def _hdf_extract_sequence(self,hdf_path,dataset = None, l_slc = None, r_slc= None, resampling_factor=None, fs_idx =None,dt_idx =None,fast_resample=True):
        '''
        extracts a sequence with the shape [seq_len x num_features]

        hdf_path: file path of hdf file, may be a string or path type
        clms: list of dataset names of sequences in hdf file
        dataset: dataset root for clms. Useful for multiples sequences stored in one file.
        l_slc: left boundary for extraction of a window of the whole sequence
        r_slc: right boundary for extraction of a window of the whole sequence
        resampling_factor: scaling factor for the sequence length, uses 'resample_interp' for resampling
        fs_idx: clms list idx of fs entry in sequence. Will be scaled by resampling_factor after resampling
        dt_idx: clms list idx of dt entry in sequence. Will be scaled by resampling_factor after resampling
        fast_resample: if True, uses linear interpolation with anti-aliasing filter for faster resampling. Is less accurate than fft based resampling
        '''

        if resampling_factor is not None:
            seq_len = r_slc-l_slc if l_slc is not None and r_slc is not None else None #calculate seq_len for later slicing, necesary because of rounding errors in resampling
            if l_slc is not None: 
                l_slc= math.floor(l_slc/resampling_factor)
            if r_slc is not None: 
                r_slc= math.ceil(r_slc/resampling_factor)

        with h5py.File(hdf_path,'r') as f:
            ds = f if dataset is None else f[dataset]
            l_array = [(ds[n][l_slc:r_slc]) for n in self.clm_names]
            seq = np.stack(l_array,axis=-1)

        if resampling_factor is not None:
            if fast_resample:
                res_seq = resample_interp(seq,resampling_factor)
            else:
                res_seq = resample(seq,int(seq.shape[0]*resampling_factor),window=('kaiser', 14.0))
            
            if fs_idx is not None: 
                res_seq[:,fs_idx] = seq[0,fs_idx] * resampling_factor
            if dt_idx is not None: 
                res_seq[:,dt_idx] = seq[0,dt_idx] / resampling_factor
            seq = res_seq

            if seq_len is not None: 
                seq = seq[:seq_len] #cut the part of the sequence that is too long because of resampling rounding errors

        return seq
    
    def _extract_dict_sequence(self,item):
        if hasattr(item,'keys'):
            path = item['path']
            dataset = item['dataset'] if 'dataset' in item else None
            l_slc = item['l_slc'] if 'l_slc' in item else None
            r_slc = item['r_slc'] if 'r_slc' in item else None
            resampling_factor = item['resampling_factor'] if 'resampling_factor' in item else None

            if self.cached:
                seq = self._exseq(path,dataset,None,None,resampling_factor,self.fs_idx,self.dt_idx,self.fast_resample)[l_slc:r_slc]
            else:
                seq = self._exseq(path,dataset,l_slc,r_slc,resampling_factor,self.fs_idx,self.dt_idx,self.fast_resample)
        else:
            seq = self._exseq(str(item),None,None,None,None,None)

        #shift clms of result by given value 
        if self.clm_shift is not None:
            l_seq = seq.shape[0]
            seq = np.stack([seq[self.l_shift[i]:l_seq+self.r_shift[i],i] for i in range(seq.shape[1])],axis=-1)
            
        if self.truncate_sz is not None:
            seq = seq[self.truncate_sz:]
        
        #it is important to slice first and then do the class conversion
#         return self.to_cls(seq.astype('f8'))#workaround for random bug, that mitigates convergence if the numpy array is an f4 array. Seems to make no sense because the result does not change. 
        return self.to_cls(seq)

    def encodes(self, item)->None: 
        return self._extract_dict_sequence(item)

# %% ../../nbs/00_data/00_core.ipynb 82
def hdf2scalars(hdf_path,c_names,dataset = None):
    with h5py.File(hdf_path,'r') as f:
        ds = f if dataset is None else f[dataset]
        l_array = [ds.attrs[n] for n in c_names]
        scalars = np.stack(l_array,axis=-1)
#         import pdb; pdb.set_trace()
#         l_array = [f[n][:][:,None] for n in c_names]
#         seq = np.concatenate(l_array,axis=1)
        return scalars

# %% ../../nbs/00_data/00_core.ipynb 84
class HDF2Scalars(Transform):
    
    def __init__(self, clm_names,to_cls=noop):
        store_attr('clm_names,to_cls')
    
    def _extract_dict_scalars(self,item):
        if isinstance(item,dict):
            path = item['path']
            dataset = item['dataset'] if 'dataset' in item else None

            seq = hdf2scalars(path,self.clm_names,dataset)
        else:
            seq = hdf2scalars(str(item),self.clm_names)
        return self.to_cls(seq)

    def encodes(self, item)->None: 
        return self._extract_dict_scalars(item)

# %% ../../nbs/00_data/00_core.ipynb 87
class ScalarSequenceElement(Transform):
    
    def __init__(self, idx,to_cls=noop):
        store_attr('idx,to_cls')

    def encodes(self, item)->None: 
        return self.to_cls(item[self.idx])

# %% ../../nbs/00_data/00_core.ipynb 90
class TensorSequences(TensorBase):#TensorBase
#     def __init__(self,x,c_names=None, **kwargs):
#         super().__init__()
#         self.c_names = c_names
    
    def show(self, ctx=None, **kwargs):
#         import pdb; pdb.set_trace()
        ax = ctx
        if ax is None: 
            _,ax = plt.subplots()
        ax.plot(self)
#         if title is not None: ax.set_title(title)
        return ax

    @classmethod
    @delegates(HDF2Sequence, keep=True)
    def from_hdf(cls,clm_names,**kwargs):
        return HDF2Sequence(clm_names,**kwargs)
    
class TensorSequencesInput(TensorSequences): 
    pass

class TensorSequencesOutput(TensorSequences): 
    pass

# %% ../../nbs/00_data/00_core.ipynb 93
@Transform
def toTensorSequencesInput(o): return TensorSequencesInput(o)
@Transform
def toTensorSequencesOutput(o): return TensorSequencesOutput(o)

# %% ../../nbs/00_data/00_core.ipynb 94
class TensorScalars(TensorBase):
    @classmethod
    @delegates(HDF2Scalars, keep=True)
    def from_hdf(cls,clm_names,**kwargs):
        return HDF2Scalars(clm_names,**kwargs)
    
    
class TensorScalarsInput(TensorScalars): 
    pass

class TensorScalarsOutput(TensorScalars): 
    pass


# %% ../../nbs/00_data/00_core.ipynb 96
for f in torch.nn.functional.mse_loss,torch.nn.functional.huber_loss, Tensor.__getitem__, Tensor.__ne__,Tensor.__eq__,Tensor.add,Tensor.sub,Tensor.mul,Tensor.div,Tensor.__rsub__,Tensor.__radd__,Tensor.matmul,Tensor.bmm:
    TensorBase.register_func(f,TensorSequences)
    TensorBase.register_func(f,TensorScalars)

# %% ../../nbs/00_data/00_core.ipynb 102
def plot_sequence(axs,in_sig,targ_sig,out_sig=None,**kwargs):
    for j,ax in  enumerate(axs[:-1]):
        ax.plot(targ_sig[:,j])
        if out_sig is not None: 
            ax.plot(out_sig[:,j])
            ax.legend(['y','ŷ'])
            if 'ref' in kwargs:
                ax.plot(kwargs['ref'][:,j]) 
        ax.label_outer()
    axs[-1].plot(in_sig)

# %% ../../nbs/00_data/00_core.ipynb 103
def plot_seqs_single_figure(n_samples,n_targ,samples,plot_func,outs=None,**kwargs):
    rows=max(1,((n_samples-1) // 3)+1)
    cols=min(3,n_samples)
    fig = plt.figure(figsize=(9,2*cols))
    outer_grid = fig.add_gridspec(rows, cols)
#     import pdb; pdb.set_trace()
    for i in range(n_samples):
        in_sig = samples[i][0]
        targ_sig = samples[i][1]
        if outs is not None: 
            out_sig = outs[i][0]
        inner_grid = outer_grid[i].subgridspec(n_targ+1, 1)
        axs = [fig.add_subplot(inner_grid[j]) for j in range(n_targ+1)]
        plot_func(axs,in_sig,targ_sig,out_sig=out_sig if outs is not None else None,**kwargs)
    plt.tight_layout()

# %% ../../nbs/00_data/00_core.ipynb 104
def plot_seqs_multi_figures(n_samples,n_targ,samples,plot_func,outs=None,**kwargs):
    for i in range(n_samples):
        fig = plt.figure(figsize=(9,3))
        axs = fig.subplots(nrows=n_targ+1,sharex=True)
        in_sig = samples[i][0]
        targ_sig = samples[i][1]
        if outs is not None:  
            out_sig = outs[i][0]
            
        plot_func(axs,in_sig,targ_sig,out_sig=out_sig if outs is not None else None,**kwargs)
        
        plt.tight_layout()

# %% ../../nbs/00_data/00_core.ipynb 105
from plum import dispatch


@dispatch
def show_batch(x:TensorSequences, y:TensorSequences, samples, ctxs=None, max_n=6, **kwargs):
    n_samples = min(len(samples), max_n)
    n_targ = samples[0][1].shape[1]
    if n_samples > 3:
        #if there are more then 3 samples to plot then put them in a single figure
        plot_seqs_single_figure(n_samples,n_targ,samples,plot_sequence, **kwargs)
    else:
        #if there are less then 3 samples to plot then put each in its own figure
        plot_seqs_multi_figures(n_samples,n_targ,samples,plot_sequence, **kwargs)
    return ctxs

# %% ../../nbs/00_data/00_core.ipynb 106
@dispatch
def show_results(x:TensorSequences, y:TensorSequences, samples, outs, ctxs=None, max_n=2, **kwargs):
    n_samples = min(len(samples), max_n)
    n_targ = samples[0][1].shape[1]
    if n_samples > 3:
        #if there are more then 3 samples to plot then put them in a single figure
        plot_seqs_single_figure(n_samples,n_targ,samples,plot_sequence,outs, **kwargs)
    else:
        #if there are less then 3 samples to plot then put each in its own figure
        plot_seqs_multi_figures(n_samples,n_targ,samples,plot_sequence,outs, **kwargs)
    return ctxs
