#!/usr/bin/env python3
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np

#throttle

raw = [[920,725, 300,255, 300,255, 300,270, 915,215, 350,205, 920,220, 345,210, 345,210, 925,715, 925,215, 920,215, 345,210, 925,220, 915,215, 345,215, 925,220, 350,210, 345,715, 295,780, 925,205, 920,210, 345],
       [925,725, 295,265, 295,265, 305,265, 920,210, 345,210, 925,215, 350,210, 915,720, 920,215, 350,215, 920,210, 345,210, 915,225, 920,210, 345,210, 915,225, 350,205, 350,715, 295,775, 300,760, 925,715, 925],
       [920,725, 300,255, 300,255, 300,265, 920,215, 350,205, 920,220, 345,715, 295,770, 925,715, 295,270, 925,210, 350,210, 915,225, 920,210, 350,210, 915,225, 350,205, 350,710, 300,270, 925,710, 925,205, 920],
       [920,725, 295,265, 300,255, 300,265, 920,215, 340,215, 920,220, 870,265, 870,765, 295,770, 300,265, 925,205, 350,205, 920,220, 865,265, 350,210, 870,270, 345,215, 345,720, 300,775, 295,770, 300,255, 925],
       [925,720, 300,260, 295,260, 295,270, 920,210, 345,715, 295,775, 300,255, 300,260, 295,765, 920,225, 875,255, 350,715, 295,270, 920,210, 345,210, 865,275, 350,205, 350,715, 295,775, 920,215, 340,215, 350],
       [895,920, 310,320, 365,270, 365,265, 945,270, 365,265, 945,265, 365,265, 945,265, 945,870, 305,325, 895,320, 305,325, 900,315, 895,320, 305,325, 895,320, 305,325, 310,920, 310,920, 310,320, 920,295, 890]] # last fake
'''
CAA BAB AABD BAB BABAAC DBA
CAA BAB ABDA BAB BABAAC CDD
CAA BAB ACDC BAB BABAAC BDB
CAA BAB BBCC BAB BABAAC CCB
CAA BAC CAAD BAC BABAAC DAA
'''
#throttle mittel, rechts links
'''
raw = [[925,720, 300,260, 295,260, 350,215, 350,205, 350,210, 350,220, 350,715, 295,265, 920,215, 875,270, 870,265, 350,210, 870,270, 875,260, 345,210, 865,275, 350,210, 350,715, 295,275, 300,255, 870,765, 920],
       [925,725, 305,255, 300,255, 345,220, 925,210, 350,210, 350,220, 345,715, 295,265, 920,215, 875,265, 870,265, 350,205, 870,270, 875,255, 345,210, 865,275, 350,210, 350,715, 295,270, 920,210, 865,770, 925],
       [930,720, 300,260, 295,260, 355,720, 300,260, 295,765, 920,225, 350,710, 300,260, 875,260, 865,275, 870,260, 345,210, 875,270, 870,265, 350,205, 870,270, 345,215, 350,710, 300,775, 305,255, 870,265, 350]]
CAA AAA ACBB BAB BABAAC ABD
CAA BAA ACBB BAB BABAAC BBD
CAA CAD ACBB BAB BABAAC CBA
'''

# throtthle mittel, vorwärts -> rückwärts
'''
raw = [[925,725, 295,265, 300,255, 350,220, 870,265, 350,205, 350,220, 350,715, 295,265, 875,260, 345,225, 350,205, 350,210, 350,220, 870,265, 350,205, 870,270, 345,215, 350,710, 300,270, 295,260, 875,765, 920],
       [920,730, 300,255, 300,260, 345,220, 875,255, 345,210, 345,220, 345,720, 300,255, 870,265, 350,215, 870,260, 340,215, 875,270, 865,265, 350,210, 870,275, 350,205, 350,715, 295,270, 875,260, 875,765, 295],
       [920,730, 300,285, 275,260, 295,270, 925,210, 300,260, 300,270, 300,765, 300,285, 900,205, 300,770, 295,265, 295,770, 925,220, 915,215, 295,265, 920,220, 295,260, 295,770, 300,775, 300,285, 900,205, 300]]
CAA BAA ACBA AAA BABAAC ABD
CAA BAA ACBA BAB BABAAC BBC
CAA BAA ACBA CAD BABAAC CBA
'''

# Trim, rechts -> links
'''

raw = [[920,225, 345,720, 300,260, 295,270, 875,260, 345,210, 345,220, 345,720, 300,260, 875,255, 350,220, 870,265, 350,210, 870,270, 875,260, 345,210, 875,265, 345,210, 345,715, 305,800, 895,720, 915,720, 300],
       [920,220, 875,760, 295,260, 875,270, 875,260, 345,210, 345,220, 345,715, 300,255, 870,265, 350,215, 870,265, 350,205, 870,270, 875,255, 345,210, 865,275, 350,210, 345,715, 305,800, 270,795, 900,715, 875],
       [925,725, 305,760, 295,795, 900,220, 870,265, 350,210, 350,220, 345,715, 295,265, 870,265, 350,215, 870,260, 345,215, 875,270, 865,265, 350,210, 870,275, 350,205, 350,715, 295,270, 875,765, 870,265, 870]]
ACA BAA ACBA BAB BABAAC DDC
BCB BAA ACBA BAB BABAAC CDD
CCD BAA ACBA BAB BABAAC BDB
'''

# Keine Taste, Linke Taste, Rechte Taste, Beide Tasten
'''
raw = [[925,220, 870,770, 295,260, 920,220, 925,210, 350,210, 345,220, 345,715, 305,255, 920,215, 350,215, 920,215, 350,205, 920,220, 925,210, 345,210, 925,220, 350,210, 345,715, 300,775, 295,770, 925,715, 920],
       [925,220, 870,770, 300,260, 355,215, 870,260, 345,215, 340,225, 355,710, 300,260, 875,260, 345,220, 865,265, 350,210, 345,220, 870,265, 350,210, 870,780, 300,260, 355,710, 300,270, 295,765, 930,710, 925],
       [920,220, 875,765, 295,265, 350,220, 915,215, 350,205, 350,220, 350,715, 295,265, 920,215, 350,215, 920,215, 345,215, 920,220, 925,210, 350,210, 915,225, 920,215, 350,710, 300,775, 930,715, 925,715, 305],
       [925,220, 925,715, 295,260, 355,215, 920,210, 345,215, 355,215, 350,740, 270,260, 920,210, 345,220, 925,210, 345,210, 345,220, 925,210, 345,210, 925,750, 895,210, 345,720, 300,265, 925,715, 925,745, 905]]

BCB BAA ACBA BAB BAB A AC CDD
BCA BAA ACBA BAA BAB C AC ADD
BCA BAA ACBA BAB BAB B AC DDC
BCA BAA ACBA BAA BAB D AC BDD
'''

# Modeltaste: A,B,C, C mit l Taste
'''
raw = [[925,220, 870,770, 300,260, 875,270, 870,265, 350,210, 870,270, 875,260, 345,210, 345,215, 875,270, 875,255, 345,210, 870,275, 870,260, 340,215, 870,270, 345,215, 355,710, 300,270, 295,765, 300,765, 920],
       [925,220, 870,770, 300,260, 350,215, 870,260, 345,215, 875,265, 870,265, 350,210, 350,210, 345,220, 875,255, 345,210, 345,220, 875,255, 350,210, 870,275, 350,210, 870,265, 870,270, 350,710, 925,210, 870],
       [920,225, 875,765, 295,265, 870,275, 875,260, 345,210, 875,265, 870,265, 350,205, 350,210, 870,275, 870,260, 345,215, 875,270, 875,255, 345,210, 870,275, 350,715, 295,265, 300,265, 345,210, 345,210, 875],
       [925,215, 920,745, 275,260, 295,270, 920,210, 300,260, 925,220, 925,205, 300,255, 300,255, 300,270, 925,210, 300,260, 925,215, 920,215, 295,265, 920,730, 300,765, 300,255, 300,775, 380,180, 300,260, 925]]
BCB BAB BAAB BAB BABAAC ACD
BCA BAB BAAA BAA BABABB ADB
BCB BAB BAAB BAB BABACA AAB
BCA BAB BAAA BAB BAB CCA CAB
'''
#TRIM R-L THRO FBW ??? BTNS CHS?

y = []
x = []
symbols = []
bitfields = []
bfMap=['00','01', '10', '11']
symbolMap=['A','B','C','D']

for i,iraw in enumerate(raw):
    yval = 1
    xc = 0
    x.append([0,0])
    y.append([0,1])
    symbol = ""
    bitfield = ""
    for j,l in enumerate(iraw):
        xc += l
        x[i].append(xc)
        y[i].append(yval)
        yval = not yval
        x[i].append(xc)
        y[i].append(int(yval))
        if j > 0 and j%2==0:
            low = iraw[j-1]
            high = iraw[j]
            if low < 450 and high < 450:
                symbol += 'A'
                bitfield += "00"
            elif low < 450 and high >= 450:
                symbol += 'B'
                bitfield += "01"
            elif low >= 450 and high < 450:
                symbol += 'C'
                bitfield += "10"
            elif low >= 450 and high >= 450:
                symbol += 'D'
                bitfield += "11"
            else:
                print("Error in encoding.")
    symbols.append(symbol)
    bitfields.append(bitfield)

xmax = max([item[-1] for item in x])
for i, iraw in enumerate(raw):
    x[i].append(xmax)
    y[i].append(0)

gs = gridspec.GridSpec(len(x),1)
fig = plt.figure()

for i in range(len(raw)):
    s = 6
    e = 10
    #print(symbols[i][s:e])
    sb = bitfields[i][s*2:e*2]
    #print(sb)
    print(int(sb, 2))

# BCA BAB BAAA BAB BABCCA CAB
#TRIM R-L THRO FBW ??? BTNS CHS?

wordsI = [(0,3), (3,6), (6, 10), (10, 13), (13, 16), (16, 19)]
# Checksum
for i in range(len(raw)):
    print("\n", i)
    s = 19
    e = 22
    #print(symbols[i][s:e])
    sb = bitfields[i][s*2:e*2]
    print(sb)
    #print(int(sb, 2))
    #noOnes = 0
    #for bit in bitfields[i]:
    #    if bit == '1':
    #        noOnes += 1
    #print(noOnes)
    words = []
    for tup in wordsI:
        words.append(int(bitfields[i][tup[0]*2:tup[1]*2],2))
    chks = 0
    for w in words:
       print(bin(w))
       chks ^= w
    print(str(bin(chks))[-6:])

for i,l in enumerate(zip(x,y)):
    print(symbols[i])
    print(bitfields[i])
    ax = fig.add_subplot(gs[i])
    zeroline = np.zeros(len(l[0]))
    ax.plot(l[0],l[1])
    ax.fill_between(l[0],l[1],zeroline)
    ax.set_ylabel(str(i), size =16)
    ax.get_yaxis().set_label_coords(-0.1,0.5)
    ax.set_ylim((0, 1.1))
gs.tight_layout(fig)

#plt.show()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File              : /home/wang/processor/s1p/ma_data_proc_leeflt_3.0/python_scripts/global_production.py
# Date              : 10.04.2020 00:18:06
# Last Modified Date: 10.04.2020 00:18:06
# Last Modified By  : Yuanyuan Wang <y.wang@tum.de>
"""
Created on Tue Jun 29 09:06:08 2018
@author: hu jingliang
"""
# Last modified: 19.03.2020 21:01:25 Yuanyuan Wang
# modified GPU index

# Last modified: 10.04.2020 00:18:26 Yuanyuan Wang
# commented out printing messages

# Last modified: 10.04.2020 00:18:48 Yuanyuan Wang
# changed

import os
import sys
import time
import logging
# import tensorflow as tf
import numpy as np
import sen1_cnntrain_uil as sen1cnn
import sen1_production_uil as gp
# import se1Processing_uil as pul

# supress warning messages like SSE4.1, AVX, etc.
os.environ['TF_CPP_MIN_LOG_LEVEL']='1'
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"
os.environ["TF_ENABLE_AUTO_MIXED_PRECISION"] = "1"

# config = tf.config.experimental()
# config.gpu_options.per_process_gpu_memory_fraction = 0.5
# config.gpu_options.allow_growth = True
# session = tf.Session(config=config)

import net_resnet_large_window
import resnet_v2
from keras.optimizers import Nadam
from keras.callbacks import EarlyStopping, ModelCheckpoint

def predict_classes(x):
#generate class prediction from the input samples
    y=x.argmax(axis=1)
    return y

#path2DataOfCity = '/data/hu/test/00005_21206_Mumbai'
path2DataOfCity = sys.argv[1]

#path2NetModel = '/data/hu/so2sat_CNN/RESNET20/models/S1_RESNET20_BS32_LR1e-4_IN32-32-7_PRO52A-R10-GLOBAL_2019_06_22.h5'
path2NetModel = sys.argv[2]

# 1. parameter setting
# feature selection
featSelection = [1,1,1,0,0,0,0,1,1,1,1]
# patch size
patchsize = 32
# patch shape
patch_shape = (32, 32, np.sum(featSelection))
# batch size
batch_size = 64

# 2. read data of a given city, feature extraction, patch cutting, and initial classification map tiff file
# read data path
dataPathTime = sen1cnn.getPathOfTime(path2DataOfCity)
# dataPathTif  = sen1cnn.getPath2Data(dataPathTime[0])
dataPathTif = dataPathTime[0] +'/mosaic.tif'

# initial classification map tiff file
LCZPaths = gp.initialLCZGrids_demo(dataPathTif)
# get patch coordinate
coordCell = gp.getCoordLCZGrid(LCZPaths[0])
coordImage = gp.getImageCoordByXYCoord(coordCell,dataPathTif)
# read data and feature extraction
dataFeature,_,_  = sen1cnn.getDataNormMStd(dataPathTif, featSelection)
# cutting patches
dataPatches = gp.getPatch(dataFeature,coordImage,patchsize)

# 3. loading trained CNN and predict label
#model = net_resnet_large_window.build_model(patch_shape, nb_classes = 17 , start_ch = 16, depth = 4, inc_rate = 2, nb_skipped = 2, activation = 'relu')
model = resnet_v2.resnet_v2_s1(input_shape=patch_shape,depth=20, num_classes = 17)
timeStart = time.time()
model.load_weights(path2NetModel)
timeEnd = time.time()
#print('INFO:    TIME SPENT FOR LOADING TRAINED CNN MODEL IS: '+str(timeEnd-timeStart))
timeStart = time.time()
probPred = model.predict(dataPatches,batch_size = batch_size)
timeEnd = time.time()
print('INFO:    Time for infereencing '+str(timeEnd-timeStart))
#print(np.min(probPred))
#print(np.median(probPred))
#print(np.max(probPred))

labelPred = predict_classes(probPred)

# 4. save predicted probability and label
probPath = gp.saveProbabilityPrediction(probPred,LCZPaths[1])
labelPath = gp.saveLabelPrediction(labelPred,LCZPaths[0])

#print('9')
# 5. write ok file
try:
    #print('Try writing OK files')
    #os.system("printf '%s' \"0\" > OK.lcz_classification")
    with open('OK.lcz_classification', 'w') as out:
        out.write("0")
    #print('Write OK file finished')

except RuntimeError:
    print('ERROR:   Error in writing OK.lcz_classification.')
    print(e)

#print('10')


from base_lot import *
import numpy as np
import manage_xyz as mx
from units import *
import os
import re

class Molpro(Lot):

    # designed to do multiple multiplicities at once... maybe not such a good idea, that feature is currently broken
    #TODO
    def run(self,geom):
        if self.has_nelectrons==False:
            for i in self.states:
                self.get_nelec(geom,i[0])
            self.has_nelectrons==True

        #TODO gopro needs a number
        tempfilename = 'scratch/gopro.com'
        tempfile = open(tempfilename,'w')
        tempfile.write(' file,2,mp_0000_{:03d}\n'.format(self.node_id))
        tempfile.write(' memory,800,m\n')
        tempfile.write(' symmetry,nosym\n')
        tempfile.write(' orient,noorient\n\n')
        tempfile.write(' geometry={\n')
        for coord in geom:
            for i in coord:
                tempfile.write(' '+str(i))
            tempfile.write('\n')
        tempfile.write('}\n\n')
        singlets=self.search_tuple(self.states,1)
        len_singlets=len(singlets)
        len_E_singlets=singlets[-1][1] +1 #final adiabat +1 because 0 indexed
        triplets=self.search_tuple(self.states,3)
        len_triplets=len(triplets)

        if self.lot_inp_file == False:
            tempfile.write(' basis={}\n\n'.format(self.basis))
            if len_singlets is not 0:
                tempfile.write(' {multi\n')
                nclosed = self.nocc
                nocc = nclosed+self.nactive
                tempfile.write(' direct\n')
                tempfile.write(' closed,{}\n'.format(nclosed))
                tempfile.write(' occ,{}\n'.format(nocc))
                tempfile.write(' wf,{},1,0\n'.format(self.n_electrons))
                #this can be made the len of singlets
                tempfile.write(' state,{}\n'.format(len_singlets))

                for state in singlets:
                    s=state[1]
                    grad_name="510"+str(s)+".1"
                    tempfile.write(' CPMCSCF,GRAD,{}.1,record={}\n'.format(s+1,grad_name))

                #TODO this can only do coupling if states is 2, want to generalize to 3 states
                if self.do_coupling==True and len(singlets)==2:
                    tempfile.write('CPMCSCF,NACM,{}.1,{}.1,record=5200.1\n'.format(singlets[0][1]+1,singlets[1][1]+1))
                tempfile.write(' }\n')

                for state in singlets:
                    s=state[1]
                    grad_name="510"+str(s)+".1"
                    tempfile.write('Force;SAMC,{};varsav\n'.format(grad_name))
                if self.do_coupling==True and len(singlets)==2:
                    tempfile.write('Force;SAMC,5200.1;varsav\n')
            if len_triplets is not 0:
                tempfile.write(' {multi\n')
                nclosed = self.nocc
                nocc = nclosed+self.nactive
                tempfile.write(' closed,{}\n'.format(nclosed))
                tempfile.write(' occ,{}\n'.format(nocc))
                tempfile.write(' wf,{},1,2\n'.format(self.n_electrons))
                nstates = len(self.states)
                tempfile.write(' state,{}\n'.format(len_triplets))

                for state in triplets:
                    s=state[1]
                    grad_name="511"+str(s)+".1"
                    tempfile.write(' CPMCSCF,GRAD,{}.1,record={}\n'.format(s+1,grad_name))
                tempfile.write(' }\n')

                for s in triplets:
                    s=state[1]
                    grad_name="511"+str(s)+".1"
                    tempfile.write('Force;SAMC,{};varsav\n'.format(grad_name))
        else:
            with open(self.lot_inp_file) as lot_inp:
                lot_inp_lines = lot_inp.readlines()
            for line in lot_inp_lines:
                #print line
                tempfile.write(line)

        tempfile.close()

        cmd = "molpro -W scratch -n {} {} --no-xml-output".format(self.nproc,tempfilename)
        os.system(cmd)

        tempfileout='scratch/gopro.out'
        pattern = re.compile(r'MCSCF STATE \d.1 Energy \s* ([-+]?[0-9]*\.?[0-9]+)')
        self.E = []
        tmp =[]
        for i,line in enumerate(open(tempfileout)):
            for match in re.finditer(pattern,line):
                tmp.append(float(match.group(1)))

        for i in range(len_E_singlets):
            self.E.append((1,i,tmp[i]))
        for i in range(len_triplets): #triplets broken here
            self.E.append((3,i,tmp[len_singlets+i]))

        tmpgrada=[]
        tmpgrad=[]
        self.grada=[]
        self.coup=[]
        with open(tempfileout,"r") as f:
            for line in f:
                if line.startswith("GRADIENT FOR STATE",7): #will work for SA-MC and RSPT2 HF
                    for i in range(3):
                        next(f)
                    for i in range(len(geom)):
                        findline = next(f,'').strip()
                        mobj = re.match(r'^\s*(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s*$', findline)
                        tmpgrad.append([
                            float(mobj.group(2)),
                            float(mobj.group(3)),
                            float(mobj.group(4)),
                            ])
                    tmpgrada.append(tmpgrad)
                    tmpgrad = []
                if line.startswith(" SA-MC NACME FOR STATES"):
                    for i in range(3):
                        next(f)
                    for i in range(len(geom)):
                        findline = next(f,'').strip()
                        mobj = re.match(r'^\s*(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s*$', findline)
                        self.coup.append([
                            float(mobj.group(2)),
                            float(mobj.group(3)),
                            float(mobj.group(4)),
                            ])

        for count,i in enumerate(self.states):
            if i[0]==1:
                self.grada.append((1,i[1],tmpgrada[count]))
            if i[0]==3:
                self.grada.append((3,i[1],tmpgrada[count]))
        self.hasRanForCurrentCoords=True
        return

    def get_energy(self,coords,multiplicity,state):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).any():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        return self.search_PES_tuple(self.E,multiplicity,state)[0][2]*KCAL_MOL_PER_AU

    def get_gradient(self,coords,multiplicity,state):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).any():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        tmp = self.search_PES_tuple(self.grada,multiplicity,state)[0][2]
        return np.asarray(tmp)*ANGSTROM_TO_AU

    def get_coupling(self,coords,multiplicity,state1,state2):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).any():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        return np.reshape(self.coup,(3*len(self.coup),1))*ANGSTROM_TO_AU

    #TODO molpro requires extra things when copying. . . can this be done in the base_lot?
    # e.g. if cls=="Molpro": #do molpro stuff?
    @classmethod
    def copy(cls,lot,options):
         """ create a copy of this lot object"""
        #print(" creating copy, new node id =",node_id)
        #print(" old node id = ",self.node_id)
        if node_id != self.node_id:
            cmd = "cp scratch/mp_0000_{:03d} scratch/mp_0000_{:03d}".format(self.node_id,node_id)
            print(" ",cmd)
            os.system(cmd)
        return cls(lot.options.copy().set_values(options))


import logging
import copy
import numpy as np
import pandas as pd
import os
import time

from supervised.algorithms.algorithm import BaseAlgorithm
from supervised.algorithms.registry import AlgorithmsRegistry
from supervised.algorithms.registry import (
    BINARY_CLASSIFICATION,
    MULTICLASS_CLASSIFICATION,
    REGRESSION,
)
from supervised.preprocessing.preprocessing_utils import PreprocessingUtils
from supervised.utils.metric import (
    CatBoostEvalMetricSpearman,
    CatBoostEvalMetricPearson,
    CatBoostEvalMetricAveragePrecision,
    CatBoostEvalMetricMSE,
    CatBoostEvalMetricUserDefined,
)

from supervised.utils.config import LOG_LEVEL

logger = logging.getLogger(__name__)
logger.setLevel(LOG_LEVEL)

from catboost import CatBoostClassifier, CatBoostRegressor, CatBoost, Pool
import catboost

def catboost_eval_metric(ml_task, eval_metric):
    if eval_metric == "user_defined_metric":
        return eval_metric
    metric_name_mapping = {
        BINARY_CLASSIFICATION: {
            "auc": "AUC",
            "logloss": "Logloss",
            "f1": "F1",
            "average_precision": "average_precision",
            "accuracy": "Accuracy",
        },
        MULTICLASS_CLASSIFICATION: {
            "logloss": "MultiClass",
            "f1": "TotalF1:average=Micro",
            "accuracy": "Accuracy",
        },
        REGRESSION: {
            "rmse": "RMSE",
            "mse": "mse",
            "mae": "MAE",
            "mape": "MAPE",
            "r2": "R2",
            "spearman": "spearman",
            "pearson": "pearson",
        },
    }
    return metric_name_mapping[ml_task][eval_metric]

def catboost_objective(ml_task, eval_metric):
    objective = "RMSE"
    if ml_task == BINARY_CLASSIFICATION:
        objective = "Logloss"
    elif ml_task == MULTICLASS_CLASSIFICATION:
        objective = "MultiClass"
    else:  # ml_task == REGRESSION
        objective = catboost_eval_metric(REGRESSION, eval_metric)
        if objective in [
            "mse",
            "R2",
            "spearman",
            "pearson",
            "user_defined_metric",
        ]:  # cant optimize them directly
            objective = "RMSE"
    return objective

class CatBoostAlgorithm(BaseAlgorithm):

    algorithm_name = "CatBoost"
    algorithm_short_name = "CatBoost"
    warmup_iterations = 20

    def __init__(self, params):
        super(CatBoostAlgorithm, self).__init__(params)
        self.library_version = catboost.__version__
        self.snapshot_file_path = "training_snapshot"

        self.explain_level = params.get("explain_level", 0)
        self.rounds = additional.get("max_rounds", 10000)
        self.max_iters = 1
        self.early_stopping_rounds = additional.get("early_stopping_rounds", 50)

        Algo = CatBoostClassifier
        loss_function = "Logloss"
        if self.params["ml_task"] == BINARY_CLASSIFICATION:
            loss_function = self.params.get("loss_function", "Logloss")
        elif self.params["ml_task"] == MULTICLASS_CLASSIFICATION:
            loss_function = self.params.get("loss_function", "MultiClass")
        elif self.params["ml_task"] == REGRESSION:
            loss_function = self.params.get("loss_function", "RMSE")
            Algo = CatBoostRegressor

        cat_params = {
            "iterations": self.params.get("num_boost_round", self.rounds),
            "learning_rate": self.params.get("learning_rate", 0.1),
            "depth": self.params.get("depth", 3),
            "rsm": self.params.get("rsm", 1.0),
            "l2_leaf_reg": self.params.get("l2_leaf_reg", 3.0),
            "random_strength": self.params.get("random_strength", 1.0),
            "loss_function": loss_function,
            "eval_metric": self.params.get("eval_metric", loss_function),
            # "custom_metric": self.params.get("eval_metric", loss_function),
            "thread_count": self.params.get("n_jobs", -1),
            "verbose": False,
            "allow_writing_files": False,
            "random_seed": self.params.get("seed", 1),
        }

        for extra_param in [
            "min_data_in_leaf",
            "bootstrap_type",
            "bagging_temperature",
            "subsample",
            "border_count",
        ]:
            if extra_param in self.params:
                cat_params[extra_param] = self.params[extra_param]

        self.log_metric_name = cat_params["eval_metric"]
        if cat_params["eval_metric"] == "spearman":
            cat_params["eval_metric"] = CatBoostEvalMetricSpearman()
            self.log_metric_name = "CatBoostEvalMetricSpearman"
        elif cat_params["eval_metric"] == "pearson":
            cat_params["eval_metric"] = CatBoostEvalMetricPearson()
            self.log_metric_name = "CatBoostEvalMetricPearson"
        elif cat_params["eval_metric"] == "average_precision":
            cat_params["eval_metric"] = CatBoostEvalMetricAveragePrecision()
            self.log_metric_name = "CatBoostEvalMetricAveragePrecision"
        elif cat_params["eval_metric"] == "mse":
            cat_params["eval_metric"] = CatBoostEvalMetricMSE()
            self.log_metric_name = "CatBoostEvalMetricMSE"
        elif cat_params["eval_metric"] == "user_defined_metric":
            cat_params["eval_metric"] = CatBoostEvalMetricUserDefined()
            self.log_metric_name = "CatBoostEvalMetricUserDefined"

        self.model = Algo(**cat_params)
        self.cat_features = None
        self.best_ntree_limit = 0

        logger.debug("CatBoostAlgorithm.__init__")

    def _assess_iterations(self, X, y, sample_weight, eval_set, max_time=None):
        if max_time is None:
            max_time = 3600
        try:
            model = copy.deepcopy(self.model)
            model.set_params(iterations=self.warmup_iterations)
            start_time = time.time()
            model.fit(
                X,
                y,
                sample_weight=sample_weight,
                cat_features=self.cat_features,
                init_model=None if self.model.tree_count_ is None else self.model,
                eval_set=eval_set,
                early_stopping_rounds=self.early_stopping_rounds,
                verbose_eval=False,
            )
            elapsed_time = (time.time() - start_time) / float(self.warmup_iterations)
            # print(max_time, elapsed_time, max_time / elapsed_time, np.round(time.time() - start_time, 2))
            new_rounds = int(min(10000, max_time / elapsed_time))
            new_rounds = max(new_rounds, 10)
            return model, new_rounds
        except Exception as e:
            # print(str(e))
            return None, 1000

    def fit(
        self,
        X,
        y,
        sample_weight=None,
        X_validation=None,
        y_validation=None,
        sample_weight_validation=None,
        log_to_file=None,
        max_time=None,
    ):
        if self.is_fitted():
            print("CatBoost model already fitted. Skip fit().")
            return

        if self.cat_features is None:
            self.cat_features = []
            for i in range(X.shape[1]):
                if PreprocessingUtils.is_categorical(X.iloc[:, i]):
                    self.cat_features += [i]
                    X.iloc[:, i] = X.iloc[:, i].astype(str)
                    if X_validation is not None:
                        X_validation.iloc[:, i] = X_validation.iloc[:, i].astype(str)

        eval_set = None
        if X_validation is not None and y_validation is not None:
            eval_set = Pool(
                data=X_validation,
                label=y_validation,
                cat_features=self.cat_features,
                weight=sample_weight_validation,
            )

        if self.params.get("num_boost_round") is None:
            model_init, new_iterations = self._assess_iterations(
                X, y, sample_weight, eval_set, max_time
            )
            self.model.set_params(iterations=new_iterations)
        else:
            model_init = None
            self.model.set_params(iterations=self.params.get("num_boost_round"))
            self.early_stopping_rounds = self.params.get("early_stopping_rounds", 50)

        self.model.fit(
            X,
            y,
            sample_weight=sample_weight,
            cat_features=self.cat_features,
            init_model=model_init,
            eval_set=eval_set,
            early_stopping_rounds=self.early_stopping_rounds,
            verbose_eval=False,
        )

        if self.model.best_iteration_ is not None:
            if model_init is not None:
                self.best_ntree_limit = (
                    self.model.best_iteration_ + model_init.tree_count_ + 1
                )
            else:
                self.best_ntree_limit = self.model.best_iteration_ + 1

        else:
            # just take all the trees
            # the warm-up trees are already included
            # dont need to add +1
            self.best_ntree_limit = self.model.tree_count_

        if log_to_file is not None:
            train_scores = self.model.evals_result_["learn"].get(self.log_metric_name)
            validation_scores = self.model.evals_result_["validation"].get(
                self.log_metric_name
            )
            if model_init is not None:
                if train_scores is not None:
                    train_scores = (
                        model_init.evals_result_["learn"].get(self.log_metric_name)
                        + train_scores
                    )
                if validation_scores is not None:
                    validation_scores = (
                        model_init.evals_result_["validation"].get(self.log_metric_name)
                        + validation_scores
                    )
            iteration = None
            if train_scores is not None:
                iteration = range(len(validation_scores))
            elif validation_scores is not None:
                iteration = range(len(validation_scores))

            result = pd.DataFrame(
                {
                    "iteration": iteration,
                    "train": train_scores,
                    "validation": validation_scores,
                }
            )
            result.to_csv(log_to_file, index=False, header=False)

    def is_fitted(self):
        return self.model is not None and self.model.tree_count_ is not None

    def predict(self, X):
        self.reload()
        if self.params["ml_task"] == BINARY_CLASSIFICATION:
            return self.model.predict_proba(X, ntree_end=self.best_ntree_limit)[:, 1]
        elif self.params["ml_task"] == MULTICLASS_CLASSIFICATION:
            return self.model.predict_proba(X, ntree_end=self.best_ntree_limit)

        return self.model.predict(X, ntree_end=self.best_ntree_limit)

    def copy(self):
        return copy.deepcopy(self)

    def save(self, model_file_path):
        self.model.save_model(model_file_path)
        self.model_file_path = model_file_path
        logger.debug("CatBoostAlgorithm save model to %s" % model_file_path)

    def load(self, model_file_path):
        logger.debug("CatBoostLearner load model from %s" % model_file_path)

        # waiting for fix https://github.com/catboost/catboost/issues/696
        Algo = CatBoostClassifier
        if self.params["ml_task"] == REGRESSION:
            Algo = CatBoostRegressor

        # loading might throw warnings in the case of custom eval_metric
        # check https://github.com/catboost/catboost/issues/1169
        self.model = Algo().load_model(model_file_path)
        self.model_file_path = model_file_path

    def file_extension(self):
        return "catboost"

    def get_metric_name(self):
        metric = self.params.get("eval_metric")
        if metric is None:
            return None
        if metric == "Logloss":
            return "logloss"
        elif metric == "AUC":
            return "auc"
        elif metric == "MultiClass":
            return "logloss"
        elif metric == "RMSE":
            return "rmse"
        elif metric == "MSE":
            return "mse"
        elif metric == "MAE":
            return "mae"
        elif metric == "MAPE":
            return "mape"
        elif metric in ["F1", "TotalF1:average=Micro"]:
            return "f1"
        elif metric == "Accuracy":
            return "accuracy"
        return metric

classification_params = {
    "learning_rate": [0.025, 0.05, 0.1, 0.2],
    "depth": [4, 5, 6, 7, 8, 9],
    "rsm": [0.7, 0.8, 0.9, 1],  # random subspace method
    "loss_function": ["Logloss"],
}

classification_default_params = {
    "learning_rate": 0.1,
    "depth": 6,
    "rsm": 1,
    "loss_function": "Logloss",
}

additional = {
    "max_rounds": 10000,
    "early_stopping_rounds": 50,
    "max_rows_limit": None,
    "max_cols_limit": None,
}
required_preprocessing = [
    "missing_values_inputation",
    "datetime_transform",
    "text_transform",
    "target_as_integer",
]

AlgorithmsRegistry.add(
    BINARY_CLASSIFICATION,
    CatBoostAlgorithm,
    classification_params,
    required_preprocessing,
    additional,
    classification_default_params,
)

multiclass_classification_params = copy.deepcopy(classification_params)
multiclass_classification_params["loss_function"] = ["MultiClass"]
multiclass_classification_params["depth"] = [3, 4, 5, 6]
multiclass_classification_params["learning_rate"] = [0.1, 0.15, 0.2]

multiclass_classification_default_params = copy.deepcopy(classification_default_params)
multiclass_classification_default_params["loss_function"] = "MultiClass"
multiclass_classification_default_params["depth"] = 5
multiclass_classification_default_params["learning_rate"] = 0.15

AlgorithmsRegistry.add(
    MULTICLASS_CLASSIFICATION,
    CatBoostAlgorithm,
    multiclass_classification_params,
    required_preprocessing,
    additional,
    multiclass_classification_default_params,
)

regression_params = copy.deepcopy(classification_params)
regression_params["loss_function"] = ["RMSE", "MAE", "MAPE"]

regression_required_preprocessing = [
    "missing_values_inputation",
    "datetime_transform",
    "text_transform",
    "target_scale",
]

regression_default_params = {
    "learning_rate": 0.1,
    "depth": 6,
    "rsm": 1,
    "loss_function": "RMSE",
}

AlgorithmsRegistry.add(
    REGRESSION,
    CatBoostAlgorithm,
    regression_params,
    regression_required_preprocessing,
    additional,
    regression_default_params,
)


# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
# usage:
# python GEN_WFLW.py
# dataset details in https://wywu.github.io/projects/LAB/WFLW.html
####################################################################
import os, sys, glob, torch, copy
from os import path as osp
from collections import OrderedDict, defaultdict
from scipy.io import loadmat
import numpy as np
from tqdm import tqdm
from pathlib import Path
lib_dir = (Path(__file__).parent / '..' / 'lib').resolve()
if str(lib_dir) not in sys.path: sys.path.insert(0, str(lib_dir))
assert sys.version_info.major == 3, 'Please upgrade from {:} to Python 3.x'.format(sys.version_info)
print ('lib-dir : {:}'.format(lib_dir))
import datasets
from utils.file_utils import load_list_from_folders, load_txt_file
from xvision import to_tensor
from xvision import draw_points, normalize_L

class FACE:
  def __init__(self, string):
    values = string.split(' ')
    assert len(values) == 196 + 4 + 6 + 1, 'invalid length : {:}'.format(len(values))
    coordinates = [float(x) for x in values[:196]]
    coordinates = np.array( coordinates )
    X, Y = coordinates[::2], coordinates[1::2]
    xbox = [float(x) for x in values[196:200]]
    [pose, expression, illumination, makeup, occlusion, blur] = [int(x) for x in values[200:206]]
    name = values[-1]
    self.landmarks = np.concatenate((np.expand_dims(X,axis=1), np.expand_dims(Y,axis=1), np.ones((98,1))), axis=1)
    self.xbox = xbox
    self.name = name

def return_box(landmarks, np_boxes, USE_BOX):
  if USE_BOX == 'GTL':
    box = datasets.dataset_utils.PTSconvert2box(landmarks.copy().T)
    box = (max(0, box[0]-1), max(0, box[1]-1), box[2]+1, box[3]+1)
  elif USE_BOX == 'GTB':
    box = (float(np_boxes[0]), float(np_boxes[1]), float(np_boxes[2]), float(np_boxes[3]))
  else:
    raise ValueError('The box indicator not find : {:}'.format(USE_BOX))
  return box

def Generate_WFLW_LIST(root, save_dir):
  assert osp.isdir(root), '{:} is not dir'.format(root)
  #assert osp.isdir(save_dir), '{} is not dir'.format(save_dir)
  if not osp.isdir(save_dir): os.makedirs(save_dir)
  image_dir = osp.join(root, 'WFLW_images')

  train_list = osp.join(root, 'WFLW_annotations', 'list_98pt_rect_attr_train_test', 'list_98pt_rect_attr_train.txt')
  with open(train_list) as f:
    content = f.readlines()
    train_faces = [FACE( x.strip() ) for x in content]
  assert len(train_faces) == 7500, 'invalid number of training faces : {:}'.format( len(train_faces) )

  SUFFIXS = ['GTB', 'GTL']
  mean_landmark = {SUFFIX : [[]for i in range(98)] for SUFFIX in SUFFIXS}

  trainData = [] # OrderedDict()
  for face in train_faces:
    landmarks = face.landmarks
    data = {'points': landmarks.T, 'name'  : face.name}
    for SUFFIX in SUFFIXS:
      box = return_box(landmarks, face.xbox, SUFFIX)
      data['box-{:}'.format(SUFFIX)] = box
      for idx in range(98):
        if int(landmarks[idx,2] + 0.5) == 0: continue
        x, y = float(landmarks[idx,0]-box[0]), float(landmarks[idx,1]-box[1])
        x, y = normalize_L(x, box[2]-box[0]), normalize_L(y, box[3]-box[1])
        mean_landmark[SUFFIX][idx].append( (x,y) )
    data['box-default']    = data['box-GTB']
    data['previous_frame'] = None
    data['current_frame']  = osp.join(image_dir, face.name)
    data['next_frame']     = None
    assert osp.isfile( data['current_frame'] )
    trainData.append( data )
  torch.save(trainData, osp.join(save_dir, 'train.pth'))
  for SUFFIX in SUFFIXS:
    allp = mean_landmark[SUFFIX]
    allp = np.array(allp)
    mean_landmark[SUFFIX] = np.mean(allp, axis=1)
    mean_landmark[SUFFIX] = mean_landmark[SUFFIX] * 0.9
    image = draw_points(mean_landmark[SUFFIX], 600, 500, True)
    image.save(osp.join(save_dir, 'train-{:}.png'.format(SUFFIX)))
  mean_landmark['default'] = mean_landmark['GTB']
  torch.save(mean_landmark, osp.join(save_dir, 'train-mean.pth'))
  print ('Training Set   : {:5d} facial images.'.format(len(trainData)))

  test_list = osp.join(root, 'WFLW_annotations', 'list_98pt_rect_attr_train_test', 'list_98pt_rect_attr_test.txt')
  with open(test_list) as f:
    content = f.readlines()
    test_faces = [FACE( x.strip() ) for x in content]
  assert len(test_faces) == 2500, 'invalid number of training faces : {:}'.format( len(train_faces) )
  testData = []
  for face in test_faces:
    landmarks = face.landmarks
    data = {'points': landmarks.T, 'name'  : face.name}
    for SUFFIX in SUFFIXS:
      box = return_box(landmarks, face.xbox, SUFFIX)
      data['box-{:}'.format(SUFFIX)] = box
    data['box-default']    = data['box-GTB']
    data['previous_frame'] = None
    data['current_frame']  = osp.join(image_dir, face.name)
    data['next_frame']     = None
    assert osp.isfile( data['current_frame'] )
    testData.append( data )
  torch.save(testData, osp.join(save_dir, 'test.pth'))
  print ('Test Set   : {:5d} facial images.'.format(len(testData)))
  print ('Save all dataset files into {:}'.format(save_dir))

if __name__ == '__main__':
  HOME_STR = 'DOME_HOME'
  if HOME_STR not in os.environ: HOME_STR = 'HOME'
  assert HOME_STR in os.environ, 'Doest not find the HOME dir : {}'.format(HOME_STR)
  this_dir = osp.dirname(os.path.abspath(__file__))
  SAVE_DIR = osp.join(this_dir, 'lists', 'WFLW')
  print ('This dir : {:}, HOME : [{:}] : {:}'.format(this_dir, HOME_STR, os.environ[HOME_STR]))
  root_path = osp.join( os.environ[HOME_STR], 'datasets', 'landmark-datasets', 'WFLW')

  Generate_WFLW_LIST(root_path, SAVE_DIR)


# -*- coding: utf-8 -*-
from csv import excel
import shutil
from functools import lru_cache
from typing import List, Tuple, Union

import numpy as np
import wrapt_timeout_decorator
from ase import Atoms
from ase.calculators.bond_polarizability import BondPolarizability
from ase.vibrations import Infrared
from ase.vibrations.placzek import PlaczekStatic
from ase.vibrations.raman import StaticRamanCalculator
from fastapi.logger import logger
from rdkit import Chem
from scipy import spatial
from xtb.ase.calculator import XTB
from math import pi, sqrt, log
from .cache import ir_cache, ir_from_molfile_cache, ir_from_smiles_cache
from .models import IRResult
from .optimize import run_xtb_opt
from .settings import IMAGINARY_FREQ_THRESHOLD, MAX_ATOMS_FF, MAX_ATOMS_XTB, TIMEOUT
from .utils import (
    get_moments_of_inertia,
    hash_atoms,
    hash_object,
    molfile2ase,
    smiles2ase,
)

def get_max_atoms(method):
    if method == "GFNFF":
        return MAX_ATOMS_FF
    elif method == "GFN2xTB":
        return MAX_ATOMS_XTB
    elif method == "GFN1xTB":
        return MAX_ATOMS_XTB

def ir_hash(atoms, method):
    return hash_object(str(hash_atoms(atoms)) + method)

def get_raman_spectrum(
    pz,
    modes,
    start=0,
    end=4000,
    npts=None,
    width=4,
    type="Gaussian",
    method="standard",
    direction="central",
    normalize=False,
):
    """Get spectrum.

    The method returns wavenumbers in cm^-1 with corresponding
    absolute infrared intensity.
    Start and end point, and width of the Gaussian/Lorentzian should
    be given in cm^-1.
    normalize=True ensures the integral over the peaks to give the
    intensity.
    """
    frequencies = pz.vibrations.get_frequencies(method, direction)
    intensities = [mode['ramanIntensity'] for mode in modes]
    return get_spectrum(modes, frequencies, intensities, start, end, npts, width, type, normalize)

def get_ir_spectrum(ir, modes, start=0,
    end=4000,
    npts=None,
    width=4,
    type="Gaussian",
    method="standard",
    direction="central",
    normalize=False,):
    """Get spectrum.

    The method returns wavenumbers in cm^-1 with corresponding
    absolute infrared intensity.
    Start and end point, and width of the Gaussian/Lorentzian should
    be given in cm^-1.
    normalize=True ensures the integral over the peaks to give the
    intensity.
    """
    frequencies = ir.get_frequencies(method, direction)
    intensities = [mode['intensity'] for mode in modes]
    return get_spectrum(modes, frequencies, intensities, start, end, npts, width, type, normalize)

def fold(frequencies, intensities,
            start=800.0, end=4000.0, npts=None, width=4.0,
            type='Gaussian', normalize=False):
    """Fold frequencies and intensities within the given range
    and folding method (Gaussian/Lorentzian).
    The energy unit is cm^-1.
    normalize=True ensures the integral over the peaks to give the
    intensity.
    """

    lctype = type.lower()
    assert lctype in ['gaussian', 'lorentzian']
    if not npts:
        npts = int((end - start) / width * 10 + 1)
    prefactor = 1
    if lctype == 'lorentzian':
        intensities = intensities * width * pi / 2.
        if normalize:
            prefactor = 2. / width / pi
    else:
        sigma = width / 2. / sqrt(2. * log(2.))
        if normalize:
            prefactor = 1. / sigma / sqrt(2 * pi)

    # Make array with spectrum data
    spectrum = np.empty(npts)
    energies = np.linspace(start, end, npts)
    for i, energy in enumerate(energies):
        energies[i] = energy
        if lctype == 'lorentzian':
            spectrum[i] = (intensities * 0.5 * width / pi /
                            ((frequencies - energy)**2 +
                            0.25 * width**2)).sum()
        else:
            spectrum[i] = (intensities *
                            np.exp(-(frequencies - energy)**2 /
                                    2. / sigma**2)).sum()
    return [energies, prefactor * spectrum]

def get_spectrum(modes, frequencies, intensities,     start=0,
    end=4000,
    npts=None,
    width=4,
    type="Gaussian",

    normalize=False):

    filtered_frequencies, filtered_intensities = [], []

    for freq, int, mode in zip(frequencies, intensities, modes):
        if mode['modeType'] == 'vibration':
            filtered_frequencies.append(freq.real)
            filtered_intensities.append(int)
    return fold(
        filtered_frequencies, filtered_intensities, start, end, npts, width, type, normalize
    )

def run_xtb_ir(
    atoms: Atoms, method: str = "GFNFF", mol: Union[None, Chem.Mol] = None
) -> IRResult:
    if mol is None:
        raise Exception

    this_hash = ir_hash(atoms, method)
    logger.debug(f"Running IR for {this_hash}")
    moi = atoms.get_moments_of_inertia()
    linear = sum(moi > 0.01) == 2
    result = ir_cache.get(this_hash)

    if result is None:
        logger.debug(f"IR not in cache for {this_hash}, running")
        atoms.pbc = False
        atoms.calc = XTB(method=method)

        try:
            rm = StaticRamanCalculator(atoms, BondPolarizability, name=str(this_hash))
            rm.ir = True
            rm.run()
            pz = PlaczekStatic(atoms, name=str(this_hash))
            raman_intensities = pz.get_absolute_intensities()

        except Exception as e:
            print(e)
            shutil.rmtree(str(this_hash))
            raman_intensities = None
            raman_spectrum = None

        ir = Infrared(atoms, name=str(this_hash))
        ir.run()

        zpe = ir.get_zero_point_energy()
        most_relevant_mode_for_bond = None
        bond_displacements = None
        if mol is not None:
            bond_displacements = compile_all_bond_displacements(mol, atoms, ir)
            mask = np.zeros_like(bond_displacements)
            if len(atoms) > 2:
                if linear:
                    mask[:5, :] = 1
                else:
                    mask[:6, :] = 1
                masked_bond_displacements = np.ma.masked_array(bond_displacements, mask)
            else:
                masked_bond_displacements = bond_displacements
            most_relevant_mode_for_bond_ = masked_bond_displacements.argmax(axis=0)
            most_relevant_mode_for_bond = []
            bonds = get_bonds_from_mol(mol)
            for i, mode in enumerate(most_relevant_mode_for_bond_):
                most_relevant_mode_for_bond.append(
                    {
                        "startAtom": bonds[i][0],
                        "endAtom": bonds[i][1],
                        "mode": int(mode),
                        "displacement": bond_displacements[mode][i],
                    }
                )
        displacement_alignments = [
            get_alignment(ir, n) for n in range(3 * len(ir.indices))
        ]

        mode_info, has_imaginary, has_large_imaginary = compile_modes_info(
            ir,
            linear,
            displacement_alignments,
            bond_displacements,
            bonds,
            raman_intensities,
        )

        spectrum = get_ir_spectrum(ir, mode_info)
        try:
            raman_spectrum = list(get_raman_spectrum(pz, mode_info)[1])
        except Exception as e:
            raman_spectrum = None
        if raman_spectrum is not None:
            assert len(spectrum[0]) == len(raman_spectrum)

        result = IRResult(
            wavenumbers=list(spectrum[0]),
            intensities=list(spectrum[1]),
            ramanIntensities=raman_spectrum,
            zeroPointEnergy=zpe,
            modes=mode_info,
            hasImaginaryFrequency=has_imaginary,
            mostRelevantModesOfAtoms=get_max_displacements(ir, linear),
            mostRelevantModesOfBonds=most_relevant_mode_for_bond,
            isLinear=linear,
            momentsOfInertia=[float(i) for i in moi],
            hasLargeImaginaryFrequency=has_large_imaginary,
        )
        ir_cache.set(this_hash, result)

        shutil.rmtree(ir.cache.directory)
        ir.clean()
    return result

@wrapt_timeout_decorator.timeout(TIMEOUT, use_signals=False)
def calculate_from_smiles(smiles, method, myhash):
    atoms, mol = smiles2ase(smiles, get_max_atoms(method))
    opt_result = run_xtb_opt(atoms, method=method)
    result = run_xtb_ir(opt_result.atoms, method=method, mol=mol)
    ir_from_smiles_cache.set(myhash, result, expire=None)
    return result

def ir_from_smiles(smiles, method):
    myhash = hash_object(smiles + method)
    result = ir_from_smiles_cache.get(myhash)
    if result is None:
        result = calculate_from_smiles(smiles, method, myhash)
    return result

@wrapt_timeout_decorator.timeout(TIMEOUT, use_signals=False)
def calculate_from_molfile(molfile, method, myhash):
    atoms, mol = molfile2ase(molfile, get_max_atoms(method))
    opt_result = run_xtb_opt(atoms, method=method)
    result = run_xtb_ir(opt_result.atoms, method=method, mol=mol)
    ir_from_molfile_cache.set(myhash, result, expire=None)
    return result

def ir_from_molfile(molfile, method):
    myhash = hash_object(molfile + method)

    result = ir_from_molfile_cache.get(myhash)

    if result is None:
        result = calculate_from_molfile(molfile, method, myhash)
    return result

def compile_all_bond_displacements(mol, atoms, ir):
    bond_displacements = []
    for mode_number in range(3 * len(ir.indices)):
        bond_displacements.append(
            get_bond_displacements(mol, atoms, ir.get_mode(mode_number))
        )

    return np.vstack(bond_displacements)

def clean_frequency(frequencies, n):
    if frequencies[n].imag != 0:
        c = "i"
        freq = frequencies[n].imag

    else:
        freq = frequencies[n].real
        c = " "
    return freq, c

def compile_modes_info(
    ir, linear, alignments, bond_displacements=None, bonds=None, raman_intensities=None
):
    frequencies = ir.get_frequencies()
    symbols = ir.atoms.get_chemical_symbols()
    modes = []
    sorted_alignments = sorted(alignments, reverse=True)
    mapping = dict(zip(np.arange(len(frequencies)), np.argsort(frequencies)))
    third_best_alignment = sorted_alignments[2]
    has_imaginary = False
    has_large_imaginary = False
    num_modes = 3 * len(ir.indices)
    if raman_intensities is None:
        raman_intensities = [None] * num_modes
    for n in range(num_modes):
        n = int(mapping[n])
        if n < 3:
            # print("below 5", alignments[n])
            # if alignments[n] >= third_best_alignment:
            modeType = "translation"
            # else:
            #     modeType = "rotation"
        elif n < 5:
            modeType = "rotation"
        elif n == 5:
            if linear:
                modeType = "vibration"
            else:
                if alignments[n] >= third_best_alignment:
                    modeType = "translation"
                else:
                    modeType = "rotation"
        else:
            modeType = "vibration"

        f, c = clean_frequency(frequencies, n)
        if c == "i":
            has_imaginary = True
            if f > IMAGINARY_FREQ_THRESHOLD:
                has_large_imaginary = True
        mostContributingBonds = None
        if bond_displacements is not None:
            mostContributingBonds = select_most_contributing_bonds(
                bond_displacements[n, :]
            )
            mostContributingBonds = [bonds[i] for i in mostContributingBonds]
            mode = ir.get_mode(n)

        ramanIntensity = float(raman_intensities[n]) if raman_intensities[n] is not None else None
        modes.append(
            {
                "number": n,
                "displacements": get_displacement_xyz_for_mode(
                    ir, frequencies, symbols, n
                ),
                "intensity": float(ir.intensities[n]),
                "ramanIntensity": ramanIntensity,
                "wavenumber": float(f),
                "imaginary": True if c == "i" else False,
                "mostDisplacedAtoms": [
                    int(i)
                    for i in np.argsort(np.linalg.norm(mode - mode.sum(axis=0), axis=1))
                ][::-1],
                "mostContributingAtoms": [
                    int(i) for i in select_most_contributing_atoms(ir, n)
                ],
                "mostContributingBonds": mostContributingBonds,
                "modeType": modeType,
                "centerOfMassDisplacement": float(
                    np.linalg.norm(ir.get_mode(n).sum(axis=0))
                ),
                "totalChangeOfMomentOfInteria": get_change_in_moi(ir.atoms, ir, n),
                "displacementAlignment": alignments[n],
            }
        )

    return modes, has_imaginary, has_large_imaginary

def get_max_displacements(ir, linear):
    mode_abs_displacements = []

    for n in range(3 * len(ir.indices)):
        mode_abs_displacements.append(np.linalg.norm(ir.get_mode(n), axis=1))

    mode_abs_displacements = np.stack(mode_abs_displacements)
    if linear:
        mode_abs_displacements[:5, :] = 0
    else:
        mode_abs_displacements[:6, :] = 0

    return dict(
        zip(
            ir.indices,
            [list(a)[::-1] for a in mode_abs_displacements.argsort(axis=0).T],
        )
    )

def get_alignment(ir, mode_number):
    dot_result = []

    displacements = ir.get_mode(mode_number)

    for i, displ_i in enumerate(displacements):
        for j, displ_j in enumerate(displacements):
            if i < j:
                dot_result.append(spatial.distance.cosine(displ_i, displ_j))

    return np.mean(dot_result)

def get_displacement_xyz_for_mode(ir, frequencies, symbols, n):
    xyz_file = []
    xyz_file.append("%6d\n" % len(ir.atoms))

    f, c = clean_frequency(frequencies, n)

    xyz_file.append("Mode #%d, f = %.1f%s cm^-1" % (n, float(f.real), c))

    if ir.ir:
        xyz_file.append(", I = %.4f (D/Å)^2 amu^-1.\n" % ir.intensities[n])
    else:
        xyz_file.append(".\n")

    # dict_label = xyz_file[-1] + xyz_file[-2]
    # dict_label = dict_label.strip('\n')

    mode = ir.get_mode(n)
    for i, pos in enumerate(ir.atoms.positions):
        xyz_file.append(
            "%2s %12.5f %12.5f %12.5f %12.5f %12.5f %12.5f\n"
            % (symbols[i], pos[0], pos[1], pos[2], mode[i, 0], mode[i, 1], mode[i, 2],)
        )

    xyz_file_string = "".join(xyz_file)
    return xyz_file_string

def get_displacement_xyz_dict(ir):
    symbols = ir.atoms.get_chemical_symbols()
    frequencies = ir.get_frequencies()
    modes = {}

    for n in range(3 * len(ir.indices)):
        modes[n] = get_displacement_xyz_for_mode(ir, frequencies, symbols, n)

    return modes

def select_most_contributing_atoms(ir, mode, threshold: float = 0.4):
    displacements = ir.get_mode(mode)
    relative_contribution = (
        np.linalg.norm(displacements, axis=1)
        / np.linalg.norm(displacements, axis=1).max()
    )
    res = np.where(
        relative_contribution
        > threshold * np.max(np.abs(np.diff(relative_contribution)))
    )[0]

    return res

def select_most_contributing_bonds(displacements, threshold: float = 0.4):

    if len(displacements) > 1:
        relative_contribution = displacements / displacements.sum()
        return np.where(
            relative_contribution
            > threshold * np.max(np.abs(np.diff(relative_contribution)))
        )[0]
    else:
        return np.array([0])

@lru_cache()
def get_bonds_from_mol(mol) -> List[Tuple[int, int]]:
    all_bonds = []
    for i in range(mol.GetNumBonds()):
        bond = mol.GetBondWithIdx(i)
        all_bonds.append((bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()))

    return all_bonds

def get_bond_vector(positions, bond):
    return positions[bond[1]] - positions[bond[0]]

def get_displaced_positions(positions, mode):
    return positions + mode

def get_bond_displacements(mol, atoms, mode):
    bonds = get_bonds_from_mol(mol)
    positions = atoms.positions
    displaced_positions = get_displaced_positions(positions, mode) - mode.sum(axis=0)
    changes = []

    for bond in bonds:
        bond_displacements = np.linalg.norm(
            get_bond_vector(positions, bond)
        ) - np.linalg.norm(get_bond_vector(displaced_positions, bond))

        changes.append(np.linalg.norm(bond_displacements))

    return changes

def get_change_in_moi(atoms, ir, mode_number):
    return np.linalg.norm(
        np.linalg.norm(
            get_moments_of_inertia(
                get_displaced_positions(atoms.positions, ir.get_mode(mode_number)),
                atoms.get_masses(),
            )
        )
        - np.linalg.norm(get_moments_of_inertia(atoms.positions, atoms.get_masses()))
    )


import datetime
import logging
import os
from pprint import pprint

import numpy as np
from hyperopt import hp, fmin, tpe, Trials, STATUS_OK, STATUS_FAIL
from hyperopt.pyll.base import scope
from traitlets import Int, Unicode, Enum

from citeomatic import file_util
from citeomatic.common import PAPER_EMBEDDING_MODEL, CITATION_RANKER_MODEL, DatasetPaths
from citeomatic.config import App
from citeomatic.candidate_selectors import ANNCandidateSelector, BM25CandidateSelector
from citeomatic.models.options import ModelOptions
from citeomatic.serialization import model_from_directory
from citeomatic.neighbors import EmbeddingModel, ANN
from citeomatic.ranker import Ranker, NoneRanker
from citeomatic.training import end_to_end_training
from citeomatic.training import eval_text_model, EVAL_DATASET_KEYS
import pickle
import tensorflow as tf
import keras.backend as K

class TrainCiteomatic(App, ModelOptions):
    dataset_type = Enum(('dblp', 'pubmed', 'oc'), default_value='dblp')

    # Whether to train based on given options or to run a hyperopt
    mode = Enum(('train', 'hyperopt'))

    # Training parameters
    hyperopts_results_pkl = Unicode(default_value=None, allow_none=True)
    options_json = Unicode(default_value=None, allow_none=True)

    # hyperopt parameters
    max_evals_initial = Int(default_value=25)
    max_evals_secondary = Int(default_value=5)
    total_samples_initial = Int(default_value=5000000)
    total_samples_secondary = Int(default_value=50000000)
    n_eval = Int(default_value=500, allow_none=True)
    models_ann_dir = Unicode(default_value=None, allow_none=True)
    models_dir_base = Unicode(
        default_value='data/models/'
    )
    run_identifier = Unicode(default_value=None, allow_none=True)
    version = Unicode(default_value='v0')

    # to be filled in later
    models_dir = Unicode(default_value=None, allow_none=True)
    ann = None

    def main(self, args):

        if self.mode == 'hyperopt':
            self.run_hyperopt()
        elif self.mode == 'train':
            self.run_train()
        else:
            assert False

    def run_train(self):
        eval_params = {}
        if self.hyperopts_results_pkl is not None:
            params = pickle.load(open(self.hyperopts_results_pkl, "rb"))
            for k, v in params[1][0]['result']['params'].items():
                eval_params[k] = v
        if self.options_json is not None:
            obj = file_util.read_json(self.options_json)
            eval_params.update(obj)

        if self.model_name == PAPER_EMBEDDING_MODEL:
            self.models_ann_dir = None
            self.models_dir = os.path.join(self.models_dir_base, PAPER_EMBEDDING_MODEL)

        self.train_and_evaluate(eval_params)

    @staticmethod
    def _hyperopt_space(model_name, total_samples):

        use_pretrained = hp.choice('use_pretrained', [True, False])

        common_param_space = {
            'total_samples':
                total_samples,
            'lr':
                hp.choice('lr', [0.1, 0.01, 0.001, 0.0001, 0.00001]),
            'l1_lambda':
                hp.choice('l1_lambda', np.append(np.logspace(-7, -2, 6), 0)),
            'dropout_p':
                hp.quniform('dropout_p', 0.0, 0.75, 0.05),
            'margin_multiplier':
                hp.choice('margin_multiplier', [0.5, 0.75, 1.0, 1.25, 1.5])
        }

        pre_trained_params = {
            True: {
                'use_pretrained': True,
                'l2_lambda': 0,
                'dense_dim': 300,
                'enable_fine_tune': hp.choice('enable_fine_tune', [True, False])
            },
            False: {
                'use_pretrained': False,
                'l2_lambda': hp.choice('l2_lambda', np.append(np.logspace(-7, -2, 6), 0)),
                'dense_dim': scope.int(hp.quniform('dense_dim', 25, 325, 25)),
                'enable_fine_tune': True # doesn't matter what goes here
            }
        }

        # the search space
        # note that the scope.int code is a hack to get integers out of the sampler
        if model_name == CITATION_RANKER_MODEL:
            ranker_model_params = {
                'embedding_type':
                    hp.choice('embedding_type', ['sum']),
                'metadata_dim':
                    scope.int(hp.quniform('metadata_dim', 5, 55, 5)),
            }

            space = scope.switch(
                scope.int(use_pretrained),
                {**pre_trained_params[False],
                 **common_param_space,
                 **ranker_model_params,
                 },
                {**pre_trained_params[True],
                 **common_param_space,
                 **ranker_model_params
                 }
            )
        elif model_name == PAPER_EMBEDDING_MODEL:
            space = scope.switch(
                scope.int(use_pretrained),
                {**pre_trained_params[False],
                 **common_param_space
                 },
                {**pre_trained_params[True],
                 **common_param_space
                 }
            )
        else:
            # Should not come here. Adding this to make pycharm happy.
            assert False

        return space

    def run_hyperopt(self):
        # run identifier
        if self.run_identifier is None:
            self.run_identifier = '_'.join(
                [
                    'citeomatic_hyperopt',
                    self.model_name,
                    self.dataset_type,
                    datetime.datetime.now().strftime("%Y-%m-%d"),
                    self.version
                ]
            )

        self.models_dir = os.path.join(self.models_dir_base, self.run_identifier)
        if self.model_name == PAPER_EMBEDDING_MODEL:
            self.models_ann_dir = None

        space = self._hyperopt_space(self.model_name, self.total_samples_initial)

        # stage 1: run hyperopt for max_evals_initial
        # using a max of total_samples_initial samples
        trials = Trials()
        _ = fmin(
            fn=self.train_and_evaluate,
            space=space,
            algo=tpe.suggest,
            max_evals=self.max_evals_initial,
            trials=trials
        )

        sorted_results_stage_1 = sorted(
            trials.trials, key=lambda x: x['result']['loss']
        )

        # stage 2: run the top max_evals_seconadry from stage 1
        # using a max of total_samples_secondary samples
        results_stage_2 = []
        for result in sorted_results_stage_1[:self.max_evals_secondary]:
            params = result['result']['params']
            params['total_samples'] = self.total_samples_secondary
            out = self.train_and_evaluate(params)
            results_stage_2.append({'params': params, 'result': out})

        sorted_results_stage_2 = sorted(
            results_stage_2, key=lambda x: x['result']['loss']
        )

        # save and display results
        results_save_file = 'hyperopt_results.pickle'
        file_util.write_pickle(
            os.path.join(self.models_dir, results_save_file),
            (sorted_results_stage_1, sorted_results_stage_2)
        )
        pprint(sorted_results_stage_2[0])

    def train_and_evaluate(self, eval_params):
        # Needed especially for hyperopt runs
        K.clear_session()

        # eval_params为args传递参数
        model_kw = {name: getattr(self, name) for name in ModelOptions.class_traits().keys()}
        model_kw.update(eval_params)
        model_options = ModelOptions(**model_kw)

        if model_options.use_metadata:
            model_options.use_keyphrases = True
            model_options.use_authors = True
            model_options.use_venue = True

        print("====== OPTIONS =====")
        print(model_options)
        print("======")

        if model_options.train_for_test_set:
            logging.info("\n\n============== TRAINING FOR TEST SET =============\n\n")
        # Todo train-1 训练主函数
        training_outputs = end_to_end_training(
            model_options,
            self.dataset_type,
            self.models_dir,
            self.models_ann_dir
        )
        corpus, featurizer, model_options, citeomatic_model, embedding_model = training_outputs

        if self.candidate_selector_type == 'ann':
            # if no ann_dir is provided, then we use the model that was just trained
            # and have to rebuild the ANN
            if self.models_ann_dir is None:
                print(
                    'Using embedding model that was just trained for eval. Building...')
                paper_embedding_model = EmbeddingModel(
                    featurizer,
                    embedding_model
                )
                self.ann = ANN.build(paper_embedding_model, corpus)
            # if a dir is provided, then go ahead and load it
            else:
                featurizer_for_ann, ann_models = model_from_directory(
                    self.models_ann_dir, on_cpu=True
                )
                paper_embedding_model = EmbeddingModel(
                    featurizer_for_ann,
                    ann_models['embedding']
                )
                # the ANN itself needs to be only built once
                if self.ann is None:
                    if corpus.corpus_type == 'oc' and os.path.exists(DatasetPaths.OC_ANN_FILE + ".pickle"):
                        self.ann = ANN.load(DatasetPaths.OC_ANN_FILE)
                    else:
                        self.ann = ANN.build(paper_embedding_model, corpus)

            candidate_selector = ANNCandidateSelector(
                corpus=corpus,
                ann=self.ann,
                paper_embedding_model=paper_embedding_model,
                top_k=model_options.num_ann_nbrs_to_fetch,
                extend_candidate_citations=model_options.extend_candidate_citations
            )
        elif self.candidate_selector_type == 'bm25':
            dp = DatasetPaths()
            candidate_selector = BM25CandidateSelector(
                corpus=corpus,
                index_path=dp.get_bm25_index_path(self.dataset_type),
                top_k=model_options.num_ann_nbrs_to_fetch,
                extend_candidate_citations=model_options.extend_candidate_citations
            )
        else:
            # Should not come here. Adding this to make pycharm happy.
            assert False

        if self.citation_ranker_type == 'neural':
            ranker = Ranker(
                corpus=corpus,
                featurizer=featurizer,
                citation_ranker=citeomatic_model,
                num_candidates_to_rank=model_options.num_candidates_to_rank
            )
        elif self.citation_ranker_type == 'none':
            ranker = NoneRanker()
        else:
            # Should not come here. Adding this to make pycharm happy.
            assert False

        if self.mode != 'hyperopt' or model_options.total_samples == self.total_samples_secondary:
            results_training = eval_text_model(
                corpus,
                candidate_selector,
                ranker,
                papers_source='train',
                n_eval=self.n_eval
            )
        else:
            results_training = {}

        results_validation = eval_text_model(
            corpus,
            candidate_selector,
            ranker,
            papers_source='valid',
            n_eval=self.n_eval
        )

        logging.info("===== Validation Results ===== ")
        logging.info("Validation Precision\n\n{}".format(results_validation['precision_1']))
        logging.info("Validation Recall\n\n{}".format(results_validation['recall_1']))

        p = results_validation['precision_1'][EVAL_DATASET_KEYS[self.dataset_type]]
        r = results_validation['recall_1'][EVAL_DATASET_KEYS[self.dataset_type]]
        f1 = results_validation['f1_1'][EVAL_DATASET_KEYS[self.dataset_type]]

        if self.model_name == PAPER_EMBEDDING_MODEL:
            # optimizing for recall
            l = -r
        else:
            # optimizing for F1
            l = -f1

        out = {
            'loss': l,  # have to negate since we're minimizing
            'losses_training': results_training,
            'losses_validation': results_validation,
            'status': STATUS_FAIL if np.isnan(f1) else STATUS_OK,
            'params': eval_params
        }

        return out

TrainCiteomatic.run(__name__)


"""Convert raw data into suitable inputs for Keras."""

import pickle

import numpy as np
from keras.utils import np_utils
from sklearn.model_selection import train_test_split

from KanaRecognizer.trainer.preprocessing.data_utils import get_ETL_data
from util.util import PROCESSED_DATA_PATH, HIRAGANA, KATAKANA

def get_data(writers_per_char=160, mode='all', get_scripts=False, test_size=0.2):
    """
    Load the characters into a format suitable for Keras
    Args:
        writers_per_char (int): number of samples per Japanese character
        mode (string): specify the type of Japanese characters: 'all', 'hiragana', 'kanji', or 'katakana'
        get_scripts (bool): 'True' returns a label for the type of script corresponding to each Japanese character
        test_size (float): fraction of data to use for obtaining a test error
    Returns:
        X_train (np.array): training data
        Y_train (np.array): training labels
        X_test (np.array): test data
        Y_test (np.array): test labels
        labels_dict: label str list in the order of label index
    """
    characters = np.empty((0, 64, 64))
    labels = []
    scripts = []
    if mode in ('kanji', 'all'):
        for i in range(1, 4):
            if i == 3:
                max_records = 315
            else:
                max_records = 319

            if i != 1:
                start_record = 0
            else:
                start_record = 75

            chars, labs, spts = get_ETL_data(
                i, range(start_record, max_records), writers_per_char)

            characters = np.concatenate((characters, chars), axis=0)
            labels = np.concatenate((labels, labs), axis=0)
            scripts = np.concatenate((scripts, spts), axis=0)

    if mode in (HIRAGANA, 'all', 'kana'):
        max_records = 75
        chars, labs, spts = get_ETL_data(
            1, range(0, max_records), writers_per_char)
        print("hiragana data num:", chars.shape[0])
        characters = np.concatenate((characters, chars), axis=0)
        labels = np.concatenate((labels, labs), axis=0)
        scripts = np.concatenate((scripts, spts), axis=0)

    if mode in (KATAKANA, 'all', 'kana'):
        katakana_num = 0
        for i in range(7, 14):
            if i < 10:
                filename = '0' + str(i)
            else:
                filename = str(i)

            chars, labs, spts = get_ETL_data(filename, range(0, 8 if i < 13 else 3),
                                             writers_per_char, database='ETL1C')
            katakana_num += chars.shape[0]
            characters = np.concatenate((characters, chars), axis=0)
            labels = np.concatenate((labels, labs), axis=0)
            scripts = np.concatenate((scripts, spts), axis=0)

        print("katakana data num:", katakana_num)

    # rename labels from 0 to n_labels-1
    unique_labels = sorted(list(set(labels)))
    print('unique_labels num =', len(unique_labels))
    labels_dict = {unique_labels[i]: i for i in range(len(unique_labels))}
    new_labels = np.array([labels_dict[l] for l in labels], dtype=np.int32)

    if get_scripts:
        x_train, x_test, y_train, y_test = train_test_split(characters,
                                                            scripts,
                                                            test_size=test_size,
                                                            random_state=42)
    elif mode in ('all', 'kanji', HIRAGANA, KATAKANA, 'kana'):
        x_train, x_test, y_train, y_test = train_test_split(characters,
                                                            new_labels,
                                                            test_size=test_size,
                                                            random_state=42)
    print("X_train:", x_train.shape)
    print("y_train:", y_train.shape)
    print("X_test:", x_test.shape)
    print("y_test:", y_test.shape)

    # reshape to (1, 64, 64)
    X_train = x_train.reshape(
        (x_train.shape[0], 1, x_train.shape[1], x_train.shape[2]))
    X_test = x_test.reshape(
        (x_test.shape[0], 1, x_test.shape[1], x_test.shape[2]))

    # convert class vectors to binary class matrices
    nb_classes = len(unique_labels)
    Y_train = np_utils.to_categorical(y_train, nb_classes)
    Y_test = np_utils.to_categorical(y_test, nb_classes)
    return X_train, Y_train, X_test, Y_test, unique_labels

if __name__ == '__main__':
    X_train, Y_train, X_test, Y_test, unique_labels = \
        get_data(writers_per_char=160, mode='hiragana', get_scripts=False, test_size=0.2)
    pickle.dump(X_train, open(PROCESSED_DATA_PATH + '/hiragana_x_train.p', 'wb'))
    pickle.dump(Y_train, open(PROCESSED_DATA_PATH + '/hiragana_y_train.p', 'wb'))
    pickle.dump(X_test, open(PROCESSED_DATA_PATH + '/hiragana_x_test.p', 'wb'))
    pickle.dump(Y_test, open(PROCESSED_DATA_PATH + '/hiragana_y_test.p', 'wb'))
    print(unique_labels)

    X_train, Y_train, X_test, Y_test, unique_labels = \
        get_data(writers_per_char=300, mode='katakana', get_scripts=False, test_size=0.2)
    pickle.dump(X_train, open(PROCESSED_DATA_PATH + '/katakana_x_train.p', 'wb'))
    pickle.dump(Y_train, open(PROCESSED_DATA_PATH + '/katakana_y_train.p', 'wb'))
    pickle.dump(X_test, open(PROCESSED_DATA_PATH + '/katakana_x_test.p', 'wb'))
    pickle.dump(Y_test, open(PROCESSED_DATA_PATH + '/katakana_y_test.p', 'wb'))
    print(unique_labels)

    # X_train, Y_train, X_test, Y_test, unique_labels = \
    #     get_data(writers_per_char=160, mode='kana', get_scripts=False, test_size=0.2)
    # pickle.dump(X_train, open(PROCESSED_DATA_PATH + '/x_train.p', 'wb'))
    # pickle.dump(Y_train, open(PROCESSED_DATA_PATH + '/y_train.p', 'wb'))
    # pickle.dump(X_test, open(PROCESSED_DATA_PATH + '/x_test.p', 'wb'))
    # pickle.dump(Y_test, open(PROCESSED_DATA_PATH + '/y_test.p', 'wb'))
    # pickle.dump(unique_labels, open(PROCESSED_DATA_PATH + '/unique_labels.p', 'wb'))
    # for i in range(0, len(unique_labels)//10 + 1):
    #     print(unique_labels[i*10: i*10 + 10])


# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import random
import time
import math
from functools import partial

import numpy as np
import paddle
from paddle.io import DataLoader

import paddlenlp as ppnlp
from paddlenlp.transformers import LinearDecayWithWarmup
from paddlenlp.metrics import ChunkEvaluator
from datasets import load_dataset
from paddlenlp.transformers import BertForTokenClassification, BertTokenizer
from paddlenlp.transformers import ErnieForTokenClassification, ErnieTokenizer
from paddlenlp.transformers import ErnieCtmForTokenClassification, ErnieCtmTokenizer
from paddlenlp.data import DataCollatorForTokenClassification
from paddlenlp.utils.log import logger

MODEL_CLASSES = {
    "bert": (BertForTokenClassification, BertTokenizer),
    "ernie": (ErnieForTokenClassification, ErnieTokenizer),
    "ernie-ctm": (ErnieCtmForTokenClassification, ErnieCtmTokenizer)
}

parser = argparse.ArgumentParser()

# yapf: disable
parser.add_argument("--model_type", default="bert", type=str, help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()), )
parser.add_argument("--model_name_or_path", default=None, type=str, required=True, help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join( sum([ list(classes[-1].pretrained_init_configuration.keys())  for classes in MODEL_CLASSES.values() ], [])), )
parser.add_argument("--dataset", default="msra_ner", type=str, choices=["msra_ner", "peoples_daily_ner"] ,help="The named entity recognition datasets.")
parser.add_argument("--output_dir", default=None, type=str, required=True, help="The output directory where the model predictions and checkpoints will be written.")
parser.add_argument("--max_seq_length", default=128, type=int, help="The maximum total input sequence length after tokenization. Sequences longer than this will be truncated, sequences shorter will be padded.")
parser.add_argument("--batch_size", default=8, type=int, help="Batch size per GPU/CPU for training.")
parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
parser.add_argument("--num_train_epochs", default=3, type=int, help="Total number of training epochs to perform.", )
parser.add_argument("--max_steps", default=-1, type=int, help="If > 0: set total number of training steps to perform. Override num_train_epochs.",)
parser.add_argument("--warmup_steps", default=0, type=int, help="Linear warmup over warmup_steps.")
parser.add_argument("--logging_steps", type=int, default=1, help="Log every X updates steps.")
parser.add_argument("--save_steps", type=int, default=100, help="Save checkpoint every X updates steps.")
parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
parser.add_argument("--device", default="gpu", type=str, choices=["cpu", "gpu", "xpu"] ,help="The device to select to train the model, is must be cpu/gpu/xpu.")
# yapf: enable

@paddle.no_grad()
def evaluate(model, loss_fct, metric, data_loader, label_num, mode="valid"):
    model.eval()
    metric.reset()
    avg_loss, precision, recall, f1_score = 0, 0, 0, 0
    for batch in data_loader:
        logits = model(batch['input_ids'], batch['token_type_ids'])
        loss = loss_fct(logits, batch['labels'])
        avg_loss = paddle.mean(loss)
        preds = logits.argmax(axis=2)
        num_infer_chunks, num_label_chunks, num_correct_chunks = metric.compute(
            batch['seq_len'], preds, batch['labels'])
        metric.update(num_infer_chunks.numpy(),
                      num_label_chunks.numpy(), num_correct_chunks.numpy())
        precision, recall, f1_score = metric.accumulate()
    print("%s: eval loss: %f, precision: %f, recall: %f, f1: %f" %
          (mode, avg_loss, precision, recall, f1_score))
    model.train()

def do_train(args):
    paddle.set_device(args.device)
    if paddle.distributed.get_world_size() > 1:
        paddle.distributed.init_parallel_env()

    # Create dataset, tokenizer and dataloader.
    if args.dataset == "peoples_daily_ner":
        raw_datasets = load_dataset(args.dataset)
    else:
        raw_datasets = load_dataset(args.dataset)

    AutoForTokenClassification, AutoTokenizer = MODEL_CLASSES[args.model_type]
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
    train_ds = raw_datasets['train']

    label_list = train_ds.features['ner_tags'].feature.names
    label_num = len(label_list)
    no_entity_id = 0

    def tokenize_and_align_labels(examples):
        tokenized_inputs = tokenizer(
            examples['tokens'],
            max_seq_len=args.max_seq_length,
            # We use this argument because the texts in our dataset are lists of words (with a label for each word).
            is_split_into_words=True,
            return_length=True)
        labels = []

        for i, label in enumerate(examples['ner_tags']):
            label_ids = label
            if len(tokenized_inputs['input_ids'][i]) - 2 < len(label_ids):
                label_ids = label_ids[:len(tokenized_inputs['input_ids'][i]) -
                                      2]
            label_ids = [no_entity_id] + label_ids + [no_entity_id]
            label_ids += [no_entity_id] * (
                len(tokenized_inputs['input_ids'][i]) - len(label_ids))

            labels.append(label_ids)
        tokenized_inputs["labels"] = labels
        return tokenized_inputs

    train_ds = train_ds.select(range(len(train_ds) - 1))
    column_names = train_ds.column_names
    train_ds = train_ds.map(tokenize_and_align_labels,
                            batched=True,
                            remove_columns=column_names)

    ignore_label = -100

    batchify_fn = DataCollatorForTokenClassification(
        tokenizer=tokenizer, label_pad_token_id=ignore_label)

    train_batch_sampler = paddle.io.DistributedBatchSampler(
        train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True)

    train_data_loader = DataLoader(
        dataset=train_ds,
        collate_fn=batchify_fn,
        num_workers=0,
        batch_sampler=train_batch_sampler,
        return_list=True)

    test_ds = raw_datasets['test']
    test_ds = test_ds.select(range(len(test_ds) - 1))
    test_ds = test_ds.map(tokenize_and_align_labels,
                          batched=True,
                          remove_columns=column_names)

    test_data_loader = DataLoader(
        dataset=test_ds,
        collate_fn=batchify_fn,
        num_workers=0,
        batch_size=args.batch_size,
        return_list=True)

    if args.dataset == "peoples_daily_ner":
        dev_ds = raw_datasets['validation']
        dev_ds = dev_ds.select(range(len(dev_ds) - 1))
        dev_ds = dev_ds.map(tokenize_and_align_labels,
                            batched=True,
                            remove_columns=column_names)

        dev_data_loader = DataLoader(
            dataset=dev_ds,
            collate_fn=batchify_fn,
            num_workers=0,
            batch_size=args.batch_size,
            return_list=True)

    # Define the model netword and its loss
    model = AutoForTokenClassification.from_pretrained(
        args.model_name_or_path, num_classes=label_num)

    if paddle.distributed.get_world_size() > 1:
        model = paddle.DataParallel(model)

    num_training_steps = args.max_steps if args.max_steps > 0 else len(
        train_data_loader) * args.num_train_epochs

    lr_scheduler = LinearDecayWithWarmup(args.learning_rate, num_training_steps,
                                         args.warmup_steps)

    # Generate parameter names needed to perform weight decay.
    # All bias and LayerNorm parameters are excluded.
    decay_params = [
        p.name for n, p in model.named_parameters()
        if not any(nd in n for nd in ["bias", "norm"])
    ]
    optimizer = paddle.optimizer.AdamW(
        learning_rate=lr_scheduler,
        epsilon=args.adam_epsilon,
        parameters=model.parameters(),
        weight_decay=args.weight_decay,
        apply_decay_param_fun=lambda x: x in decay_params)

    loss_fct = paddle.nn.loss.CrossEntropyLoss(ignore_index=ignore_label)

    metric = ChunkEvaluator(label_list=label_list)

    global_step = 0
    last_step = args.num_train_epochs * len(train_data_loader)
    tic_train = time.time()
    for epoch in range(args.num_train_epochs):
        for step, batch in enumerate(train_data_loader):
            global_step += 1
            logits = model(batch['input_ids'], batch['token_type_ids'])
            loss = loss_fct(logits, batch['labels'])
            avg_loss = paddle.mean(loss)
            if global_step % args.logging_steps == 0:
                print(
                    "global step %d, epoch: %d, batch: %d, loss: %f, speed: %.2f step/s"
                    % (global_step, epoch, step, avg_loss,
                       args.logging_steps / (time.time() - tic_train)))
                tic_train = time.time()
            avg_loss.backward()
            optimizer.step()
            lr_scheduler.step()
            optimizer.clear_grad()
            if global_step % args.save_steps == 0 or global_step == num_training_steps:
                if paddle.distributed.get_rank() == 0:
                    if args.dataset == "peoples_daily_ner":
                        evaluate(model, loss_fct, metric, dev_data_loader,
                                 label_num, "valid")
                    evaluate(model, loss_fct, metric, test_data_loader,
                             label_num, "test")

                    paddle.save(model.state_dict(),
                                os.path.join(args.output_dir,
                                             "model_%d.pdparams" % global_step))
            if global_step >= num_training_steps:
                return

if __name__ == "__main__":
    args = parser.parse_args()
    for arg in vars(args):
        logger.info('{:20}:{}'.format(arg, getattr(args, arg)))

    do_train(args)


"""
Generates 16-measure waltzes based on Mozart's "Musical Table", a collection
of 176 measures. Measures are selected according to the rules of Mozart
himself.
"""

import time
from numpy.random import randint
from pyknon.genmidi import Midi
from pyknon.music import Note, NoteSeq
import pygame as pg

def get_measures():
    """
    The text file holds note data. Each line is in the form

    <note> <beat of actuation> <note duration>

    . For example,

    C3 6 2

    represents holding C3 for two beats, starting on beat 6.

    This function reads the text file and separates the notes into measures,
    each 3 beats long.
    """
    notes = [line.rstrip("\n").split(" ") for line in open("notes.txt")]

    measures = []
    current = []
    last_note = notes[-1]
    last_multiple = 0

    for note in notes:
        if (float(note[1]) % 3 == 0) and (float(note[1])//3 != last_multiple):
            last_multiple = float(note[1])//3
            measures.append(current)
            current = []
            current.append([note[0], float(note[1]), float(note[2])])
        else:
            current.append([note[0], float(note[1]), float(note[2])])
        if note == last_note:
            measures.append(current)

    return(measures)

def roll():
    """
    Simulates a roll of two dice and returns the total.
    """
    dice_1 = randint(1, 7)
    dice_2 = randint(1, 7)
    return(dice_1 + dice_2)

def measures_to_notes(measures):
    """
    measures: List[List[List[str, float, float]]]
         -> List[List[str, float, float]]

    Breaks up a list of measures of notes into a list of notes.
    """
    notes = []
    for i in range(len(measures)):
        notes.extend(measures[i])
    return(notes)

def note_to_index(note_str):
    """
    note_str: str -> int

    Returns the integer value of the note in note_str, given as an index in
    the list below.
    """
    notes = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
    return(notes.index(note_str))

def create_song(song_notes, tempo):
    """
    song_notes: List[List[str, float, float]], tempo: int -> None

    Creates a MIDI file with notes given in song_notes and saves it at
    tmp/song.mid.
    """
    notes = []
    beat = 0

    while beat < 48:
        chord = []
        for song_note in song_notes:
            if song_note[1] == beat:
                if beat == 47:  # make last note longer
                    dur = 0.5
                else:
                    dur = song_note[2]/4
                note = Note(
                    value=note_to_index(song_note[0][:-1]),
                    octave=int(song_note[0][-1]) + 1,
                    dur=dur)
                chord.append(note)
        if chord != []:  # beat isn't a rest
            notes.append([NoteSeq(chord), beat])
        beat += 0.5

    midi = Midi(tempo=tempo)
    for note in notes:
        midi.seq_chords([note[0]], time=note[1])

    midi.write("tmp/song.mid")

def play_song(tempo):
    """
    Plays the created MIDI file located at tmp/song.mid.
    """
    pg.init()
    pg.mixer.music.load("tmp/song.mid")
    pg.mixer.music.play()
    time.sleep(48*60/tempo)

def main():
    """
    Simulates 16 dice rolls and makes the appropriate measure selections.
    Plays the created waltz afterwards.
    """
    song = []
    measures = get_measures()
    measure_selections = [
        [96, 32, 69, 40, 148, 104, 152, 119, 98, 3, 54],
        [22, 6, 95, 17, 74, 157, 60, 84, 142, 87, 130],
        [141, 128, 158, 113, 163, 27, 171, 114, 42, 165, 10],
        [41, 63, 13, 85, 45, 167, 53, 50, 156, 61, 103],
        [105, 146, 153, 161, 80, 154, 99, 140, 75, 135, 28],
        [122, 46, 55, 2, 97, 68, 133, 86, 129, 47, 37],
        [11, 134, 110, 159, 36, 118, 21, 169, 62, 147, 106],
        [30, 81, 24, 100, 107, 91, 127, 94, 123, 33, 5],
        [70, 117, 66, 90, 25, 138, 16, 120, 65, 102, 35],
        [121, 39, 136, 176, 143, 71, 155, 88, 77, 4, 20],
        [26, 126, 15, 7, 64, 150, 57, 48, 19, 31, 108],
        [9, 56, 132, 34, 125, 29, 175, 166, 82, 164, 92],
        [112, 174, 73, 67, 76, 101, 43, 51, 137, 144, 12],
        [49, 18, 58, 160, 136, 162, 168, 115, 38, 59, 124],
        [109, 116, 145, 52, 1, 23, 89, 72, 149, 173, 44],
        [14, 83, 79, 170, 93, 151, 172, 111, 8, 78, 131]]

    for i in range(16):
        roll_ = roll()
        measure_selection = measures[measure_selections[i][roll_ - 2] - 1][:]
        offset = measure_selection[0][1] - 3*i
        for j in range(len(measure_selection)):
            measure_selection[j][1] -= offset
        song.append(measure_selection)

    tempo = 120
    create_song(measures_to_notes(song), tempo)
    play_song(tempo)

if __name__ == "__main__":
    main()


# Helper module for creating masks and masked tensors for RFold input

import torch
import numpy as np
#from icecream import ic
from dj_util import SampledMask

import sys, os, copy

# absolute path to folder containing this file
script_dir = os.path.dirname(os.path.realpath(__file__))

sys.path.append(script_dir+'/../hallucination/util/')
import contigs

# hyperparams for model
model_param = {'SE3_param': {'div': 2,
                            'l0_in_features': 32,
                            'l0_out_features': 8,
                            'l1_in_features': 3,
                            'l1_out_features': 2,
                            'n_heads': 4,
                            'num_channels': 16,
                            'num_degrees': 2,
                            'num_edge_features': 32,
                            'num_layers': 3},
              'd_hidden': 64,
              'd_msa': 384,
              'd_msa_full': 64,
              'd_pair': 288,
              'd_templ': 64,
              'n_head_msa': 12,
              'n_head_pair': 8,
              'n_head_templ': 4,
              'n_layer': 1,
              'n_module': 8,
              'n_module_str': 4,
              'n_resblock': 1,
              'p_drop': 0.15,
              'performer_L_opts': {'feature_redraw_interval': 10000, 'nb_features': 64},
              'performer_N_opts': {'feature_redraw_interval': 10000, 'nb_features': 64},
              'r_ff': 4}
model_param['use_templ'] = True

def get_mask(args, parsed_pdb, chain_idx_gap=500):
    """
    Gets 1-d boolean mask corresponding to residue positions to replace/inpaint.

    Parameters
    ----------
        args : argparse.Namespace
            Parsed command-line arguments, in particular the variables: window, contigs, len, keep_order, min_gap, receptor_chain.
        parsed_pdb : dict
            Output of parse_pdb command, in particular the key pdb_idx.

    Returns
    -------
        mask : np.array of bool (L,)
            Boolean array of length L of designed protein indicating positions to be replaced/inpainted.
        mappings : dict
            Dictionary with various indices describing the mapping of kept regions:
                con_ref_pdb_idx: kept regions, original pdb numbering [(chain, num),...]
                con_ref_idx0:    kept regions, 0-indexed continguous "pytorch tensor numbering"
                con_hal_pdb_idx: kept regions, output pdb numbering, 1-indexed contiguous by chain [(chain, num),...]
                con_hal_idx0:    kept regions, 0-indexed output numbering
                idx_rf:          indices for RoseTTAFold input (with large offsets added between chains)
                hal_idx1:        all regions (masked and kept), output pdb numbering, 1-indexed contiguous by chain
    """
    # contig/inverse mask mode: specify residues to keep
    if args.contigs is not None:
        # contig mode: chain letter starting each range
        if all([x[0].isalpha() for x in args.contigs.split(',')]):
            if args.len is None:
                sys.exit('ERROR: Either --len must be given or --contigs argument must have gap lengths.')
            _, mappings = contigs.scatter_contigs(contigs=args.contigs,
                                                  pdb_out=parsed_pdb,
                                                  L_range=args.len,
                                                  keep_order=args.keep_order,
                                                  min_gap=args.min_gap)

        # inverse mask mode: some ranges don't have chain letter
        else:
            _, mappings = contigs.apply_mask(args.contigs, parsed_pdb)

        mask_str = mappings['sampled_mask']
        sm = SampledMask(mask_str, ref_pdb_idxs=parsed_pdb['pdb_idx'])

        # add a receptor
        if args.receptor_chain:
            pdb_idx_rec = [x for x in parsed_pdb['pdb_idx'] if x[0]==args.receptor_chain]
            rec_str = SampledMask.contract(pdb_idx_rec)
            sm.add_receptor(rec_str, location='second')

    else:
        # add entire protein
        # assumes masking only happens on chain A
        pdb_idx = parsed_pdb['pdb_idx']
        pdb_idx_A = [idx for idx in pdb_idx if idx[0]=='A']
        mask_str = SampledMask.contract(pdb_idx_A)
        sm = SampledMask(mask_str, ref_pdb_idxs=pdb_idx_A)

        # window mode: specify residues to replace using DJ's syntax ("A,1,10:A,15,15")
        if args.window:
            window = args.window.split(':')           # list of A,B where A is start, B is end
            window = [i.split(',') for i in window]   # list of lists. [[chain1, a1, b1],[chain2, a2,b2]]
            window = [[ch,int(r1), int(r2)] for ch,r1,r2 in window]
            for ch, s, e in window:
                sm.add_inpaint_range(f'{ch}{s}-{e}')

        # mask mode: specify residues to replace using Doug/Jue syntax ("A1-10,A15")
        elif args.mask:
            for region in args.mask.split(','):
                sm.add_inpaint_range(region)

        # add a receptor
        if args.receptor_chain:
            pdb_idx_rec = [x for x in parsed_pdb['pdb_idx'] if x[0]==args.receptor_chain]
            rec_str = SampledMask.contract(pdb_idx_rec)
            sm.add_receptor(rec_str, location='second')

    mask_str = sm.inpaint
    mappings = sm.mappings

    # mask sequence separately from structure
    if args.mask_seq:
        sm_seq = copy.deepcopy(sm)
        for region in args.mask_seq.split(','):
            sm_seq.add_inpaint_range(region)
    else:
        sm_seq = sm

    mask_seq = sm_seq.inpaint

    return sm, mappings, mask_str, mask_seq

def mask_inputs(args, mask_str, seq, msa, msa_full, xyz_t, f1d_t, mask_seq=None):

    if mask_seq is None:
        mask_seq = mask_str

    if args.task == 'den':
        # do denoising preprocessing
        # Ideally, here we would do mutation of amino acids and noising of structure, etc
        pass

    elif args.task == 'des':

        # remove sequence information within the window (seq --> 20, msa --> 21)
        msa[:,mask_seq]      = args.msa_mask_token
        msa_full[:,mask_seq] = args.msa_mask_token
        seq[mask_seq]      = 20
        if args.mask_query:
            seq[:] = 20
            msa[0,:] = args.msa_mask_token
            msa_full[0,:] = args.msa_mask_token

    elif args.task == 'str':
        # do structure prediction preprocessing
        xyz_t[:,:,mask_str] = np.nan

    elif args.task == 'hal':
        # do hallucination preprocessing (combination of des + str)

        if args.init_seq == None:
            msa[:,mask_seq]      = args.msa_mask_token
            msa_full[:,mask_seq] = args.msa_mask_token
            seq[mask_seq]        = 20

        elif args.init_seq == 'poly':
            # poly AA initialization
            msa[:,mask_seq]      = util.aa2num[args.poly]
            msa_full[:,mask_seq] = util.aa2num[args.poly]
            seq[mask_seq]        = util.aa2num[args.poly]

        else:
            raise NotImplementedError('Shouldnt be here :/')

        f1d_t[:,:,mask_str] = 0 # zero confidence in the masked region(s)

        # remove structure information within the window (change template coords to nan)
        xyz_t[:,:,mask_str] = float('nan')

    else:
        sys.exit(f'Invalid task {args.task}')

    return seq, msa, msa_full, xyz_t, f1d_t

def rw_ins_input(args,mask_orig,xyz_t,seq,msa,msa_full,mask,DEVICE):
    idx_ins      = np.random.choice(np.where(mask_orig)[0], 1, replace=False)[0]
    xyz_t_new    = torch.full((1,xyz_t.shape[1],1,3,3),np.nan).to(DEVICE) #torch.full((1,1,L,3,3),np.nan)
    xyz_t        = torch.cat([xyz_t[:,:,:idx_ins], xyz_t_new, xyz_t[:,:,idx_ins:]], 2)
    seq_new      = torch.full((1,), 20)[None].to(DEVICE)
    seq          = torch.cat([seq[:,:idx_ins],seq_new,seq[:,idx_ins:]],1)
    msa_new      = torch.full((msa.shape[0],1), args.msa_mask_token)[None].to(DEVICE)
    msa          = torch.cat([msa[:,:,:idx_ins],msa_new,msa[:,:,idx_ins:]],2)
    msa_full_new = torch.full((msa_full.shape[0],1), args.msa_mask_token)[None].to(DEVICE)
    msa_full     = torch.cat([msa_full[:,:,:idx_ins],msa_full_new,msa_full[:,:,idx_ins:]],2)
    mask_orig    = np.insert(mask_orig, idx_ins, True, axis=0)
    mask         = np.insert(mask, idx_ins, True, axis=0)
    return mask_orig,xyz_t,seq,msa,msa_full,mask,idx_ins

def rw_del_input(args,mask_orig,xyz_t,seq,msa,msa_full,mask,DEVICE):
    idx_del      = np.random.choice(np.where(mask_orig)[0], 1, replace=False)[0]
    xyz_t        = torch.cat([xyz_t[:,:,:idx_del], xyz_t[:,:,idx_del+1:]], 2)
    seq          = torch.cat([seq[:,:idx_del],seq[:,idx_del+1:]],1)
    msa          = torch.cat([msa[:,:,:idx_del],msa[:,:,idx_del+1:]],2)
    msa_full     = torch.cat([msa_full[:,:,:idx_del],msa_full[:,:,idx_del+1:]],2)
    mask_orig    = np.delete(mask_orig, idx_del, axis=0)
    mask         = np.delete(mask, idx_del, axis=0)
    return mask_orig,xyz_t,seq,msa,msa_full,mask,idx_del