import numpy as np
import pickle
import os

def GenerateFeature_alpha(ligand_name, working_dir):
    Cut = 12.0

    LIGELE = ['C','N','O','S','CN','CO','CS','NO','NS','OS','CCl','CBr','CP','CF','CNO','CNS','COS','NOS','CNOS','CNOSPFClBrI','H','CH','NH','OH','SH','CNH','COH','CSH','NOH','NSH','OSH','CNOH','CNSH','COSH','NOSH','CNOSH','CNOSPFClBrIH','CClH','CBrH','CPH','CFH']

    Feature_i = []
    pdb = ligand_name
    InFile = open(working_dir+'/'+ligand_name+'_alpha.pkl')
    BarCollection = pickle.load(InFile)
    for el in LIGELE:
        if 'lig_'+el in BarCollection.keys():
            Bars = BarCollection['lig_'+el]
            Bar0Birth = []; Bar0Death = []; Bar1Birth = []; Bar1Death = []; Bar2Birth = []; Bar2Death = [];
            for Bar in Bars:
                if Bar[2] < Bar[1]:
                    continue
                if Bar[2] > 12.0 and Bar[0] == 0: continue
                if Bar[2] > 12.0 and Bar[0] > 0: Bar[2] = 12.0
                if Bar[0] == 0:
                    Bar0Birth.append(Bar[1])
                    Bar0Death.append(Bar[2])
                if Bar[0] == 1:
                    Bar1Birth.append(Bar[1])
                    Bar1Death.append(Bar[2])
                if Bar[0] == 2:
                    Bar2Birth.append(Bar[1])
                    Bar2Death.append(Bar[2])
            if len(Bar0Birth) > 0:
                Bar0Birth = np.asarray(Bar0Birth, float)
                Bar0Death = np.asarray(Bar0Death, float)
            if len(Bar1Birth) > 0:
                Bar1Birth = np.asarray(Bar1Birth, float)
                Bar1Death = np.asarray(Bar1Death, float)
            if len(Bar2Birth) > 0:
                Bar2Birth = np.asarray(Bar2Birth, float)
                Bar2Death = np.asarray(Bar2Death, float)
            if len(Bar0Death) > 0:
                Feature_i.append(np.mean(Bar0Death[:]))
                Feature_i.append(np.std(Bar0Death[:]))
                Feature_i.append(np.max(Bar0Death[:]))
                Feature_i.append(np.min(Bar0Death[:]))
                Feature_i.append(np.sum(Bar0Death[:]))
                Feature_i.append(len(Bar0Death))
            else:
                Feature_i.extend([0.]*6)
            if len(Bar1Death) > 0:
                Feature_i.append(np.mean(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.std(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.max(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.min(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.sum(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(Bar1Birth[np.argmax(Bar1Death[:] - Bar1Birth[:])])
                Feature_i.append(Bar1Death[np.argmax(Bar1Death[:] - Bar1Birth[:])])
                Feature_i.append(np.mean(Bar1Birth[:]))
                Feature_i.append(np.std(Bar1Birth[:]))
                Feature_i.append(np.max(Bar1Birth[:]))
                Feature_i.append(np.min(Bar1Birth[:]))
                Feature_i.append(np.sum(Bar1Birth[:]))
                Feature_i.append(np.mean(Bar1Death[:]))
                Feature_i.append(np.std(Bar1Death[:]))
                Feature_i.append(np.max(Bar1Death[:]))
                Feature_i.append(np.min(Bar1Death[:]))
                Feature_i.append(np.sum(Bar1Death[:]))
                Feature_i.append(len(Bar1Death))
            else:
                Feature_i.extend([0.]*18)
            if len(Bar2Death) > 0:
                Feature_i.append(np.mean(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.std(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.max(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.min(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.sum(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(Bar2Birth[np.argmax(Bar2Death[:] - Bar2Birth[:])])
                Feature_i.append(Bar2Death[np.argmax(Bar2Death[:] - Bar2Birth[:])])
                Feature_i.append(np.mean(Bar2Birth[:]))
                Feature_i.append(np.std(Bar2Birth[:]))
                Feature_i.append(np.max(Bar2Birth[:]))
                Feature_i.append(np.min(Bar2Birth[:]))
                Feature_i.append(np.sum(Bar2Birth[:]))
                Feature_i.append(np.mean(Bar2Death[:]))
                Feature_i.append(np.std(Bar2Death[:]))
                Feature_i.append(np.max(Bar2Death[:]))
                Feature_i.append(np.min(Bar2Death[:]))
                Feature_i.append(np.sum(Bar2Death[:]))
                Feature_i.append(len(Bar2Death))
            else:
                Feature_i.extend([0.]*18)
        else:
            Feature_i.extend([0.]*42)
    Feature_i = np.asarray(Feature_i, float)

    outfile = open(working_dir+'/'+ligand_name+'_feature_alpha_handcrafted.npy', 'w')
    np.save(outfile, Feature_i)
    outfile.close()

def GenerateFeature_level1(ligand_name, working_dir):
    small = 0.01
    Feature_i = []
    Cut = 12.0

    LIGELE = ['C','N','O','S','CN','CO','CS','NO','NS','OS','CCl','CBr','CP','CF','CNO','CNS','COS','NOS','CNOS','CNOSPFClBrI','H','CH','NH','OH','SH','CNH','COH','CSH','NOH','NSH','OSH','CNOH','CNSH','COSH','NOSH','CNOSH','CNOSPFClBrIH','CClH','CBrH','CPH','CFH']

    pdb = ligand_name

    for el in LIGELE:
        if os.path.exists(working_dir+'/'+ligand_name+'_'+el+'_level1.PH'):
            InFile = open(working_dir+'/'+ligand_name+'_'+el+'_level1.PH')
            lines = InFile.read().splitlines()
            Bars = []
            for line in lines:
                a,b,c = line.split()
                Bars.append([int(a), float(b), float(c)])
            InFile.close()
            Bar0Birth = []; Bar0Death = []; Bar1Birth = []; Bar1Death = []; Bar2Birth = []; Bar2Death = [];
            for Bar in Bars:
                if Bar[2] < Bar[1]:
                    continue
                if Bar[2] > 12.0 and Bar[0] == 0: continue
                if Bar[2] > 12.0 and Bar[0] > 0: Bar[2] = 12.0
                if Bar[0] == 0 and Bar[2]-Bar[1] >= small:
                    Bar0Birth.append(Bar[1])
                    Bar0Death.append(Bar[2])
                if Bar[0] == 1 and Bar[2]-Bar[1] >= small:
                    Bar1Birth.append(Bar[1])
                    Bar1Death.append(Bar[2])
                if Bar[0] == 2 and Bar[2]-Bar[1] >= small:
                    Bar2Birth.append(Bar[1])
                    Bar2Death.append(Bar[2])
            if len(Bar0Birth) > 0:
                Bar0Birth = np.asarray(Bar0Birth, float)
                Bar0Death = np.asarray(Bar0Death, float)
            if len(Bar1Birth) > 0:
                Bar1Birth = np.asarray(Bar1Birth, float)
                Bar1Death = np.asarray(Bar1Death, float)
            if len(Bar2Birth) > 0:
                Bar2Birth = np.asarray(Bar2Birth, float)
                Bar2Death = np.asarray(Bar2Death, float)
            if len(Bar0Death) > 0:
                Feature_i.append(np.mean(Bar0Death[:]))
                Feature_i.append(np.std(Bar0Death[:]))
                Feature_i.append(np.max(Bar0Death[:]))
                Feature_i.append(np.min(Bar0Death[:]))
                Feature_i.append(np.sum(Bar0Death[:]))
                Feature_i.append(len(Bar0Death))
            else:
                Feature_i.extend([0.]*6)
            if len(Bar1Death) > 0:
                Feature_i.append(np.mean(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.std(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.max(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.min(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(np.sum(Bar1Death[:] - Bar1Birth[:]))
                Feature_i.append(Bar1Birth[np.argmax(Bar1Death[:] - Bar1Birth[:])])
                Feature_i.append(Bar1Death[np.argmax(Bar1Death[:] - Bar1Birth[:])])
                Feature_i.append(np.mean(Bar1Birth[:]))
                Feature_i.append(np.std(Bar1Birth[:]))
                Feature_i.append(np.max(Bar1Birth[:]))
                Feature_i.append(np.min(Bar1Birth[:]))
                Feature_i.append(np.sum(Bar1Birth[:]))
                Feature_i.append(np.mean(Bar1Death[:]))
                Feature_i.append(np.std(Bar1Death[:]))
                Feature_i.append(np.max(Bar1Death[:]))
                Feature_i.append(np.min(Bar1Death[:]))
                Feature_i.append(np.sum(Bar1Death[:]))
                Feature_i.append(len(Bar1Death))
            else:
                Feature_i.extend([0.]*18)
            if len(Bar2Death) > 0:
                Feature_i.append(np.mean(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.std(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.max(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.min(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(np.sum(Bar2Death[:] - Bar2Birth[:]))
                Feature_i.append(Bar2Birth[np.argmax(Bar2Death[:] - Bar2Birth[:])])
                Feature_i.append(Bar2Death[np.argmax(Bar2Death[:] - Bar2Birth[:])])
                Feature_i.append(np.mean(Bar2Birth[:]))
                Feature_i.append(np.std(Bar2Birth[:]))
                Feature_i.append(np.max(Bar2Birth[:]))
                Feature_i.append(np.min(Bar2Birth[:]))
                Feature_i.append(np.sum(Bar2Birth[:]))
                Feature_i.append(np.mean(Bar2Death[:]))
                Feature_i.append(np.std(Bar2Death[:]))
                Feature_i.append(np.max(Bar2Death[:]))
                Feature_i.append(np.min(Bar2Death[:]))
                Feature_i.append(np.sum(Bar2Death[:]))
                Feature_i.append(len(Bar2Death))
            else:
                Feature_i.extend([0.]*18)
        else:
            Feature_i.extend([0.]*42)
    Feature_i = np.asarray(Feature_i, float)

    outfile = open(working_dir+'/'+ligand_name+'_feature_ligand_level1_handcrafted.npy', 'w')
    np.save(outfile, Feature_i)
    outfile.close()

[(13, 'pickle.load', 'pickle.load', (['InFile'], {}), False, 'import pickle\n'), (94, 'numpy.asarray', 'np.asarray', (['Feature_i', 'float'], {}), True, 'import numpy as np\n'), (97, 'numpy.save', 'np.save', (['outfile', 'Feature_i'], {}), True, 'import numpy as np\n'), (195, 'numpy.asarray', 'np.asarray', (['Feature_i', 'float'], {}), True, 'import numpy as np\n'), (198, 'numpy.save', 'np.save', (['outfile', 'Feature_i'], {}), True, 'import numpy as np\n'), (110, 'os.path.exists', 'os.path.exists', (["(working_dir + '/' + ligand_name + '_' + el + '_level1.PH')"], {}), False, 'import os\n'), (33, 'numpy.asarray', 'np.asarray', (['Bar0Birth', 'float'], {}), True, 'import numpy as np\n'), (34, 'numpy.asarray', 'np.asarray', (['Bar0Death', 'float'], {}), True, 'import numpy as np\n'), (36, 'numpy.asarray', 'np.asarray', (['Bar1Birth', 'float'], {}), True, 'import numpy as np\n'), (37, 'numpy.asarray', 'np.asarray', (['Bar1Death', 'float'], {}), True, 'import numpy as np\n'), (39, 'numpy.asarray', 'np.asarray', (['Bar2Birth', 'float'], {}), True, 'import numpy as np\n'), (40, 'numpy.asarray', 'np.asarray', (['Bar2Death', 'float'], {}), True, 'import numpy as np\n'), (134, 'numpy.asarray', 'np.asarray', (['Bar0Birth', 'float'], {}), True, 'import numpy as np\n'), (135, 'numpy.asarray', 'np.asarray', (['Bar0Death', 'float'], {}), True, 'import numpy as np\n'), (137, 'numpy.asarray', 'np.asarray', (['Bar1Birth', 'float'], {}), True, 'import numpy as np\n'), (138, 'numpy.asarray', 'np.asarray', (['Bar1Death', 'float'], {}), True, 'import numpy as np\n'), (140, 'numpy.asarray', 'np.asarray', (['Bar2Birth', 'float'], {}), True, 'import numpy as np\n'), (141, 'numpy.asarray', 'np.asarray', (['Bar2Death', 'float'], {}), True, 'import numpy as np\n'), (42, 'numpy.mean', 'np.mean', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (43, 'numpy.std', 'np.std', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (44, 'numpy.max', 'np.max', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (45, 'numpy.min', 'np.min', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (46, 'numpy.sum', 'np.sum', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (51, 'numpy.mean', 'np.mean', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (52, 'numpy.std', 'np.std', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (53, 'numpy.max', 'np.max', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (54, 'numpy.min', 'np.min', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (55, 'numpy.sum', 'np.sum', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (58, 'numpy.mean', 'np.mean', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (59, 'numpy.std', 'np.std', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (60, 'numpy.max', 'np.max', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (61, 'numpy.min', 'np.min', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (62, 'numpy.sum', 'np.sum', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (63, 'numpy.mean', 'np.mean', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (64, 'numpy.std', 'np.std', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (65, 'numpy.max', 'np.max', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (66, 'numpy.min', 'np.min', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (67, 'numpy.sum', 'np.sum', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (72, 'numpy.mean', 'np.mean', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (73, 'numpy.std', 'np.std', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (74, 'numpy.max', 'np.max', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (75, 'numpy.min', 'np.min', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (76, 'numpy.sum', 'np.sum', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (79, 'numpy.mean', 'np.mean', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (80, 'numpy.std', 'np.std', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (81, 'numpy.max', 'np.max', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (82, 'numpy.min', 'np.min', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (83, 'numpy.sum', 'np.sum', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (84, 'numpy.mean', 'np.mean', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (85, 'numpy.std', 'np.std', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (86, 'numpy.max', 'np.max', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (87, 'numpy.min', 'np.min', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (88, 'numpy.sum', 'np.sum', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (143, 'numpy.mean', 'np.mean', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (144, 'numpy.std', 'np.std', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (145, 'numpy.max', 'np.max', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (146, 'numpy.min', 'np.min', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (147, 'numpy.sum', 'np.sum', (['Bar0Death[:]'], {}), True, 'import numpy as np\n'), (152, 'numpy.mean', 'np.mean', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (153, 'numpy.std', 'np.std', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (154, 'numpy.max', 'np.max', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (155, 'numpy.min', 'np.min', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (156, 'numpy.sum', 'np.sum', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (159, 'numpy.mean', 'np.mean', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (160, 'numpy.std', 'np.std', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (161, 'numpy.max', 'np.max', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (162, 'numpy.min', 'np.min', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (163, 'numpy.sum', 'np.sum', (['Bar1Birth[:]'], {}), True, 'import numpy as np\n'), (164, 'numpy.mean', 'np.mean', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (165, 'numpy.std', 'np.std', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (166, 'numpy.max', 'np.max', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (167, 'numpy.min', 'np.min', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (168, 'numpy.sum', 'np.sum', (['Bar1Death[:]'], {}), True, 'import numpy as np\n'), (173, 'numpy.mean', 'np.mean', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (174, 'numpy.std', 'np.std', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (175, 'numpy.max', 'np.max', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (176, 'numpy.min', 'np.min', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (177, 'numpy.sum', 'np.sum', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (180, 'numpy.mean', 'np.mean', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (181, 'numpy.std', 'np.std', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (182, 'numpy.max', 'np.max', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (183, 'numpy.min', 'np.min', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (184, 'numpy.sum', 'np.sum', (['Bar2Birth[:]'], {}), True, 'import numpy as np\n'), (185, 'numpy.mean', 'np.mean', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (186, 'numpy.std', 'np.std', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (187, 'numpy.max', 'np.max', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (188, 'numpy.min', 'np.min', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (189, 'numpy.sum', 'np.sum', (['Bar2Death[:]'], {}), True, 'import numpy as np\n'), (56, 'numpy.argmax', 'np.argmax', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (57, 'numpy.argmax', 'np.argmax', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (77, 'numpy.argmax', 'np.argmax', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (78, 'numpy.argmax', 'np.argmax', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (157, 'numpy.argmax', 'np.argmax', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (158, 'numpy.argmax', 'np.argmax', (['(Bar1Death[:] - Bar1Birth[:])'], {}), True, 'import numpy as np\n'), (178, 'numpy.argmax', 'np.argmax', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n'), (179, 'numpy.argmax', 'np.argmax', (['(Bar2Death[:] - Bar2Birth[:])'], {}), True, 'import numpy as np\n')]