import shutil


def combine_ligands_gro(in_file_A, in_file_B, out_file, ligand_A='MOL', ligand_B='MOL'):
    """Add ligand B coordinates to coordinate (.gro) file of ligand A in complex with protein
    Parameters
    ----------
    in_file_A : .gro file
        Complex A coordinate file
    in_file_B : .gro file
        Complex B coordinate file
    out_file : .gro file
        Complex A + ligand B coordinate file
    ligand_A: str
        Three letter code for ligand A
    ligand_B: str
        Three letter code for ligand B
    """
    # make copy of in_file_A ---> new outfile
    shutil.copy(in_file_A, out_file)
    file = open(in_file_B, 'r')
    text_B = file.readlines()
    file.close()
    file = open(out_file, 'r')
    text_A = file.readlines()
    file = open(out_file, 'w')
    lig2 = []
    #Iterate over complex B, store ligand B lines
    for idx, line in enumerate(text_B):
        if ligand_B in line:
            lig2.append(line)

    lig1 = []
    # Iterate over complex A, store ligand A lines
    for idx, line in enumerate(text_A):
        if ligand_A in line:
            lig1.append(line)
    # Iterate over complex A, insert ligand B lines
    for idx, line in enumerate(text_A):
        # Add number of atoms of ligandB to total number of atoms systemA
        if idx == 1:
            line = '%i\n' % (int(line) + len(lig2))
            # Insert ligandB coordinates right after ligandA
        if line == lig1[-1]:
            count = 0
            for i in lig2:
                index = idx + count + 1
                text_A.insert(index, i)
                count += 1
        file.write(line)
    file.close()

    return out_file