Topology Handling

# Access topology attributes from atoms
atoms = u.atoms

# Basic atomic properties
names = atoms.names        # Atom names (e.g., 'CA', 'CB', 'N')
types = atoms.types        # Atom types (e.g., 'CT', 'NH1', 'O')  
masses = atoms.masses      # Atomic masses in amu
charges = atoms.charges    # Partial charges in elementary charge units
radii = atoms.radii        # Van der Waals radii in Angstrom

# Residue-level properties  
resnames = atoms.resnames  # Residue names (e.g., 'ALA', 'GLY', 'PRO')
resids = atoms.resids      # Residue IDs/numbers
resnums = atoms.resnums    # Sequential residue numbers

# Segment-level properties
segids = atoms.segids      # Segment identifiers

def add_TopologyAttr(universe, topologyattr, values=None):
    """
    Add a topology attribute to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    topologyattr : str or TopologyAttr
        Name of topology attribute or TopologyAttr class instance.
    values : array-like, optional
        Values for the new attribute. If None, default values are used.
        
    Examples
    --------
    >>> # Add masses for all atoms
    >>> u.add_TopologyAttr('masses', [12.0, 14.0, 16.0] * (u.atoms.n_atoms // 3))
    
    >>> # Add custom attribute
    >>> u.add_TopologyAttr('occupancy', np.ones(u.atoms.n_atoms))
    """

def del_TopologyAttr(universe, topologyattr):
    """
    Remove a topology attribute from the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    topologyattr : str
        Name of topology attribute to remove.
        
    Examples
    --------
    >>> u.del_TopologyAttr('masses')
    >>> # Now u.atoms.masses will raise AttributeError
    """

# Available topology attributes
class AtomAttr:
    """Base class for atom-level topology attributes."""
    
    # Standard atom attributes
    names = "Atom names"
    types = "Atom types" 
    masses = "Atomic masses"
    charges = "Partial charges"
    radii = "Van der Waals radii"
    bfactors = "B-factors from crystallography"
    occupancies = "Occupancies from crystallography"
    altlocs = "Alternate locations from crystallography"
    elements = "Chemical elements"
    ids = "Atom IDs/indices"

class ResidueAttr:
    """Base class for residue-level topology attributes."""
    
    resnames = "Residue names"
    resids = "Residue IDs"
    resnums = "Sequential residue numbers" 
    resnames = "Residue names"
    icodes = "Insertion codes"

class SegmentAttr:
    """Base class for segment-level topology attributes."""
    
    segids = "Segment identifiers"

def add_bonds(universe, values, types=None, guessed=False, order=None):
    """
    Add bond information to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    values : array-like
        Bond connectivity as array of shape (n_bonds, 2) containing
        atom indices for each bond.
    types : array-like, optional
        Bond types for each bond (e.g., 'single', 'double', 'aromatic').
    guessed : bool, optional
        Whether bonds were guessed from distances (default False).
    order : array-like, optional
        Bond orders as numerical values (1.0, 2.0, 1.5 for aromatic).
        
    Examples
    --------
    >>> # Add bonds between consecutive atoms
    >>> bond_list = [(i, i+1) for i in range(u.atoms.n_atoms-1)]
    >>> u.add_bonds(bond_list)
    
    >>> # Add typed bonds
    >>> u.add_bonds([(0, 1), (1, 2)], types=['single', 'double'])
    
    >>> # Add bonds with orders
    >>> u.add_bonds([(0, 1), (2, 3)], order=[1.0, 2.0])
    """

def delete_bonds(universe, values):
    """
    Remove specific bonds from the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    values : array-like
        Bonds to remove as array of shape (n_bonds, 2).
        
    Examples
    --------
    >>> # Remove specific bonds
    >>> bonds_to_remove = [(10, 11), (15, 16)]
    >>> u.delete_bonds(bonds_to_remove)
    """

def guess_bonds(atomgroup, vdwradii=None, fudge_factor=0.55, lower_bound=0.1):
    """
    Guess bonds based on inter-atomic distances.
    
    Parameters
    ----------
    atomgroup : AtomGroup
        Atoms for bond guessing.
    vdwradii : dict, optional
        Van der Waals radii for elements. If None, uses default values.
    fudge_factor : float, optional
        Tolerance factor for bond distance (default 0.55).
    lower_bound : float, optional
        Minimum distance for bond formation (default 0.1).
        
    Returns
    -------
    numpy.ndarray
        Array of shape (n_bonds, 2) containing guessed bonds.
        
    Examples
    --------
    >>> # Guess all bonds in system
    >>> bonds = u.atoms.guess_bonds()
    >>> u.add_bonds(bonds, guessed=True)
    
    >>> # Custom radii for specific elements
    >>> custom_radii = {'C': 0.77, 'N': 0.70, 'O': 0.66}
    >>> bonds = u.atoms.guess_bonds(vdwradii=custom_radii)
    """

# Access existing bonds
bonds = u.atoms.bonds  # TopologyGroup containing all bonds
bond_indices = bonds.indices  # Array of shape (n_bonds, 2)
bond_types = bonds.types  # Bond types if available

def add_angles(universe, values, types=None, guessed=False):
    """
    Add angle information to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    values : array-like
        Angle connectivity as array of shape (n_angles, 3) containing
        atom indices (i, j, k) where j is the central atom.
    types : array-like, optional
        Angle types for each angle.
    guessed : bool, optional
        Whether angles were guessed from connectivity (default False).
        
    Examples
    --------
    >>> # Add angles from bond connectivity
    >>> angles = [(0, 1, 2), (1, 2, 3), (2, 3, 4)]
    >>> u.add_angles(angles)
    """

def add_dihedrals(universe, values, types=None, guessed=False):
    """
    Add dihedral angle information to the Universe.
    
    Parameters
    ----------  
    universe : Universe
        Universe to modify.
    values : array-like
        Dihedral connectivity as array of shape (n_dihedrals, 4) containing
        atom indices (i, j, k, l) defining the dihedral.
    types : array-like, optional
        Dihedral types for each dihedral.
    guessed : bool, optional
        Whether dihedrals were guessed from connectivity (default False).
        
    Examples
    --------
    >>> # Add backbone dihedrals
    >>> dihedrals = [(0, 1, 2, 3), (1, 2, 3, 4)]
    >>> u.add_dihedrals(dihedrals)
    """

def add_impropers(universe, values, types=None, guessed=False):
    """
    Add improper dihedral information to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    values : array-like  
        Improper connectivity as array of shape (n_impropers, 4).
    types : array-like, optional
        Improper types for each improper.
    guessed : bool, optional
        Whether impropers were guessed (default False).
        
    Examples
    --------
    >>> # Add improper dihedrals for planarity  
    >>> impropers = [(0, 1, 2, 3), (4, 5, 6, 7)]
    >>> u.add_impropers(impropers)
    """

# Access connectivity information
angles = u.atoms.angles        # All angles
dihedrals = u.atoms.dihedrals  # All dihedrals  
impropers = u.atoms.impropers  # All impropers

# Get indices for connectivity
angle_indices = angles.indices      # Shape: (n_angles, 3)
dihedral_indices = dihedrals.indices  # Shape: (n_dihedrals, 4)

def add_Residue(universe, segment=None, **attrs):
    """
    Add a new residue to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    segment : Segment, optional
        Segment to add residue to. If None, added to first segment.
    **attrs
        Attributes for the new residue (resname, resid, etc.).
        
    Returns
    -------
    Residue
        New residue object.
        
    Examples
    --------
    >>> # Add new residue  
    >>> new_res = u.add_Residue(resname='LIG', resid=999)
    
    >>> # Add to specific segment
    >>> seg = u.segments[0] 
    >>> new_res = u.add_Residue(segment=seg, resname='WAT', resid=1000)
    """

def add_Segment(universe, **attrs):
    """
    Add a new segment to the Universe.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    **attrs
        Attributes for the new segment (segid, etc.).
        
    Returns
    -------
    Segment
        New segment object.
        
    Examples
    --------
    >>> # Add new segment
    >>> new_seg = u.add_Segment(segid='LIG')
    """

# Adding atoms requires more complex topology manipulation
# Usually done through specialized functions or parsers

def guess_TopologyAttrs(universe, to_guess=None, skip_fields=None):
    """
    Guess missing topology attributes from available information.
    
    Parameters
    ----------
    universe : Universe
        Universe to modify.
    to_guess : list, optional
        List of attributes to guess. If None, guesses all possible.
    skip_fields : list, optional
        List of attributes to skip during guessing.
        
    Examples
    --------
    >>> # Guess all missing attributes
    >>> u.guess_TopologyAttrs()
    
    >>> # Guess only specific attributes
    >>> u.guess_TopologyAttrs(to_guess=['masses', 'bonds'])
    
    >>> # Skip certain fields
    >>> u.guess_TopologyAttrs(skip_fields=['charges'])
    """

# Available guessers
from MDAnalysis.guesser import (
    guess_bonds,      # Guess bonds from distances
    guess_angles,     # Guess angles from bonds  
    guess_dihedrals,  # Guess dihedrals from angles
    guess_atom_type,  # Guess atom types from names
    guess_atom_mass,  # Guess masses from types/elements
    guess_atom_charge # Guess charges from types
)

class TopologyGroup:
    """
    Base class for topology objects like bonds, angles, dihedrals.
    
    Provides array-like access to connectivity information with
    additional methods for analysis.
    """
    
    @property
    def indices(self):
        """
        Connectivity indices for topology objects.
        
        Returns
        -------
        numpy.ndarray
            Array of atom indices defining connectivity.
            Shape depends on topology type:
            - Bonds: (n_bonds, 2)  
            - Angles: (n_angles, 3)
            - Dihedrals: (n_dihedrals, 4)
        """
    
    @property
    def types(self):
        """
        Types for each topology object.
        
        Returns
        -------
        numpy.ndarray or None
            Array of type strings, None if no type information.
        """
    
    def __len__(self):
        """Number of topology objects."""
    
    def __getitem__(self, key):
        """Access specific topology objects by index."""

# Specific topology groups
class BondGroup(TopologyGroup):
    """Group of bonds with bond-specific methods."""
    
    def length(self, pbc=False):
        """
        Calculate bond lengths.
        
        Parameters
        ----------
        pbc : bool, optional
            Whether to consider periodic boundary conditions.
            
        Returns
        -------
        numpy.ndarray
            Array of bond lengths in Angstrom.
        """
    
    def partner(self, atom):
        """
        Get bonding partners for specified atom.
        
        Parameters
        ----------
        atom : Atom or int
            Atom or atom index.
            
        Returns
        -------
        AtomGroup
            Atoms bonded to specified atom.
        """

class AngleGroup(TopologyGroup):
    """Group of angles with angle-specific methods."""
    
    def angle(self, pbc=False):
        """
        Calculate angle values.
        
        Parameters
        ----------
        pbc : bool, optional
            Whether to consider periodic boundary conditions.
            
        Returns
        -------
        numpy.ndarray
            Array of angle values in degrees.
        """

class DihedralGroup(TopologyGroup):
    """Group of dihedrals with dihedral-specific methods."""
    
    def dihedral(self, pbc=False):
        """
        Calculate dihedral angle values.
        
        Parameters
        ----------
        pbc : bool, optional
            Whether to consider periodic boundary conditions.
            
        Returns
        -------
        numpy.ndarray
            Array of dihedral angle values in degrees.
        """

# Available parsers for different formats
from MDAnalysis.topology import (
    PSFParser,      # CHARMM PSF format
    PDBParser,      # Protein Data Bank format  
    GROParser,      # GROMACS GRO format
    TPRParser,      # GROMACS TPR binary format
    PrmtopParser,   # AMBER topology format
    MOL2Parser,     # Tripos MOL2 format
    PQRParser,      # PQR format (PDB with charges/radii)
    CRDParser,      # CHARMM CRD format
    XYZParser       # Generic XYZ format
)

class TopologyReaderBase:
    """
    Base class for topology parsers.
    
    All topology parsers inherit from this class and implement
    the parse() method to extract topology information.
    """
    
    def parse(self, **kwargs):
        """
        Parse topology information from file.
        
        Returns
        -------
        dict
            Dictionary containing parsed topology data with keys:
            - 'atoms': atom information
            - 'bonds': bond connectivity  
            - 'residues': residue information
            - 'segments': segment information
            - Additional format-specific data
        """

# Example: Custom topology parsing
class CustomParser(TopologyReaderBase):
    """Example custom topology parser."""
    
    format = 'CUSTOM'
    
    def parse(self, **kwargs):
        """Parse custom topology format."""
        # Implementation details
        atoms = self._parse_atoms()
        bonds = self._parse_bonds()
        residues = self._parse_residues()
        
        return {
            'atoms': atoms,
            'bonds': bonds, 
            'residues': residues
        }

# PSF files contain complete topology information
u = mda.Universe("system.psf")

# Access PSF-specific information
print(u.atoms.types)      # CHARMM atom types
print(u.atoms.charges)    # Partial charges
print(u.atoms.bonds)      # Complete bond connectivity
print(u.atoms.angles)     # All angles
print(u.atoms.dihedrals)  # All dihedrals  
print(u.atoms.impropers)  # Improper dihedrals

# TPR files are binary and contain full force field information
u = mda.Universe("topol.tpr")

# TPR provides extensive topology data
print(u.atoms.masses)     # Atomic masses
print(u.atoms.charges)    # Partial charges
print(u.atoms.types)      # GROMACS atom types
print(u.atoms.bonds)      # Bond connectivity

# PRMTOP files contain AMBER force field topology
u = mda.Universe("system.prmtop")

# AMBER-specific topology information
print(u.atoms.types)      # AMBER atom types  
print(u.atoms.charges)    # Partial charges
print(u.atoms.masses)     # Atomic masses
print(u.atoms.bonds)      # Bond information

# Create universe and examine topology
u = mda.Universe("topology.psf", "trajectory.dcd")

# Check available topology attributes
print("Available attributes:")
for attr in ['names', 'types', 'masses', 'charges']:
    if hasattr(u.atoms, attr):
        print(f"  {attr}: {getattr(u.atoms, attr)[:5]}...")  # First 5 values

# Connectivity information
print(f"Number of bonds: {len(u.atoms.bonds)}")
print(f"Number of angles: {len(u.atoms.angles)}")  
print(f"Number of dihedrals: {len(u.atoms.dihedrals)}")

# Bond analysis
bond_lengths = u.atoms.bonds.length()
print(f"Average bond length: {np.mean(bond_lengths):.2f} Å")

# Start with minimal topology (e.g., from PDB)
u = mda.Universe("structure.pdb")

# Add missing information
if not hasattr(u.atoms, 'masses'):
    print("Guessing atomic masses...")
    u.guess_TopologyAttrs(to_guess=['masses'])

if not hasattr(u.atoms, 'bonds'):
    print("Guessing bond connectivity...")
    u.guess_TopologyAttrs(to_guess=['bonds'])

# Add custom attributes
occupancies = np.random.random(u.atoms.n_atoms)
u.add_TopologyAttr('occupancy', occupancies)

# Now access new attributes
print(f"First 5 masses: {u.atoms.masses[:5]}")
print(f"Number of bonds: {len(u.atoms.bonds)}")
print(f"Average occupancy: {np.mean(u.atoms.occupancy):.2f}")

def analyze_connectivity(universe):
    """
    Analyze connectivity patterns in molecular system.
    """
    atoms = universe.atoms
    
    # Bond statistics
    if hasattr(atoms, 'bonds'):
        bonds = atoms.bonds
        bond_lengths = bonds.length()
        
        print("Bond Analysis:")
        print(f"  Total bonds: {len(bonds)}")
        print(f"  Average length: {np.mean(bond_lengths):.2f} ± {np.std(bond_lengths):.2f} Å")
        print(f"  Length range: {np.min(bond_lengths):.2f} - {np.max(bond_lengths):.2f} Å")
    
    # Coordination numbers
    coordination = {}
    for atom in atoms:
        n_bonds = len(atom.bonds)
        coordination[n_bonds] = coordination.get(n_bonds, 0) + 1
    
    print("\nCoordination Analysis:")
    for coord, count in sorted(coordination.items()):
        print(f"  {count} atoms with {coord} bonds")
    
    # Fragment analysis  
    if hasattr(atoms, 'bonds'):
        fragments = atoms.split('fragments')
        print(f"\nFragment Analysis:")
        print(f"  Number of fragments: {len(fragments)}")
        fragment_sizes = [len(frag) for frag in fragments]
        print(f"  Largest fragment: {max(fragment_sizes)} atoms")
        print(f"  Average fragment size: {np.mean(fragment_sizes):.1f} atoms")

# Run connectivity analysis
analyze_connectivity(u)

def modify_topology_example(universe):
    """
    Example of dynamic topology modification during simulation.
    """
    # Monitor and modify bonds based on distance criteria
    cutoff = 1.8  # Bond formation cutoff in Angstrom
    
    for ts in universe.trajectory:
        # Get current positions
        positions = universe.atoms.positions
        
        # Calculate all pairwise distances
        from MDAnalysis.lib.distances import distance_array
        dist_matrix = distance_array(positions, positions, 
                                   box=universe.dimensions)
        
        # Find potential new bonds (close atoms not already bonded)
        existing_bonds = set()
        if hasattr(universe.atoms, 'bonds'):
            for bond in universe.atoms.bonds:
                i, j = bond.indices[0]
                existing_bonds.add((min(i, j), max(i, j)))
        
        # Identify new bonds
        new_bonds = []
        for i in range(len(positions)):
            for j in range(i + 1, len(positions)):
                if dist_matrix[i, j] < cutoff and (i, j) not in existing_bonds:
                    new_bonds.append((i, j))
        
        # Add new bonds if found
        if new_bonds:
            print(f"Frame {ts.frame}: Adding {len(new_bonds)} new bonds")
            universe.add_bonds(new_bonds, guessed=True)
            
        # Break bonds that are too long
        if hasattr(universe.atoms, 'bonds'):
            bonds_to_remove = []
            for bond in universe.atoms.bonds:
                i, j = bond.indices[0]
                if dist_matrix[i, j] > 2.5:  # Breaking threshold
                    bonds_to_remove.append((i, j))
            
            if bonds_to_remove:
                print(f"Frame {ts.frame}: Removing {len(bonds_to_remove)} broken bonds")
                universe.delete_bonds(bonds_to_remove)

# Example usage (be careful with trajectory modification!)
# modify_topology_example(u)

def validate_topology(universe):
    """
    Validate topology consistency and completeness.
    """
    atoms = universe.atoms
    issues = []
    
    # Check for missing critical attributes
    required_attrs = ['names', 'resnames', 'segids']
    for attr in required_attrs:
        if not hasattr(atoms, attr):
            issues.append(f"Missing required attribute: {attr}")
    
    # Check for reasonable masses
    if hasattr(atoms, 'masses'):
        masses = atoms.masses
        if np.any(masses <= 0):
            issues.append("Found zero or negative masses")
        if np.any(masses > 300):  # Unusually heavy atoms
            issues.append("Found unusually heavy atoms (>300 amu)")
    
    # Check bond lengths
    if hasattr(atoms, 'bonds'):
        bond_lengths = atoms.bonds.length()
        if np.any(bond_lengths < 0.5):
            issues.append("Found unusually short bonds (<0.5 Å)")
        if np.any(bond_lengths > 3.0):
            issues.append("Found unusually long bonds (>3.0 Å)")
    
    # Check for disconnected residues
    if hasattr(atoms, 'bonds'):
        for residue in atoms.residues:
            res_atoms = residue.atoms
            res_bonds = [bond for bond in atoms.bonds 
                        if all(idx in res_atoms.indices for idx in bond.indices)]
            if len(res_bonds) == 0 and len(res_atoms) > 1:
                issues.append(f"Residue {residue.resname}{residue.resid} has no internal bonds")
    
    # Report results
    if issues:
        print("Topology validation found issues:")
        for issue in issues:
            print(f"  - {issue}")
    else:
        print("Topology validation passed - no issues found")
    
    return len(issues) == 0

# Validate topology
is_valid = validate_topology(u)