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', False, 'import pickle\n'), (94, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (97, 'numpy.save', 'np.save', True, 'import numpy as np\n'), (195, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (198, 'numpy.save', 'np.save', True, 'import numpy as np\n'), (110, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (33, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (34, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (36, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (37, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (39, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (40, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (134, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (135, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (137, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (138, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (140, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (141, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (42, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (43, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (44, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (45, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (46, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (51, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (52, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (53, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (54, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (55, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (58, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (59, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (60, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (61, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (62, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (63, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (64, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (65, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (66, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (67, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (72, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (73, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (74, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (75, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (76, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (79, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (80, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (81, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (82, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (83, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (84, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (85, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (86, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (87, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (88, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (143, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (144, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (145, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (146, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (147, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (152, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (153, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (154, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (155, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (156, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (159, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (160, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (161, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (162, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (163, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (164, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (165, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (166, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (167, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (168, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (173, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (174, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (175, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (176, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (177, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (180, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (181, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (182, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (183, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (184, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (185, 'numpy.mean', 'np.mean', True, 'import numpy as np\n'), (186, 'numpy.std', 'np.std', True, 'import numpy as np\n'), (187, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (188, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (189, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (56, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (57, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (77, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (78, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (157, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (158, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (178, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n'), (179, 'numpy.argmax', 'np.argmax', True, 'import numpy as np\n')]