# -*- coding: utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# fluctmatch --- https://github.com/tclick/python-fluctmatch
# Copyright (c) 2013-2017 The fluctmatch Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the New BSD license.
#
# Please cite your use of fluctmatch in published work:
#
# <NAME>, <NAME>, and <NAME>.
# Calculation of Enzyme Fluctuograms from All-Atom Molecular Dynamics
# Simulation. Meth Enzymology. 578 (2016), 327-342,
# doi:10.1016/bs.mie.2016.05.024.
#
from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
)
from future.builtins import (
    super, )

import numpy as np
from MDAnalysis.core import selection


class BioIonSelection(selection.Selection):
    """Contains atoms commonly found in proteins.
    """
    token = "bioion"
    ion_atoms = np.array(["MG", "CAL", "MN", "FE", "CU", "ZN", "AG"])

    def __init__(self, parser, tokens):
        pass

    def apply(self, group):
        mask = np.in1d(group.names, self.ion_atoms)
        return group[mask].unique


class WaterSelection(selection.Selection):
    """Contains atoms commonly found in water.
    """
    token = "water"
    water_atoms = np.array(["OW", "HW1", "HW2", "MW"])

    def __init__(self, parser, tokens):
        pass

    def apply(self, group):
        mask = np.in1d(group.names, self.water_atoms)
        return group[mask].unique


class BackboneSelection(selection.BackboneSelection):
    """Contains all heavy atoms within a protein backbone including the terminal carboxyl oxygens.
    """
    token = "backbone"
    oxy_atoms = ["OXT", "OT1", "OT2"]

    def apply(self, group):
        mask = np.in1d(group.names,
                       np.concatenate([self.bb_atoms, self.oxy_atoms]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HBackboneSelection(BackboneSelection):
    """Includes all atoms found within a protein backbone including hydrogens.
    """
    token = "hbackbone"
    hbb_atoms = np.array([
        "H", "HN", "H1", "H2", "H3", "HT1", "HT2", "HT3", "HA", "HA1", "HA2",
        "1HA", "2HA"
    ])

    def apply(self, group):
        mask = np.in1d(group.names,
                       np.concatenate(
                           [self.bb_atoms, self.oxy_atoms, self.hbb_atoms]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CalphaSelection(selection.ProteinSelection):
    """Contains only the alpha-carbon of a protein.
    """
    token = "calpha"
    calpha = np.array(["CA"])

    def apply(self, group):
        mask = np.in1d(group.names, self.calpha)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HCalphaSelection(CalphaSelection):
    """Contains the alpha-carbon and alpha-hydrogens of a protein.
    """
    token = "hcalpha"
    hcalpha = np.array(["HA", "HA1", "HA2", "1HA", "2HA"])

    def apply(self, group):
        mask = np.in1d(group.names, np.concatenate([self.calpha,
                                                    self.hcalpha]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CbetaSelection(selection.ProteinSelection):
    """Contains only the beta-carbon of a protein.
    """
    token = "cbeta"
    cbeta = np.array(["CB"])

    def apply(self, group):
        mask = np.in1d(group.names, self.cbeta)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class AmineSelection(selection.ProteinSelection):
    """Contains atoms within the amine group of a protein.
    """
    token = "amine"
    amine = np.array(["N", "HN", "H", "H1", "H2", "H3", "HT1", "HT2", "HT3"])

    def apply(self, group):
        mask = np.in1d(group.names, self.amine)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CarboxylSelection(selection.ProteinSelection):
    """Contains atoms within the carboxyl group of a protein.
    """
    token = "carboxyl"
    carboxyl = np.array(["C", "O", "OXT", "OT1", "OT2"])

    def apply(self, group):
        mask = np.in1d(group.names, self.carboxyl)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HSidechainSelection(HBackboneSelection):
    """Includes hydrogens on the protein sidechain.
    """
    token = "hsidechain"

    def apply(self, group):
        mask = np.in1d(
            group.names,
            np.concatenate([self.bb_atoms, self.oxy_atoms, self.hbb_atoms]),
            invert=True)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class AdditionalNucleicSelection(selection.NucleicSelection):
    """Contains additional nucleic acid residues."""
    token = "nucleic"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.nucl_res = np.concatenate((self.nucl_res, ["OXG", "HPX", "DC35"]), axis=0)

    def apply(self, group):
        mask = np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class HNucleicSugarSelection(AdditionalNucleicSelection,
                             selection.NucleicSugarSelection):
    """Contains the additional atoms definitions for the sugar.
    """
    token = "hnucleicsugar"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.sug_atoms = np.concatenate(
            (self.sug_atoms,
             np.array([
                 "H1'", "O1'", "O2'", "H2'", "H2''", "O3'", "H3'", "H3T", "H4'"
             ])),
            axis=0)

    def apply(self, group):
        mask = np.in1d(group.names, self.sug_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class HBaseSelection(AdditionalNucleicSelection, selection.BaseSelection):
    """Contains additional atoms on the base region of the nucleic acids.
    """
    token = "hnucleicbase"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.base_atoms = np.concatenate(
            (self.base_atoms, [
                "O8", "H8", "H21", "H22", "H2", "O6", "H6", "H61", "H62",
                "H41", "H42", "H5", "H51", "H52", "H53", "H3", "H7"
            ]),
            axis=0)

    def apply(self, group):
        mask = np.in1d(group.names, self.base_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicPhosphateSelection(AdditionalNucleicSelection):
    """Contains the nucleic phosphate group including the C5'.
    """
    token = "nucleicphosphate"
    phos_atoms = np.array(
        ["P", "O1P", "O2P", "O5'", "C5'", "H5'", "H5''", "O5T", "H5T"])

    def apply(self, group):
        mask = np.in1d(group.names, self.phos_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicC2Selection(AdditionalNucleicSelection):
    """Contains the definition for the C3' region.
    """
    token = "sugarC2"
    c3_atoms = np.array([
        "C1'",
        "H1'",
        "C2'",
        "O2'",
        "H2'",
        "H2''",
    ])

    def apply(self, group):
        mask = np.in1d(group.names, self.c3_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicC4Selection(AdditionalNucleicSelection):
    """Contains the definition for the C4' region.
    """
    token = "sugarC4"
    c3_atoms = np.array([
        "C3'",
        "O3'",
        "H3'",
        "H3T",
        "C4'",
        "O4'",
        "H4'",
    ])

    def apply(self, group):
        mask = np.in1d(group.names, self.c3_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class BaseCenterSelection(AdditionalNucleicSelection):
    """Contains the central atoms (C4 and C5) on the base of the nuleic acid.
    """
    token = "nucleiccenter"
    center_atoms = np.array(["C4", "C5"])

    def apply(self, group):
        mask = np.in1d(group.names, self.center_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique

# -*- coding: utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# fluctmatch --- https://github.com/tclick/python-fluctmatch
# Copyright (c) 2013-2017 The fluctmatch Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the New BSD license.
#
# Please cite your use of fluctmatch in published work:
#
# <NAME>, <NAME>, and <NAME>.
# Calculation of Enzyme Fluctuograms from All-Atom Molecular Dynamics
# Simulation. Meth Enzymology. 578 (2016), 327-342,
# doi:10.1016/bs.mie.2016.05.024.
#
from __future__ import (
    absolute_import,
    division,
    print_function,
    unicode_literals,
)
from future.builtins import (
    super, )

import numpy as np
from MDAnalysis.core import selection


class BioIonSelection(selection.Selection):
    """Contains atoms commonly found in proteins.
    """
    token = "bioion"
    ion_atoms = np.array(["MG", "CAL", "MN", "FE", "CU", "ZN", "AG"])

    def __init__(self, parser, tokens):
        pass

    def apply(self, group):
        mask = np.in1d(group.names, self.ion_atoms)
        return group[mask].unique


class WaterSelection(selection.Selection):
    """Contains atoms commonly found in water.
    """
    token = "water"
    water_atoms = np.array(["OW", "HW1", "HW2", "MW"])

    def __init__(self, parser, tokens):
        pass

    def apply(self, group):
        mask = np.in1d(group.names, self.water_atoms)
        return group[mask].unique


class BackboneSelection(selection.BackboneSelection):
    """Contains all heavy atoms within a protein backbone including the terminal carboxyl oxygens.
    """
    token = "backbone"
    oxy_atoms = ["OXT", "OT1", "OT2"]

    def apply(self, group):
        mask = np.in1d(group.names,
                       np.concatenate([self.bb_atoms, self.oxy_atoms]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HBackboneSelection(BackboneSelection):
    """Includes all atoms found within a protein backbone including hydrogens.
    """
    token = "hbackbone"
    hbb_atoms = np.array([
        "H", "HN", "H1", "H2", "H3", "HT1", "HT2", "HT3", "HA", "HA1", "HA2",
        "1HA", "2HA"
    ])

    def apply(self, group):
        mask = np.in1d(group.names,
                       np.concatenate(
                           [self.bb_atoms, self.oxy_atoms, self.hbb_atoms]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CalphaSelection(selection.ProteinSelection):
    """Contains only the alpha-carbon of a protein.
    """
    token = "calpha"
    calpha = np.array(["CA"])

    def apply(self, group):
        mask = np.in1d(group.names, self.calpha)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HCalphaSelection(CalphaSelection):
    """Contains the alpha-carbon and alpha-hydrogens of a protein.
    """
    token = "hcalpha"
    hcalpha = np.array(["HA", "HA1", "HA2", "1HA", "2HA"])

    def apply(self, group):
        mask = np.in1d(group.names, np.concatenate([self.calpha,
                                                    self.hcalpha]))
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CbetaSelection(selection.ProteinSelection):
    """Contains only the beta-carbon of a protein.
    """
    token = "cbeta"
    cbeta = np.array(["CB"])

    def apply(self, group):
        mask = np.in1d(group.names, self.cbeta)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class AmineSelection(selection.ProteinSelection):
    """Contains atoms within the amine group of a protein.
    """
    token = "amine"
    amine = np.array(["N", "HN", "H", "H1", "H2", "H3", "HT1", "HT2", "HT3"])

    def apply(self, group):
        mask = np.in1d(group.names, self.amine)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class CarboxylSelection(selection.ProteinSelection):
    """Contains atoms within the carboxyl group of a protein.
    """
    token = "carboxyl"
    carboxyl = np.array(["C", "O", "OXT", "OT1", "OT2"])

    def apply(self, group):
        mask = np.in1d(group.names, self.carboxyl)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class HSidechainSelection(HBackboneSelection):
    """Includes hydrogens on the protein sidechain.
    """
    token = "hsidechain"

    def apply(self, group):
        mask = np.in1d(
            group.names,
            np.concatenate([self.bb_atoms, self.oxy_atoms, self.hbb_atoms]),
            invert=True)
        mask &= np.in1d(group.resnames, self.prot_res)
        return group[mask].unique


class AdditionalNucleicSelection(selection.NucleicSelection):
    """Contains additional nucleic acid residues."""
    token = "nucleic"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.nucl_res = np.concatenate((self.nucl_res, ["OXG", "HPX", "DC35"]), axis=0)

    def apply(self, group):
        mask = np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class HNucleicSugarSelection(AdditionalNucleicSelection,
                             selection.NucleicSugarSelection):
    """Contains the additional atoms definitions for the sugar.
    """
    token = "hnucleicsugar"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.sug_atoms = np.concatenate(
            (self.sug_atoms,
             np.array([
                 "H1'", "O1'", "O2'", "H2'", "H2''", "O3'", "H3'", "H3T", "H4'"
             ])),
            axis=0)

    def apply(self, group):
        mask = np.in1d(group.names, self.sug_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class HBaseSelection(AdditionalNucleicSelection, selection.BaseSelection):
    """Contains additional atoms on the base region of the nucleic acids.
    """
    token = "hnucleicbase"

    def __init__(self, parser, tokens):
        super().__init__(parser, tokens)
        self.base_atoms = np.concatenate(
            (self.base_atoms, [
                "O8", "H8", "H21", "H22", "H2", "O6", "H6", "H61", "H62",
                "H41", "H42", "H5", "H51", "H52", "H53", "H3", "H7"
            ]),
            axis=0)

    def apply(self, group):
        mask = np.in1d(group.names, self.base_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicPhosphateSelection(AdditionalNucleicSelection):
    """Contains the nucleic phosphate group including the C5'.
    """
    token = "nucleicphosphate"
    phos_atoms = np.array(
        ["P", "O1P", "O2P", "O5'", "C5'", "H5'", "H5''", "O5T", "H5T"])

    def apply(self, group):
        mask = np.in1d(group.names, self.phos_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicC2Selection(AdditionalNucleicSelection):
    """Contains the definition for the C3' region.
    """
    token = "sugarC2"
    c3_atoms = np.array([
        "C1'",
        "H1'",
        "C2'",
        "O2'",
        "H2'",
        "H2''",
    ])

    def apply(self, group):
        mask = np.in1d(group.names, self.c3_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class NucleicC4Selection(AdditionalNucleicSelection):
    """Contains the definition for the C4' region.
    """
    token = "sugarC4"
    c3_atoms = np.array([
        "C3'",
        "O3'",
        "H3'",
        "H3T",
        "C4'",
        "O4'",
        "H4'",
    ])

    def apply(self, group):
        mask = np.in1d(group.names, self.c3_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique


class BaseCenterSelection(AdditionalNucleicSelection):
    """Contains the central atoms (C4 and C5) on the base of the nuleic acid.
    """
    token = "nucleiccenter"
    center_atoms = np.array(["C4", "C5"])

    def apply(self, group):
        mask = np.in1d(group.names, self.center_atoms)
        mask &= np.in1d(group.resnames, self.nucl_res)
        return group[mask].unique

# -*- coding: utf-8 -*- # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 # # fluctmatch --- https://github.com/tclick/python-fluctmatch # Copyright (c) 2013-2017 The fluctmatch Development Team and contributors # (see the file AUTHORS for the full list of names) # # Released under the New BSD license. # # Please cite your use of fluctmatch in published work: # # <NAME>, <NAME>, and <NAME>. # Calculation of Enzyme Fluctuograms from All-Atom Molecular Dynamics # Simulation. Meth Enzymology. 578 (2016), 327-342, # doi:10.1016/bs.mie.2016.05.024. # Contains atoms commonly found in proteins. Contains atoms commonly found in water. Contains all heavy atoms within a protein backbone including the terminal carboxyl oxygens. Includes all atoms found within a protein backbone including hydrogens. Contains only the alpha-carbon of a protein. Contains the alpha-carbon and alpha-hydrogens of a protein. Contains only the beta-carbon of a protein. Contains atoms within the amine group of a protein. Contains atoms within the carboxyl group of a protein. Includes hydrogens on the protein sidechain. Contains additional nucleic acid residues. Contains the additional atoms definitions for the sugar. Contains additional atoms on the base region of the nucleic acids. Contains the nucleic phosphate group including the C5'. Contains the definition for the C3' region. Contains the definition for the C4' region. Contains the central atoms (C4 and C5) on the base of the nuleic acid.