# -*- coding: utf-8 -*-
"""
Подсчёт ДНК-белковых контактов в 1I3J.
Белок - цепь A, ДНК - цепи B и C.

Критерии:
    полярный контакт    - атомы O/N белка и ДНК на расстоянии <= 3.5 A
    неполярный контакт   - атомы C/P/S белка и ДНК на расстоянии <= 4.5 A

ДНК разбита на 4 категории атомов:
    2'-дезоксирибоза   : C1',C2',C3',C4',C5',O4'
    фосфатная группа   : P,OP1,OP2,O5',O3'
    основание, большая бороздка (major_atoms)
    основание, малая бороздка (minor_atoms)

атомы большой/малой бороздки для A/G/T/C ниже заданы по стандартной
геометрической классификации Сеемана

"""
from pymol import cmd

PDB_FILE = "X:/2/1i3j.pdb"
PROTEIN = "chain A"
DNA = "(chain B or chain C)"

# атомы оснований, обращённые в большую / малую бороздку (по типу основания)
MAJOR_ATOMS = {
    "DA": "N6+N7+C5+C6+C8",
    "DG": "O6+N7+C5+C6+C8",
    "DT": "O4+C5+C6+C7",
    "DC": "N4+C5+C6",
}
MINOR_ATOMS = {
    "DA": "N1+N3+C2+C4",
    "DG": "N2+N3+N9+C2+C4",
    "DT": "O2+N3",
    "DC": "O2+N3",
}

def _base_sel(atom_map):
    parts = []
    for resn, atoms in atom_map.items():
        parts.append(f"(resn {resn} and name {atoms})")
    return " or ".join(parts)

def count_contacts():
    cmd.reinitialize()
    cmd.load(PDB_FILE, "s")

    deoxyribose = f"s and {DNA} and name C1'+C2'+C3'+C4'+C5'+O4'"
    phosphate   = f"s and {DNA} and name P+OP1+OP2+O5'+O3'"
    major       = f"s and {DNA} and ({_base_sel(MAJOR_ATOMS)})"
    minor       = f"s and {DNA} and ({_base_sel(MINOR_ATOMS)})"

    categories = {
        "2'-дезоксирибоза": deoxyribose,
        "фосфатная группа": phosphate,
        "основания (большая бороздка)": major,
        "основания (малая бороздка)": minor,
    }

    protein_polar    = f"s and {PROTEIN} and (elem N+O)"
    protein_nonpolar = f"s and {PROTEIN} and (elem C+P+S)"

    print(f"{'Категория':35s} {'Полярные':>10s} {'Неполярные':>12s} {'Всего':>8s}")
    for label, sel in categories.items():
        dna_polar    = f"({sel}) and (elem N+O)"
        dna_nonpolar = f"({sel}) and (elem C+P+S)"

        polar_pairs = cmd.find_pairs(dna_polar, protein_polar, cutoff=3.5)
        nonpolar_pairs = cmd.find_pairs(dna_nonpolar, protein_nonpolar, cutoff=4.5)

        n_polar = len(polar_pairs)
        n_nonpolar = len(nonpolar_pairs)
        print(f"{label:35s} {n_polar:10d} {n_nonpolar:12d} {n_polar+n_nonpolar:8d}")


cmd.extend("count_contacts", count_contacts)
count_contacts()