from pymol import cmd

def shortest_permol(mol_1, mol_2):
    '''
    ref: mol_1
    compute rmsd of mol_2
    '''
    mol1_model = cmd.get_model(mol_1)
    mol2_model = cmd.get_model(mol_2)
    num = 0
    lst_atom_dis = []
    dis_shorest = 1000
    for a in mol2_model.atom:
        # num += 1
        for b in mol1_model.atom:
            # print("model %s and index %d"%(mol_2, a.index),"model %s and index %d"%(mol_1, b.index))
            dis = cmd.get_distance("model %s and index %d"%(mol_2, a.index), "model %s and index %d"%(mol_1, b.index))
            if dis_shorest > dis:
                dis_shorest = dis
        lst_atom_dis.append(dis_shorest)
        num += dis_shorest
    return num/len(lst_atom_dis)

def rmsd_three(mol_1, mol_2):
    refmol1_commol2 = shortest_permol(mol_1, mol_2)
    refmol2_commol1 = shortest_permol(mol_2, mol_1)
    rmsd = round((refmol1_commol2 + refmol2_commol1)/2,2)
    # print(refmol1_commol2, refmol2_commol1)
    print("The rmsd of " + mol_1 + " between " + mol_2 + " is " + str(rmsd))
    return rmsd

cmd.extend("rmsd_three", rmsd_three)