#!/usr/bin/env python3
"""
Tool: pdb_rmsd_contact_map.py
Category: Structural Bioinformatics & Biophysics
Description: Parses 3D atomic coordinates from Protein Data Bank (PDB) structures,
             calculates Root Mean Square Deviation (RMSD) between protein conformations,
             and generates C-alpha Euclidean distance contact maps.
Usage:
    python pdb_rmsd_contact_map.py -p structure.pdb [--contact-cutoff 8.0]
"""

import argparse
import math
import sys

def parse_ca_coordinates(filepath):
    """Extracts C-alpha (CA) atom coordinates from PDB file format."""
    coords = []
    residues = []
    with open(filepath, 'r') as f:
        for line in f:
            if line.startswith('ATOM') and line[12:16].strip() == 'CA':
                res_name = line[17:20].strip()
                res_seq = int(line[22:26].strip())
                x = float(line[30:38].strip())
                y = float(line[38:46].strip())
                z = float(line[46:54].strip())
                coords.append((x, y, z))
                residues.append(f"{res_name}{res_seq}")
    return residues, coords

def calculate_rmsd(coords1, coords2):
    """Calculates coordinate Root Mean Square Deviation (RMSD) in Angstroms."""
    if len(coords1) != len(coords2) or len(coords1) == 0:
        raise ValueError("Coordinate sets must have identical non-zero dimensions.")
    sq_dist_sum = sum(
        (c1[0] - c2[0])**2 + (c1[1] - c2[1])**2 + (c1[2] - c2[2])**2
        for c1, c2 in zip(coords1, coords2)
    )
    return math.sqrt(sq_dist_sum / len(coords1))

def contact_map_matrix(coords, cutoff=8.0):
    """Identifies residue pairs within the Euclidean contact distance cutoff."""
    n = len(coords)
    contacts = []
    for i in range(n):
        for j in range(i + 1, n):
            c1, c2 = coords[i], coords[j]
            d = math.sqrt((c1[0] - c2[0])**2 + (c1[1] - c2[1])**2 + (c1[2] - c2[2])**2)
            if d <= cutoff:
                contacts.append((i, j, round(d, 2)))
    return contacts

def main():
    print("=== Protein PDB Structural Analysis & Contact Map ===")
    
    # Create demo mini PDB file (Alpha-helix fragment)
    demo_pdb = "/tmp/demo_helix.pdb"
    with open(demo_pdb, "w") as f:
        f.write("ATOM      1  CA  MET A   1      20.154  14.221   8.312  1.00 20.00           C\n")
        f.write("ATOM      8  CA  ALA A   2      22.450  16.890  10.120  1.00 20.00           C\n")
        f.write("ATOM     15  CA  LEU A   3      25.600  15.100  11.450  1.00 20.00           C\n")
        f.write("ATOM     22  CA  GLU A   4      24.120  12.300  13.200  1.00 20.00           C\n")
        f.write("ATOM     29  CA  LYS A   5      21.300  14.200  15.400  1.00 20.00           C\n")

    res, coords = parse_ca_coordinates(demo_pdb)
    print(f"Extracted {len(coords)} C-alpha atoms:")
    for r, c in zip(res, coords):
        print(f"  - {r:6}: ({c[0]:6.2f}, {c[1]:6.2f}, {c[2]:6.2f}) Å")

    contacts = contact_map_matrix(coords, cutoff=6.5)
    print(f"\nResidue-Residue Spatial Contacts (<= 6.5 Å):")
    for i, j, dist in contacts:
        print(f"  - {res[i]} <---> {res[j]}: {dist} Å")

if __name__ == "__main__":
    main()