from psfdataset import transforms, PSFDataset, PSFZippedDataset
import numpy as np
import tensorflow as tf
from tensorflow.keras.utils import to_categorical
from src.Autoencoder import SigAutoencoder
from sklearn.preprocessing import MinMaxScaler
import os
import pickle
from src.data_grabbing.data_grabber import DataGrabber
from sklearn.model_selection import StratifiedKFold
from src.algos.utils import SkeletonUtils

class LinConvData:
    SINGLE_LABELS=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,
    23,24,25,26,27,28,29,30,31,32,33,34,35,
    36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,62,63,64,65,66,67,68,69,70,71,
    72,73,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,
    117,118,119,120,121,122,123,124,125,126,127,128,129,136,137,138,139,140,141,
    142,143,144,145,147,149,150,151,152,153,154]
    DOUBLE_LABELS=[74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,
    93,94,95,96,97,98,130,131,132,133,134,135,146,148]
    # JOINT_LABELS=[75,76,78,88,89,93,95,97,98,130,132,133]

    DEFAULT_DATA_PATH="/scratch/gale/UAVHuman/newData/"

    def __init__(self,pre_normaliser=None,transform='landmarks',classes=155, debug=False, bodies='both',
                autoencoded_latent_dim=50,load_autoencoded_data=False,add_center_joint=False,with_center_joint=True,
                load_dejanked_data=False, no_head=False,
                data_path=DEFAULT_DATA_PATH):
        '''
        Args:
        input_shape: (N_TIMESTEPS, N_AXES, N_JOINTS, N_PERSONS)
        pre_normaliser: object is defined if data is not to be loaded. The path signature features are
            ... directly created using 'self.sig_transform' and saved in 'self.sig_data' (a pair (trainingset, valset)
            ... of class PSFDataset or PSFZippedDataset.
        transform: shape of extracted data
        '''
        self.no_head=no_head
        self.N_TIMESTEPS = 305
        self.N_AXES =2
        self.N_JOINTS = SkeletonUtils.UAVHuman_Joints(no_head=self.no_head,with_center=with_center_joint)
        self.N_PERSONS = 2
        self.N_CLASSES=classes
        self.transform = transform
        self._debug = debug
        self.pre_normaliser = pre_normaliser
        self.path=data_path
        self.bodies=bodies
        self.load_autoencoded_data=load_autoencoded_data
        self.datagrabber=DataGrabber()
        self.fold_idx_dict=self._cross_validation_fold()
        self.autoencoded_latent_dim=autoencoded_latent_dim
        self.add_center_joint=add_center_joint
        self.load_with_center_joint=with_center_joint
        self.load_dejanked_data=load_dejanked_data

        if self.add_center_joint or self.load_with_center_joint:
            self.N_JOINTS = SkeletonUtils.UAVHuman_Joints(with_center=True,no_head=self.no_head)
            self.joints_to_bones_matrix=SkeletonUtils.joints_to_bones_matrix(SkeletonUtils.UAVHuman_Joints_graph_edges(no_head=self.no_head),self.N_JOINTS)
            self.N_BONES = len(SkeletonUtils.UAVHuman_Joints_graph_edges(no_head=self.no_head))
        if pre_normaliser is None:
            self.load_transform=True
        else:
            self.load_transform=False

        # if self.transform=='landmarks':
        #      self.train_data, self.train_label, self.val_data, self.val_label=self.load_landmarks()
    def _allowed_labels(self):
        if self.bodies=='both':
            return range(self.N_CLASSES)
        elif self.bodies=='single':
            return LinConvData.SINGLE_LABELS
        elif self.bodies=='double':
            return LinConvData.DOUBLE_LABELS

    def _label_translator(self,label):
        ''' takes original label and gives new label, i.e. the position of the
        original label in the current allowed label list. throws ValueError if
        not present'''
        return self._allowed_labels().index(label)

    def _inverse_label_translator(self,new_label):
        return self._allowed_labels()[new_label]

    def _cross_validation_fold(self,data_input=None,cross_val_fold_num=5,verbose=0,process_labels=True):
        if data_input is None:
            data_uav, label_uav = self.datagrabber.train_data, self.datagrabber.train_label
        elif data_input=='load':
            data_uav, label_uav = self.datagrabber._load_uav_data()
        else:
            data_uav, label_uav = data_input
        # length_uav = lengths   ER: ???
        fold_num=cross_val_fold_num

        if self.pre_normaliser is not None:
            if self.pre_normaliser.debug:
                data_uav=data_uav[:self.pre_normaliser.debug]
                label_uav=label_uav[:self.pre_normaliser.debug]
                fold_num=2
        if process_labels:
            label_uav=self.get_body_type_labels(label_uav)
        data_uav=data_uav[:label_uav.shape[0]]

        def verboseprint(*args,**kwargs):
            if verbose==0:
                return
            else:
                print(*args,**kwargs)

        # LABELS=self._allowed_labels()
        # class_num=len(LABELS)
        # verboseprint('Check data imbalance...')
        # class_cnt = np.zeros((class_num,)))
        # for l in label_uav:
        #     class_cnt[l] += 1
        # verboseprint(class_cnt)
        # verboseprint('Avg sample num: ',class_cnt.mean())
        # verboseprint('Max sample num: ',class_cnt.max())
        # verboseprint('Min sample num: ',class_cnt.min())

        k_fold = StratifiedKFold(cross_val_fold_num)
        k_fold.get_n_splits(data_uav,label_uav)
        k_fold_idx_dict = dict()

        verboseprint('Create {}-fold for cross validation...'.format(cross_val_fold_num))
        for k, (train_idx, val_idx) in enumerate(k_fold.split(data_uav,label_uav)):
            k_fold_idx_dict.update({str(k):{'train':train_idx, 'val':val_idx}})
            verboseprint(k+1,'- fold:','Trainset size: ',len(train_idx),' Valset size: ',len(val_idx))
        return k_fold_idx_dict

    def get_body_type(self,data,labels,by_label=False,translate_labels=True):
        '''get data and labels for single,double or both bodies from preNormaliser
            shape of input data = Batch x Coords x Time x Joints x People
            shape of output data = Batch x Time x Coords x (Joints x People)
        '''
        if labels is not None and len(labels.shape)>1:
            labels=np.argmax(labels, axis=1)
        B, C, T, J, P = data.shape

        if by_label:
            if self.bodies=='both':
                return data.reshape((B,C,T,J*P)).transpose([0,2,1,3]), labels
            elif self.bodies=='single':
                single_indices=[]
                for i in range(labels.shape[0]):
                    if labels[body] in self._allowed_labels():
                        single_indices.append(i)
                if translate_labels:
                    print('Translating single body labels')
                    for i in single_indices:
                        new_label=self._label_translator(labels[body])
                        np.put(labels,i,new_label)
                return data[single_indices,:,:,:,0].transpose([0,2,1,3]),labels[single_indices]
            elif self.bodies=='double':
                double_indices=[]
                for i in range(labels.shape[0]):
                    if labels[body] in self._allowed_labels():
                        double_indices.append(i)
                if translate_labels:
                    print('Translating double body labels')
                    for i in double_indices:
                        new_label=self._label_translator(labels[body])
                        np.put(labels,i,new_label)
                return data[double_indices].reshape((-1,C,T,J*P)).transpose([0,2,1,3]),labels[double_indices]
        else:
            if self.bodies=='both':
                if labels is None:
                    return data.reshape((B,C,T,J*P)).transpose([0,2,1,3])
                else:
                    return data.reshape((B,C,T,J*P)).transpose([0,2,1,3]), labels
            elif self.bodies=='single':
                single_indices=[]
                for i in range(data.shape[0]):
                    if np.array_equal(data[i,:,:,:,1],np.zeros_like(data[i,:,:,:,1])):
                        single_indices.append(i)
                if translate_labels and labels is not None:
                    for i in single_indices:
                        new_label=self._label_translator(labels[body])
                        np.put(labels,i,new_label)
                if labels is None:
                    return data[single_indices,:,:,:,0].transpose([0,2,1,3])
                else:
                    return data[single_indices,:,:,:,0].transpose([0,2,1,3]),labels[single_indices]
            elif self.bodies=='double':
                double_indices=[]
                for i in range(data.shape[0]):
                    if not np.array_equal(data[i,:,:,:,1],np.zeros_like(data[i,:,:,:,1])):
                        double_indices.append(i)
                if translate_labels and labels is not None:
                    for i in double_indices:
                        new_label=self._label_translator(labels[body])
                        np.put(labels,i,new_label)
                if labels is None:
                    return data[double_indices].reshape((-1,C,T,J*P)).transpose([0,2,1,3])
                else:
                    return data[double_indices].reshape((-1,C,T,J*P)).transpose([0,2,1,3]),labels[double_indices]

    def get_body_type_labels(self,labels,translate_labels=True):

        LABELS=self._allowed_labels()

        label_list=[]
        for i in range(labels.shape[0]):
            if labels[i] in LABELS:
                label_list.append(i)
        if translate_labels:
            for i in label_list:
                new_label=self._label_translator(labels[i])
                np.put(labels,i,new_label)
        return labels[label_list]

    def _load_data(self,fold_idx=0,full_data=False):
        '''
            Will only be applied if data is to be loaded. Loads the data and the fold indices from the prenormaliser.

        '''
        isFastTestSubset = self._debug # --- for preliminary testing purposes only

        data_uav,label_uav = self.load_full_landmarks('train')

        print('Prenormalized data shape:',data_uav.shape)
        print('Prenormalized label shape',label_uav.shape)

        train_index = self.fold_idx_dict[str(fold_idx)]['train'] if not isFastTestSubset else fold_idx_dict[str(fold_idx)]['train'][::100]
        val_index = self.fold_idx_dict[str(fold_idx)]['val'] if not isFastTestSubset else fold_idx_dict[str(fold_idx)]['val'][::100]

        train_data = data_uav[train_index]
        train_label = label_uav[train_index]
        # train_length = length_uav[train_index]   # see Weixin's notebook: can be used to add the temporal length of a movie as a feature; then self._get_iter(...) has to be adapted

        val_data = data_uav[val_index]
        val_label = label_uav[val_index]
        # val_length = length_uav[val_index]   # see above

        print('Test split data shape', train_data.shape, val_data.shape)

        return train_data, train_label, val_data, val_label

    def _get_iter(self,data,label):
        """
        Returns iterators over the requested
        subset of the data.
        """
        return iter(list(zip(data, label)))

    def load_autoencoded_signature(self,fold_idx=0):
        path=self.path
        train_data = np.load(os.path.join(path,'{}_body_{}_sig_autoencoded.npy'.format(self.bodies,'train')))
        val_data = np.load(os.path.join(path,'{}_body_{}_sig_autoencoded.npy'.format(self.bodies,'val')))
        train_label=np.load(path+self.bodies+'_body_uav_train_landmarks_labels.npy')
        val_label=np.load(path+self.bodies+'_body_uav_val_landmarks_labels.npy')

        return train_data,train_label,val_data,val_label

    def load_full_autoencoded_signature(self,flag):
        path=self.path

        label=np.load(os.path.join(path,'{}_body_full_{}_landmarks_labels.npy'.format(self.bodies,flag)))
        data = np.load(os.path.join(path,'{}_body_full_{}_sig_autoencoded.npy'.format(self.bodies,flag)))
        return data, label

    def load_autoencoded(self,fold_idx=0):
        path = self.path
        latent_dim=self.autoencoded_latent_dim

        if self.load_autoencoded==True:
            print('Load autoencoded trial train data')
            train_data=np.load(os.path.join(path,self.bodies+'_body_trial_{}_autoencoded_dim_{}.npy'.format('train',latent_dim)))
            val_data=np.load(os.path.join(path,self.bodies+'_body_trial_{}_autoencoded_dim_{}.npy'.format('val',latent_dim)))

            train_label=np.load(path+self.bodies+'_body_uav_train_landmarks_labels.npy')
            val_label=np.load(path+self.bodies+'_body_uav_val_landmarks_labels.npy')

            return train_data, train_label, val_data, val_label
        else:
            dimV=2
            train_data, train_label, val_data, val_label = self.load_landmarks(fold_idx)
            N_JOINTS=train_data.shape[3]

            if self.bodies=='single':
                autoencoder=SigAutoencoder(latent_dim,N_JOINTS=N_JOINTS,load_weights=True,save_file_name_suffix='_single_body')
            else:
                autoencoder=SigAutoencoder(latent_dim,N_JOINTS=N_JOINTS,load_weights=True,save_file_name_suffix='_double_body')

            print('Encoding data using autoencoder for',self.bodies,'body type')
            def encode_data(data):
                '''input shape B x T X C X (JxP)
                    output shape B x T x LD'''
                B,T,C,JP=data.shape
                LD=latent_dim
                data=data.transpose([0,1,3,2])
                print('encode_data debug 1:',data.shape)
                data=data.reshape((-1, JP,C))

                data=autoencoder.encode(data[:,:,:dimV])
                return data.reshape((B,T,LD))

            train_data=encode_data(train_data)
            val_data=encode_data(val_data)

            print('Saving autoencoded trial train data')
            np.save(os.path.join(path,self.bodies+'_body_trial_{}_autoencoded_dim_{}.npy'.format('train',latent_dim)),train_data)
            np.save(os.path.join(path,self.bodies+'_body_trial_{}_autoencoded_dim_{}.npy'.format('val',latent_dim)),val_data)

        print('autoencoded data shape',train_data.shape, val_data.shape)

        return train_data, train_label, val_data, val_label

    def load_full_autoencoded(self,flag='train'):
        path = self.path
        latent_dim=self.autoencoded_latent_dim

        if self.load_autoencoded==True:
            print('Load autoencoded trial train data')
            data=np.load(os.path.join(path,self.bodies+'_body_full_{}_autoencoded_dim_{}.npy'.format(flag,latent_dim)))

            label=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_labels.npy'.format(flag)))

            return data, label
        else:
            dimV=2
            data, label = self.load_full_landmarks(flag)
            N_JOINTS=data.shape[3]

            if self.bodies=='single':
                autoencoder=SigAutoencoder(latent_dim,N_JOINTS=N_JOINTS,load_weights=True,save_file_name_suffix='_single_body')
            else:
                autoencoder=SigAutoencoder(latent_dim,N_JOINTS=N_JOINTS,load_weights=True,save_file_name_suffix='_double_body')

            print('Encoding data using autoencoder for',self.bodies,'body type')
            def encode_data(data):
                '''input shape B x T x C x (JxP)
                    output shape B x T x LD'''
                B,T,C,JP=data.shape
                LD=latent_dim
                data=data.transpose([0,1,3,2])
                print('encode_data debug 2:',data.shape)
                data=data.reshape((-1, JP,C))

                data=autoencoder.encode(data[:,:,:dimV])
                return data.reshape((B,T,LD))

            data=encode_data(data)

            print('Saving autoencoded full',flag,'data')
            np.save(os.path.join(path,self.bodies+'_body_trial_{}_autoencoded_dim_{}.npy'.format(flag,latent_dim)),data)

        print('autoencoded data shape',data.shape)

        return data,label

    def load_landmarks(self,fold_idx=0):
        print('load_landmarks(',fold_idx,')')
        data, labels = self.load_full_landmarks('train')

        train_index = self.fold_idx_dict[str(fold_idx)]['train']
        val_index = self.fold_idx_dict[str(fold_idx)]['val']

        train_data = data[train_index]
        train_label = labels[train_index]

        val_data = data[val_index]
        val_label = labels[val_index]

        return train_data,train_label,val_data,val_label

    def load_sampled_landmarks(self,fold_idx=0):
        '''
            Either loads the landmark data from path, or creates from pre_normaliser and saves to file

        '''

        path = self.path
        if self.load_transform==True:
            print('Load data')
            train_data=np.load(path+self.bodies+'_body_sampled_train_landmarks.npy')
            train_label_trans=np.load(path+self.bodies+'_body_sampled_train_landmarks_labels.npy')

            val_data=np.load(path+self.bodies+'_body_sampled_val_landmarks.npy')
            val_label_trans=np.load(path+self.bodies+'_body_sampled_val_landmarks_labels.npy')

            print('reshaped data',train_data.shape, val_data.shape)
            return train_data,train_label_trans, val_data, val_label_trans

        else:
            print('Create data')

            train_data, train_label, val_data, val_label = self._load_data(fold_idx)
            mask = np.random.choice(self.N_TIMESTEPS,100,replace=False)
            # transforming data [N,C,T,V,M] to [N,T,V,M,C] to [N,T,C,V*M] -- changed AG
            # train_data = np.transpose(train_data,(0,2,1,3))
            # train_data=train_data.reshape(train_data.shape[:3] + (self.N_JOINTS*self.N_PERSONS,))
            # val_data = np.transpose(val_data,(0,2,1,3))
            # val_data=val_data.reshape(val_data.shape[:3]+(self.N_JOINTS*self.N_PERSONS,))
            train_label_trans = to_categorical(train_label,len(self._allowed_labels()))
            val_label_trans =to_categorical(val_label,len(self._allowed_labels()))

            print('reshaped data',train_data.shape, val_data.shape)

            print('Save landmark data')
            np.save(path+self.bodies+'_body_sampled_train_landmarks.npy', train_data[:,mask])
            np.save(path+self.bodies+'_body_sampled_train_landmarks_labels.npy', train_label_trans)

            np.save(path+self.bodies+'_body_sampled_val_landmarks.npy', val_data[:,mask])
            np.save(path+self.bodies+'_body_sampled_val_landmarks_labels.npy', val_label_trans)

            return train_data[:,mask],train_label_trans[:,mask], val_data[:,mask], val_label_trans[:,mask]

    def load_test_train(self,fold_idx=0):
        '''
            Either loads the landmark data from path, or creates from pre_normaliser and saves to file
            The path is ~/scr/preprocessing/LinConvData.
        '''
        path = self.path
        print('Load test and train data')

        data=np.concatenate((np.load(os.path.join(path,self.bodies+'_body_full_train_landmarks.npy')),
            np.load(os.path.join(path,self.bodies+'_body_full_test_landmarks.npy'))),
            axis=0)
        label=np.concatenate((np.load(os.path.join(path,self.bodies+'_body_full_train_landmarks_labels.npy')),
            np.load(os.path.join(path,self.bodies+'_body_full_test_landmarks_labels.npy'))),
            axis=0)
        print("full data shape", data.shape, label.shape)

        def shuffle_in_unison(a, b):
            rng_state = np.random.get_state()
            np.random.shuffle(a)
            np.random.set_state(rng_state)
            np.random.shuffle(b)

        shuffle_in_unison(data,label)

        fold_idx_dict=self._cross_validation_fold(data_input=(data,np.argmax(label,axis=1)),process_labels=False)
        train_index = fold_idx_dict[str(fold_idx)]['train']
        val_index = fold_idx_dict[str(fold_idx)]['val']

        train_data = data[train_index]
        train_label = label[train_index]

        val_data = data[val_index]
        val_label = label[val_index]

        print('Test split data shape', train_data.shape, val_data.shape)

        return train_data,train_label, val_data, val_label

    def load_pairwise_landmarks(self, classes):
        X_train, y_train=self.load_full_landmarks(flag='train')
        X_test, y_test=self.load_full_landmarks(flag='test')

        ind_classes_tr=np.nonzero(y_train[:,classes])
        ind_classes_test=np.nonzero(y_test[:,classes])

        X_train_cut=X_train[ind_classes_tr[0]]
        X_test_cut=X_test[ind_classes_test[0]]
        y_train_cut=y_train[ind_classes_tr[0]]
        y_train_cut=y_train_cut[:,classes]
        y_train_cut=np.argmax(y_train_cut,1)

        y_test_cut=y_test[ind_classes_test[0]]
        y_test_cut=y_test_cut[:,classes]
        y_test_cut=np.argmax(y_test_cut,1)

        y_train_cut=y_train_cut.reshape(-1,1)
        y_test_cut=y_test_cut.reshape(-1,1)

        return X_train_cut, y_train_cut, X_test_cut, y_test_cut

    def load_full_landmarks(self, flag='train',raw=False):

        path = self.path
        if self.load_transform==True:
            print('Load full',flag,'data')
            if self.load_with_center_joint:
                if self.load_dejanked_data:
                    data=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_dejanked.npy'.format(flag)))
                    label_trans=np.load(os.path.join(path,self.bodies+'_body_full_{}_label_dejanked.npy'.format(flag)))

                    print('loading dejanked',flag,'data')
                else:
                    if raw:
                        label_trans=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_labels.npy'.format(flag)))
                        data=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_with_center_raw.npy'.format(flag)))
                    else:
                        label_trans=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_labels.npy'.format(flag)))
                        data=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_with_center.npy'.format(flag)))
            else:
                if self.load_dejanked_data:
                    data=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_dejanked.npy'.format(flag)))
                    label_trans=np.load(os.path.join(path,self.bodies+'_body_full_{}_label_dejanked.npy'.format(flag)))
                    data=data.reshape((-1,305,2,18,2))[:,:,:,:17,:].reshape((-1,305,2,34))
                    print('loading dejanked',flag,'data')
                else:
                    if raw:
                        data=np.load(os.path.join(path,self.bodies+'_body_full_raw_{}_landmarks.npy'.format(flag)))
                        label_trans=np.load(os.path.join(path,self.bodies+'_body_full_raw_{}_landmarks_labels.npy'.format(flag)))
                        print('raw')
                    else:
                        data=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks.npy'.format(flag)))
                        label_trans=np.load(os.path.join(path,self.bodies+'_body_full_{}_landmarks_labels.npy'.format(flag)))
                        print('scaled')
            # data=np.load(os.path.join(path,'{}_data.npy'.format(flag)))
            # with open(os.path.join(path,'{}_label.pkl'.format(flag)), 'rb') as f:
            #         print('loading labels')
            #         sample_name, label_uav = pickle.load(f)
            #         label = np.array(label_uav)
            # data = np.transpose(data,(0,2,1,3,4))
            # data=data.reshape(data.shape[:3] + (self.N_JOINTS*self.N_PERSONS,))
            # label_trans = to_categorical(label,self.N_CLASSES)

        else:
            print('Create full data')
            if flag=='train':
                data = self.pre_normaliser.train_prenorm_data.copy()
                label = self.pre_normaliser.train_prenorm_label.copy()
            elif flag=='test':
                data = self.pre_normaliser.test_prenorm_data.copy()
                label = self.pre_normaliser.test_prenorm_label.copy()

            print('Prenormalized data shape:',data.shape)
            print('Prenormalized label shape',label.shape)
            # transforming data [N,C,T,V,M] to [N,T,V,M,C] to [N,T,C,V*M]
            # data = np.transpose(data,(0,2,1,3,4))
            # print('transposed data shape', data.shape)
            # data=data.reshape(data.shape[:3] + (self.N_JOINTS*self.N_PERSONS,))
            # print('reshapen data shape', data.shape)
            # done in get_body_type AG

            print('Get',self.bodies,'body data')
            data,label = self.get_body_type(data,label,by_label=True)
            print(data.shape,label.shape)
            label_trans = to_categorical(label,len(self._allowed_labels()))
            print('Save landmark data')

            np.save(os.path.join(path,self.bodies+'_body_full_{}_landmarks.npy'.format(flag)), data)
            np.save(os.path.join(path,self.bodies+'_body_full_{}_landmarks_labels.npy'.format(flag)), label_trans)

        if self.add_center_joint:
            print('Adding center joint as joint 17')
            data=self._add_center_joint(data)
            np.save(os.path.join(path,self.bodies+'_body_full_{}_landmarks_with_center.npy'.format(flag)), data)
        if self.no_head:
            shape=data.shape
            data=data.reshape(shape[:3]+(int(shape[3]/2),2))
            data=data[:,:,:,5:]
            data=data.reshape(shape[:3]+(self.N_JOINTS*2,))
        return data, label_trans

    def _add_center_joint(self,data):
        if self.bodies=='single':
            bodies=1
        else:
            bodies=2
        data=data.reshape((-1,305,3,17,bodies))
        new_data=np.zeros((data.shape[0],305,3,18,bodies))
        new_data[:,:,:,:17,:]=data

        for i in range(len(data)):
            print(i,'added',end='\r')
            for frame in range(305):
                for body in range(bodies):
                    x_center=(data[i,frame,0,5,body]+data[i,frame,0,6,body]+data[i,frame,0,11,body]+data[i,frame,0,12,body])/4
                    y_center=(data[i,frame,1,5,body]+data[i,frame,1,6,body]+data[i,frame,1,11,body]+data[i,frame,1,12,body])/4
                    c_center=(data[i,frame,2,5,body]+data[i,frame,2,6,body]+data[i,frame,2,11,body]+data[i,frame,2,12,body])/4
                    new_data[i,frame,:,17,body]=np.array([x_center,y_center,c_center])

        return new_data.reshape((-1,305,3,18*bodies))

    def load_test_landmarks(self,load=False):

        path = self.path
        if load:
            print('Load test data from file')
            test_data=np.load(path+self.bodies+'_body_full_test_landmarks.npy')
        else:
            print('Create full test data')
            test_data = self.pre_normaliser.test_prenorm_data.copy()
            print('Prenormalized test data shape:',test_data.shape)
            # transforming data [N,C,T,V,M] to [N,T,V,M,C] to [N,T,C,V*M]
            test_data = np.transpose(test_data,(0,2,1,3,4))
            test_data=test_data.reshape(test_data.shape[:3] + (self.N_JOINTS*self.N_PERSONS,))
            print('Get',self.bodies,'body data')
            test_data = get_body_type(data,labels=None,by_label=False)
            print('Save test landmark data')
            np.save(path+self.bodies+'_body_full_test_landmarks.npy', test_data)

        return test_data

    def sig_transform(self):
        '''
            Either loads the data from path, or creates the desired PSFDataset object (trainingset, validation set). These
            ... objects contain the desired PSF transform. Then trainingset and validation set are filled with the data so that
            ... they are ready to be used by torch for training and validation. In the second case, data will be saved
            ... in the path. The path is ~/scr/preprocessing/PSF.
        '''

        path = os.path.dirname(__file__) + "/LinConvData/"
        if self.load_transform==True:
            print('Load data')
            trainingset = PSFDataset()
            valset = PSFDataset()
            trainingset.load(path+"uav_train_"+self.transform)
            valset.load(path+"uav_val_"+self.transform)
        else:
            print('Create data')
            if self.transform=='pairs':
                tr = transforms.Compose([
                        transforms.spatial.Tuples(2)

                ])
                trainingset = PSFDataset(transform=tr)
                valset = PSFDataset(transform=tr)

                train_data, train_label, val_data, val_label = self._load_data()
                # transforming data [N,C,T,V,M] to [N,T,V,M,C] to [N,T,V*M,C]
                train_data = np.transpose(train_data,(0,2,3,4,1))
                train_data = train_data.reshape(*train_data.shape[:2],-1,*train_data.shape[4:])
                val_data = np.transpose(val_data,(0,2,3,4,1))
                val_data = val_data.reshape(*val_data.shape[:2],-1,*val_data.shape[4:])
                if self._debug:
                # consider only the 5 first joints of the first person
                    train_data = train_data[:,:,:5,:]
                    val_data = val_data[:,:,:5,:]
                print('Training data shape:',train_data.shape)
                print('Validation data shape:',val_data.shape)
                trainingset.fill_from_iterator(self._get_iter(train_data, train_label))
                valset.fill_from_iterator(self._get_iter(val_data, val_label))
            else:
                raise NotImplementedError

            print('Save PSF.')
            trainingset.save(path+"uav_train_"+self.transform)
            valset.save(path+"uav_val_"+self.transform)

        print("Number of trainingset elements:", len(trainingset))
        print("Number of validation set elements", len(valset))
        print("Dimension of feature vector:", trainingset.get_data_dimension())

        return trainingset, valset


import os
import shutil
from gi.repository import Gdk
from gi.repository import GLib
from gi.repository import Gtk
from gi.repository import Pango

from kuwo import Config
from kuwo import Widgets

_ = Config._

MARGIN_LEFT = 15
MARGIN_TOP = 20
ShortcutMode = Config.ShortcutMode

class NoteTab(Gtk.Box):
    def __init__(self):
        super().__init__(orientation=Gtk.Orientation.VERTICAL)
        self.set_border_width(10)

class ColorButton(Gtk.ColorButton):
    def __init__(self, color):
        super().__init__()

class ColorBox(Gtk.Box):
    def __init__(self, label, conf, color_name, use_margin=False):
        super().__init__()
        self.conf = conf
        self.color_name = color_name
        left_label = Gtk.Label(label)
        self.pack_start(left_label, False, True, 0)

        color_button = Gtk.ColorButton()
        color_button.set_use_alpha(True)
        color_rgba = Gdk.RGBA()
        color_rgba.parse(conf[color_name])
        color_button.set_rgba(color_rgba)
        color_button.connect('color-set', self.on_color_set)
        self.pack_end(color_button, False, True, 0)

        if use_margin:
            self.props.margin_left = 20

    def on_color_set(self, color_button):
        color_rgba = color_button.get_rgba()
        if color_rgba.alpha == 1:
            color_rgba.alpha = 0.999
        self.conf[self.color_name] = color_rgba.to_string()

class FontBox(Gtk.Box):
    def __init__(self, label, conf, font_name, use_margin=True):
        super().__init__()
        self.conf = conf
        self.font_name = font_name
        left_label = Gtk.Label(label)
        self.pack_start(left_label, False, True, 0)

        font_button = Gtk.SpinButton()
        adjustment = Gtk.Adjustment(conf[font_name], 4, 72, 1, 10)
        adjustment.connect('value-changed', self.on_font_set)
        font_button.set_adjustment(adjustment)
        font_button.set_value(conf[font_name])
        self.pack_end(font_button, False, True, 0)

        if use_margin:
            self.props.margin_left = 20

    def on_font_set(self, adjustment):
        self.conf[self.font_name] = adjustment.get_value()

class ChooseFolder(Gtk.Box):
    def __init__(self, parent, conf_name, toggle_label):
        super().__init__(orientation=Gtk.Orientation.VERTICAL)
        self.props.margin_left = MARGIN_LEFT
        self.props.margin_top = 5
        self.parent = parent
        self.app = parent.app
        self.conf_name = conf_name
        self.old_dir = self.app.conf[conf_name]

        hbox = Gtk.Box(spacing=5)
        self.pack_start(hbox, False, True, 0)

        self.dir_entry = Gtk.Entry()
        self.dir_entry.set_text(self.old_dir)
        self.dir_entry.props.editable = False
        self.dir_entry.props.can_focus = False
        self.dir_entry.props.width_chars = 20
        hbox.pack_start(self.dir_entry, True, True, 0)

        choose_button = Gtk.Button('...')
        choose_button.connect('clicked', self.on_choose_button_clicked)
        hbox.pack_start(choose_button, False, False, 0)

    def on_choose_button_clicked(self, button):
        def on_dialog_file_activated(dialog):
            new_dir = dialog.get_filename()
            dialog.destroy()
            self.dir_entry.set_text(new_dir)
            if new_dir != self.app.conf[self.conf_name]:
                self.app.conf[self.conf_name] = new_dir
            return

        dialog = Gtk.FileChooserDialog(
                _('Choose a Folder'), self.parent,
                Gtk.FileChooserAction.SELECT_FOLDER,
                (Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL,
                 Gtk.STOCK_OK, Gtk.ResponseType.OK))

        dialog.connect('file-activated', on_dialog_file_activated)
        response = dialog.run()
        if response == Gtk.ResponseType.OK:
            on_dialog_file_activated(dialog)
            return
        dialog.destroy()

class Preferences(Gtk.Dialog):
    def __init__(self, app):
        self.app = app
        super().__init__(
                _('Preferences'), app.window, 0,
                (Gtk.STOCK_CLOSE, Gtk.ResponseType.CLOSE,))
        self.set_modal(True)
        self.set_transient_for(app.window)
        self.set_default_size(600, 320)
        self.set_border_width(5)
        box = self.get_content_area()
        #box.props.margin_left = MARGIN_LEFT

        notebook = Gtk.Notebook()
        box.pack_start(notebook, True, True, 0)

        # generic tab
        generic_box = NoteTab()
        notebook.append_page(generic_box, Gtk.Label(_('Generic')))

        status_button = Gtk.CheckButton(_('Close to system tray'))
        status_button.set_active(app.conf['use-status-icon'])
        status_button.connect('toggled', self.on_status_button_toggled)
        generic_box.pack_start(status_button, False, False, 0)

        notify_button = Gtk.CheckButton(_('Show kwplayer on lock screen'))
        notify_button.set_tooltip_text(
            _('Only works with gdm3/gnome3.8+\n') +
            _('Please disable it on other desktop environments (like KDE)'))
        notify_button.set_active(app.conf['use-notify'])
        notify_button.connect('toggled', self.on_notify_button_toggled)
        generic_box.pack_start(notify_button, False, False, 0)

        # format tab
        format_box = NoteTab()
        notebook.append_page(format_box, Gtk.Label(_('Format')))

        audio_label = Widgets.BoldLabel(_('Prefered Audio Format'))
        format_box.pack_start(audio_label, False, False, 0)
        radio_mp3 = Gtk.RadioButton(_('MP3 (faster)'))
        radio_mp3.props.margin_left = MARGIN_LEFT
        radio_mp3.connect('toggled', self.on_audio_toggled)
        format_box.pack_start(radio_mp3, False, False, 0)
        radio_ape = Gtk.RadioButton(_('APE (better)'))
        radio_ape.join_group(radio_mp3)
        radio_ape.props.margin_left = MARGIN_LEFT
        radio_ape.set_active(app.conf['use-ape'])
        radio_ape.connect('toggled', self.on_audio_toggled)
        format_box.pack_start(radio_ape, False, False, 0)

        video_label = Widgets.BoldLabel(_('Prefered Video Format'))
        video_label.props.margin_top = MARGIN_TOP
        format_box.pack_start(video_label, False, False, 0)
        radio_mp4 = Gtk.RadioButton(_('MP4 (faster)'))
        radio_mp4.props.margin_left = MARGIN_LEFT
        radio_mp4.connect('toggled', self.on_video_toggled)
        format_box.pack_start(radio_mp4, False, False, 0)
        radio_mkv = Gtk.RadioButton(_('MKV (better)'))
        radio_mkv.props.margin_left = MARGIN_LEFT
        radio_mkv.join_group(radio_mp4)
        radio_mkv.set_active(app.conf['use-mkv'])
        radio_mkv.connect('toggled', self.on_video_toggled)
        format_box.pack_start(radio_mkv, False, False, 0)

        # lyrics tab
        lrc_box = NoteTab()
        notebook.append_page(lrc_box, Gtk.Label(_('Lyrics')))

        lrc_normal_text_label = Widgets.BoldLabel(_('Normal Text'))
        lrc_box.pack_start(lrc_normal_text_label, False, True, 0)

        lrc_normal_text_size = FontBox(
                _('text size'), app.conf, 'lrc-text-size', use_margin=True)
        lrc_box.pack_start(lrc_normal_text_size, False, True, 0)

        lrc_normal_text_color = ColorBox(
                _('text color'), app.conf, 'lrc-text-color', use_margin=True)
        lrc_box.pack_start(lrc_normal_text_color, False, True, 0)
        lrc_normal_text_color.props.margin_bottom = 10

        lrc_highlighted_text_label = Widgets.BoldLabel(
                _('Highlighted Text'))
        lrc_box.pack_start(lrc_highlighted_text_label, False, True, 0)

        lrc_highlighted_text_size = FontBox(
                _('text size'), app.conf, 'lrc-highlighted-text-size',
                use_margin=True)
        lrc_box.pack_start(lrc_highlighted_text_size, False, True, 0)

        lrc_highlighted_text_color = ColorBox(
                _('text color'), app.conf, 'lrc-highlighted-text-color',
                use_margin=True)
        lrc_highlighted_text_color.props.margin_bottom = 10
        lrc_box.pack_start(lrc_highlighted_text_color, False, True, 0)

        lrc_word_back_color = ColorBox(
                _('Lyrics Text Background color'), app.conf, 'lrc-back-color')
        lrc_box.pack_start(lrc_word_back_color, False, True, 0)

        # folders tab
        folder_box = NoteTab()
        notebook.append_page(folder_box, Gtk.Label(_('Folders')))

        song_folder_label = Widgets.BoldLabel(_('Place to store sogns'))
        folder_box.pack_start(song_folder_label, False, False, 0)
        song_folder = ChooseFolder(
                self, 'song-dir', _('Moving cached songs to new folder'))
        folder_box.pack_start(song_folder, False, False, 0)

        mv_folder_label = Widgets.BoldLabel(_('Place to store MVs'))
        mv_folder_label.props.margin_top = MARGIN_TOP
        folder_box.pack_start(mv_folder_label, False, False, 0)
        mv_folder = ChooseFolder(
                self, 'mv-dir', _('Moving cached MVs to new folder'))
        folder_box.pack_start(mv_folder, False, False, 0)

        self.notebook = notebook

        # shortcut tab
        self.init_shortcut_tab()

    def init_shortcut_tab(self):
        curr_mode = self.app.conf['shortcut-mode']

        box = NoteTab()
        self.notebook.append_page(box, Gtk.Label(_('Shortcut')))

        self.shortcut_win = Gtk.ScrolledWindow()

        disable_btn = Gtk.RadioButton(_('Disable Keyboard Shortcut'))
        disable_btn.connect(
                'toggled', self.on_shortcut_btn_toggled, ShortcutMode.NONE)
        disable_btn.set_active(curr_mode == ShortcutMode.NONE)
        box.pack_start(disable_btn, False, False, 0)

        default_btn = Gtk.RadioButton(_('Use Default MultiMedia Key'))
        default_btn.connect(
                'toggled', self.on_shortcut_btn_toggled, ShortcutMode.DEFAULT)
        default_btn.join_group(disable_btn)
        default_btn.set_active(curr_mode == ShortcutMode.DEFAULT)
        box.pack_start(default_btn, False, False, 0)

        custom_btn = Gtk.RadioButton(_('Use Custom Keyboard Shortcut'))
        custom_btn.connect(
                'toggled', self.on_shortcut_btn_toggled, ShortcutMode.CUSTOM)
        custom_btn.join_group(default_btn)
        custom_btn.set_active(curr_mode == ShortcutMode.CUSTOM)
        box.pack_start(custom_btn, False, False, 0)

        self.shortcut_win.props.margin_left = 10
        self.shortcut_win.set_sensitive(curr_mode == ShortcutMode.CUSTOM)
        box.pack_start(self.shortcut_win, True, True, 0)

        # name, shortct key, shortcut modifiers
        self.shortcut_liststore = Gtk.ListStore(str, int, int)
        tv = Gtk.TreeView(model=self.shortcut_liststore)
        self.shortcut_win.add(tv)

        name_cell = Gtk.CellRendererText()
        name_col = Gtk.TreeViewColumn('Action', name_cell, text=0)
        tv.append_column(name_col)

        key_cell = Gtk.CellRendererAccel(editable=True)
        key_cell.connect('accel-edited', self.on_shortcut_key_cell_edited)
        key_col = Gtk.TreeViewColumn(
                'Shortcut Key', key_cell, accel_key=1, accel_mods=2)
        tv.append_column(key_col)

        for name in self.app.conf['custom-shortcut']:
            key = self.app.conf['custom-shortcut'][name]
            i18n_name = Config.SHORT_CUT_I18N[name]
            k, m = Gtk.accelerator_parse(key)
            self.shortcut_liststore.append([i18n_name, k, m, ])

    def run(self):
        self.get_content_area().show_all()
        super().run()

    def on_destroy(self):
        print('dialog.on_destroy()')
        Config.dump_conf(self.app.conf)

    # generic tab signal handlers
    def on_status_button_toggled(self, button):
        self.app.conf['use-status-icon'] = button.get_active()

    def on_notify_button_toggled(self, button):
        self.app.conf['use-notify'] = button.get_active()

    # format tab signal handlers
    def on_audio_toggled(self, radiobtn):
        self.app.conf['use-ape'] = radiobtn.get_group()[0].get_active()

    def on_video_toggled(self, radiobtn):
        # radio_group[0] is MKV
        self.app.conf['use-mkv'] = radiobtn.get_group()[0].get_active()

    def on_shortcut_btn_toggled(self, button, mode):
        if button.get_active() is False:
            return
        self.app.conf['shortcut-mode'] = mode
        self.shortcut_win.set_sensitive(mode == ShortcutMode.CUSTOM)

    def on_shortcut_key_cell_edited(self, accel, path, key, mod,
                                    hardware_keycode):
        accel_key = Gtk.accelerator_name(key, mod)
        name = self.shortcut_liststore[path][0]
        self.shortcut_liststore[path][1] = key
        self.shortcut_liststore[path][2] = int(mod)
        self.app.conf['custom-shortcut'][name] = accel_key


import uuid
import queue

from rnnvis.utils.tree import TreeNode, Tree

class XTree(object):
    """"This class is just a wrap of class SuffixTree which provide both prefix and suffix
    """
    def __init__(self, data, sentence_delimiter='\n'):
        self._data = data
        self._prefix_data, self._suffix_data = self.data_preprocessing(
            self._data, sentence_delimiter)

        self._prefix_tree = SuffixTree(self._prefix_data)
        self._suffix_tree = SuffixTree(self._suffix_data)

    def fetch_prefix_suffix(self, gram):
        if isinstance(gram, str):
            gram = gram.split()
        prefix_result = self._prefix_tree.search_suffix(list(reversed(gram)))
        suffix_result = self._suffix_tree.search_suffix(gram)

        q = queue.Queue()
        q.put(prefix_result)
        while not q.empty():
            node = q.get()
            prefix_position = (node['position'][0],
                               len(self._prefix_data[node['position'][0]]) -
                               node['position'][1] - len(node['data']))
            node['position'] = prefix_position
            node['data'] = ' '.join(list(reversed(node['data'])))
            for child in node['children']:
                q.put(child)

        q.put(suffix_result)
        while not q.empty():
            node = q.get()
            node['data'] = ' '.join(node['data'])
            for child in node['children']:
                q.put(child)

        return prefix_result, suffix_result

    def data_preprocessing(self, data, sentence_delimiter='\n'):
        """ Preprocess the raw data to produce prefix and suffix data
        """
        lines = data.split(sentence_delimiter)
        prefix_data = []
        for i, line in enumerate(lines):
            words = line.split()
            words = list(reversed(words))
            words.append('<bos_' + str(i) + '>')
            prefix_data.append(words)
        suffix_data = []
        for i, line in enumerate(lines):
            words = line.split()
            words.append('<eos_' + str(i) + '>')
            suffix_data.append(words)

        return prefix_data, suffix_data

class SuffixTree(Tree):
    """This class is used to construct suffix tree
    """

    def __init__(self, data):
        Tree.__init__(self)
        root_node = SuffixTreeNode(data=[], position=(-1, -1))
        self.add_node(root_node)
        self._root = self.get_root()
        self._data = data

        self._leaves = set()
        self._status = None
        self._remainder = 0
        self._previous_node_added = None
        self._reverse_suffix_links = {}
        self._tmp_str = None

        self._words = None
        self._sentence_num = -1

        self.construct(self._data)
        self._tree_dict = self.to_dict(self.get_root())

    def construct(self, data):
        for i, line in enumerate(data):
            self._status = {'active_node': self._root,
                            'active_edge': None,
                            'active_length': 0
                            }
            self._leaves = set()
            self._remainder = 0
            self._words = line
            self._sentence_num = i

            for j, word in enumerate(line):
                for leaf in self._leaves:
                    self.get_node(leaf).append_suffix(word)
                self._remainder += 1
                self.add_to_tree(j)
                self._previous_node_added = None
                # self._tmp_str = self.tree_str(self.get_root())
            # print(self.to_dict(self.get_root()))

    def search_suffix(self, gram):
        if isinstance(gram, str):
            gram = gram.split()
        gram_node, position = self.find_node(self.get_root(), gram)
        if position < 0:
            return None
        tree_dict = self.to_dict(gram_node)
        tree_dict['data'] = ' '.join(gram_node.data[position:])
        return tree_dict

    def to_dict(self, node):
        node_info = {}
        node_info['data'] = node.data if not node.is_root() else ['#root#']
        node_info['position'] = node.position
        node_info['children'] = []
        for child_id in node.children_id:
            child_node = self.get_node(child_id)
            node_info['children'].append(self.to_dict(child_node))
        return node_info

    def find_node(self, start_node, gram):
        if not start_node.children_prefix(gram[0]):
            return start_node, -1

        edge = self.get_node(start_node.children_prefix(gram[0]))
        for i in range(min(len(gram), len(edge.data))):
            if gram[i] != edge.data[i]:
                return start_node, -1
        i += 1
        if i == len(gram):
            return edge, i
        elif i == len(edge.data):
            return self.find_node(edge, gram[i:])

    def tree_str(self, node):
        if len(node.children_id) < 1:
            return ' '.join(node.data)
        strs = []
        for child_id in node.children_id:
            strs.append(self.tree_str(self.get_node(child_id)))
        key_ = ' '.join(node.data) if not node.is_root() else 'root'
        return {key_: strs}

    def add_node(self, node, parent=None):
        if not isinstance(node, TreeNode):
            raise TypeError("arg node should be of type TreeNode, but it's of type {:s}".format(type(node)))
        if node.id in self._dict:
            raise ValueError("node with id {:s} already exists!".format(node.id))
        if parent is None:
            if self.root_id is not None:
                raise ValueError("Attempt to add a multiple root! The tree already has a root.")
            self.root_id = node.id
        else:
            parent = self.get_node(parent)
            if parent is None:
                raise ValueError("the given parent is not in the tree!")
            parent.children_id.add(node.id)
            parent.add_child(node)
            node.parent_id = parent.id
        self._dict[node.id] = node

    def add_to_tree(self, i):
        if self._remainder == 0:
            return

        # word = words[i]
        # if self._status['active_node'].is_root():
        #     tmp_word = words[i - self._remainder + 1]
        # else:
        #     tmp_word = words[i - self._status['active_length']]
        # word = self._words[i - self._remainder + 1] if self._status['active_node'].is_root() else \
        #     self._words[i - self._status['active_length']]
        word_pos = i - self._remainder + 1 if self._status['active_node'].is_root() else \
            i - self._status['active_length']
        word = self._words[word_pos]
        if self._status['active_edge'] is None:
            if not self._status['active_node'].children_prefix(word):
                node = SuffixTreeNode(data=[word], position=(self._sentence_num, word_pos))
                self.add_node(node, parent=self._status['active_node'])
                self._leaves.add(node.id)
                self._remainder -= 1
                if not self._status['active_node'].is_root():
                    self._status['active_node'] = self.get_node(
                        self._status['active_node'].suffix_link) if \
                        self._status['active_node'].suffix_link else self.get_root()
                    # if self._previous_node_added is not None:
                    #     self._previous_node_added.suffix_link = node.id
                    # self._previous_node_added = node
                    self._previous_node_added = None
                    self.add_to_tree(i)

            else:
                edge_node_id = self._status['active_node'].children_prefix(word)
                edge_node = self.get_node(edge_node_id)
                self._status['active_edge'] = edge_node
                self._status['active_length'] = 0
                counter = 0
                while counter+1 <= self._remainder and self._status['active_edge'] and self._status['active_edge'].data[
                    self._status['active_length']] == self._words[i - self._remainder + counter + 1]:

                    self._status['active_length'] += 1
                    counter += 1
                    if len(self._status['active_edge'].data) == self._status['active_length']:
                        self._status['active_node'] = self._status['active_edge']
                        self._status['active_edge'] = self._status['active_node'].children_prefix(
                            self._words[i - self._remainder + counter + 2]) if counter+2 <= self._remainder else None
                        self._status['active_length'] = 0

                if counter == self._remainder:
                    # contain all infomation
                    return
                self.add_to_tree(i)

        elif self._status['active_edge'].data[self._status['active_length']] == self._words[i]:
            self._status['active_length'] += 1
            if len(self._status['active_edge'].data) == self._status['active_length']:
                self._status['active_node'] = self._status['active_edge']
                self._status['active_edge'] = None
                self._status['active_length'] = 0

        elif self._status['active_edge'].data[self._status['active_length']] != self._words[i]:
            new_node = SuffixTreeNode(data=[self._words[i]], position=(self._sentence_num, i))
            other_node = self.split_node(self._status['active_edge'], self._status['active_length'])
            self.add_node(new_node, self._status['active_edge'])
            self.add_node(other_node, self._status['active_edge'])
            self._leaves.add(new_node.id)

            if self._previous_node_added is not None:
                self._previous_node_added.suffix_link = self._status['active_edge'].id
            self._previous_node_added = self._status['active_edge']

            self._remainder -= 1

            # update active node and active edge
            if self._status['active_node'].is_root():
                self._status['active_length'] -= 1
                edge_node_id = self._status['active_node'].children_prefix(self._words[i - self._status['active_length']])
                self._status['active_edge'] = self.get_node(edge_node_id) if edge_node_id else None
                while self._status['active_edge'] and self._status['active_length'] >= len(self._status['active_edge'].data):
                    self._status['active_node'] = self._status['active_edge']
                    self._status['active_length'] -= len(self._status['active_edge'].data)
                    edge_node_id = self._status['active_node'].children_prefix(self._words[i - self._status['active_length']])
                    self._status['active_edge'] = self.get_node(edge_node_id) if edge_node_id else None

                self.add_to_tree(i)
            elif self._status['active_node'].suffix_link:
                self._status['active_node'] = self.get_node(
                    self._status['active_node'].suffix_link)
                edge_node_id = self._status['active_node'].children_prefix(self._words[i - self._status['active_length']])
                self._status['active_edge'] = self.get_node(edge_node_id)
                self.add_to_tree(i)
            else:
                self._status['active_node'] = self.get_root()
                self._status['active_edge'] = None
                self.add_to_tree(i)

    def split_node(self, node, position_):
        new_node = SuffixTreeNode(data=node.data[position_:],
                                  position=(node.position[0], node.position[1]+position_))
        node.data = node.data[:position_]
        for child_id in node.children_id:
            child_node = self.get_node(child_id)
            child_node.parent_id = new_node.id
        new_node.children_id = node.children_id
        new_node.child_prefix_dict = node.child_prefix_dict
        node.delete_all_children()
        if node.id in self._leaves:
            self._leaves.remove(node.id)
            self._leaves.add(new_node.id)
        return new_node

class SuffixTreeNode(TreeNode):
    '''Suffix tree node class, inherit from TreeNode
    '''

    def __init__(self, data, position, id_=None, parent_id=None):
        if id_ is None:
            id_ = uuid.uuid1().hex
        TreeNode.__init__(self, id_=id_, parent_id=parent_id)
        self._data = data
        self._child_prefix_dict = {}
        self._suffix_link = None
        self._position = position

    def first_element(self):
        """Return the first element of node data
        """
        if not isinstance(self._data, list):
            raise ValueError('The node data should be a list of words')
        if len(self._data) < 1:
            raise ValueError('The length of node data should at least be 1')
        return self._data[0]

    def children_prefix(self, word):
        return self._child_prefix_dict.get(word)

    def append_suffix(self, word):
        self._data.append(word)

    def add_child(self, child_node):
        self._child_prefix_dict[child_node.first_element()] = child_node.id
        self._children_id.add(child_node.id)

    def delete_all_children(self):
        self._child_prefix_dict = {}
        self._children_id = set()

    @property
    def position(self):
        return self._position

    @position.setter
    def position(self, new_position):
        self._position = new_position

    @property
    def child_prefix_dict(self):
        return self._child_prefix_dict

    @child_prefix_dict.setter
    def child_prefix_dict(self, dict_):
        self._child_prefix_dict = dict_

    @property
    def children_id(self):
        return self._children_id

    @children_id.setter
    def children_id(self, children_id_):
        self._children_id = children_id_

    @property
    def child_prefix(self):
        return self._child_prefix_dict

    @property
    def suffix_link(self):
        '''Suffix link getter
        '''
        return self._suffix_link

    @suffix_link.setter
    def suffix_link(self, node_id):
        '''Suffix link setter
        '''
        self._suffix_link = node_id

    @property
    def data(self):
        '''Node data getter
        '''
        return self._data

    @data.setter
    def data(self, data_):
        '''Node data setter
        '''
        self._data = data_


#!/usr/bin/env python3

def psf_study(quadrant_path, ztfname, filtercode, logger, args):
    from utils import ListTable, match_pixel_space, quadrant_size_px
    import numpy as np
    from saunerie.plottools import binplot
    from scipy.stats import norm
    import matplotlib
    import matplotlib.pyplot as plt
    import pickle
    from matplotlib.ticker import MultipleLocator, FormatStrFormatter, AutoMinorLocator
    matplotlib.use('Agg')

    if not quadrant_path.joinpath("psfstars.list").exists():
        return True

    psf_resid = ListTable.from_filename(quadrant_path.joinpath("psf_resid_tuple.dat"))
    psf_tuple = ListTable.from_filename(quadrant_path.joinpath("psftuple.list"))
    psf_stars = ListTable.from_filename(quadrant_path.joinpath("psfstars.list"))
    stand = ListTable.from_filename(quadrant_path.joinpath("standalone_stars.list"))

    stand.df['mag'] = -2.5*np.log10(stand.df['flux'])
    psf_stars.df['mag'] = -2.5*np.log10(psf_stars.df['flux'])
    psf_residuals_px = (psf_resid.df['fimg'] - psf_resid.df['fpsf']).to_numpy()
    mag = np.linspace(np.min(stand.df['mag']), np.max(stand.df['mag']), 100)

    subx = 4
    suby = 1

    # Fit order 2 polynomials on skewness/mag relation
    def _fit_skewness_subquadrants(axis, sub):
        ranges = np.linspace(0., quadrant_size_px[axis], sub+1)
        skew_polynomials = []
        for i in range(sub):
            range_mask = (stand.df['y'] >= ranges[i]) & (stand.df['y'] < ranges[i+1])
            skew_polynomials.append(np.polynomial.Polynomial.fit(stand.df.loc[range_mask]['mag'], stand.df.loc[range_mask]['gm{}3'.format(axis)], 1))

        return skew_polynomials

    skew_polynomial_x = np.polynomial.Polynomial.fit(stand.df['mag'], stand.df['gmx3'], 1)
    skew_polynomial_y = np.polynomial.Polynomial.fit(stand.df['mag'], stand.df['gmy3'], 1)

    skew_polynomial_subx = _fit_skewness_subquadrants('x', subx)
    skew_polynomial_suby = _fit_skewness_subquadrants('y', suby)

    with open(quadrant_path.joinpath("stamp_skewness.pickle"), 'wb') as f:
        pickle.dump({'poly_x': skew_polynomial_x, 'poly_y': skew_polynomial_y,
                     'poly_x_sub': skew_polynomial_subx, 'poly_y_sub': skew_polynomial_suby}, f)

    plt.figure(figsize=(7., 7.))
    ax = plt.axes()
    ax.tick_params(which='both', direction='in')
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.grid(linestyle='--', color='xkcd:sky blue')
    plt.suptitle("Skewness plane for standalone stars stamps in \n {}".format(quadrant_path.name))
    plt.scatter(stand.df['gmx3'], stand.df['gmy3'], s=2., color='black')
    plt.xlabel("$M^g_{xxx}$")
    plt.ylabel("$M^g_{yyy}$")
    plt.axvline(0.)
    plt.axhline(0.)
    plt.axvline(skew_polynomial_x.coef[0], ls='--')
    plt.axhline(skew_polynomial_y.coef[0], ls='--')
    plt.axis('equal')
    plt.savefig(quadrant_path.joinpath("{}_mx3_my3_plane.png".format(quadrant_path.name)), dpi=300.)
    plt.close()

    plt.subplots(ncols=1, nrows=2, sharex=True, figsize=(10., 5.))
    plt.suptitle("Skewness vs magnitude for standalone stars stamps in {}".format(quadrant_path.name))

    ax = plt.subplot(2, 1, 1)
    ax.tick_params(which='both', direction='in')
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.grid(linestyle='--', color='xkcd:sky blue')
    plt.scatter(stand.df['mag'], stand.df['gmx3'], s=2.)
    plt.plot(mag, skew_polynomial_x(mag), color='black')
    plt.axhline(0.)
    plt.ylim(-0.5, 0.3)
    plt.xlim(-15., -7)
    plt.ylabel("$M^g_{xxx}$")

    ax = plt.subplot(2, 1, 2)
    ax.tick_params(which='both', direction='in')
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.grid(linestyle='--', color='xkcd:sky blue')
    plt.scatter(stand.df['mag'], stand.df['gmy3'], s=2.)
    plt.plot(mag, skew_polynomial_y(mag), color='black')
    plt.axhline(0.)
    plt.xlabel("$m$")
    plt.ylim(-0.3, 0.3)
    plt.xlim(-15., -6)
    plt.ylabel("$M^g_{yyy}$")

    plt.savefig(quadrant_path.joinpath("{}_skewness_magnitude.png".format(quadrant_path.name)), dpi=300.)
    plt.close()
    return

    star_indices = match_pixel_space(psf_stars.df[['x', 'y']].to_records(), psf_resid.df[['xc', 'yc']].rename(mapper={'xc': 'x', 'yc': 'y'}, axis='columns').to_records())
    psf_size = int(np.sqrt(sum(star_indices==0)))
    psf_residuals = np.zeros([len(psf_stars.df), psf_size*psf_size])
    for star_index in range(len(np.bincount(star_indices))):
        star_mask = (star_indices==star_index)
        ij = (psf_resid.df.iloc[star_mask]['j']*psf_size + psf_resid.df.iloc[star_mask]['i'] + int(np.floor(psf_size**2/2))).to_numpy()
        np.put_along_axis(psf_residuals[star_index], ij, psf_residuals_px[star_mask], 0)

    psf_residuals = psf_residuals.reshape(len(psf_stars.df), psf_size, psf_size)

    bins_count = 5
    mag_range = (psf_stars.df['mag'].min(), psf_stars.df['mag'].max())
    bins = np.linspace(*mag_range, bins_count+1)
    psf_residual_means = []
    psf_residual_count = []
    for i in range(bins_count):
        lower_bound = bins[i]
        upper_bound = bins[i+1]

        binned_stars_mask = np.all([psf_stars.df['mag'] < upper_bound, psf_stars.df['mag'] >= lower_bound], axis=0)
        psf_residual_means.append(np.mean(psf_residuals[binned_stars_mask], axis=0))
        psf_residual_count.append(sum(binned_stars_mask))

    plt.subplots(ncols=bins_count, nrows=1, figsize=(5.*bins_count, 5.))
    plt.suptitle("Binned PSF residuals for {}".format(quadrant_path.name))
    for i, psf_residual_mean in enumerate(psf_residual_means):
        plt.subplot(1, bins_count,  1+i)
        plt.imshow(psf_residual_mean)
        plt.axis('off')
        plt.title("${0:.2f} \leq m < {1:.2f}$\n$N={2}$".format(bins[i], bins[i+1], psf_residual_count[i]))

    plt.savefig(quadrant_path.joinpath("{}_psf_residuals.png".format(quadrant_path.name)))
    plt.close()
    # plt.subplots(ncols=2, nrows=1, sharex=True, sharey=True)
    # # plt.subplot(1, 2, 1)
    # # plt.scatter(stand.df['gmx3'], stand.df['gmy3'], c=stand.df['x'], s=2.)
    # # plt.axvline(0.)
    # # plt.axhline(0.)

    # # plt.subplot(1, 2, 2)
    # # plt.scatter(stand.df['gmx3'], stand.df['gmy3'], c=stand.df['y'], s=2.)
    # # plt.axvline(0.)
    # # plt.axhline(0.)
    # # plt.show()

    # plt.subplot(2, 2, 1)
    # plt.scatter(stand.df['mag'], stand.df['gmx3'], c=stand.df['x'], s=1.)
    # plt.axhline(0.)
    # plt.subplot(2, 2, 2)
    # plt.scatter(stand.df['mag'], stand.df['gmy3'], c=stand.df['x'])
    # plt.axhline(0.)
    # plt.subplot(2, 2, 3)
    # plt.scatter(stand.df['mag'], stand.df['gmx3'], c=stand.df['y'])
    # plt.axhline(0.)
    # plt.subplot(2, 2, 4)
    # plt.scatter(stand.df['mag'], stand.df['gmy3'], c=stand.df['y'])
    # plt.axhline(0.)
    # plt.show()

    # psf_resid.df.to_csv("psf_resid_tuple.csv", sep=",")
    # psf_tuple.df.to_csv("psftuple.csv", sep=",")
    # psf_stars.df.to_csv("psfstars.csv", sep=",")

    # psf_resid.df['flux'] = psf_stars.df['flux'].iloc[idx].reset_index(drop=True)
    # psf_resid.df['eflux'] = psf_stars.df['eflux'].iloc[idx].reset_index(drop=True)

    # limits = [np.min(res), np.max(res)]
    # f = np.linspace(*limits, 200)
    # s = np.var(res)
    # m = np.mean(res)
    # print(np.sqrt(s))
    # print(m)

    # plt.hist(res, bins=1000, density=True)
    # plt.plot(f, norm.pdf(f, loc=m, scale=np.sqrt(s)))
    # plt.xlim(limits)
    # plt.show()

    # plt.plot(psf_resid.df['flux'], psf_resid.df['eflux']/psf_resid.df['flux'], '.')
    # plt.show()

    # plt.subplots(nrows=2, ncols=1, figsize=(10., 5.), sharex=True)
    # plt.subplot(2, 1, 1)
    # binned_mag, plot_res, res_dispersion = binplot(psf_resid.df['flux'].to_numpy(), psf_resid.df['fimg'] - psf_resid.df['fpsf'], nbins=10, data=True, rms=True, scale=False)

    # plt.ylabel("res")
    # plt.grid()

    # plt.subplot(2, 1, 2)
    # plt.plot(binned_mag, res_dispersion, color='black')
    # plt.xlabel("Flux [ADU]")
    # plt.ylabel("$\sigma_\\mathrm{res}$")
    # plt.grid()

    # plt.show()

    return True

def psf_study_reduce(band_path, ztfname, filtercode, logger, args):
    import numpy as np
    import pandas as pd
    from deppol_utils import quadrants_from_band_path
    from utils import get_header_from_quadrant_path, ListTable
    import matplotlib
    import matplotlib.pyplot as plt
    from matplotlib.ticker import MultipleLocator, FormatStrFormatter, AutoMinorLocator
    from utils import idx2markerstyle, plot_ztf_focal_plan_values
    import itertools
    import pickle
    from astropy import time
    from matplotlib.cm import ScalarMappable
    from matplotlib.colorbar import Colorbar
    # matplotlib.use('Agg')

    output_path = band_path.joinpath("psf_study")
    output_path.mkdir(exist_ok=True)

    # Compare seeing computed by sextractor vs the one computed by the ZTF pipeline
    quadrant_paths = quadrants_from_band_path(band_path, logger, check_files="psfstars.list", ignore_noprocess=False)

    #chi2_df = pd.read_csv(band_path.joinpath("astrometry/ref2px_plots/chi2_quadrants.csv"), index_col=0)

    # def _extract_seeing(quadrant_path):
    #     hdr = get_header_from_quadrant_path(quadrant_path)
    #     if quadrant_path.name in chi2_df.index.tolist():
    #         chi2 = chi2_df.at[quadrant_path.name, 'chi2']
    #     else:
    #         chi2 = float('nan')

    #     psfstars_list = ListTable.from_filename(quadrant_path.joinpath("psfstars.list"))
    #     return 2.355*hdr['gfseeing'], hdr['seeing'], len(psfstars_list.df), chi2

    # seeings = np.array(list(map(_extract_seeing, quadrant_paths))).T

    def _extract_skewness(quadrant_path):
        hdr = get_header_from_quadrant_path(quadrant_path)
        with open(quadrant_path.joinpath("stamp_skewness.pickle"), 'rb') as f:
            poly = pickle.load(f)
            poly_x = poly['poly_x']
            poly_y = poly['poly_y']
            poly_x_sub = poly['poly_x_sub']
            poly_y_sub = poly['poly_y_sub']

        #out_dict = {'mjd': hdr['obsmjd'], 'quadrant_id': "{}_{}".format(hdr['ccdid'], hdr['qid'])}
        out_dict = {'mjd': hdr['obsmjd'], 'ccdid': int(hdr['ccdid']), 'qid': int(hdr['qid']), 'rcid': int(hdr['rcid']),
                    'sky': hdr['sexsky'], 'skysigma': hdr['sexsigma'], 'gfseeing': hdr['gfseeing'], 'seeing': hdr['seeing']}

        out_dict.update(dict(('x{}'.format(i), coef) for i, coef in enumerate(poly_x.coef)))
        out_dict.update(dict(('y{}'.format(i), coef) for i, coef in enumerate(poly_y.coef)))
        out_dict.update({'x1_sub': [poly.coef[1] for poly in poly_x_sub],
                         'y1_sub': [poly.coef[1] for poly in poly_y_sub]})

        return out_dict

    skewness_df = pd.DataFrame(list(map(_extract_skewness, quadrant_paths)))
    cmap = 'coolwarm'

    unique_mjds = list(set(skewness_df['mjd']))
    xskewness = dict([(mjd, dict([(i+1, dict([(j, None) for j in range(4)])) for i in range(16)])) for mjd in unique_mjds])
    yskewness = dict([(mjd, dict([(i+1, dict([(j, None) for j in range(4)])) for i in range(16)])) for mjd in unique_mjds])

    for mjd in unique_mjds:
        skewness_today_df = skewness_df.loc[skewness_df['mjd']==mjd]
        for row in skewness_today_df.iterrows():
            # Dirty
            row = row[1]
            ccdid = int(row['ccdid'])
            qid = int(row['qid'])

            xskewness[mjd][ccdid][qid-1] = np.tile(np.array(row['x1_sub']), (len(row['x1_sub']), 1))
            yskewness[mjd][ccdid][qid-1] = np.tile(np.array(row['y1_sub']), (len(row['y1_sub']), 1))

    def _plot_focal_plane_skewness(x, vmin, vmax):
        cm = ScalarMappable(cmap=cmap)
        cm.set_clim(vmin=vmin, vmax=vmax)

        if x == 'x1':
            skewness = xskewness
        elif x == 'y1':
            skewness = yskewness

        for mjd in unique_mjds:
            t = time.Time(mjd, format='mjd')
            fig = plt.figure(figsize=(5., 6.), constrained_layout=True)
            f1, f2 = fig.subfigures(ncols=1, nrows=2, height_ratios=[8., 1.])
            plot_ztf_focal_plan_values(f1, skewness[mjd], vmin=vmin, vmax=vmax, cmap=cmap)
            ax = f2.add_subplot()
            Colorbar(ax, cm, orientation='horizontal', label="${}_{}$".format(x[0], x[1]))
            fig.suptitle("Focal plane skewness in ${}$ direction the {}".format(x[0], t.to_value('iso', subfmt='date')), fontsize='large')
            plt.savefig(output_path.joinpath("{}_focal_plane_{}skewness.png".format(t.to_value('iso', subfmt='date'), x[0])), dpi=300.)
            plt.close()

    sharedscale = False
    center = True
    if sharedscale:
        vmin, vmax = np.min([skewness_df['x1'], skewness_df['y1']]), np.max([skewness_df['x1'], skewness_df['y1']])
        xvmin = vmin
        xvmax = vmax
        yvmin = vmin
        yvmax = vmax
    else:
        xvmin, xvmax = np.min(skewness_df['x1']), np.max(skewness_df['x1'])
        yvmin, yvmax = np.min(skewness_df['y1']), np.max(skewness_df['y1'])

    if center:
        if sharedscale:
            m = max(np.abs(vmin), np.abs(vmax))
            xvmin = -m
            xvmax = m
            yvmin = -m
            yvmax = m
        else:
            m = max(np.abs(xvmin), np.abs(xvmax))
            xvmin = -m
            xvmax = m
            m = max(np.abs(yvmin), np.abs(yvmax))
            yvmin = -m
            yvmax = m

    _plot_focal_plane_skewness('x1', xvmin, xvmax)
    _plot_focal_plane_skewness('y1', yvmin, yvmax)

    rcids = set(skewness_df['rcid'])
    plt.subplots(nrows=2, ncols=1, figsize=(10., 5.))
    plt.suptitle("Order 1 polynomial coefficients fit on skewness/mag")
    ax = plt.subplot(2, 1, 1)
    if len(rcids) <= len(idx2markerstyle):
        for i, rcid in enumerate(rcids):
            rcid_mask = (skewness_df['rcid'] == rcid)
            plt.plot(skewness_df.loc[rcid_mask]['mjd'], skewness_df[rcid_mask]['x1'], idx2markerstyle[i], color='black', label=rcid)
        plt.legend(title="Quadrant ID")
    else:
        plt.scatter(skewness_df['mjd'], skewness_df['x1'], c=skewness_df['rcid'].tolist(), s=1.)
        plt.colorbar()
    plt.axhline(0., color='black')
    plt.ylim(vmin, vmax)
    ax.tick_params(which='both', direction='in')
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.grid(linestyle='--', color='xkcd:sky blue')
    plt.xlabel("MJD")
    plt.ylabel("$S^x_1$")

    ax = plt.subplot(2, 1, 2)
    if len(rcids) <= len(idx2markerstyle):
        for i, rcid in enumerate(rcids):
            rcid_mask = (skewness_df['rcid'] == rcid)
            plt.plot(skewness_df.loc[rcid_mask]['mjd'], skewness_df[rcid_mask]['y1'], idx2markerstyle[i], color='black', label=rcid)
        plt.legend(title="Quadrant ID")
    else:
        plt.scatter(skewness_df['mjd'], skewness_df['y1'], c=skewness_df['rcid'].tolist(), s=1.)
        plt.colorbar()
    plt.axhline(0., color='black')
    plt.ylim(vmin, vmax)
    ax.tick_params(which='both', direction='in')
    ax.xaxis.set_minor_locator(AutoMinorLocator())
    ax.yaxis.set_minor_locator(AutoMinorLocator())
    plt.grid(linestyle='--', color='xkcd:sky blue')
    plt.xlabel("MJD")
    plt.ylabel("$S^y_1$")

    plt.tight_layout()
    plt.savefig(output_path.joinpath("skewness_mjd.png"), dpi=300.)
    plt.close()

    # ax = plt.subplot(2, 1, 2)
    # if len(qids) <= len(idx2markerstyle):
    #     for i, qid in enumerate(qids):
    #         qid_mask = (skewness_df['quadrant_id'] == qid)
    #         plt.plot(skewness_df.loc[qid_mask]['mjd'], skewness_df[qid_mask]['y1'], idx2markerstyle[i], color='black', label=qid)
    #     plt.legend(title="Quadrant ID")
    # else:
    #     plt.scatter(skewness_df['mjd'], skewness_df['y1'], c=skewness_df['quadrant_id'].tolist(), s=1.)
    #     plt.colorbar()
    # plt.axhline(0., color='black')
    # plt.ylim(*ylim)
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.xlabel("MJD")
    # plt.ylabel("$S^y_1$")

    return

    # ax = plt.axes()
    # plt.plot(seeings[0], seeings[1], '.')
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.xlabel("GF seeing [pixel]")
    # plt.ylabel("ZTF seeing [pixel]")
    # plt.show()
    # plt.close()

    # ax = plt.axes()
    # plt.plot(seeings[2], seeings[3], '.')
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.xlabel("PSF stars count")
    # plt.ylabel("Chi2")
    # plt.show()
    # plt.close()

    # ax = plt.axes()
    # plt.plot(seeings[2], seeings[0]-seeings[1]-np.mean(seeings[0]-seeings[1]), '.')
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.xlabel("PSF stars count")
    # plt.ylabel("GF - ZTF - <GF-ZTF> (seeing) [pixel]")
    # plt.show()
    # plt.close()

    # ax = plt.axes()
    # plt.plot(seeings[3], seeings[0]-seeings[1]-np.mean(seeings[0]-seeings[1]), '.')
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.xlabel("Chi2")
    # plt.ylabel("GF - ZTF - <GF-ZTF> (seeing) [pixel]")
    # plt.show()
    # plt.close()

    # ax = plt.axes()
    # plt.hist(seeings[0]-seeings[1]-np.mean(seeings[0]-seeings[1]), bins='auto', histtype='step', color='black')
    # ax.tick_params(which='both', direction='in')
    # ax.xaxis.set_minor_locator(AutoMinorLocator())
    # ax.yaxis.set_minor_locator(AutoMinorLocator())
    # plt.grid(linestyle='--', color='xkcd:sky blue')
    # plt.show()

def seeing_study(band_path, ztfname, filtercode, logger, args):
    from utils import get_header_from_quadrant_path, plot_ztf_focal_plan_values, ListTable
    from deppol_utils import quadrants_from_band_path
    import pickle
    import numpy as np
    import itertools
    import matplotlib
    import matplotlib.pyplot as plt
    from matplotlib.cm import ScalarMappable
    from matplotlib.colorbar import Colorbar
    quadrant_paths = quadrants_from_band_path(band_path, logger, check_files="psfstars.list", ignore_noprocess=False)

    keys = {'seeing': 'hdr',
            'seseeing': 'hdr',
            'gfseeing': 'gfseeing',
            'momentx1': 'psfstars',
            'momenty1': 'psfstars',
            'momentx2': 'psfstars',
            'momenty2': 'psfstars',
            'momentxy': 'psfstars'}

    if not band_path.joinpath("seeings.pickle").exists() or args.from_scratch:
        vals = {}

        for key in keys:
            vals[key] = {}
            for i in range(1, 17):
                vals[key][i] = {}
                for j in range(4):
                    vals[key][i][j] = []

        for quadrant_path in quadrant_paths:
            hdr = get_header_from_quadrant_path(quadrant_path)
            psfstars_list = ListTable.from_filename(quadrant_path.joinpath("psfstars.list"))

            for key in keys.keys():
                if keys[key] == 'psfstars':
                    vals[key][hdr['ccdid']][hdr['qid']-1].append(psfstars_list.header[key])
                elif key == 'seseeing':
                    vals[key][hdr['ccdid']][hdr['qid']-1].append(2.355*hdr[key])
                elif keys[key] == 'hdr':
                    vals[key][hdr['ccdid']][hdr['qid']-1].append(hdr[key])
                elif key == 'gfseeing':
                    with open(quadrant_path.joinpath("psf.dat"), 'r') as f:
                        lines = f.readlines()
                        gfseeing = float(lines[3].split(" ")[0].strip())
                    vals[key][hdr['ccdid']][hdr['qid']-1].append(2.355*gfseeing)

        with open(band_path.joinpath("seeings.pickle"), 'wb') as f:
            pickle.dump(vals, f)
    else:
        with open(band_path.joinpath("seeings.pickle"), 'rb') as f:
            vals = pickle.load(f)

    for key in keys:
        for key_i in vals[key].keys():
            for key_j in vals[key][key_i].keys():
                vals[key][key_i][key_j] = np.median(list(map(lambda x: float(x), vals[key][key_i][key_j])))

    cmap = 'coolwarm'
    def _plot_focal_plane_seeing(key):
        val_list = [[vals[key][key_i][key_j] for key_j in vals[key][key_i].keys()] for key_i in vals[key].keys()]
        val_list = list(itertools.chain(*val_list))
        vmin = min(val_list)
        vmax = max(val_list)

        cm = ScalarMappable(cmap=cmap)
        cm.set_clim(vmin=vmin, vmax=vmax)

        fig = plt.figure(figsize=(5., 6.), constrained_layout=False)
        #f1, f2, f3 = fig.subfigures(ncols=1, nrows=3, height_ratios=[12., 1., 1.])
        f1, f2 = fig.subfigures(ncols=1, nrows=2, height_ratios=[12., 0.3], wspace=-1., hspace=-1.)
        f1.suptitle("\n{}-{} Focal plane {}".format(ztfname, filtercode, key), fontsize='large')
        plot_ztf_focal_plan_values(f1, vals[key], vmin=vmin, vmax=vmax, cmap=cmap, scalar=True)
        ax2 = f2.add_subplot(clip_on=False)
        f2.subplots_adjust(bottom=-5, top=5.)
        ax2.set_position((0.2, 0, 0.6, 1))
        # ax2 = f2.add_subplot()
        # ax3 = f3.add_subplot(axis='none')
        # Colorbar(ax, cm, orientation='horizontal')
        cb = f1.colorbar(cm, cax=ax2, orientation='horizontal')
        # ax2.autoscale(tight=True)
        ax2.set_clip_on(False)
        ax2.minorticks_on()
        ax2.tick_params(direction='inout')
        plt.savefig(band_path.joinpath("{}-{}_focal_plane_{}.png".format(ztfname, filtercode, key)), dpi=300., bbox_inches='tight', pad_inches=0.1)
        plt.close()

    for key in keys:
        # print("Plotting {}".format(key))
        _plot_focal_plane_seeing(key)

def match_gaia(quadrant_path, ztfname, filtercode, logger, args):
    from itertools import chain
    from utils import read_list, get_wcs_from_quadrant, get_mjd_from_quadrant_path, contained_in_exposure, match_pixel_space, gaiarefmjd, quadrant_width_px, quadrant_height_px
    import pandas as pd
    import numpy as np
    from astropy.coordinates import SkyCoord
    import astropy.units as u

    if not quadrant_path.joinpath("psfstars.list").exists():
        return False

    _, stars_df = read_list(quadrant_path.joinpath("psfstars.list"))
    _, aper_stars_df = read_list(quadrant_path.joinpath("standalone_stars.list"))
    wcs = get_wcs_from_quadrant(quadrant_path)
    obsmjd = get_mjd_from_quadrant_path(quadrant_path)
    center_radec = wcs.pixel_to_world_values(np.array([[quadrant_width_px/2., quadrant_height_px/2.]]))[0]
    center_skycoord = SkyCoord(ra=center_radec[0], dec=center_radec[1], unit='deg')

    # Match aperture photometry and psf photometry catalogs
    i = match_pixel_space(stars_df[['x', 'y']], aper_stars_df[['x', 'y']], radius=1.)
    stars_df = stars_df.iloc[i[i>=0]].reset_index(drop=True)
    aper_stars_df = aper_stars_df.iloc[i>=0].reset_index(drop=True)

    # Add aperture photometry
    aper_fields = [['apfl'+str(i), 'eapfl'+str(i), 'rad'+str(i)] for i in range(10)]
    aper_fields = list(chain(*aper_fields))
    stars_df[aper_fields] = aper_stars_df[aper_fields]

    gaia_stars_df = pd.read_hdf(args.lc_folder.joinpath("{}.hd5".format(ztfname)), key='gaia_cal')
    gaia_stars_df.rename({'pmde': 'pmdec'}, axis='columns', inplace=True)
    gaia_stars_df['gaiaid'] = gaia_stars_df.index

    # First order proper motion correction for indentification
    # We do it in astrometry code
    # gaia_stars_df['ra'] = gaia_stars_df['ra']+(obsmjd-gaiarefmjd)*gaia_stars_df['pmra']/np.cos(gaia_stars_df['dec']/180.*np.pi)/1000./3600./365.25
    # gaia_stars_df['dec'] = gaia_stars_df['dec']+(obsmjd-gaiarefmjd)*gaia_stars_df['pmde']/1000./3600./365.25
    # mas/year -> deg/day
    gaia_stars_df['pmra'] = gaia_stars_df['pmra']/np.cos(np.deg2rad(gaia_stars_df['dec']))/1000./3600./365.25
    gaia_stars_df['pmdec'] = gaia_stars_df['pmdec']/1000./3600./365.25
    gaia_stars_df.dropna(inplace=True)

    ra_corrected = gaia_stars_df['ra']+(obsmjd-gaiarefmjd)*gaia_stars_df['pmra']
    dec_corrected = gaia_stars_df['dec']+(obsmjd-gaiarefmjd)*gaia_stars_df['pmdec']

    # gaia_stars_radec = SkyCoord(gaia_stars_df['ra'], gaia_stars_df['dec'], unit='deg')
    gaia_stars_skycoords = SkyCoord(ra_corrected, dec_corrected, unit='deg')

    sep = center_skycoord.separation(gaia_stars_skycoords)
    idxc = (sep < 0.5*u.deg)
    gaia_stars_skycoords = gaia_stars_skycoords[idxc]
    gaia_stars_df = gaia_stars_df[idxc]

    #gaia_mask = contained_in_exposure(gaia_stars_radec, wcs, return_mask=True)
    #gaia_stars_df = gaia_stars_df.iloc[gaia_mask]
    #x, y = gaia_stars_radec[gaia_mask].to_pixel(wcs)
    x, y = gaia_stars_skycoords.to_pixel(wcs)
    gaia_stars_df['x'] = x
    gaia_stars_df['y'] = y
    gaia_stars_df = gaia_stars_df.dropna()

    try:
        i = match_pixel_space(gaia_stars_df[['x', 'y']].to_records(), stars_df[['x', 'y']].to_records(), radius=2.)
    except Exception as e:
        logger.error("Could not match any Gaia stars!")
        raise ValueError("Could not match any Gaia stars!")

    matched_gaia_stars_df = gaia_stars_df.iloc[i[i>=0]].reset_index(drop=True)
    matched_stars_df = stars_df.iloc[i>=0].reset_index(drop=True)
    logger.info("Matched {} GAIA stars".format(len(matched_gaia_stars_df)))

    with pd.HDFStore(quadrant_path.joinpath("matched_stars.hd5"), 'w') as hdfstore:
        hdfstore.put('matched_gaia_stars', matched_gaia_stars_df)
        hdfstore.put('matched_stars', matched_stars_df)

    return True

def match_gaia_reduce(band_path, ztfname, filtercode, logger, args):
    import pandas as pd
    import numpy as np
    from utils import get_header_from_quadrant_path
    from deppol_utils import quadrants_from_band_path

    quadrant_paths = quadrants_from_band_path(band_path, logger, check_files="match_gaia.success")

    logger.info("Concatenating matched stars measurements from {} quadrants.".format(len(quadrant_paths)))

    matched_stars_list = []
    quadrants_dict = {}

    for quadrant_path in quadrant_paths:
        logger.info(quadrant_path.name)
        matched_gaia_stars_df = pd.read_hdf(quadrant_path.joinpath("matched_stars.hd5"), key='matched_gaia_stars')
        matched_stars_df = pd.read_hdf(quadrant_path.joinpath("matched_stars.hd5"), key='matched_stars')

        matched_gaia_stars_df.rename(columns={'x': 'gaia_x', 'y': 'gaia_y'}, inplace=True)
        matched_stars_df['mag'] = -2.5*np.log10(matched_stars_df['flux'])
        matched_stars_df['emag'] = matched_stars_df['eflux']/matched_stars_df['flux']

        matched_stars_df = pd.concat([matched_stars_df, matched_gaia_stars_df], axis=1)

        quadrant_dict = {}
        header = get_header_from_quadrant_path(quadrant_path)
        quadrant_dict['quadrant'] = quadrant_path.name
        quadrant_dict['airmass'] = header['airmass']
        quadrant_dict['mjd'] = header['obsmjd']
        quadrant_dict['seeing'] = header['seeing']
        quadrant_dict['ha'] = header['hourangd'] #*15
        quadrant_dict['ha_15'] = 15.*header['hourangd']
        quadrant_dict['lst'] = header['oblst']
        quadrant_dict['azimuth'] = header['azimuth']
        quadrant_dict['dome_azimuth'] = header['dome_az']
        quadrant_dict['elevation'] = header['elvation']
        quadrant_dict['z'] = 90. - header['elvation']
        quadrant_dict['telra'] = header['telrad']
        quadrant_dict['teldec'] = header['teldecd']

        quadrant_dict['field'] = header['dbfield']
        quadrant_dict['ccdid'] = header['ccd_id']
        quadrant_dict['qid'] = header['amp_id']
        quadrant_dict['rcid'] = header['dbrcid']

        quadrant_dict['fid'] = header['dbfid']

        quadrant_dict['temperature'] = header['tempture']
        quadrant_dict['head_temperature'] = header['headtemp']
        quadrant_dict['ccdtemp'] = float(header['ccdtmp{}'.format(str(header['ccd_id']).zfill(2))])

        quadrant_dict['wind_speed'] = header['windspd']
        quadrant_dict['wind_dir'] = header['winddir']
        quadrant_dict['dewpoint'] = header['dewpoint']
        quadrant_dict['humidity'] = header['humidity']
        quadrant_dict['wetness'] = header['wetness']
        quadrant_dict['pressure'] = header['pressure']

        quadrant_dict['crpix1'] = header['crpix1']
        quadrant_dict['crpix2'] = header['crpix2']
        quadrant_dict['crval1'] = header['crval1']
        quadrant_dict['crval2'] = header['crval2']
        quadrant_dict['cd_11'] = header['cd1_1']
        quadrant_dict['cd_12'] = header['cd1_2']
        quadrant_dict['cd_21'] = header['cd2_1']
        quadrant_dict['cd_22'] = header['cd2_2']

        quadrant_dict['skylev'] = header['sexsky']
        quadrant_dict['sigma_skylev'] = header['sexsigma']

        for key in quadrant_dict.keys():
            matched_stars_df[key] = quadrant_dict[key]

        matched_stars_list.append(matched_stars_df)
        quadrants_dict[quadrant_path.name] = quadrant_dict

    matched_stars_df = pd.concat(matched_stars_list, axis=0, ignore_index=True)

    # Build gaia star catalog
    gaia_stars_df = pd.read_hdf(args.lc_folder.joinpath("{}.hd5".format(ztfname)), key='gaia_cal')
    gaia_stars_df.rename({'pmde': 'pmdec'}, axis='columns', inplace=True)
    gaia_stars_df['gaiaid'] = gaia_stars_df.index
    gaia_stars_df = gaia_stars_df[~gaia_stars_df.index.duplicated(keep='first')]

    gaia_stars = []
    for gaiaid in set(matched_stars_df['gaiaid']):
        #gaia_stars.append(pd.Series(gaia_stars_df.drop(['rcid', 'field'], axis='columns').loc[gaiaid], name=gaiaid))
        gaia_stars.append(pd.Series(gaia_stars_df.loc[gaiaid], name=gaiaid))

    gaia_stars_df = pd.DataFrame(data=gaia_stars)
    gaia_stars_df.rename({'pmde': 'pmdec'}, axis='columns', inplace=True)
    gaia_stars_df['pmra'] = gaia_stars_df['pmra']/np.cos(np.deg2rad(gaia_stars_df['dec']))/1000./3600./365.25
    gaia_stars_df['pmdec'] = gaia_stars_df['pmdec']/1000./3600./365.25

    # Remove measures with Nan's
    nan_mask = matched_stars_df.isna().any(axis=1)
    matched_stars_df = matched_stars_df[~nan_mask]
    logger.info("Removed {} measurements with NaN's".format(sum(nan_mask)))

    nan_mask = gaia_stars_df.isna().any(axis=1)
    gaia_stars_df = gaia_stars_df[~nan_mask]
    logger.info("Removed {} Gaia stars with NaN's".format(sum(nan_mask)))

    # Compute color
    matched_stars_df['colormag'] = matched_stars_df['bpmag'] - matched_stars_df['rpmag']
    gaia_stars_df['colormag'] = gaia_stars_df['bpmag'] - gaia_stars_df['rpmag']

    # Save to disk
    matched_stars_df.to_parquet(band_path.joinpath("matched_stars.parquet"))
    logger.info("Total matched Gaia stars: {}".format(len(matched_stars_df)))

    quadrants_df = pd.DataFrame.from_dict(quadrants_dict, orient='index')
    quadrants_df.to_parquet(band_path.joinpath("quadrants.parquet"))

    gaia_stars_df.to_parquet(band_path.joinpath("gaia_stars.parquet"))

    return True

# Extract data from standalone stars and plot several distributions
def stats(quadrant_path, ztfname, filtercode, logger, args):
    import warnings
    import pandas as pd
    from utils import read_list
    from astropy.io import fits

    warnings.filterwarnings('ignore',category=pd.io.pytables.PerformanceWarning)

    import pandas as pd

    def _extract_from_list(list_filename, hdfstore):
        list_path = quadrant_path.joinpath(list_filename).with_suffix(".list")

        if not list_path.exists():
            return False

        with open(list_path, mode='r') as f:
            global_params, df = read_list(f)

        hdfstore.put(list_path.stem, df)
        hdfstore.put("{}_globals".format(list_path.stem), pd.DataFrame([global_params]))

        return True

    with pd.HDFStore(quadrant_path.joinpath("lists.hdf5"), mode='w') as hdfstore:
        # From make_catalog
        _extract_from_list("se", hdfstore)

        # From mkcat2
        cont = _extract_from_list("standalone_stars", hdfstore)

        if not cont:
            return True

        _extract_from_list("aperse", hdfstore)

        # From calibrated.fits
        keywords = ['sexsky', 'sexsigma', 'bscale', 'bzero', 'origsatu', 'saturlev', 'backlev', 'back_sub', 'seseeing', 'gfseeing']

        calibrated = {}
        with fits.open(quadrant_path.joinpath("calibrated.fits")) as hdul:
            for keyword in keywords:
                calibrated[keyword] = hdul[0].header[keyword]

            hdfstore.put('calibrated', pd.DataFrame([calibrated]))

        # From makepsf
        cont = _extract_from_list("psfstars", hdfstore)

        if not cont:
            return True

        _extract_from_list("psftuples", hdfstore)

    return True

def stats_reduce(band_path, ztfname, filtercode, logger, args):
    import pandas as pd
    import matplotlib
    import matplotlib.pyplot as plt
    from deppol_utils import quadrants_from_band_path
    from utils import get_header_from_quadrant_path, ListTable
    matplotlib.use('Agg')

    # Seeing histogram
    folders = [folder for folder in band_path.glob("*") if folder.is_dir()]
    quadrant_paths = quadrants_from_band_path(band_path, logger, ignore_noprocess=True)

    logger.info("Plotting fitted seeing histogram")
    seeings = []
    for quadrant_path in quadrant_paths:
        if quadrant_path.joinpath("calibrated.fits").exists():
           hdr = get_header_from_quadrant_path(quadrant_path)
           seeings.append(hdr['seeing'])

    ax = plt.axes()
    plt.suptitle("Seeing distribution (computed by the ZTF pipeline)")
    #plt.hist(seeings, bins=int(len(seeings)/4), color='xkcd:dark grey', histtype='step')
    plt.hist(seeings, bins='auto', color='xkcd:dark grey', histtype='step')
    if args.max_seeing:
        removed_quadrants = len(list(filter(lambda x: x >= args.max_seeing, seeings)))
        plt.axvline(args.max_seeing)
        plt.text(0.2, 0.9, "{:.02f}% quadrants with seeing > {} ({} quadrants)".format(removed_quadrants/len(seeings)*100, args.max_seeing, len(seeings)), transform=ax.transAxes)
    plt.grid()
    plt.xlabel("Seeing FWHM [pixel]")
    plt.ylabel("Count")
    plt.savefig(band_path.joinpath("{}-{}_seeing_dist.png".format(ztfname, filtercode)), dpi=300)
    plt.close()

    quadrant_paths = quadrants_from_band_path(band_path, logger, check_files=['psfstars.list', 'calibrated.fits'], ignore_noprocess=True)
    logger.info("Plotting psfstars count vs seeing")

    seeings = []
    psfstars_counts = []
    for quadrant_path in quadrant_paths:
        hdr = get_header_from_quadrant_path(quadrant_path)
        seeings.append(hdr['seeing'])

        psfstars_list = ListTable.from_filename(quadrant_path.joinpath("psfstars.list"))
        psfstars_counts.append(len(psfstars_list.df))

    ax = plt.axes()
    plt.plot(seeings, psfstars_counts, '.')

    if args.min_psfstars:
        filtered_psfstars_count = len(list(filter(lambda x: x < args.min_psfstars, psfstars_counts)))
        plt.axhline(args.min_psfstars)
        plt.text(0.2, 0.1, "{:.02f}% quadrants ({}) with PSF stars count < {}".format(filtered_psfstars_count/len(quadrant_paths)*100., filtered_psfstars_count, args.min_psfstars), transform=ax.transAxes)

    if args.max_seeing:
        plt.axvline(args.max_seeing)

    plt.grid()
    plt.xlabel("Seeing FWHM [pixel]")
    plt.ylabel("PSF stars count")
    plt.savefig(band_path.joinpath("{}-{}_psfstars_seeing.png".format(ztfname, filtercode)), dpi=300.)
    plt.close()

    with open(band_path.joinpath("{}-{}_failures.txt".format(ztfname, filtercode)), 'w') as f:
        # Failure rates
        def _failure_rate(listname, func):
            success_count = 0
            for quadrant_path in quadrant_paths:
                if quadrant_path.joinpath("{}.list".format(listname)).exists():
                    success_count += 1

            f.writelines(["For {}:\n".format(func),
                          " Success={}/{}\n".format(success_count, len(folders)),
                          " Rate={}\n\n".format(float(success_count)/len(folders))])

        _failure_rate("se", 'make_catalog')
        _failure_rate("standalone_stars", 'mkcat2')
        _failure_rate("psfstars", 'makepsf')

    logger.info("Plotting computing time histograms")
    # Plot results_*.csv histogram
    result_paths = list(band_path.glob("results_*.csv"))
    for result_path in result_paths:
        func = "_".join(str(result_path.stem).split("_")[1:])
        result_df = pd.read_csv(result_path)
        computation_times = (result_df['time_end'] - result_df['time_start']).to_numpy()
        plt.hist(computation_times, bins=int(len(result_df)/4), histtype='step')
        plt.xlabel("Computation time (s)")
        plt.ylabel("Count")
        plt.title("Computation time for {}".format(func))
        plt.grid()
        plt.savefig(band_path.joinpath("{}-{}_{}_compute_time_dist.png".format(ztfname, filtercode, func)), dpi=300)
        plt.close()

def clean(quadrant_path, ztfname, filtercode, logger, args):
    # We want to delete all files in order to get back to the prepare_deppol stage
    files_to_keep = ["elixir.fits", "dead.fits.gz", ".dbstuff"]
    files_to_delete = list(filter(lambda f: f.name not in files_to_keep, list(quadrant_path.glob("*"))))

    for file_to_delete in files_to_delete:
            file_to_delete.unlink()

    return True

def clean_reduce(band_path, ztfname, filtercode, logger, args):
    from shutil import rmtree

    # We want to delete all files in order to get back to the prepare_deppol stage
    files_to_keep = ["prepare.log"]

    # Delete all files
    files_to_delete = list(filter(lambda f: f.is_file() and (f.name not in files_to_keep), list(band_path.glob("*"))))

    [f.unlink() for f in files_to_delete]

    # Delete output folders
    # We delete everything but the quadrant folders
    folders_to_keep = list(band_path.glob("ztf_*"))
    folders_to_delete = [folder for folder in list(band_path.glob("*")) if (folder not in folders_to_keep) or folder.is_file()]

    [rmtree(band_path.joinpath(folder), ignore_errors=True) for folder in folders_to_delete]

def filter_psfstars_count(band_path, ztfname, filtercode, logger, args):
    from deppol_utils import quadrants_from_band_path, noprocess_quadrants
    from utils import ListTable

    quadrant_paths = quadrants_from_band_path(band_path, logger, ignore_noprocess=True)
    noprocess = noprocess_quadrants(band_path)
    flagged_count = 0
    quadrants_to_flag = []

    for quadrant_path in quadrant_paths:
        if quadrant_path.joinpath("psfstars.list").exists():
            quadrant = quadrant_path.name
            psfstars_list = ListTable.from_filename(quadrant_path.joinpath("psfstars.list"))
            psfstars_count = len(psfstars_list.df)

            if psfstars_count <= args.min_psfstars:
                flagged_count += 1

                if quadrant not in noprocess:
                    quadrants_to_flag.append(quadrant)

    if quadrants_to_flag:
        with open(band_path.joinpath("noprocess"), 'a') as f:
            for quadrant_to_flag in quadrants_to_flag:
                f.write("{}\n".format(quadrant_to_flag))

    flagged_count = len(quadrants_to_flag)
    logger.info("{} quadrants flagged as having PSF stars count <= {}.".format(flagged_count, args.min_psfstars))
    logger.info("{} quadrants added to the noprocess list.".format(len(quadrants_to_flag)))

    return True

def filter_astro_chi2(band_path, ztfname, filtercode, logger, args):
    from deppol_utils import noprocess_quadrants
    import pandas as pd
    import numpy as np

    noprocess = noprocess_quadrants(band_path)

    chi2_path = band_path.joinpath("astrometry/ref2px_chi2_quadrants.csv")
    if not chi2_path.exists():
        logger.info("Could not find file {}!".format(chi2_path))
        return False

    chi2_df = pd.read_csv(chi2_path, index_col=0)

    to_filter = (np.any([chi2_df['chi2'] >= args.astro_max_chi2, chi2_df['chi2'].isna()], axis=0))
    # to_filter = (np.any([chi2_df['chi2'].isna()], axis=0))
    logger.info("{} quadrants flagged as having astrometry chi2 >= {} or NaN values.".format(sum(to_filter), args.astro_max_chi2))
    logger.info("List of filtered quadrants:")
    logger.info(chi2_df.loc[to_filter])

    quadrants_to_flag = 0
    with open(band_path.joinpath("noprocess"), 'a') as f:
        for quadrant in chi2_df.loc[to_filter].index.tolist():
            if quadrant not in noprocess:
                f.write("{}\n".format(quadrant))
                quadrants_to_flag += 1

    logger.info("{} quadrants added to the noprocess list.".format(quadrants_to_flag))

    return True

def filter_seeing(band_path, ztfname, filtercode, logger, args):
    from utils import get_header_from_quadrant_path
    from deppol_utils import quadrants_from_band_path, noprocess_quadrants, run_and_log

    quadrant_paths = quadrants_from_band_path(band_path, logger)
    flagged_count = 0
    quadrants_to_flag = []
    noprocess = noprocess_quadrants(band_path)

    for quadrant_path in quadrant_paths:
        quadrant = quadrant_path.name

        try:
            seeing = get_header_from_quadrant_path(quadrant_path)['SEEING']
        except FileNotFoundError as e:
            logger.error(e)
            logger.error("Folder content:")
            run_and_log(["ls", "-lah", quadrant_path], logger)
            continue
        except KeyError as e:
            logger.error(quadrant)
            logger.error(e)
            quadrants_to_flag.append(quadrant)
            continue

        if seeing > args.max_seeing:
            flagged_count += 1

            if quadrant not in noprocess:
                quadrants_to_flag.append(quadrant)

    if quadrants_to_flag:
        with open(band_path.joinpath("noprocess"), 'a') as f:
            for quadrant_to_flag in quadrants_to_flag:
                f.write("{}\n".format(quadrant_to_flag))

    logger.info("{} quadrants flagged as having seeing > {}.".format(flagged_count, args.max_seeing))
    logger.info("{} quadrants added to the noprocess list.".format(len(quadrants_to_flag)))

    return True

def discard_calibrated(quadrant_path, ztfname, filtercode, logger, args):
    pass

discard_calibrated_rm = ["calibrated.fits", "weight.fz"]