__all__ = ['butter_filter', 'butter_filter_frequency_response', 'apply_butter_filter', 'save_dataset', 'load_dataset',
           'dataset_add_chunk_col', 'dataset_add_length_col', 'dataset_add_labels_col', 'dataset_add_mean_col',
           'dataset_add_median_col', 'dataset_add_std_col', 'dataset_add_iqr_col', 'dataset_get_stats',
           'npys_to_memmap_batched', 'npys_to_memmap', 'reformat_as_memmap', 'TimeseriesDatasetCrops', 'RandomCrop',
           'CenterCrop', 'GaussianNoise', 'Rescale', 'ToTensor', 'Normalize', 'NormalizeBatch', 'ButterFilter',
           'ChannelFilter', 'Transform', 'TupleTransform', 'aggregate_predictions']

import numpy as np
import torch
import torch.utils.data
from torch import nn
from pathlib import Path
from scipy.stats import iqr

try:
    import pickle5 as pickle
except ImportError as e:
    import pickle

#Note: due to issues with the numpy rng for multiprocessing (https://github.com/pytorch/pytorch/issues/5059) that could be fixed by a custom worker_init_fn we use random throught for convenience
import random

#Note: multiprocessing issues with python lists and dicts (https://github.com/pytorch/pytorch/issues/13246) and pandas dfs (https://github.com/pytorch/pytorch/issues/5902)
import multiprocessing as mp

from skimage import transform

import warnings
warnings.filterwarnings("ignore", category=UserWarning)

from scipy.signal import butter, sosfilt, sosfiltfilt, sosfreqz

from tqdm.auto import tqdm

from collections import namedtuple

import pathlib
pathlib.WindowsPath = pathlib.PosixPath

tsdata_static = namedtuple("tsdata_static",("data","label","static"))
tsdata = namedtuple("tsdata",("data","label"))

# Cell
#https://stackoverflow.com/questions/12093594/how-to-implement-band-pass-butterworth-filter-with-scipy-signal-butter
def butter_filter(lowcut=10, highcut=20, fs=50, order=5, btype='band'):
    '''returns butterworth filter with given specifications'''
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq

    sos = butter(order, [low, high] if btype=="band" else (low if btype=="low" else high), analog=False, btype=btype, output='sos')
    return sos

def butter_filter_frequency_response(filter):
    '''returns frequency response of a given filter (result of call of butter_filter)'''
    w, h = sosfreqz(filter)
    #gain vs. freq(Hz)
    #plt.plot((fs * 0.5 / np.pi) * w, abs(h))
    return w,h

def apply_butter_filter(data, filter, forwardbackward=True):
    '''pass filter from call of butter_filter to data (assuming time axis at dimension 0)'''
    if(forwardbackward):
        return sosfiltfilt(filter, data, axis=0)
    else:
        data = sosfilt(filter, data, axis=0)

def save_dataset(df,lbl_itos,mean,std,target_root,filename_postfix="",protocol=4):
    target_root = Path(target_root)
    df.to_pickle(target_root/("df"+filename_postfix+".pkl"), protocol=protocol)

    if(isinstance(lbl_itos,dict)):#dict as pickle
        outfile = open(target_root/("lbl_itos"+filename_postfix+".pkl"), "wb")
        pickle.dump(lbl_itos, outfile, protocol=protocol)
        outfile.close()
    else:#array
        np.save(target_root/("lbl_itos"+filename_postfix+".npy"),lbl_itos)

    np.save(target_root/("mean"+filename_postfix+".npy"),mean)
    np.save(target_root/("std"+filename_postfix+".npy"),std)

def load_dataset(target_root,filename_postfix="",df_mapped=True):
    target_root = Path(target_root)
    # if(df_mapped):
    #     df = pd.read_pickle(target_root/("df_memmap"+filename_postfix+".pkl"))
    # else:
    #     df = pd.read_pickle(target_root/("df"+filename_postfix+".pkl")

    ### due to pickle 5 protocol error

    if(df_mapped):
        df = pickle.load(open(target_root/("df_memmap"+filename_postfix+".pkl"), "rb"))
    else:
        df = pickle.load(open(target_root/("df"+filename_postfix+".pkl"), "rb"))

    if((target_root/("lbl_itos"+filename_postfix+".pkl")).exists()):#dict as pickle
        infile = open(target_root/("lbl_itos"+filename_postfix+".pkl"), "rb")
        lbl_itos=pickle.load(infile)
        infile.close()
    else:#array
        lbl_itos = np.load(target_root/("lbl_itos"+filename_postfix+".npy"))

    mean = np.load(target_root/("mean"+filename_postfix+".npy"))
    std = np.load(target_root/("std"+filename_postfix+".npy"))
    return df, lbl_itos, mean, std

def dataset_add_chunk_col(df, col="data"):
    '''add a chunk column to the dataset df'''
    df["chunk"]=df.groupby(col).cumcount()

def dataset_add_length_col(df, col="data", data_folder=None):
    '''add a length column to the dataset df'''
    df[col+"_length"]=df[col].apply(lambda x: len(np.load(x if data_folder is None else data_folder/x, allow_pickle=True)))

def dataset_add_labels_col(df, col="label", data_folder=None):
    '''add a column with unique labels in column col'''
    df[col+"_labels"]=df[col].apply(lambda x: list(np.unique(np.load(x if data_folder is None else data_folder/x, allow_pickle=True))))

def dataset_add_mean_col(df, col="data", axis=(0), data_folder=None):
    '''adds a column with mean'''
    df[col+"_mean"]=df[col].apply(lambda x: np.mean(np.load(x if data_folder is None else data_folder/x, allow_pickle=True),axis=axis))

def dataset_add_median_col(df, col="data", axis=(0), data_folder=None):
    '''adds a column with median'''
    df[col+"_median"]=df[col].apply(lambda x: np.median(np.load(x if data_folder is None else data_folder/x, allow_pickle=True),axis=axis))

def dataset_add_std_col(df, col="data", axis=(0), data_folder=None):
    '''adds a column with mean'''
    df[col+"_std"]=df[col].apply(lambda x: np.std(np.load(x if data_folder is None else data_folder/x, allow_pickle=True),axis=axis))

def dataset_add_iqr_col(df, col="data", axis=(0), data_folder=None):
    '''adds a column with mean'''
    df[col+"_iqr"]=df[col].apply(lambda x: iqr(np.load(x if data_folder is None else data_folder/x, allow_pickle=True),axis=axis))

def dataset_get_stats(df, col="data", simple=True):
    '''creates (weighted) means and stds from mean, std and length cols of the df'''
    if(simple):
        return df[col+"_mean"].mean(), df[col+"_std"].mean()
    else:
        #https://notmatthancock.github.io/2017/03/23/simple-batch-stat-updates.html
        #or https://gist.github.com/thomasbrandon/ad5b1218fc573c10ea4e1f0c63658469
        def combine_two_means_vars(x1,x2):
            (mean1,var1,n1) = x1
            (mean2,var2,n2) = x2
            mean = mean1*n1/(n1+n2)+ mean2*n2/(n1+n2)
            var = var1*n1/(n1+n2)+ var2*n2/(n1+n2)+n1*n2/(n1+n2)/(n1+n2)*np.power(mean1-mean2,2)
            return (mean, var, (n1+n2))

        def combine_all_means_vars(means,vars,lengths):
            inputs = list(zip(means,vars,lengths))
            result = inputs[0]

            for inputs2 in inputs[1:]:
                result= combine_two_means_vars(result,inputs2)
            return result

        means = list(df[col+"_mean"])
        vars = np.power(list(df[col+"_std"]),2)
        lengths = list(df[col+"_length"])
        mean,var,length = combine_all_means_vars(means,vars,lengths)
        return mean, np.sqrt(var)

def npys_to_memmap_batched(npys, target_filename, max_len=0, delete_npys=True, batch_length=900000):
    memmap = None
    start = np.array([0])#start_idx in current memmap file (always already the next start- delete last token in the end)
    length = []#length of segment
    filenames= []#memmap files
    file_idx=[]#corresponding memmap file for sample
    shape=[]#shapes of all memmap files

    data = []
    data_lengths=[]
    dtype = None

    for idx,npy in tqdm(list(enumerate(npys))):

        data.append(np.load(npy, allow_pickle=True))
        data_lengths.append(len(data[-1]))

        if(idx==len(npys)-1 or np.sum(data_lengths)>batch_length):#flush
            data = np.concatenate(data)
            if(memmap is None or (max_len>0 and start[-1]>max_len)):#new memmap file has to be created
                if(max_len>0):
                    filenames.append(target_filename.parent/(target_filename.stem+"_"+str(len(filenames))+".npy"))
                else:
                    filenames.append(target_filename)

                shape.append([np.sum(data_lengths)]+[l for l in data.shape[1:]])#insert present shape

                if(memmap is not None):#an existing memmap exceeded max_len
                    del memmap
                #create new memmap
                start[-1] = 0
                start = np.concatenate([start,np.cumsum(data_lengths)])
                length = np.concatenate([length,data_lengths])

                memmap = np.memmap(filenames[-1], dtype=data.dtype, mode='w+', shape=data.shape)
            else:
                #append to existing memmap
                start = np.concatenate([start,start[-1]+np.cumsum(data_lengths)])
                length = np.concatenate([length,data_lengths])
                shape[-1] = [start[-1]]+[l for l in data.shape[1:]]
                memmap = np.memmap(filenames[-1], dtype=data.dtype, mode='r+', shape=tuple(shape[-1]))

            #store mapping memmap_id to memmap_file_id
            file_idx=np.concatenate([file_idx,[(len(filenames)-1)]*len(data_lengths)])
            #insert the actual data
            memmap[start[-len(data_lengths)-1]:start[-len(data_lengths)-1]+len(data)]=data[:]
            memmap.flush()
            dtype = data.dtype
            data = []#reset data storage
            data_lengths = []

    start= start[:-1]#remove the last element
    #cleanup
    for npy in npys:
        if(delete_npys is True):
            npy.unlink()
    del memmap

    #convert everything to relative paths
    filenames= [f.name for f in filenames]
    #save metadata
    np.savez(target_filename.parent/(target_filename.stem+"_meta.npz"),start=start,length=length,shape=shape,file_idx=file_idx,dtype=dtype,filenames=filenames)

def npys_to_memmap(npys, target_filename, max_len=0, delete_npys=True):
    memmap = None
    start = []#start_idx in current memmap file
    length = []#length of segment
    filenames= []#memmap files
    file_idx=[]#corresponding memmap file for sample
    shape=[]

    for _,npy in tqdm(list(enumerate(npys))):
        data = np.load(npy, allow_pickle=True)
        if(memmap is None or (max_len>0 and start[-1]+length[-1]>max_len)):
            if(max_len>0):
                filenames.append(target_filename.parent/(target_filename.stem+"_"+str(len(filenames))+".npy"))
            else:
                filenames.append(target_filename)

            if(memmap is not None):#an existing memmap exceeded max_len
                shape.append([start[-1]+length[-1]]+[l for l in data.shape[1:]])
                del memmap
            #create new memmap
            start.append(0)
            length.append(data.shape[0])
            memmap = np.memmap(filenames[-1], dtype=data.dtype, mode='w+', shape=data.shape)
        else:
            #append to existing memmap
            start.append(start[-1]+length[-1])
            length.append(data.shape[0])
            memmap = np.memmap(filenames[-1], dtype=data.dtype, mode='r+', shape=tuple([start[-1]+length[-1]]+[l for l in data.shape[1:]]))

        #store mapping memmap_id to memmap_file_id
        file_idx.append(len(filenames)-1)
        #insert the actual data
        memmap[start[-1]:start[-1]+length[-1]]=data[:]
        memmap.flush()
        if(delete_npys is True):
            npy.unlink()
    del memmap

    #append final shape if necessary
    if(len(shape)<len(filenames)):
        shape.append([start[-1]+length[-1]]+[l for l in data.shape[1:]])
    #convert everything to relative paths
    filenames= [f.name for f in filenames]
    #save metadata
    np.savez(target_filename.parent/(target_filename.stem+"_meta.npz"),start=start,length=length,shape=shape,file_idx=file_idx,dtype=data.dtype,filenames=filenames)

def reformat_as_memmap(df, target_filename, data_folder=None, annotation=False, max_len=0, delete_npys=True,col_data="data",col_label="label", batch_length=0):
    target_filename = Path(target_filename)
    data_folder = Path(data_folder)

    npys_data = []
    npys_label = []

    for _,row in df.iterrows():
        npys_data.append(data_folder/row[col_data] if data_folder is not None else row[col_data])
        if(annotation):
            npys_label.append(data_folder/row[col_label] if data_folder is not None else row[col_label])
    if(batch_length==0):
        npys_to_memmap(npys_data, target_filename, max_len=max_len, delete_npys=delete_npys)
    else:
        npys_to_memmap_batched(npys_data, target_filename, max_len=max_len, delete_npys=delete_npys,batch_length=batch_length)
    if(annotation):
        if(batch_length==0):
            npys_to_memmap(npys_label, target_filename.parent/(target_filename.stem+"_label.npy"), max_len=max_len, delete_npys=delete_npys)
        else:
            npys_to_memmap_batched(npys_label, target_filename.parent/(target_filename.stem+"_label.npy"), max_len=max_len, delete_npys=delete_npys, batch_length=batch_length)

    #replace data(filename) by integer
    df_mapped = df.copy()
    df_mapped[col_data+"_original"]=df_mapped.data
    df_mapped[col_data]=np.arange(len(df_mapped))

    df_mapped.to_pickle(target_filename.parent/("df_"+target_filename.stem+".pkl"))
    return df_mapped

class ConcatDatasetTimeseriesDatasetCrops(torch.utils.data.ConcatDataset):
    '''ConcatDataset that handles id mapping correctly (to allow to aggregate predictions)'''
    def __init__(self, datasets):
        super(ConcatDatasetTimeseriesCrops, self).__init__(datasets)
        idmaps = []
        for dataset_idx,ds in enumerate(self.datasets):
            idmap = ds.get_id_mapping()
            remap_dict = {x:j+(self.cumulative_sizes[dataset_idx-1] if dataset_idx>0 else 0) for j,x in enumerate(np.unique(idmap))}
            idmaps.append(np.array([remap_dict[x] for x in idmap]))
        self.df_idx_mapping = np.concatenate(idmaps)

    def get_id_mapping(self):
        return self.df_idx_mapping

class TimeseriesDatasetCrops(torch.utils.data.Dataset):
    """timeseries dataset with partial crops."""

    def __init__(self, df, output_size, chunk_length, min_chunk_length, memmap_filename=None, npy_data=None, random_crop=True, data_folder=None, num_classes=2, copies=0, col_lbl="label", cols_static=None, stride=None, start_idx=0, annotation=False, transforms=None, sample_items_per_record=1):
        """
        accepts three kinds of input:
        1) filenames pointing to aligned numpy arrays [timesteps,channels,...] for data and either integer labels or filename pointing to numpy arrays[timesteps,...] e.g. for annotations
        2) memmap_filename to memmap file (same argument that was passed to reformat_as_memmap) for data [concatenated,...] and labels- data column in df corresponds to index in this memmap
        3) npy_data [samples,ts,...] (either path or np.array directly- also supporting variable length input) - data column in df corresponds to sampleid

        transforms: list of callables (transformations) or (preferred) single instance e.g. from torchvision.transforms.Compose (applied in the specified order i.e. leftmost element first)

        col_lbl = None: return dummy label 0 (e.g. for unsupervised pretraining)
        cols_static: (optional) list of cols with extra static information
        """
        assert not((memmap_filename is not None) and (npy_data is not None))
        # require integer entries if using memmap or npy
        # keys (in column data) have to be unique
        assert(len(df["data"].unique())==len(df))

        self.timeseries_df_data = np.array(df["data"])
        if(self.timeseries_df_data.dtype not in [np.int16, np.int32, np.int64]):
            assert(memmap_filename is None and npy_data is None) #only for filenames in mode files
            self.timeseries_df_data = np.array(df["data"].astype(str)).astype(np.string_)

        if(col_lbl is None):# use dummy labels
            self.timeseries_df_label = np.zeros(len(df))
        else: # use actual labels
            if(isinstance(df[col_lbl].iloc[0],list) or isinstance(df[col_lbl].iloc[0],np.ndarray)):#stack arrays/lists for proper batching
                self.timeseries_df_label = np.stack(df[col_lbl])
            else: # single integers/floats
                self.timeseries_df_label = np.array(df[col_lbl])

            if(self.timeseries_df_label.dtype not in [np.int16, np.int32, np.int64, np.float32, np.float64]): #everything else cannot be batched anyway mp.Manager().list(self.timeseries_df_label)
                assert(annotation and memmap_filename is None and npy_data is None)#only for filenames in mode files
                self.timeseries_df_label = np.array(df[col_lbl].apply(lambda x:str(x))).astype(np.string_)

        if(cols_static is not None):
            self.timeseries_df_static = np.array(df[cols_static]).astype(np.float32)
            self.static = True
        else:
            self.static = False

        self.output_size = output_size
        self.data_folder = data_folder
        self.transforms = transforms
        if(isinstance(self.transforms,list) or isinstance(self.transforms,np.ndarray)):
            print("Warning: the use of lists as arguments for transforms is discouraged")
        self.annotation = annotation
        self.col_lbl = col_lbl

        self.c = num_classes

        self.mode="files"

        if(memmap_filename is not None):
            self.memmap_meta_filename = memmap_filename.parent/(memmap_filename.stem+"_meta.npz")
            self.mode="memmap"
            memmap_meta = np.load(self.memmap_meta_filename, allow_pickle=True)
            self.memmap_start = memmap_meta["start"]
            self.memmap_shape = memmap_meta["shape"]
            self.memmap_length = memmap_meta["length"]
            self.memmap_file_idx = memmap_meta["file_idx"]
            self.memmap_dtype = np.dtype(str(memmap_meta["dtype"]))
            self.memmap_filenames = np.array(memmap_meta["filenames"]).astype(np.string_)#save as byte to avoid issue with mp
            if(annotation):
                memmap_meta_label = np.load(self.memmap_meta_filename.parent/("_".join(self.memmap_meta_filename.stem.split("_")[:-1])+"_label_meta.npz"), allow_pickle=True)
                self.memmap_shape_label = memmap_meta_label["shape"]
                self.memmap_filenames_label = np.array(memmap_meta_label["filenames"]).astype(np.string_)
                self.memmap_dtype_label = np.dtype(str(memmap_meta_label["dtype"]))
        elif(npy_data is not None):
            self.mode="npy"
            if(isinstance(npy_data,np.ndarray) or isinstance(npy_data,list)):
                self.npy_data = np.array(npy_data)
                assert(annotation is False)
            else:
                self.npy_data = np.load(npy_data, allow_pickle=True)
            if(annotation):
                self.npy_data_label = np.load(npy_data.parent/(npy_data.stem+"_label.npy"), allow_pickle=True)

        self.random_crop = random_crop
        self.sample_items_per_record = sample_items_per_record

        self.df_idx_mapping=[]
        self.start_idx_mapping=[]
        self.end_idx_mapping=[]

        for df_idx,(id,row) in enumerate(df.iterrows()):
            if(self.mode=="files"):
                data_length = row["data_length"]
            elif(self.mode=="memmap"):
                data_length= self.memmap_length[row["data"]]
            else: #npy
                data_length = len(self.npy_data[row["data"]])

            if(chunk_length == 0):#do not split
                idx_start = [start_idx]
                idx_end = [data_length]
            else:
                idx_start = list(range(start_idx,data_length,chunk_length if stride is None else stride))
                idx_end = [min(l+chunk_length, data_length) for l in idx_start]

            #remove final chunk(s) if too short
            for i in range(len(idx_start)):
                if(idx_end[i]-idx_start[i]< min_chunk_length):
                    del idx_start[i:]
                    del idx_end[i:]
                    break
            #append to lists
            for _ in range(copies+1):
                for i_s,i_e in zip(idx_start,idx_end):
                    self.df_idx_mapping.append(df_idx)
                    self.start_idx_mapping.append(i_s)
                    self.end_idx_mapping.append(i_e)
        #convert to np.array to avoid mp issues with python lists
        self.df_idx_mapping = np.array(self.df_idx_mapping)
        self.start_idx_mapping = np.array(self.start_idx_mapping)
        self.end_idx_mapping = np.array(self.end_idx_mapping)

    def __len__(self):
        return len(self.df_idx_mapping)

    @property
    def is_empty(self):
        return len(self.df_idx_mapping)==0

    def __getitem__(self, idx):
        lst=[]
        for _ in range(self.sample_items_per_record):
            #determine crop idxs
            timesteps= self.get_sample_length(idx)

            if(self.random_crop):#random crop
                if(timesteps==self.output_size):
                    start_idx_rel = 0
                else:
                    start_idx_rel = random.randint(0, timesteps - self.output_size -1)#np.random.randint(0, timesteps - self.output_size)
            else:
                start_idx_rel =  (timesteps - self.output_size)//2
            if(self.sample_items_per_record==1):
                return self._getitem(idx,start_idx_rel)
            else:
                lst.append(self._getitem(idx,start_idx_rel))
        return tuple(lst)

    def _getitem(self, idx,start_idx_rel):
        #low-level function that actually fetches the data
        df_idx = self.df_idx_mapping[idx]
        start_idx = self.start_idx_mapping[idx]
        end_idx = self.end_idx_mapping[idx]
        #determine crop idxs
        timesteps= end_idx - start_idx
        assert(timesteps>=self.output_size)
        start_idx_crop = start_idx + start_idx_rel
        end_idx_crop = start_idx_crop+self.output_size

        #print(idx,start_idx,end_idx,start_idx_crop,end_idx_crop)
        #load the actual data
        if(self.mode=="files"):#from separate files
            data_filename = str(self.timeseries_df_data[df_idx],encoding='utf-8') #todo: fix potential issues here
            if self.data_folder is not None:
                data_filename = self.data_folder/data_filename
            data = np.load(data_filename, allow_pickle=True)[start_idx_crop:end_idx_crop] #data type has to be adjusted when saving to npy

            ID = data_filename.stem

            if(self.annotation is True):
                label_filename = str(self.timeseries_df_label[df_idx],encoding='utf-8')
                if self.data_folder is not None:
                    label_filename = self.data_folder/label_filename
                label = np.load(label_filename, allow_pickle=True)[start_idx_crop:end_idx_crop] #data type has to be adjusted when saving to npy
            else:
                label = self.timeseries_df_label[df_idx] #input type has to be adjusted in the dataframe
        elif(self.mode=="memmap"): #from one memmap file
            memmap_idx = self.timeseries_df_data[df_idx] #grab the actual index (Note the df to create the ds might be a subset of the original df used to create the memmap)
            memmap_file_idx = self.memmap_file_idx[memmap_idx]
            idx_offset = self.memmap_start[memmap_idx]

            #wi = torch.utils.data.get_worker_info()
            #pid = 0 if wi is None else wi.id#os.getpid()
            #print("idx",idx,"ID",ID,"idx_offset",idx_offset,"start_idx_crop",start_idx_crop,"df_idx", self.df_idx_mapping[idx],"pid",pid)
            mem_filename = str(self.memmap_filenames[memmap_file_idx],encoding='utf-8')
            mem_file = np.memmap(self.memmap_meta_filename.parent/mem_filename, self.memmap_dtype, mode='r', shape=tuple(self.memmap_shape[memmap_file_idx]))
            data = np.copy(mem_file[idx_offset + start_idx_crop: idx_offset + end_idx_crop])
            del mem_file
            #print(mem_file[idx_offset + start_idx_crop: idx_offset + end_idx_crop])
            if(self.annotation):
                mem_filename_label = str(self.memmap_filenames_label[memmap_file_idx],encoding='utf-8')
                mem_file_label = np.memmap(self.memmap_meta_filename.parent/mem_filename_label, self.memmap_dtype_label, mode='r', shape=tuple(self.memmap_shape_label[memmap_file_idx]))

                label = np.copy(mem_file_label[idx_offset + start_idx_crop: idx_offset + end_idx_crop])
                del mem_file_label
            else:
                label = self.timeseries_df_label[df_idx]
        else:#single npy array
            ID = self.timeseries_df_data[df_idx]

            data = self.npy_data[ID][start_idx_crop:end_idx_crop]

            if(self.annotation):
                label = self.npy_data_label[ID][start_idx_crop:end_idx_crop]
            else:
                label = self.timeseries_df_label[df_idx]

        sample = (data, label, self.timeseries_df_static[df_idx] if self.static else None)

        if(isinstance(self.transforms,list)):#transforms passed as list
            for t in self.transforms:
                sample = t(sample)
        elif(self.transforms is not None):#single transform e.g. from torchvision.transforms.Compose
            sample = self.transforms(sample)

        # consistency check: make sure that data and annotation lengths match
        assert(self.annotation is False or len(sample[0])==len(sample[1]))

        if(self.static is True):
            return tsdata_static(sample[0],sample[1], sample[2])
        else:
            return tsdata(sample[0], sample[1])

    def get_sampling_weights(self, class_weight_dict,length_weighting=False, timeseries_df_group_by_col=None):
        '''
        class_weight_dict: dictionary of class weights
        length_weighting: weigh samples by length
        timeseries_df_group_by_col: column of the pandas df used to create the object'''
        assert(self.annotation is False)
        assert(length_weighting is False or timeseries_df_group_by_col is None)
        weights = np.zeros(len(self.df_idx_mapping),dtype=np.float32)
        length_per_class = {}
        length_per_group = {}
        for iw,(i,s,e) in enumerate(zip(self.df_idx_mapping,self.start_idx_mapping,self.end_idx_mapping)):
            label = self.timeseries_df_label[i]
            weight = class_weight_dict[label]
            if(length_weighting):
                if label in length_per_class.keys():
                    length_per_class[label] += e-s
                else:
                    length_per_class[label] = e-s
            if(timeseries_df_group_by_col is not None):
                group = timeseries_df_group_by_col[i]
                if group in length_per_group.keys():
                    length_per_group[group] += e-s
                else:
                    length_per_group[group] = e-s
            weights[iw] = weight

        if(length_weighting):#need second pass to properly take into account the total length per class
            for iw,(i,s,e) in enumerate(zip(self.df_idx_mapping,self.start_idx_mapping,self.end_idx_mapping)):
                label = self.timeseries_df_label[i]
                weights[iw]= (e-s)/length_per_class[label]*weights[iw]
        if(timeseries_df_group_by_col is not None):
            for iw,(i,s,e) in enumerate(zip(self.df_idx_mapping,self.start_idx_mapping,self.end_idx_mapping)):
                group = timeseries_df_group_by_col[i]
                weights[iw]= (e-s)/length_per_group[group]*weights[iw]

        weights = weights/np.min(weights)#normalize smallest weight to 1
        return weights

    def get_id_mapping(self):
        return self.df_idx_mapping

    def get_sample_id(self,idx):
        return self.df_idx_mapping[idx]

    def get_sample_length(self,idx):
        return self.end_idx_mapping[idx]-self.start_idx_mapping[idx]

    def get_sample_start(self,idx):
        return self.start_idx_mapping[idx]

class RandomCrop(object):
    """Crop randomly the image in a sample.
    """

    def __init__(self, output_size,annotation=False):
        self.output_size = output_size
        self.annotation = annotation

    def __call__(self, sample):
        data, label, static = sample

        timesteps= len(data)
        assert(timesteps>=self.output_size)
        if(timesteps==self.output_size):
            start=0
        else:
            start = random.randint(0, timesteps - self.output_size-1) #np.random.randint(0, timesteps - self.output_size)

        data = data[start: start + self.output_size]
        if(self.annotation):
            label = label[start: start + self.output_size]

        return data, label, static

class CenterCrop(object):
    """Center crop the image in a sample.
    """

    def __init__(self, output_size, annotation=False):
        self.output_size = output_size
        self.annotation = annotation

    def __call__(self, sample):
        data, label, static = sample

        timesteps= len(data)
        start = (timesteps - self.output_size)//2

        data = data[start: start + self.output_size]
        if(self.annotation):
            label = label[start: start + self.output_size]

        return data, label, static

class GaussianNoise(object):
    """Add gaussian noise to sample.
    """

    def __init__(self, scale=0.1):
        self.scale = scale

    def __call__(self, sample):
        if self.scale ==0:
            return sample
        else:
            data, label, static = sample
            data = data + np.reshape(np.array([random.gauss(0,self.scale) for _ in range(np.prod(data.shape))]),data.shape)#np.random.normal(scale=self.scale,size=data.shape).astype(np.float32)
            return data, label, static

class Rescale(object):
    """Rescale by factor.
    """

    def __init__(self, scale=0.5,interpolation_order=3):
        self.scale = scale
        self.interpolation_order = interpolation_order

    def __call__(self, sample):
        if self.scale ==1:
            return sample
        else:
            data, label, static = sample
            timesteps_new = int(self.scale * len(data))
            data = transform.resize(data,(timesteps_new,data.shape[1]),order=self.interpolation_order).astype(np.float32)
            return data, label, static

class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""
    def __init__(self, transpose_data=True, transpose_label=False):
        #swap channel and time axis for direct application of pytorch's convs
        self.transpose_data=transpose_data
        self.transpose_label=transpose_label

    def __call__(self, sample):

        def _to_tensor(data,transpose=False):
            if(isinstance(data,np.ndarray)):
                if(transpose):#seq,[x,y,]ch
                    return torch.from_numpy(np.moveaxis(data,-1,0))
                else:
                    return torch.from_numpy(data)
            else:#default_collate will take care of it
                return data

        data, label, static = sample

        if not isinstance(data,tuple):
            data = _to_tensor(data,self.transpose_data)
        else:
            data = tuple(_to_tensor(x,self.transpose_data) for x in data)

        if not isinstance(label,tuple):
            label = _to_tensor(label,self.transpose_label)
        else:
            label = tuple(_to_tensor(x,self.transpose_label) for x in label)

        return data, label, static #returning as a tuple (potentially of lists)

class Normalize(object):
    """Normalize using given stats.
    """
    def __init__(self, stats_mean, stats_std, input=True, channels=[]):
        self.stats_mean=stats_mean.astype(np.float32) if stats_mean is not None else None
        self.stats_std=stats_std.astype(np.float32)+1e-8 if stats_std is not None else None
        self.input = input
        if(len(channels)>0):
            for i in range(len(stats_mean)):
                if(not(i in channels)):
                    self.stats_mean[:,i]=0
                    self.stats_std[:,i]=1

    def __call__(self, sample):
        datax, labelx, static = sample
        data = datax if self.input else labelx
        #assuming channel last
        if(self.stats_mean is not None):
            data = data - self.stats_mean
        if(self.stats_std is not None):
            data = data/self.stats_std

        if(self.input):
            return (data, labelx, static)
        else:
            return (datax, data, static)

class NormalizeBatch(object):
    """Normalize using batch statistics.
    axis: tuple of integers of axis numbers to be normalized over (by default everything but the last)
    """
    def __init__(self, input=True, channels=[],axis=None):
        self.channels = channels
        self.channels_keep = None
        self.input = input
        self.axis = axis

    def __call__(self, sample):
        datax, labelx, static = sample
        data = datax if self.input else labelx
        #assuming channel last
        #batch_mean = np.mean(data,axis=tuple(range(0,len(data)-1)))
        #batch_std = np.std(data,axis=tuple(range(0,len(data)-1)))+1e-8
        batch_mean = np.mean(data,axis=self.axis if self.axis is not None else tuple(range(0,len(data.shape)-1)))
        batch_std = np.std(data,axis=self.axis if self.axis is not None else tuple(range(0,len(data.shape)-1)))+1e-8

        if(len(self.channels)>0):
            if(self.channels_keep is None):
                self.channels_keep = np.setdiff(range(data.shape[-1]),self.channels)

            batch_mean[self.channels_keep]=0
            batch_std[self.channels_keep]=1

        data = (data - batch_mean)/batch_std

        if(self.input):
            return (data, labelx, static)
        else:
            return (datax, data, static)

class ButterFilter(object):
    """Apply filter
    """

    def __init__(self, lowcut=50, highcut=50, fs=100, order=5, btype='band', forwardbackward=True, input=True):
        self.filter = butter_filter(lowcut,highcut,fs,order,btype)
        self.input = input
        self.forwardbackward = forwardbackward

    def __call__(self, sample):
        datax, labelx, static = sample
        data = datax if self.input else labelx

        if(self.forwardbackward):
            data = sosfiltfilt(self.filter, data, axis=0)
        else:
            data = sosfilt(self.filter, data, axis=0)

        if(self.input):
            return (data, labelx, static)
        else:
            return (datax, data, static)

class ChannelFilter(object):
    """Select certain channels.
    """

    def __init__(self, channels=[0], input=True):
        self.channels = channels
        self.input = input

    def __call__(self, sample):
        data, label, static = sample
        if(self.input):
            return (data[...,self.channels], label, static)
        else:
            return (data, label[...,self.channels], static)

class Transform(object):
    """Transforms data using a given function i.e. data_new = func(data) for input is True else label_new = func(label)
    """

    def __init__(self, func, input=False):
        self.func = func
        self.input = input

    def __call__(self, sample):
        data, label, static = sample
        if(self.input):
            return (self.func(data), label, static)
        else:
            return (data, self.func(label), static)

class TupleTransform(object):
    """Transforms data using a given function (operating on both data and label and return a tuple) i.e. data_new, label_new = func(data_old, label_old)
    """

    def __init__(self, func):
        self.func = func

    def __call__(self, sample):
        data,label,static = sample
        return  self.func(data,label,static)

def aggregate_predictions(preds,targs=None,idmap=None,aggregate_fn = np.mean,verbose=False):
    '''
    aggregates potentially multiple predictions per sample (can also pass targs for convenience)
    idmap: idmap as returned by TimeSeriesCropsDataset's get_id_mapping
    preds: ordered predictions as returned by learn.get_preds()
    aggregate_fn: function that is used to aggregate multiple predictions per sample (most commonly np.amax or np.mean)
    '''
    if(idmap is not None and len(idmap)!=len(np.unique(idmap))):
        if(verbose):
            print("aggregating predictions...")
        preds_aggregated = []
        targs_aggregated = []
        for i in np.unique(idmap):
            preds_local = preds[np.where(idmap==i)[0]]
            preds_aggregated.append(aggregate_fn(preds_local,axis=0))
            if targs is not None:
                targs_local = targs[np.where(idmap==i)[0]]
                assert(np.all(targs_local==targs_local[0])) #all labels have to agree
                targs_aggregated.append(targs_local[0])
        if(targs is None):
            return np.array(preds_aggregated)
        else:
            return np.array(preds_aggregated),np.array(targs_aggregated)
    else:
        if(targs is None):
            return preds
        else:
            return preds,targs


places_simple_template = [
    lambda c: f'a photo of a {c}.',
]

places_main_template = [
    lambda c: f'a photo of a {c}',
]

places_factor_templates = {
    "background": {
        "others": [""],
        "natural": ["in a natural setting", "outdoors in nature", "in the wild"],
        "urban": ["in an urban setting", "city background"],
        "indoor": ["indoors", "inside a building", "enclosed space"],
    },
    "quality": {
        "normal": [""],
        "good": ["good", "looks good"],
        "bad": ["bad", "poor"],
        "low_res": ["low resolution"],
        "pixelated": ["pixelated"],
        "jpeg_corrupted": ["jpeg corrupted", "jpeg"],
        "blurry": ["blurry"],
        "clean": ["clean"],
        "dirty": ["dirty"],
    },
    "condition": {
        "normal": [""],
        "cool": ["cool", "stylish", "trendy"],
        "nice": ["nice", "pleasant", "good-looking"],
        "weird": ["weird", "odd", "unusual"],
    },
}


from ansible.module_utils.basic import *
try:
    from hpOneView.oneview_client import OneViewClient
    from hpOneView.exceptions import HPOneViewException

    HAS_HPE_ONEVIEW = True
except ImportError:
    HAS_HPE_ONEVIEW = False

DOCUMENTATION = '''
module: oneview_task_facts
short_description: Retrieve facts about the OneView Tasks.
description:
    - Retrieve facts about the OneView Tasks.
requirements:
    - "python >= 2.7.9"
    - "hpOneView >= 2.0.1"
author: "Bruno Souza (@bsouza)"
options:
    config:
      description:
        - Path to a .json configuration file containing the OneView client configuration.
          The configuration file is optional. If the file path is not provided, the configuration will be loaded from
          environment variables.
      required: false
    params:
      description:
        - "List with parameters to help filter the tasks.
          Params allowed: count, fields, filter, query, sort, start, and view."
      required: false
notes:
    - "A sample configuration file for the config parameter can be found at:
       https://github.com/HewlettPackard/oneview-ansible/blob/master/examples/oneview_config-rename.json"
    - "Check how to use environment variables for configuration at:
       https://github.com/HewlettPackard/oneview-ansible#environment-variables"
'''

EXAMPLES = '''
- name: Gather facts about the last 2 tasks
  oneview_task_facts:
    config: "{{ config }}"
    params:
      count: 2

- debug: var=tasks

- name: Gather facts about the last 2 tasks associated to Server Profile templates
  oneview_task_facts:
    config: "{{ config }}"
    params:
      count: 2
      filter: "associatedResource.resourceCategory='server-profile-templates'"

- debug: var=tasks
'''

RETURN = '''
tasks:
    description: The list of tasks.
    returned: Always, but can be null.
    type: list
'''
HPE_ONEVIEW_SDK_REQUIRED = 'HPE OneView Python SDK is required for this module.'

class TaskFactsModule(object):
    argument_spec = dict(
        config=dict(required=False, type='str'),
        params=dict(required=False, type='dict')
    )

    def __init__(self):
        self.module = AnsibleModule(argument_spec=self.argument_spec, supports_check_mode=False)
        if not HAS_HPE_ONEVIEW:
            self.module.fail_json(msg=HPE_ONEVIEW_SDK_REQUIRED)

        if not self.module.params['config']:
            oneview_client = OneViewClient.from_environment_variables()
        else:
            oneview_client = OneViewClient.from_json_file(self.module.params['config'])

        self.resource_client = oneview_client.tasks

    def run(self):
        try:
            params = self.module.params.get('params') or {}
            facts = self.resource_client.get_all(**params)
            self.module.exit_json(changed=False, ansible_facts=dict(tasks=facts))

        except HPOneViewException as exception:
            self.module.fail_json(msg='; '.join(str(e) for e in exception.args))

def main():
    TaskFactsModule().run()

if __name__ == '__main__':
    main()


def func(s):
    res = ''
    if '0' <= s[0] <= '9':
        res = '*' + s[0]
    else:
        res += s[0]
    for cur in range(1, len(s)):
        if '0' <= s[cur] <= '9':
            if '0' <= s[cur - 1] <= '9':
                res += s[cur]
            else:
                res = res + '*' + s[cur]
        else:
            if '0' <= s[cur - 1] <= '9':
                res = res + '*' + s[cur]
            else:
                res += s[cur]
    if '0' <= s[-1] <= '9':
        res += '*'
    return res

while 1:
    try:
        s = str(input())
        print(func(s))
    except:
        break


"""

Construir um algoritmo para gerar a seguinte série:

s = 1/1 + 1/2 + 1/3 + ... + 1/n para os 50 primeiros termos.

"""

soma = 0
contador = 0
for c in range(1, 51):
    contador += 1
    soma += 1/contador
print(soma)


class Solution:
    def maximumBeauty(self, items: List[List[int]], queries: List[int]) -> List[int]:
        n = len(queries)
        items.sort()
        res = [0] * n
        items.sort()
        queries = sorted([(queries[i], i) for i in range(n)])
        i = 0
        maxbeauty = 0
        for q, j in queries:
            while i < len(items) and items[i][0] <= q:
                maxbeauty = max(maxbeauty, items[i][1])
                i += 1
            res[j] = maxbeauty
        return res


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

# Taking input
word = input()

# Defining the vowels characters with space
v = "aeiou "

# List comprehension
words_list1 = [i if i.lower() in v else i+'o'+i for i in word]

# printing the translated word
print("Translated word:", "".join(words_list1))


import numpy as np
import sys
from cmdinfo import CmdInfo

class Calculator:
    answer = 0
    saved_answer = 0
    cur_operator = ''
    show_help_msg = True
    show_power_off_msg = True

    help_info = CmdInfo(
        'Help',
        ['help', 'help (page #)'],
        "get help on commands"
    )
    power_off_info = CmdInfo(
        'Power Off',
        ['off', 'stop', 'exit'],
        "powers off calculator"
    )
    number_info = CmdInfo(
        'Number',
        ['(number)'],
        "sets number to calculator or use in calculation"
    )
    addition_info = CmdInfo(
        'Addition',
        ['+', 'add'],
        "adds two numbers together"
    )
    subtraction_info = CmdInfo(
        'Subtraction',
        ['-', 'subtract'],
        "subtracts two numbers"
    )
    multiplication_info = CmdInfo(
        'Multiplication',
        ['*', 'multiply'],
        "multiplies two numbers together"
    )
    division_info = CmdInfo(
        'Division',
        ['/', 'divide'],
        "divides two numbers."
    )
    exponent_info = CmdInfo(
        'Exponent',
        ['**', '^', 'exponent', 'power'],
        "set one number to the power of another"
    )
    sin_info = CmdInfo(
        'Sin',
        ['sin'],
        "get the trigonometric sin of a number"
    )
    cos_info = CmdInfo(
        'Cosine',
        ['cos', 'cosine'],
        "get the trigonometric cosine of a number"
    )
    tan_info = CmdInfo(
        'Tangent',
        ['tan', 'tangent'],
        "get the trigonometric tangent of a number"
    )
    show_help_msg_info = CmdInfo(
        'Show Help Message',
        ['show help msg'],
        "shows help message on calculator"
    )
    hide_help_msg_info = CmdInfo(
        'Hide Help Message',
        ['hide help msg'],
        "hides help message on calculator"
    )
    show_power_off_msg_info = CmdInfo(
        'Show Power Off Message',
        ['show poweroff msg'],
        "shows power off message on calculator"
    )
    hide_power_off_msg_info = CmdInfo(
        'Hide Power Off Message',
        ['hide poweroff msg'],
        "hides power off message on calculator"
    )

    cmdinfos = [
        help_info,
        power_off_info,
        number_info,
        addition_info,
        subtraction_info,
        multiplication_info,
        division_info,
        exponent_info,
        sin_info,
        cos_info,
        tan_info,
        show_help_msg_info,
        show_power_off_msg_info,
        hide_help_msg_info,
        hide_power_off_msg_info
    ]

    def __init__(self):
        pass

    def calc_screen(self):
        return f"""
____________________________
|                          |
|  {' ' * (22-len(str(self.answer))) + str(self.answer)}  |
|__________________________|""" + (f"""
|                          |
|  {' ' * ((22-len(self.cur_operator)) // 2) + self.cur_operator + ' ' * ((22-len(self.cur_operator)) // 2 + (22-len(self.cur_operator)) % 2)}  |
|__________________________|
""" if self.cur_operator != '' else '\n') + \
("""|                          |
|     Don't know how to    |
|        perform a         |
|       calculation?       |
|__________________________|
|                          |
|    enter: help (pg #)    |
|__________________________|""" + \
('\n'  if self.show_power_off_msg else '') if self.show_help_msg else '') + \
("""|                          |
|   Power off calculator   |
|        enter: off        |
|__________________________|""" if self.show_power_off_msg else '')

    def check_enter(self, enter: str):
        if enter in self.power_off_info.enter_options:
            return False
        elif enter[:4] == 'help':
            page = enter[5:]
            self.get_help(int(page) if page.isdigit() else 1)
        elif enter.isnumeric() or \
            (enter.count('.')==1 and enter.replace('.', '').isnumeric()) or \
            (enter[0]=='-' and (enter[1:].isnumeric() or \
            enter[1:].count('.')==1 and enter[1:].replace('.', '').isnumeric())):
            if self.cur_operator != '':
                self.answer = self.operate(self.saved_answer, self.cur_operator, float(enter))
            else:
                self.answer = float(enter)
            self.cur_operator = ''
        elif 'hide' in enter or 'show' in enter:
            if enter in self.show_help_msg_info.enter_options:
                self.show_help_msg = True
            elif enter in self.hide_help_msg_info.enter_options:
                self.show_help_msg = False
            elif enter in self.show_power_off_msg_info.enter_options:
                self.show_power_off_msg = True
            elif enter in self.hide_power_off_msg_info.enter_options:
                self.show_power_off_msg = False
        else:
            if enter in self.addition_info.enter_options:
                self.cur_operator = 'add'
                self.saved_answer = self.answer
            elif enter in self.subtraction_info.enter_options:
                self.cur_operator = 'subtract'
                self.saved_answer = self.answer
            elif enter in self.multiplication_info.enter_options:
                self.cur_operator = 'multiply'
                self.saved_answer = self.answer
            elif enter in self.division_info.enter_options:
                self.cur_operator = 'divide'
                self.saved_answer = self.answer
            elif enter in self.exponent_info.enter_options:
                self.cur_operator = 'exponent'
                self.saved_answer = self.answer
            elif enter in self.sin_info.enter_options:
                self.cur_operator = 'sin'
                self.saved_answer = self.answer
                self.answer = self.operate(self.saved_answer, self.cur_operator)
            elif enter in self.cos_info.enter_options:
                self.cur_operator = 'cos'
                self.saved_answer = self.answer
                self.answer = self.operate(self.saved_answer, self.cur_operator)
            elif enter in self.tan_info.enter_options:
                self.cur_operator = 'tan'
                self.saved_answer = self.answer
                self.answer = self.operate(self.saved_answer, self.cur_operator)
        return True

    def operate(self, num1, operator, num2=0):
        try:
            self.cur_operator = ''
            if operator == 'add':
                return num1 + num2
            elif operator == 'subtract':
                return num1 - num2
            elif operator == 'multiply':
                return num1 * num2
            elif operator == 'divide':
                return num1 / num2
            elif operator == 'exponent':
                return num1 ** num2
            elif operator == 'sin':
                return np.sin(num1)
            elif operator == 'cos':
                return np.cos(num1)
            elif operator == 'tan':
                return np.tan(num1)
        except OverflowError:
            input("\nOverflow Error: The resulting value is too large.\nPress enter to continue...")
            return 0
        except ZeroDivisionError:
            input("\nZero Division Error: I think that would be infinity...don't do that...\nPress enter to continue...")
        except:
            print("Unexpected error:", sys.exc_info()[0])
            return 0

    def get_help(self, page):
        if page > int(np.ceil(len(self.cmdinfos)/3)):
            print(f"""
___________________________________________________________________________

                           Page does not exist
                          (Total Page Count: {int(np.ceil(len(self.cmdinfos)/3))})
___________________________________________________________________________
""")
        else:
            print(f"""
___________________________________________________________________________

                              Help Page {page} of {int(np.ceil(len(self.cmdinfos)/3))}
___________________________________________________________________________
""")
        for i in range((page-1)*3, len(self.cmdinfos) if page*3>len(self.cmdinfos) else page*3):
                print(self.cmdinfos[i])
        input("Press enter to continue...")


# mFLIM dataset synthetic generator
import logging as log
import numpy as np
import scipy.io
from scipy import linalg

def update_args(kwargs: dict) -> dict:
    args_given = list(kwargs.keys())
    default_args = {'n_order': 2, 'n_samples': 60, 'ts': 0.25e-9, 'snr': 0, 'psnr': 0}
    for arg in args_given:
        if arg in default_args.keys():
            default_args[arg] = kwargs[arg]
    return default_args

# Check type-hint
def mflimsynth(**kwargs: [int, float]) -> [np.ndarray, np.ndarray, np.ndarray, dict]:
    """
    INPUTS:
    n_order --> Order of multi-exponential model. Only n = 2,3,4.
    n_samples --> Number of pixels in x & y axes. Size of square image.
    ts --> Sampling Time.
    snr --> Level of Gaussian noise (dB), SNR = 45,50,55.
    psnr --> Level for Shot noise (dB), PSNR = 15,20,25.
    OUTPUTS:
    y_matrix --> matrix of fluorescence decays of size [186 x (n_samples*n_samples)].
    a_matrix --> matrix of abundances of [n x (n_samples*n_samples)].
    p_matrix --> matrix of endmembers [186 x n_order].
    - Daniel U. Campos-Delgado (July/2020).
    """
    try:
        # Synthetic FLIM Dataset
        num_input_args = len(kwargs)
        if num_input_args == 0:
            args = {'n_order': 2, 'n_samples': 60, 'ts': 0.25e-9, 'snr': 40, 'psnr': 20}
        elif num_input_args in [1, 2, 3, 4, 5]:
            args = update_args(kwargs)

        if args['n_order'] > 4:
            args['n_order'] = 4
            log.info("The maximum number of components is 4!!")

        if args['snr'] or args['psnr']:
            noise_measurement = True
            if args['snr']:
                noise_gaussian = True
            else:
                noise_gaussian = False
            if args['psnr']:
                noise_shot = True
            else:
                noise_shot = False
        else:
            noise_measurement = noise_gaussian = noise_shot = False
    except Exception as error:
        log.error(f"ERROR [Input arguments step]: {error}")
    try:
        num_samples = args['n_samples']
        x = y = np.arange(1, args['n_samples'] + 1)
        xx, yy = np.meshgrid(x, y)
        k = args['n_samples'] * args['n_samples']
        dataexp1 = scipy.io.loadmat("./files/dataexp1")
        u, y = dataexp1["u"], dataexp1["y"]
        del dataexp1
        u_len = len(u)
        u = u/sum(u)  # vector of laser input
        n = np.arange(u_len).reshape((1, u_len))
        t = n*args['ts']  # vector of time samples
        zeros = np.zeros((1, u_len))
        zeros[0][0] = u[0][0]
        u_toe = linalg.toeplitz(u.T, zeros)
        tau11, tau12, tau13, tau14 = [10, 20, 5, 15]
        tau21, tau22, tau23, tau24 = [25, 10, 15, 5]
        tau31, tau32, tau33, tau34 = [5, 7, 35, 15]

        if args['n_order'] == 2:
            aa1 = 7 * np.exp(-0.001 * (xx - num_samples / 2) ** 2 - 0.001 * (yy - num_samples / 2) ** 2) + 0.5
            aa2 = 2.5 * np.exp(-0.001 * (xx - num_samples) ** 2 - 0.001 * yy ** 2) + 2.5 * np.exp(
                -0.0001 * xx ** 2 - 0.001 * (yy - num_samples) ** 2) + \
                  2.5 * np.exp(-0.001 * xx ** 2 - 0.0001 * (yy - num_samples) ** 2) + 2.5 * \
                  np.exp(-0.0001 * (xx - num_samples) ** 2 - 0.001 * (yy - num_samples) ** 2)
            a1 = np.zeros((num_samples, num_samples))
            a2 = np.zeros((num_samples, num_samples))
            for i in np.arange(num_samples):
                for j in np.arange(num_samples):
                    a1[i, j] = aa1[i, j] / (aa1[i, j] + aa2[i, j])
                    a2[i, j] = aa2[i, j] / (aa1[i, j] + aa2[i, j])

        elif args['n_order'] == 3:
            aa1 = 7 * np.exp(-0.005 * (xx - num_samples / 2) ** 2 - 0.005 * (yy - num_samples / 2) ** 2) + 0.5
            aa2 = 2.5 * np.exp(-0.001 * (xx - num_samples) ** 2 - 0.001 * yy ** 2) + \
                  2.5 * np.exp(-0.0001 * xx ** 2 - 0.001 * (yy - num_samples) ** 2)
            aa3 = 3.5 * np.exp(-0.001 * xx ** 2 - 0.0001 * (yy - num_samples) ** 2) + 2.5 * \
                  np.exp(-0.0001 * (xx - num_samples) ** 2 - 0.001 * (yy - num_samples) ** 2)
            a1 = np.zeros((num_samples, num_samples))
            a2 = np.zeros((num_samples, num_samples))
            a3 = np.zeros((num_samples, num_samples))
            for i in np.arange(num_samples):
                for j in np.arange(num_samples):
                    a1[i, j] = aa1[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j])
                    a2[i, j] = aa2[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j])
                    a3[i, j] = aa3[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j])

        elif args['n_order'] == 4:
            aa1 = 2.5 * np.exp(-0.005 * (xx - num_samples / 2) ** 2 - 0.0005 * (yy - num_samples / 2) ** 2) + 0
            # \ + 2.5 * np.exp(-0.0001 * (xx)**2 - 0.001 * (yy - num_samples)**2)
            aa2 = 2.5 * np.exp(-0.001 * (xx - num_samples) ** 2 - 0.00025 * yy ** 2)
            aa3 = 2.5 * np.exp(-0.001 * xx ** 2 - 0.0002 * (yy - num_samples) ** 2)
            aa4 = 2.5 * np.exp(-0.001 * (xx - 8 * (num_samples / 9)) ** 2 - 0.001 * (yy - 8 * (num_samples / 9)) ** 2)\
                  + 2.5 * np.exp(-0.001 * (xx - (num_samples / 9)) ** 2 - 0.001 * (yy - 8 * (num_samples / 9)) ** 2)
            a1 = np.zeros((num_samples, num_samples))
            a2 = np.zeros((num_samples, num_samples))
            a3 = np.zeros((num_samples, num_samples))
            a4 = np.zeros((num_samples, num_samples))
            for i in np.arange(num_samples):
                for j in np.arange(num_samples):
                    a1[i, j] = aa1[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j] + aa4[i, j])
                    a2[i, j] = aa2[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j] + aa4[i, j])
                    a3[i, j] = aa3[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j] + aa4[i, j])
                    a4[i, j] = aa4[i, j] / (aa1[i, j] + aa2[i, j] + aa3[i, j] + aa4[i, j])

        p11 = u_toe @ np.exp(-n.T / tau11) * 0.6
        p21 = u_toe @ np.exp(-n.T / tau21) * 0.2
        p31 = u_toe @ np.exp(-n.T / tau31) * 0.2
        p1 = np.r_[p11, p21, p31]

        if args['n_order'] >= 2:
            p12 = u_toe @ np.exp(-n.T / tau12) * 0.2
            p22 = u_toe @ np.exp(-n.T / tau22) * 0.6
            p32 = u_toe @ np.exp(-n.T / tau32) * 0.2
            p2 = np.r_[p12, p22, p32]

            if args['n_order'] >= 3:
                p13 = u_toe @ np.exp(-n.T / tau13) * 0.15
                p23 = u_toe @ np.exp(-n.T / tau23) * 0.15
                p33 = u_toe @ np.exp(-n.T / tau33) * 0.70
                p3 = np.r_[p13, p23, p33]

                if args['n_order'] >= 4:
                    p14 = u_toe @ np.exp(-n.T / tau14) * 0.0
                    p24 = u_toe @ np.exp(-n.T / tau24) * 0.4
                    p34 = u_toe @ np.exp(-n.T / tau34) * 0.6
                    p4 = np.r_[p14, p24, p34]

        y_y = np.zeros((num_samples, num_samples, 3 * u_len))

        for i in range(num_samples):
            for j in range(num_samples):
                if args['n_order'] == 2:
                    y = a1[i, j] * p1 + a2[i, j] * p2
                elif args['n_order'] == 3:
                    y = a1[i, j] * p1 + a2[i, j] * p2 + a3[i, j] * p3
                elif args['n_order'] == 4:
                    y = a1[i, j] * p1 + a2[i, j] * p2 + a3[i, j] * p3 + a4[i, j] * p4

                if noise_measurement and noise_gaussian:
                    sigma_y = np.sqrt((y.T @ y) / (10 ** (args['snr'] / 10)))
                    yy1 = sigma_y * np.random.normal(size=(3 * u_len, 1))
                else:
                    yy1 = np.zeros((3 * u_len, 1))

                if noise_measurement and noise_shot:
                    sigma_y = np.sqrt(np.max(y) / (10 ** (args['psnr'] / 10)))
                    yy2 = sigma_y * np.random.normal(size=(3 * u_len, 1)) * np.sqrt(abs(y))
                else:
                    yy2 = np.zeros((3 * u_len, 1))

                y_y[i, j, :] = y.T + yy1.T + yy2.T

        if args['n_order'] == 2:
            p_matrix = np.c_[p1, p2]
            a_matrix = np.r_[a1.reshape((1, k), order="F"), a2.reshape((1, k), order="F")]
        elif args['n_order'] == 3:
            p_matrix = np.c_[p1, p2, p3]
            a_matrix = np.r_[a1.reshape((1, k), order="F"), a2.reshape(
                (1, k), order="F"), a3.reshape((1, k), order="F")]
        elif args['n_order'] == 4:
            p_matrix = np.c_[p1, p2, p3, p4]
            a_matrix = np.r_[a1.reshape((1, k), order="F"), a2.reshape(
                (1, k), order="F"), a3.reshape((1, k), order="F"), a4.reshape((1, k), order="F")]
        p_sum = np.sum(p_matrix, 0)
        p_matrix = p_matrix / np.tile(p_sum, [p_matrix.shape[0], 1])
        y_matrix_origin = y_y.reshape((k, 3 * u_len), order="F").T
        return y_matrix_origin, p_matrix, a_matrix, args
    except Exception as error:
        log.error(f"ERROR [mFLIM matrices calculation]: {error}")


import numpy as np
import os
import sys
import re
import json
from scipy.spatial.transform import Rotation as R
import shutil
import struct

#################################
#   General Helper Functions    #
#################################

def rotate_positions_around_z(positions):
    """
    Given a list of coordinates, rotate them all by 90 degrees around the z-axis. This is used to
        convert particle coordinates from the raw tiltseries to the final reconstruction's
        coordinate system for simulated data.

    Args:
        positions: A list of [x, y, z] coordinates

    Returns: None

    """
    rot = R.from_euler("zxz", (90, 0, 0), degrees=True)
    for i, point in enumerate(positions):
        positions[i] = np.dot(rot.as_matrix(), np.array(point))

    return positions

def convert_slicer_to_motl(orientations):
    """
    Convert a set of Slicer angles to the reference-to-particle ZXZ, external Euler angles for Artiatomi
    Args:
        orientations: The list of Slicer angles

    Returns: The list of Euler angles

    """
    for i, point in enumerate(orientations):
        slicer = orientations[i]
        # Rotate by 90 around the x-axis so that the membrane is in the XY
        # plane (Phi gives in-plane rotation)
        slicer_rot = R.from_euler(
            "zyx", [slicer[2], slicer[1], slicer[0]], degrees=True
        )
        orientation_mat = np.dot(
            R.from_euler("zxz", [0, 90, 0], degrees=True).as_matrix(),
            slicer_rot.as_matrix(),
        )
        rotation = R.from_matrix(orientation_mat)
        ref_to_part = rotation.inv()
        eulers = ref_to_part.as_euler("zxz", degrees=True)
        # Like PEET, Artiatomi takes Z1, Z2, X
        orientations[i] = [eulers[0], eulers[2], eulers[1]]

    return orientations

def get_slicer_info(mod_file):
    """
    Open an IMOD .mod file and retrieve the Slicer information

    Args:
        mod_file: The .mod file path

    Returns: A list of Slicer point objects with keys {"angles", "coords"}

    """
    results = []
    with open(mod_file, "rb") as file:
        token = file.read(4)
        if token != b"IMOD":
            print(
                "ID of .mod file is not 'IMOD'. This does not seem to be an IMOD MOD file!"
            )
            exit(1)

        # Read past rest of ID and file header
        file.read(236)

        while token != b"IEOF":
            file.seek(-3, 1)
            token = file.read(4)
            if token == b"SLAN":
                # Read past SLAN object size and time
                file.read(8)

                angles = struct.unpack(">" + ("f" * 3), file.read(4 * 3))
                xyz = struct.unpack(">" + ("f" * 3), file.read(4 * 3))

                results.append({"angles": angles, "coords": xyz})
                file.read(32)

                # Read forward a little so next iteration of while loop starts at end of SLAN object
                file.read(3)
            elif token == b"OBJT":
                # Objects are 176 bytes; skip path to make reading faster
                file.read(176)
                # Read forward a little so next iteration of while loop starts at end of object
                file.read(3)

    if len(results) == 0:
        print("Reached end of MOD file without finding slicer angles!")
        exit(1)

    return results

#######################
#   IMOD Functions    #
#######################

def imod_setup_sta(artia_root, dirs_start_with):
    """
    Write out the info file necessary to run the MATLAB set up script for sub-tomogram averaging in Artiatomi
    Args:
        artia_root: The Artiatomi project directory
        dirs_start_with: A prefix for the data subdirectories within the Artiatomi project

    Returns: Text filepath containing the parsed info

    """

    lines = ["Directory Tomonum Motl\n"]
    for subdir in os.listdir(artia_root):
        if subdir.startswith(dirs_start_with):
            motl_name = ""
            full_subdir = os.path.join(artia_root, subdir)
            for file in os.listdir(full_subdir):
                if file.endswith(".st"):
                    basename = file.split(".")[0]
                    motl_name = f"{basename}_motl.em"
                    break

            motl = os.path.join(artia_root, subdir, motl_name)
            tomo_num = int(subdir.split("_")[-1]) + 1

            lines.append(f"{full_subdir} {tomo_num} {motl}\n")

    info_file = os.path.join(artia_root, "tomo_motls.txt")

    with open(info_file, "w") as f:
        f.writelines(lines)

    return info_file

def shift_coordinates_bottom_left(coords, size, binning=1):
    """
    Given an XYZ tuple of particle coordinates and the reconstruction they came from, shift the
        coordinates so that the origin is at the bottom-left of the tomogram

    Args:
        coords: the (x, y, z) coordinates for the particle
        size: the reconstruction MRC half-dimensions in (nx/2, ny/2, nz/2) form
        binning: the bin factor from the original stack to the final reconstruction, to be used if
            you are using coordinates based on the original unbinned coordinate system

    Returns: the new coordinates as a (x, y, z) tuple

    """
    return (
        float(coords[0]) / binning + size[0],
        float(coords[1]) / binning + size[1],
        float(coords[2]) / binning + size[2],
    )

def write_out_motl_files_simulated(root, name, xyz_name, motl_name, size, binning):
    """
    Iterate through an IMOD Processor project and set up text files for Artiatomi to read in for its marker files

    Args:
        root: The ETSimulations project root
        name: The project/particle name
        xyz_name: The file name to give each XYZ coordinates file
        motl_name: The file name to give each Euler angles file
        size: Final tomogram size (for shifting origin of coordinates)
        binning: The binning to apply to the raw coordinates

    Returns: None
    """
    metadata_file = os.path.join(root, "sim_metadata.json")

    with open(metadata_file, "r") as f:
        metadata = json.loads(f.read())

        # -------------------------------------
        # Retrieve parameters to write to files
        # -------------------------------------
        for num, tomogram in enumerate(metadata):
            basename = "%s_%d" % (name, tomogram["global_stack_no"])
            tomogram_dir = os.path.join(root, "processed_data/IMOD", basename)

            artia_dir = os.path.join(root, "processed_data/Artiatomi", basename)

            # Positions for TEM-Simulator are in nm, need to convert to pixels
            positions = np.array(tomogram["positions"]) / tomogram["apix"]
            # During reconstruction, there is a 90 degree rotation around the z-axis, so correct for
            # that with the positions
            positions = rotate_positions_around_z(positions)

            slicer_angles_csv = os.path.join(tomogram_dir, "%s_slicerAngles.csv" % name)
            orientations = np.loadtxt(slicer_angles_csv, delimiter=",")
            eulers = convert_slicer_to_motl(orientations)
            # Compile Slicer infos
            slicer_info = []
            for i, coords in enumerate(positions):
                angles = eulers[i]
                slicer_info.append({"coords": coords, "angles": angles})

            # Write out MOTL files for each tomogram
            xyz_motl = os.path.join(artia_dir, xyz_name)
            eulers_motl = os.path.join(artia_dir, motl_name)
            xyz_motl_file = open(xyz_motl, "w")
            eulers_motl_file = open(eulers_motl, "w")
            for info in slicer_info:
                coords = shift_coordinates_bottom_left(info["coords"], size, binning)
                angles = info["angles"]
                xyz_line = "{:f} {:f} {:f}\n".format(coords[0], coords[1], coords[2])
                xyz_motl_file.write(xyz_line)
                motl_line = "{:f} {:f} {:f}\n".format(angles[0], angles[1], angles[2])
                eulers_motl_file.write(motl_line)

            xyz_motl_file.close()
            eulers_motl_file.close()

def write_out_motl_files_real(artia_root, xyz_name, motl_name):
    """
    Iterate through Artiatomi files copied over from IMOD real data and set up text files for Artiatomi to read in for
        its marker files based on slicer .mod files

    Args:
        artia_root: The Artiatomi project root
        xyz_name: The file name to give each XYZ coordinates file
        motl_name: The file name to give each Euler angles file

    Returns: None
    """
    for subdir in os.listdir(artia_root):
        # Find the mod file
        mod_file = ""
        for file in os.listdir(os.path.join(artia_root, subdir)):
            if file.endswith(".mod"):
                mod_file = os.path.join(artia_root, subdir, file)
                break

        if mod_file == "":
            print("ERROR: No mod was found for sub-directory: %s" % subdir)
            exit(1)

        slicer_info = get_slicer_info(mod_file)

        # Write out MOTL files for each tomogram
        xyz_motl = os.path.join(artia_root, subdir, xyz_name)
        eulers_motl = os.path.join(artia_root, subdir, motl_name)
        xyz_motl_file = open(xyz_motl, "w")
        eulers_motl_file = open(eulers_motl, "w")
        for info in slicer_info:
            coords = info["coords"]
            angles = info["angles"]
            xyz_line = "{:f} {:f} {:f}\n".format(coords[0], coords[1], coords[2])
            xyz_motl_file.write(xyz_line)
            motl_line = "{:f} {:f} {:f}\n".format(angles[0], angles[2], angles[1])
            eulers_motl_file.write(motl_line)

        xyz_motl_file.close()
        eulers_motl_file.close()

def copy_over_imod_files(
    imod_root, artia_root, dir_starts_with, real_data_mode=False, mod_contains=None
):
    for subdir in os.listdir(imod_root):
        if subdir.startswith(dir_starts_with):

            # Look for the necessary IMOD files
            xf = ""
            tlt = ""
            stack = ""
            mod = ""
            if not real_data_mode:
                mod = "none"
            for file in os.listdir(os.path.join(imod_root, subdir)):
                if file.endswith(".tlt") and not file.endswith("_fid.tlt"):
                    tlt = file
                elif file.endswith(".st"):
                    stack = file
                elif file.endswith(".xf") and not file.endswith("_fid.xf"):
                    xf = file
                if real_data_mode and mod_contains in file and file.endswith(".mod"):
                    mod = file

                # Break out of loop once all three relevant files have been found
                if tlt != "" and stack != "" and xf != "" and mod != "":
                    break

            artia_stack_dir = os.path.join(artia_root, subdir)
            if not os.path.exists(artia_stack_dir):
                os.mkdir(artia_stack_dir)

            # Copy over the stack to the Artiatomi folder
            if stack != "":
                shutil.copyfile(
                    os.path.join(imod_root, subdir, stack),
                    os.path.join(artia_stack_dir, stack),
                )
            else:
                print("ERROR: No tiltseries was found for sub-directory: %s" % subdir)
                exit(1)

            # Copy over the tlt file to the Artiatomi folder
            if tlt != "":
                shutil.copyfile(
                    os.path.join(imod_root, subdir, tlt),
                    os.path.join(artia_stack_dir, tlt),
                )
            else:
                print("WARNING: No .tlt file was found for sub-directory: %s" % subdir)

            # Copy over the tlt file to the Artiatomi folder
            if xf != "":
                shutil.copyfile(
                    os.path.join(imod_root, subdir, xf),
                    os.path.join(artia_stack_dir, xf),
                )
            else:
                print("WARNING: No .xf file was found for sub-directory: %s" % subdir)

            if real_data_mode:
                # Copy over the mod file to the Artiatomi folder
                if mod != "":
                    shutil.copyfile(
                        os.path.join(imod_root, subdir, mod),
                        os.path.join(artia_stack_dir, mod),
                    )
                else:
                    print(
                        "WARNING: No .mod file was found for sub-directory: %s" % subdir
                    )

def setup_reconstructions_script(root, name, artia_args):
    """
    Function to generate the Artiatomi reconstruction setup script

    Args:
        root: The ETSimulations project root (only used for coming from simulated data)
        name: The project/particle name (only used for coming from simulated data)
        artia_args: The Artiatomi Processor arguments

    Returns: None

    """
    xyz_motl = "%s_xyz.txt" % name
    eulers_motl = "%s_motl.txt" % name

    # Get input parameters based on mode
    if artia_args["real_data_mode"]:
        artia_root = artia_args["artia_dir"]
        imod_root = artia_args["imod_dir"]
        dir_starts_with = artia_args["dir_starts_with"]
        mod_contains = artia_args["mod_contains"]
    else:
        artia_root = os.path.join(root, "processed_data", "Artiatomi")
        imod_root = os.path.join(root, "processed_data", "IMOD")
        dir_starts_with = name
        mod_contains = ""

    # -------------------------------------
    # Set up Artiatomi project directory structure
    # -------------------------------------
    print("Creating Artiatomi project directories and copying over relevant IMOD files")
    if not os.path.exists(artia_root):
        os.mkdir(artia_root)

    copy_over_imod_files(
        imod_root,
        artia_root,
        dir_starts_with,
        real_data_mode=artia_args["real_data_mode"],
        mod_contains=mod_contains,
    )

    print("Writing out MOTL-related files")
    if artia_args["real_data_mode"]:
        write_out_motl_files_real(artia_root, xyz_motl, eulers_motl)
    else:
        binning = 1
        if "position_binning" in artia_args:
            binning = artia_args["position_binning"]

        size = (
            artia_args["tomogram_size_x"] / 2,
            artia_args["tomogram_size_y"] / 2,
            artia_args["tomogram_size_z"] / 2,
        )
        write_out_motl_files_simulated(root, name, xyz_motl, eulers_motl, size, binning)

    # Use template file to create Matlab script to run the remaining steps
    current_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    template = current_dir + "../../templates/artiatomi/setup_artia_reconstructions.m"
    template_path = os.path.realpath(template)
    new_script = os.path.join(artia_root, "setup_artia_reconstructions.m")
    print("")
    print("Creating processing script at: %s" % new_script)

    with open(new_script, "w") as new_file:
        with open(template_path, "r") as base_file:
            # First look for the input params section
            while True:
                line = base_file.readline()
                if re.match(r"^%% Input parameters", line):
                    break
                else:
                    new_file.write(line)

            # Now start replacing input params
            while True:
                line = base_file.readline()
                # Break once we reach the end of the segment
                if re.match(r"^%% Process", line):
                    break

                # If we are at an assignment line
                elif re.match(r".+ =", line):
                    line = line.strip()
                    tokens = line.split(" ")
                    variable_name = tokens[0]

                    value_to_write_out = ""
                    if variable_name == "project_root":
                        value_to_write_out = f"'{artia_root}';"
                    elif variable_name == "dir_starts_with":
                        value_to_write_out = f"'{dir_starts_with}';"
                    elif variable_name == "xyz_motl":
                        value_to_write_out = f"'{xyz_motl}';"
                    elif variable_name == "eulers_motl":
                        value_to_write_out = f"'{eulers_motl}';"
                    elif variable_name in artia_args:
                        if type(artia_args[variable_name]) == str:
                            value_to_write_out = f"'{artia_args[variable_name]}';"
                        else:
                            value_to_write_out = str(artia_args[variable_name]) + ";"
                    else:
                        print(
                            "Missing Artiatomi processing parameter: %s!"
                            % variable_name
                        )
                        exit(1)

                    new_line = " ".join([variable_name, "=", value_to_write_out, "\n"])

                    new_file.write(new_line)

                # Other lines in the segment - probably just comments
                else:
                    new_file.write(line)

            # For the rest of the code, just write it out
            while True:
                line = base_file.readline()
                if len(line) == 0:
                    break
                else:
                    new_file.write(line)

def generate_reconstructions_script(root, name, artia_args):
    """
    Generate a bash script to run the reconstructions.
    Args:
        root: ETSimulations root (only used if coming from simulated data)
        name: Project/particle name (only used if coming from simulated data)
        artia_args: The Artiatomi processor arguments

    Returns: None

    """
    if artia_args["real_data_mode"]:
        artia_root = artia_args["artia_dir"]
        dir_starts_with = artia_args["dir_starts_with"]
    else:
        artia_root = os.path.join(root, "processed_data", "Artiatomi")
        dir_starts_with = name

    # Use template file to create a bash script to run EmSART on a data set
    current_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    template = current_dir + "../../templates/artiatomi/emsart_reconstruct.sh"
    template_path = os.path.realpath(template)
    new_script = os.path.join(artia_root, "emsart_reconstruct.sh")
    print("")
    print("Creating processing script at: %s" % new_script)

    with open(new_script, "w") as new_file:
        with open(template_path, "r") as base_file:
            while True:
                line = base_file.readline()
                if len(line) == 0:
                    break
                elif line.startswith("for f in"):
                    dir_pattern = f"{artia_root}/{dir_starts_with}*"
                    new_line = f"for f in {dir_pattern}\n"
                    new_file.write(new_line)
                elif line.startswith("config_file"):
                    config_file = os.path.basename(
                        artia_args["reconstruction_template_config"]
                    )
                    new_line = f'config_file="{config_file}"\n'
                    new_file.write(new_line)
                elif line.startswith("emsart_path"):
                    emsart = artia_args["emsart_path"]
                    new_line = f'emsart_path="{emsart}"\n'
                    new_file.write(new_line)
                else:
                    new_file.write(line)

def generate_sta_script(artia_root, info_file, artia_args):
    """
    Generate the MATLAB script for setting up the Artiatomi sub-tomogram averaging
    Args:
        artia_root: Artiatomi project root directory
        info_file: The text file listing the tomograms, tomogram numbers, and their motivelists
        artia_args: The Artiatomi processor arguments

    Returns: None

    """
    sta_folder = os.path.join(artia_root, "sta")
    mask_file = os.path.join(sta_folder, "other", "mask.em")
    wedge_file = os.path.join(sta_folder, "other", "wedge.em")
    maskCC_file = os.path.join(sta_folder, "other", "maskCC.em")
    global_motl_file = os.path.join(sta_folder, "motls", "motl_1.em")
    motl_file_pre = os.path.join(sta_folder, "motls", "motl_")
    particles_folder = os.path.join(sta_folder, "parts")
    part_file_pre = os.path.join(sta_folder, "parts", "part_")
    cfg_file = os.path.join(sta_folder, "sta.cfg")

    # Use template file to create Matlab script to run the remaining steps
    current_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    template = current_dir + "../../templates/artiatomi/setup_artia_sta.m"
    template_path = os.path.realpath(template)
    new_script = os.path.join(artia_root, "setup_artia_sta.m")
    print("")
    print("Creating STA processing script at: %s" % new_script)

    with open(new_script, "w") as new_file:
        with open(template_path, "r") as base_file:
            # First look for the input params section
            while True:
                line = base_file.readline()
                if re.match(r"^%% Input parameters", line):
                    break
                else:
                    new_file.write(line)

            # Now start replacing input params
            while True:
                line = base_file.readline()
                # Break once we reach the end of the segment
                if re.match(r"^%% Load motl", line):
                    break

                # If we are at an assignment line
                elif re.match(r".+ =", line):
                    line = line.strip()
                    tokens = line.split(" ")
                    variable_name = tokens[0]

                    value_to_write_out = ""
                    if variable_name == "info_file":
                        value_to_write_out = f"'{info_file}';"
                    elif variable_name == "sta_folder":
                        value_to_write_out = f"'{sta_folder}';"
                    elif variable_name == "maskFile":
                        value_to_write_out = f"'{mask_file}';"
                    elif variable_name == "wedgeFile":
                        value_to_write_out = f"'{wedge_file}';"
                    elif variable_name == "maskCCFile":
                        value_to_write_out = f"'{maskCC_file}';"
                    elif variable_name == "motlFile":
                        value_to_write_out = f"'{global_motl_file}';"
                    elif variable_name == "motlFilePre":
                        value_to_write_out = f"'{motl_file_pre}';"
                    elif variable_name == "particles_folder":
                        value_to_write_out = f"'{particles_folder}';"
                    elif variable_name == "partFilePre":
                        value_to_write_out = f"'{part_file_pre}';"
                    elif variable_name == "avgCfgFile":
                        value_to_write_out = f"'{cfg_file}';"
                    elif variable_name in artia_args:
                        if type(artia_args[variable_name]) == str:
                            value_to_write_out = f"'{artia_args[variable_name]}';"
                        else:
                            value_to_write_out = str(artia_args[variable_name]) + ";"
                    else:
                        print(
                            "Missing Artiatomi processing parameter: %s!"
                            % variable_name
                        )
                        exit(1)

                    new_line = " ".join([variable_name, "=", value_to_write_out, "\n"])

                    new_file.write(new_line)

                # Other lines in the segment - probably just comments
                else:
                    new_file.write(line)

            # For the rest of the code, just write it out
            while True:
                line = base_file.readline()
                if len(line) == 0:
                    break
                else:
                    new_file.write(line)

def get_latest_ref(sta_dir):
    """
    Find the latest reference from an Artiatomi STA run
    Args:
        sta_dir: The sub-tomogram averaging directory, where we expect to find a ref folder

    Returns: The latest reference file path

    """
    ref_dir = os.path.join(sta_dir, "ref")
    max_num = 0
    latest_ref = None
    for ref in os.listdir(ref_dir):
        if re.match("ref[0-9]+.em", ref):
            num = int(ref.split("ref")[1].split(".")[0])
            if num > max_num:
                latest_ref = os.path.join(ref_dir, ref)

    return latest_ref

def get_latest_motl(sta_dir):
    """
    Find the latest motl from an Artiatomi STA run
    Args:
        sta_dir: The sub-tomogram averaging directory, where we expect to find a ref folder

    Returns: The latest motl file path

    """
    motls_dir = os.path.join(sta_dir, "motls")
    max_num = 0
    latest_motl = None
    for motl in os.listdir(motls_dir):
        if re.match("motl_[0-9]+.em", motl):
            num = int(motl.split("_")[1].split(".")[0])
            if num > max_num:
                latest_motl = os.path.join(motls_dir, motl)

    return latest_motl

def generate_refinement_script(artia_root, info_file, artia_args):
    """
    Generate the MATLAB script for setting up the Artiatomi tilt refinement
    Args:
        artia_root: Artiatomi project root directory
        info_file: The text file listing the tomograms, tomogram numbers, and their motivelists
        artia_args: The Artiatomi processor arguments

    Returns: None

    """

    mask_file = os.path.join(artia_root, "sta", "other", "mask.em")
    wedge_file = os.path.join(artia_root, "sta", "other", "wedge.em")
    maskCC_file = os.path.join(artia_root, "sta", "other", "maskCC.em")

    refine_dir = os.path.join(artia_root, "refine")
    refine_motls_dir = os.path.join(refine_dir, "motls")

    latest_ref = get_latest_ref(os.path.join(artia_root, "sta"))
    latest_motl = get_latest_motl(os.path.join(artia_root, "sta"))

    # Use template file to create Matlab script to run the remaining steps
    current_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    template = current_dir + "../../templates/artiatomi/refine_align.m"
    template_path = os.path.realpath(template)
    new_script = os.path.join(artia_root, "refine_align.m")
    print("")
    print("Creating refinement script at: %s" % new_script)

    with open(new_script, "w") as new_file:
        with open(template_path, "r") as base_file:
            # First look for the input params section
            while True:
                line = base_file.readline()
                if re.match(r"^%% Input parameters", line):
                    break
                else:
                    new_file.write(line)

            # Now start replacing input params
            while True:
                line = base_file.readline()
                # Break once we reach the end of the segment
                if re.match(r"^%% Split latest", line):
                    break

                # If we are at an assignment line
                elif re.match(r".+ =", line):
                    line = line.strip()
                    tokens = line.split(" ")
                    variable_name = tokens[0]

                    value_to_write_out = ""
                    if variable_name == "info_file":
                        value_to_write_out = f"'{info_file}';"
                    elif variable_name == "mask_file":
                        value_to_write_out = f"'{mask_file}';"
                    elif variable_name == "wedge_file":
                        value_to_write_out = f"'{wedge_file}';"
                    elif variable_name == "maskCC_file":
                        value_to_write_out = f"'{maskCC_file}';"
                    elif variable_name == "main_root":
                        value_to_write_out = f"'{refine_dir}';"
                    elif variable_name == "refine_motls_dir":
                        value_to_write_out = f"'{refine_motls_dir}';"
                    elif variable_name == "latest_ref":
                        value_to_write_out = f"'{latest_ref}';"
                    elif variable_name == "latest_motl":
                        value_to_write_out = f"'{latest_motl}';"
                    elif variable_name in artia_args:
                        if type(artia_args[variable_name]) == str:
                            value_to_write_out = f"'{artia_args[variable_name]}';"
                        else:
                            value_to_write_out = str(artia_args[variable_name]) + ";"
                    else:
                        print(
                            "Missing Artiatomi processing parameter: %s!"
                            % variable_name
                        )
                        exit(1)

                    new_line = " ".join([variable_name, "=", value_to_write_out, "\n"])

                    new_file.write(new_line)

                # Other lines in the segment - probably just comments
                else:
                    new_file.write(line)

            # For the rest of the code, just write it out
            while True:
                line = base_file.readline()
                if len(line) == 0:
                    break
                else:
                    new_file.write(line)

def generate_extract_script(artia_root, artia_args):
    """
    Generate the MATLAB script for extracting new particles based on local refinement
    Args:
        artia_root: Artiatomi project root directory
        artia_args: The Artiatomi processor arguments

    Returns: None

    """
    refine_dir = os.path.join(artia_root, "refine")
    subvol_pre = refine_dir + "/subvols/"
    latest_motl = get_latest_motl(os.path.join(artia_root, "sta"))

    mask_file = os.path.join(artia_root, "sta", "other", "mask.em")
    wedge_file = os.path.join(artia_root, "sta", "other", "wedge.em")
    maskCC_file = os.path.join(artia_root, "sta", "other", "maskCC.em")

    # Use template file to create Matlab script to run the remaining steps
    current_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
    template = current_dir + "../../templates/artiatomi/refine_extract.m"
    template_path = os.path.realpath(template)
    new_script = os.path.join(artia_root, "refine_extract.m")
    print("")
    print("Creating refinement extraction script at: %s" % new_script)

    with open(new_script, "w") as new_file:
        with open(template_path, "r") as base_file:
            # First look for the input params section
            while True:
                line = base_file.readline()
                if re.match(r"^%% Input parameters", line):
                    break
                else:
                    new_file.write(line)

            # Now start replacing input params
            while True:
                line = base_file.readline()
                # Break once we reach the end of the segment
                if re.match(r"^%% Run extractions", line):
                    break

                # If we are at an assignment line
                elif re.match(r".+ =", line):
                    line = line.strip()
                    tokens = line.split(" ")
                    variable_name = tokens[0]

                    value_to_write_out = ""
                    if variable_name == "refineDir":
                        value_to_write_out = f"'{refine_dir}';"
                    elif variable_name == "subVolPre":
                        value_to_write_out = f"'{subvol_pre}';"
                    elif variable_name == "latestStaMotl":
                        value_to_write_out = f"'{latest_motl}';"
                    elif variable_name == "mask_file":
                        value_to_write_out = f"'{mask_file}';"
                    elif variable_name == "wedge_file":
                        value_to_write_out = f"'{wedge_file}';"
                    elif variable_name == "maskCC_file":
                        value_to_write_out = f"'{maskCC_file}';"
                    elif variable_name in artia_args:
                        if type(artia_args[variable_name]) == str:
                            value_to_write_out = f"'{artia_args[variable_name]}';"
                        else:
                            value_to_write_out = str(artia_args[variable_name]) + ";"
                    else:
                        print(
                            "Missing Artiatomi processing parameter: %s!"
                            % variable_name
                        )
                        exit(1)

                    new_line = " ".join([variable_name, "=", value_to_write_out, "\n"])

                    new_file.write(new_line)

                # Other lines in the segment - probably just comments
                else:
                    new_file.write(line)

            # For the rest of the code, just write it out
            while True:
                line = base_file.readline()
                if len(line) == 0:
                    break
                else:
                    new_file.write(line)

def artiatomi_main(root, name, artia_args):

    if (
        "setup_reconstructions_and_motls" in artia_args
        and artia_args["setup_reconstructions_and_motls"]
    ):
        setup_reconstructions_script(root, name, artia_args)

        generate_reconstructions_script(root, name, artia_args)

    if "setup_averaging" in artia_args and artia_args["setup_averaging"]:

        # Get input parameters based on mode
        if artia_args["real_data_mode"]:
            artia_root = artia_args["artia_dir"]
            dir_starts_with = artia_args["dir_starts_with"]
        else:
            artia_root = os.path.join(root, "processed_data", "Artiatomi")
            dir_starts_with = name

        # Set up averaging directory structure
        print("Creating Artiatomi averaging directories")
        sta_dir = os.path.join(artia_root, "sta")
        if not os.path.exists(sta_dir):
            os.mkdir(sta_dir)

        parts_dir = os.path.join(sta_dir, "parts")
        if not os.path.exists(parts_dir):
            os.mkdir(parts_dir)

        motls_dir = os.path.join(sta_dir, "motls")
        if not os.path.exists(motls_dir):
            os.mkdir(motls_dir)

        ref_dir = os.path.join(sta_dir, "ref")
        if not os.path.exists(ref_dir):
            os.mkdir(ref_dir)

        others_dir = os.path.join(sta_dir, "other")
        if not os.path.exists(others_dir):
            os.mkdir(others_dir)

        info_file = imod_setup_sta(artia_root, dir_starts_with)

        generate_sta_script(artia_root, info_file, artia_args)

    if "setup_refinement" in artia_args and artia_args["setup_refinement"]:

        # Get input parameters based on mode
        if artia_args["real_data_mode"]:
            artia_root = artia_args["artia_dir"]
        else:
            artia_root = os.path.join(root, "processed_data", "Artiatomi")

        info_file = os.path.join(artia_root, "tomo_motls.txt")

        # Set up averaging directory structure
        print("Creating Artiatomi refinement directories")
        refine_dir = os.path.join(artia_root, "refine")
        if not os.path.exists(refine_dir):
            os.mkdir(refine_dir)

        motls_dir = os.path.join(refine_dir, "motls")
        if not os.path.exists(motls_dir):
            os.mkdir(motls_dir)

        generate_refinement_script(artia_root, info_file, artia_args)

        generate_extract_script(artia_root, artia_args)


import time
import os

from shinken.log import logger
from livestatus_mongo_query import LiveStatusQuery
from livestatus_mongo_response import LiveStatusResponse
from livestatus_constraints import LiveStatusConstraints
from livestatus_query_metainfo import LiveStatusQueryMetainfo

class LiveStatusWaitQuery(LiveStatusQuery):

    my_type = 'wait'

    def __init__(self, *args, **kwargs):
        super(LiveStatusWaitQuery, self).__init__(*args, **kwargs)
        self.response = LiveStatusResponse()
        self.wait_start = time.time()
        self.wait_timeout = 0
        self.wait_trigger = 'all'

    def parse_input(self, data):
        """Parse the lines of a livestatus request.

        This function looks for keywords in input lines and
        sets the attributes of the request object.
        WaitCondition statements are written into the metafilter string as if they
        were ordinary Filter:-statements. (metafilter is then used for a MetaData object)

        """
        metafilter = ""
        for line in data.splitlines():
            line = line.strip()
            # Tools like NagVis send KEYWORK:option, and we prefer to have
            # a space following the:
            if ':' in line and ' ' not in line:
                line = line.replace(':', ': ')
            keyword = line.split(' ')[0].rstrip(':')
            if keyword == 'GET':  # Get the name of the base table
                _, self.table = self.split_command(line)
                metafilter += "GET %s\n" % self.table
            elif keyword == 'WaitObject':  # Pick a specific object by name
                _, item = self.split_option(line)
                # It's like Filter: name = %s
                # Only for services it's host<blank>servicedesc
                if self.table == 'services':
                    if ';' in item:
                        host_name, service_description = item.split(';', 1)
                    else:
                        host_name, service_description = item.split(' ', 1)
                    self.add_filter(self.filters_stack, '=', 'name', host_name)
                    self.add_filter(
                        self.filters_stack,
                        '=',
                        'description',
                        service_description
                    )
                    self.datamgr.stack_filter_and(
                        self.filters_stack,
                        self.table,
                        2
                    )
                else:
                    self.add_filter(self.filters_stack, '=', 'name', item)
            elif keyword == 'WaitCondition':
                try:
                    attribute, operator, reference = self.parse_filter_line(line)
                    self.add_filter(
                        self.filters_stack,
                        operator,
                        attribute,
                        reference
                    )
                except Exception as e:
                    logger.warning("[Livestatus Query] Illegal operation: %s" % e)
                    raise

            elif keyword == 'WaitConditionAnd':
                # Take the last andnum functions from the stack
                # Construct a new function which makes a logical and
                # Put the function back onto the stack
                self.datamgr.stack_filter_and(
                    self.filters_stack,
                    self.table,
                    andnum
                )
            elif keyword == 'WaitConditionOr':
                _, ornum = self.split_option(line)
                # Take the last ornum functions from the stack
                # Construct a new function which makes a logical or
                # Put the function back onto the stack
                self.datamgr.stack_filter_or(
                    self.filters_stack,
                    self.table,
                    ornum
                )
            elif keyword == 'WaitTimeout':
                _, self.wait_timeout = self.split_option(line)
                self.wait_timeout = int(self.wait_timeout) / 1000
            else:
                # This line is not valid or not implemented
                logger.warning("[Livestatus Wait Query] Received a line of input "
                               "which i can't handle: '%s'", line)

    def condition_fulfilled(self):
        '''
        :return: True if the WaitQuery condition is fulfilled.
                 False otherwise.
        The result of launch_query is non-empty when an item matches the filter criteria
        '''
        items = self.launch_query()
        if items:
            return True
        else:
            return False