#
# Shift the nucleotides along an RNA polymer maintaining backbone and
# sugar positions while aligning bases at old location to base at new location.
#
# Rigidly move the base to align specific atoms using the 3 atoms where
# base joins sugar which should make plane of base match.
#
# I want to allow shifting modified nucleotide remdesivir.  The sugar
# as extra atoms and base 3 atoms where sugar attaches differ.
#
# Exact copy sugar and backbone atoms C1', C2', C3', C4', C5', O3', O4', P, OP1, OP2.
# Align base U/C C2,N1,C6 to A/G C4,N9,C8 to N1,C5,C7 for F86 (remdesivir).
# Need also lists of all base atom names.
# TODO: To move remdesivir extra 2 sugar atoms, use 3 closest sugar atoms to align.
#
backbone_and_sugar_atoms = ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"]
backbone_link_atoms = ["P", "O3'"]

# RNA bases
base_atoms = {
    'U': ["N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"],
    'C': ["N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],
    'A': ["N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],
    'G': ["N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],
    'F86': ["C8", "C5", "N1", "N3", "N2", "N4", "C7", "C10", "C11", "C12", "C9", "N5"], # Remdesivir
    }
base_alignment_atoms = {
    'U': ["C2", "N1", "C6"],
    'C': ["C2", "N1", "C6"],
    'A': ["C4", "N9", "C8"],
    'G': ["C4", "N9", "C8"],
    'F86': ["N1", "C5", "C7"], # Remdesivir
    }

def shift_rna(session, residues, offset = 1):
    n = len(residues)
    if offset >= n:
        offset = n-1
    elif offset <= -n:
        offset = -n+1

    a5p, a3p = backbone_link_atoms
    if offset >= 0:
        overlap_range = range(0,n-offset)
        overhang_atoms = residues[-offset:].atoms
        aend = residue_atom(residues[n-offset-1], a3p)
        aover = residue_atom(residues[n-offset], a5p)
        overhang_offset = aover.scene_coord - aend.scene_coord
    else:
        overlap_range = range(n-1,-(offset+1),-1)
        overhang_atoms = residues[:-offset].atoms
        aend = residue_atom(residues[-offset], a5p)
        aover = residue_atom(residues[-offset-1], a3p)
        overhang_offset = aover.scene_coord - aend.scene_coord

    for i in overlap_range:
        r0, r1 = residues[i], residues[i+offset]
        shift_residue(r0, r1)

    # Move residues beyond the end.
    shift = aend.scene_coord - aover.scene_coord + overhang_offset
    overhang_atoms.scene_coords += shift

def shift_residue(r0, r1):
    '''
    Shift RNA residue r0 to the position of residue r1,
    copying backbone and sugar coordinates and moving the
    nucleic base rigidly to align with the target base.
    '''
    # Copy backbone and sugar position
    bs0 = residue_atoms(r0, backbone_and_sugar_atoms)
    bs1 = residue_atoms(r1, backbone_and_sugar_atoms)
    for a0,a1 in zip(bs0, bs1):
        a0.scene_coord = a1.scene_coord

    # Rigidly move base.
    tf = align_base(r0, r1)
    batoms = residue_atoms(r0, base_atoms[r0.name])
    batoms.scene_coords = tf * batoms.scene_coords

def align_base(r0, r1):
    '''
    Find transform to move RNA base r0 to RNA base r1
    aligning the 3 atoms closest to the the sugar.
    '''
    atoms0 = residue_atoms(r0, base_alignment_atoms[r0.name])
    atoms1 = residue_atoms(r1, base_alignment_atoms[r1.name])
    from chimerax.std_commands.align import align
    tf = align(r0.structure.session, atoms0, atoms1, move = False)[4]
    return tf

def residue_atoms(r, atom_names):
    alist = [residue_atom(r, atom_name) for atom_name in atom_names]
    from chimerax.atomic import Atoms
    atoms = Atoms(alist)
    return atoms

def residue_atom(r, atom_name):
    a = r.find_atom(atom_name)
    if a is None:
        from chimerax.core.errors import UserError
        raise UserError('No atom %s in residue %s' % (atom_name, str(r)))
    return a

def register_command(session):
    from chimerax.core.commands import CmdDesc, register, IntArg
    from chimerax.atomic import ResiduesArg
    desc = CmdDesc(required = [('residues', ResiduesArg)],
                   keyword = [('offset', IntArg)],
                   synopsis = 'Shift RNA by N nucleotides')
    register('shiftrna', desc, shift_rna, logger=session.logger)

register_command(session)