import sys

def readSequence(filename):
    """Return the sequence of residues from the base pair input file.
    """
    lines=open(filename).readlines()
    inSeq=False
    import string
    seq=[]
    for line in lines:
        if (line.startswith('>')):
            inSeq=True
            pass

        if inSeq and line[0].isalpha():
            seq.append( line.split() )
            #print(line)
            pass


        pass
    #print(seq)
    #print(sum(seq, []))
    return sum(seq, [])

def readResPDB(pdb_name):
    """
    Return a list of residues from a PDB file

    """
    xyz = []
    pdb_file = open(pdb_name, 'r')
    for line in pdb_file:
        n = line[12:16].strip()
        if (line.startswith("ATOM")):
            # extract n (atom name), rn (residue name), r (residue no.),x, y, z coordinates
            rn = line[17:20].strip()
            r = int(line[22:26].strip())
            x = float(line[30:38].strip())
            y = float(line[38:46].strip())
            z = float(line[46:54].strip())
            #if line[12:16].strip() == "CA":
            xyz.append([n, rn, r, x, y, z])            
    pdb_file.close()
    # define a residue dictionary
    res_dis = {}
    first_res = xyz[0][2]
    res_dis[first_res] = xyz[0][1]
    #print(first_res)
    for item in xyz:
        if (item[2] != first_res):
            res_dis[item[2]] = item[1]
        first_res = item[2]
    #print(res_dis)
    res = [res_dis[item][1] for item in range(1, len(res_dis)+1)]
    #print(res)
    return res

def checkError(seq, res_pdb):
    if len(seq) != len(res_pdb):
       print("number of residues in sequence in basepair file and pdb file doesnot match")
       sys.exit()
    count = 0
    err_count = 0
    for r_bp, r_pdb in zip(seq, res_pdb):
       count += 1
       if r_bp != r_pdb :
         print("residue: ", count, " does not match between basepair file and pdb file")
         err_count += 1
    if err_count > 0:
         sys.exit(1)


if __name__=="__main__":
    if len(sys.argv) == 3:
      file_name = sys.argv[1]
      seq = readSequence(file_name)
      pdb_file = sys.argv[2]
      res_pdb = readResPDB(pdb_file)
      checkError(seq, res_pdb)



