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

Usage:
    python mol2_to_pdb_bfactor.py input.mol2 output.pdb [--resname MOL] [--write-types types.txt]

What it does:
 - Parses the @<TRIPOS>ATOM block of a mol2 file.
 - Extracts: atom_id, atom_name, x, y, z, atom_type, partial_charge
 - Writes a simple single-residue PDB where the B-factor column contains the partial_charge.
 - Optionally writes a text file mapping atom_index -> atom_type (AMBER atom type stored in mol2 atom_type).
"""

import sys
import argparse
import re

def parse_mol2_atoms(mol2_path):
    atoms = []
    in_atoms = False
    with open(mol2_path, 'r') as fh:
        for line in fh:
            if line.startswith('@<TRIPOS>ATOM'):
                in_atoms = True
                continue
            if line.startswith('@<TRIPOS>'):
                if in_atoms:
                    break
            if not in_atoms:
                continue
            # Typical ATOM line: id name x y z type charge [subst_id subst_name ...]
            parts = line.strip().split()
            if len(parts) < 6:
                continue
            atom_id = parts[0]
            atom_name = parts[1]
            try:
                x = float(parts[2])
                y = float(parts[3])
                z = float(parts[4])
            except ValueError:
                # skip malformed line
                continue
            atom_type = parts[5]
            # charge usually last token (but sometimes additional tokens exist)
            charge = None
            # find a token that parses as float from the end
            for tok in reversed(parts):
                try:
                    charge = float(tok)
                    break
                except:
                    continue
            if charge is None:
                charge = 0.0
            atoms.append({
                'id': atom_id,
                'name': atom_name,
                'x': x, 'y': y, 'z': z,
                'type': atom_type,
                'charge': charge
            })
    return atoms

def guess_element(atom_name, atom_type):
    # Simple heuristics: element is first char(s) of atom_name or atom_type
    # prefer explicit element-like prefixes: C, N, O, H, S, P, CL, BR, F, I
    el = None
    # from atom name
    m = re.match(r'^([A-Za-z]{1,2})', atom_name)
    if m:
        cand = m.group(1)
        cand = cand.capitalize()
        if cand in ("C","N","O","H","S","P","Cl","Br","F","I"):
            return cand.upper()
        # handle 'CA' -> C? (could be alpha carbon) but prefer single-letter if ambiguous
        if cand[0].upper() in "CNOHSPFIBrcl":
            return cand[0].upper()
    # fallback to atom_type prefix
    if atom_type:
        m2 = re.match(r'^([A-Za-z]{1,2})', atom_type)
        if m2:
            return m2.group(1).upper()
    # default C
    return 'C'

def write_pdb(atoms, out_path, resname='MOL'):
    """
    Write PDB with one residue (resSeq 1), chain A.
    B-factor column will be set to the partial charge.
    Atom name uses the mol2 atom name (padded/truncated to 4).
    """
    with open(out_path, 'w') as out:
        atom_idx = 1
        for a in atoms:
            element = guess_element(a['name'], a['type'])
            # Format PDB ATOM line (fixed-column). Keep it simple.
            atom_name = a['name'][:4].rjust(4)
            altLoc = ''
            resName = resname[:3].rjust(3)
            chainID = 'A'
            resSeq = 1
            x = a['x']
            y = a['y']
            z = a['z']
            occupancy = 1.00
            bfactor = a['charge']  # this is the important part
            # build formatted line: columns as in PDB ATOM format
            # Example format:
            # ATOM  %5d %4s %3s %1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f          %2s
            line = ("ATOM  {atom_idx:5d} {atom_name:<4s}{resName:>3s} {chainID:1s}"
                    "{resSeq:4d}    {x:8.3f}{y:8.3f}{z:8.3f}{occupancy:6.2f}{bfactor:6.2f}          {element:>2s}\n").format(
                atom_idx=atom_idx,
                atom_name=atom_name,
                resName=resName,
                chainID=chainID,
                resSeq=resSeq,
                x=x, y=y, z=z,
                occupancy=occupancy,
                bfactor=bfactor,
                element=element
            )
            out.write(line)
            atom_idx += 1
        out.write("END\n")

def write_types_file(atoms, types_path):
    with open(types_path, 'w') as fh:
        for i,a in enumerate(atoms, start=1):
            fh.write(f"{i}\t{a['name']}\t{a['type']}\t{a['charge']:.6f}\n")

def main():
    parser = argparse.ArgumentParser(description="Convert mol2 -> pdb with B-factor = partial charge")
    parser.add_argument('mol2', help='input mol2 file')
    parser.add_argument('pdb', help='output pdb file')
    parser.add_argument('--resname', default='MOL', help='Residue name to use in PDB (default MOL)')
    parser.add_argument('--write-types', help='Optional: write atom index -> atom type mapping to this file')
    args = parser.parse_args()

    atoms = parse_mol2_atoms(args.mol2)
    if not atoms:
        print("No atoms found in mol2 file or parsing failed.", file=sys.stderr)
        sys.exit(2)
    write_pdb(atoms, args.pdb, resname=args.resname)
    if args.write_types:
        write_types_file(atoms, args.write_types)
    print(f"Wrote {len(atoms)} atoms to {args.pdb}")
    if args.write_types:
        print(f"Wrote atom type mapping to {args.write_types}")

if __name__ == '__main__':
    main()