# Align a structure to specified residues in another structure.  The residues in the first structure
# to pair are determined by sequence alignments (one per chain).  This is for comparing dimer interfaces
# of alphafold predictions to homologous experimental structures.

def ialign(session, structure, to_residues, cutoff_distance = None):
    to_structs = to_residues.unique_structures
    if len(to_structs) > 1:
        snames = ', '.join(str(s) for s in to_structs)
        from chimerax.core.errors import UserError
        raise UserError(f'Can only align to a single structure, got {snames}')
    to_struct = to_structs[0]

    to_res, res = paired_residues(to_residues, structure)
    report_missing_ca(res)
    report_missing_ca(to_res)

    if len(to_res) == 0:
        from chimerax.core.errors import UserError
        raise UserError(f'No residues of the {len(to_residues)} specified were paired via sequence alignments')

    session.logger.info(f'Paired {len(to_res)} of {len(to_residues)} residues using sequence alignments')

    if len(to_res) < len(to_residues):
        mres = to_residues.subtract(to_res)
        from chimerax.atomic import concise_residue_spec
        mnames = concise_residue_spec(session, mres)
        session.logger.info(f'{len(mres)} residues were not paired: {mnames}')

    ca = res.find_existing_atoms('CA')
    to_ca = to_res.find_existing_atoms('CA')
    from chimerax.std_commands.align import align
    align(session, ca, to_ca)

    res.atoms.selected = True
    
    '''
    rnums = ','.join(str(r.number) for r in res)
    to_rnums = ','.join(str(r.number) for r in to_res)
    # TODO: Missing chain id
    cmd = f'align #{structure.id_string}:{rnums}@CA to #{to_struct.id_string}:{to_rnums}@CA'
    from chimerax.core.commands import run
    run(session, cmd)
    '''

def paired_residues(residues, structure):
    # Map chain to residues in chain
    cres = {}
    for r in residues:
        c = r.chain
        if c in cres:
            cres[c].append(r)
        else:
            cres[c] = [r]

    # Find sequence alignments containing the chains for the residues
    # and a chain of the structure and map the residues to their
    # corresponding residues in the structure.
    am = structure.session.alignments
    r1, r2 = [], []
    for c, cresidues in cres.items():
        calign = []
        for a in am:
            if c in a.associations:
                for c2 in structure.chains:
                    if c2 in a.associations:
                        calign.append((a, c2))
        if len(calign) == 0:
            from chimerax.core.errors import UserError
            raise UserError(f'No sequence alignment for {c} was found.')
        if len(calign) > 1:
            from chimerax.core.errors import UserError
            anames = ', '.join(a.ident for a,s in calign)
            raise UserError(f'More than one sequence alignment for {c} was found: {anames}.')
        a, sc = calign[0]
        aseq = a.associations[c]
        sseq = a.associations[sc]
        cmm = aseq.match_maps[c]
        smm = sseq.match_maps[sc]
        cpos = cmm.res_to_pos
        sres = smm.pos_to_res
        for r in cresidues:
            if r in cpos:
                ug_pos = cpos[r]  # Ungapped
                pos = aseq.ungapped_to_gapped(ug_pos)  # Gapped
                ug_spos = sseq.gapped_to_ungapped(pos) # Ungapped in structure seq
                if ug_spos in sres:
                    r1.append(r)
                    res2 = sres[ug_spos]
                    r2.append(res2)
                    #print (f'{r} ungapped {ug_pos}, gapped {pos}, ungapped {ug_spos}, res {res2}')

    from chimerax.atomic import Residues    
    return Residues(r1), Residues(r2)

def report_missing_ca(residues):
    rmiss = [r for r,a in zip(residues, residues.find_atoms('CA')) if a is None]
    if rmiss:
        from chimerax.atomic import concise_residue_spec
        rspec = concise_residue_spec(rmiss[0].structure.session, rmiss)
        from chimerax.core.errors import UserError
        raise UserError(f'Missing CA atoms: {rspec}')

def register_command(logger):
    from chimerax.core.commands import CmdDesc, register, FloatArg
    from chimerax.atomic import StructureArg, ResiduesArg
    desc = CmdDesc(
        required = [('structure', StructureArg)],
        keyword = [('to_residues', ResiduesArg),
                   ('cutoff_distance', FloatArg),
                   ],
        required_arguments = ['to_residues'],
        synopsis = 'Align residues of a structure to residues of another structure paired using sequence alignments'
    )
    register('ialign', desc, ialign, logger=logger)

register_command(session.logger)