#!/usr/bin/env python3
"""
pdb_bfactor_diff.py

Compute B-factor differences between two PDBs:
    DeltaB = B_factor(pdb2) - B_factor(pdb1)

Usage:
    python pdb_bfactor_diff.py ref.pdb target.pdb diff.pdb

Assumptions:
 - Same number of atoms
 - Same atom ordering (same serial numbers)
"""

import sys

def read_pdb_atoms(pdb_path):
    atoms = []
    with open(pdb_path, 'r') as fh:
        for line in fh:
            if line.startswith(("ATOM  ", "HETATM")):
                try:
                    serial = int(line[6:11])
                    bfactor = float(line[60:66])
                except ValueError:
                    continue
                atoms.append({
                    "line": line.rstrip("\n"),
                    "serial": serial,
                    "bfactor": bfactor
                })
    return atoms

def write_pdb_diff(atoms1, atoms2, out_path):
    if len(atoms1) != len(atoms2):
        raise RuntimeError("PDB files have different numbers of atoms")

    with open(out_path, 'w') as out:
        for a1, a2 in zip(atoms1, atoms2):
            if a1["serial"] != a2["serial"]:
                raise RuntimeError(
                    f"Atom serial mismatch: {a1['serial']} vs {a2['serial']}"
                )

            delta_b = a2["bfactor"] - a1["bfactor"]

            # Replace B-factor field (columns 61–66)
            line = a1["line"]
            new_line = (
                line[:60] +
                f"{delta_b:6.2f}" +
                line[66:]
            )
            out.write(new_line + "\n")

        out.write("END\n")

def main():
    if len(sys.argv) != 4:
        print("Usage: python pdb_bfactor_diff.py pdb1.pdb pdb2.pdb diff.pdb")
        sys.exit(1)

    pdb1, pdb2, out_pdb = sys.argv[1:]

    atoms1 = read_pdb_atoms(pdb1)
    atoms2 = read_pdb_atoms(pdb2)

    if not atoms1 or not atoms2:
        print("Error: failed to read atoms from one or both PDBs")
        sys.exit(2)

    write_pdb_diff(atoms1, atoms2, out_pdb)
    print(f"Wrote B-factor difference PDB: {out_pdb}")

if __name__ == "__main__":
    main()