# Read the two dimer confidence score lists and open and align
# predicted structures on the interface residues to see if the
# predictions are similar.

def dimer_alignment(session, pdb_path1, pdb_path2, resA, resB):
    from chimerax.core.commands import run
    run(session, f'open "{pdb_path1}"')
    run(session, f'open "{pdb_path2}"')
    rnA = ','.join(str(rnum) for rnum in resA)
    rnB = ','.join(str(rnum) for rnum in resB)
    from_atoms, to_atoms, rmsd, ful_rmsd, tf = run(session, f'align #2/A:{rnA}/B:{rnB}@CA to #1/A:{rnA}/B:{rnB}@CA')
    run(session, 'close')
    return rmsd
    
def align_paired_structures(session, confidence1_path, confidence2_path):
    with open(dimer_conf1, 'r') as file:
        conf1_lines = file.readlines()
        gene_lines1 = {name:line for name, line in zip(parse_gene_names(conf1_lines), conf1_lines)}
    with open(dimer_conf2, 'r') as file:
        conf2_lines = file.readlines()
        gene_lines2 = {name:line for name, line in zip(parse_gene_names(conf2_lines), conf2_lines)}

    from os.path import dirname, join
    dir1, dir2 = dirname(confidence1_path), dirname(confidence2_path)
    lines = []
    for gene_names, line1 in gene_lines1.items():
        if gene_names in gene_lines2:
            line2 = gene_lines2[gene_names]
            pdb_file1, nresA1, nresB1, niA1, niB1, npair1, nconf1, pconf1, resA1, resB1 = line1.split(',')
            pdb_file2, nresA2, nresB2, niA2, niB2, npair2, nconf2, pconf2, resA2, resB2 = line2.split(',')
            pdb_path1, pdb_path2 = join(dir1, pdb_file1), join(dir2, pdb_file2)
            resA1_set = set(int(rnum) for rnum in resA1.split())
            resB1_set = set(int(rnum) for rnum in resB1.split())
            resA2_set = set(int(rnum) for rnum in resA2.split())
            resB2_set = set(int(rnum) for rnum in resB2.split())
            resA = list(resA1_set & resA2_set)
            resA.sort()
            resB = list(resB1_set & resB2_set)
            resB.sort()
            if len(resA) + len(resB) >= 3:
                rmsd = dimer_alignment(session, pdb_path1, pdb_path2, resA, resB)
                rmsd = "%.3g" % rmsd
            else:
                rmsd = ''
            lines.append(f'{pdb_path1},{pdb_path2},{niA1},{niA2},{len(resA)},{niB1},{niB2},{len(resB)},{rmsd}')
    return lines
        
def parse_gene_names(lines):
    gene_names = []
    from os.path import basename
    for line in lines:
        pdb_path = line[:line.find(',')]
        file_name = basename(pdb_path)
        name_end = file_name.find('_unrelaxed')
        gene_names.append(file_name[:name_end])
    return gene_names

from sys import argv
dimer_conf1, dimer_conf2 = argv[1:3]
lines = align_paired_structures(session, dimer_conf1, dimer_conf2)
print ('\n'.join(lines) + '\n')