#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Dec  3 11:17:08 2020

.. codeauthor:: Daniel Esteban Palma Igor <daniel.palma.i@ug.uchile.cl>
"""

from chimerax.core.commands import run
from chimerax.atomic.struct_edit import add_atom

# if num_atoms == 1, the atom is added to the original Chain,
# but if num_atoms >= 2 a new Chain is created
num_atoms = 2

def print_chains(info, structure):
    for chain in structure.chains:
        info(str(chain))
        for residue in chain.existing_residues[-5:]:
            info(str(residue))
        info("")

def filter_backbone(atoms):
    return atoms.filter(atoms.is_backbones())

info = session.logger.info
structure = run(session, "open 2efv")[0][0]
info(f"num_atoms: {num_atoms}")
info("Before adding atoms:")
print_chains(info, structure) # print
sel_atoms = filter_backbone(structure.atoms)
last_atom = sel_atoms[-1]
i = last_atom.residue.number + 1
pbgroup = structure.pseudobond_group("missing structure")
chain_id = last_atom.residue.chain_id
for _ in range(num_atoms):
    res = structure.new_residue("UNK", chain_id, i)
    new_atom = add_atom("CA", "C", res, last_atom.coord - 4)
    pbgroup.new_pseudobond(last_atom, new_atom)
    last_atom = new_atom
    i += 1
info("After adding atoms:")
print_chains(info, structure) # print