"""
Remove Fragments not in Knowledgebase
"""

__author__ = "Michael Suarez"
__email__ = "masv@connect.ust.hk"
__copyright__ = "Copyright 2019, Hong Kong University of Science and Technology"
__license__ = "3-clause BSD"

from argparse import ArgumentParser
import numpy as np
import pickle

parser = ArgumentParser(description="Build Files")
parser.add_argument("--datadir", type=str, default="Data", help="input - XXX.YYY ")
parser.add_argument("--envNewAcronym", type=str, default="PRT.SNW", help="input - XXX.YYY ")

args = parser.parse_args()


# Check the Bound Fragments

BoundFrags = np.loadtxt("../%s/%s/%s.Homogenised.boundfrags_zeros.txt" %(args.datadir, args.envNewAcronym, args.envNewAcronym), delimiter=',')

normalDF = pickle.load(open("../%s/GrandCID.dict" %(args.datadir), "rb"))

binding = np.full(BoundFrags.shape,-1)
mlength = 0
for r, i in enumerate(BoundFrags):
    for c, j in enumerate(i[i!=0]):
        try:
            # Checks whether the Fragment can be found in the 59k Fragment Base
            binding[r,c]=normalDF.index.get_loc(int(j))
        except:
            continue
    temp = binding[r]
    if temp[temp!=-1].shape[0] > mlength:
        mlength = temp[temp!=-1].shape[0]
        print(mlength) #Finds the maximum number of Fragments per environment -> 705
        
indices = np.empty(binding.shape[0])
red_binding = np.full((binding.shape[0], mlength), -1)
for j, i in enumerate(binding):
    indices[j] = i[i!=-1].shape[0]
    red_binding[j][:int(indices[j])] = i[i!=-1]
red_binding = np.delete(red_binding, np.where(indices==0), axis=0)

pickle.dump(red_binding, open("../%s/%s/%s.binding.mtr"  %(args.datadir, args.envNewAcronym, args.envNewAcronym), "wb"))

# Removes environments without binding Fragments
Features_all = pickle.load(open("../%s/%s/%s.Homogenised.property.pvar"  %(args.datadir, args.envNewAcronym, args.envNewAcronym), "rb"))
Features_all = np.delete(Features_all, np.where(indices==0), axis=0)
pickle.dump(Features_all, open("../%s/%s/%s.Homogenised.property.pvar"  %(args.datadir, args.envNewAcronym, args.envNewAcronym), "wb"))

# Removes environment annotiation without binding fragments
with open("../%s/%s/%s.Homogenised.annotation.txt" %(args.datadir, args.envNewAcronym, args.envNewAcronym), "r+") as f:
    lines = f.readlines()
    for i in np.where(indices==0)[0][::-1]:
        del lines[i]
    f.seek(0)
    f.truncate()
    f.writelines(lines)

[((14, 9, 14, 50), 'argparse.ArgumentParser', 'ArgumentParser', (), '', False, 'from argparse import ArgumentParser\n'), ((23, 13, 23, 142), 'numpy.loadtxt', 'np.loadtxt', (), '', True, 'import numpy as np\n'), ((27, 10, 27, 38), 'numpy.full', 'np.full', ({(27, 18, 27, 34): 'BoundFrags.shape', (27, 35, 27, 37): '-1'}, {}), '(BoundFrags.shape, -1)', True, 'import numpy as np\n'), ((41, 10, 41, 36), 'numpy.empty', 'np.empty', ({(41, 19, 41, 35): 'binding.shape[0]'}, {}), '(binding.shape[0])', True, 'import numpy as np\n'), ((42, 14, 42, 54), 'numpy.full', 'np.full', ({(42, 22, 42, 49): '(binding.shape[0], mlength)', (42, 51, 42, 53): '-1'}, {}), '((binding.shape[0], mlength), -1)', True, 'import numpy as np\n'), ((46, 37, 46, 57), 'numpy.where', 'np.where', ({(46, 46, 46, 56): 'indices == 0'}, {}), '(indices == 0)', True, 'import numpy as np\n'), ((52, 39, 52, 59), 'numpy.where', 'np.where', ({(52, 48, 52, 58): 'indices == 0'}, {}), '(indices == 0)', True, 'import numpy as np\n'), ((58, 13, 58, 33), 'numpy.where', 'np.where', ({(58, 22, 58, 32): '(indices == 0)'}, {}), '(indices == 0)', True, 'import numpy as np\n')]