import os
import re
import struct
import argparse
import numpy as np

ATOM_RE = re.compile(r'^(ATOM|HETATM)\s+')

def ref_lines_and_atom_count(pdbqt_path: str):
    lines = []
    atom_lines_idx = []
    with open(pdbqt_path, "r", encoding="utf-8", errors="ignore") as f:
        for i, line in enumerate(f):
            lines.append(line)
            if ATOM_RE.match(line):
                atom_lines_idx.append(i)
    if not atom_lines_idx:
        raise ValueError(f"No ATOM/HETATM lines in template: {pdbqt_path}")
    return lines, atom_lines_idx

def dpb_read_header(dpb_path: str):
    with open(dpb_path, "rb") as f:
        hdr = f.read(64)
        if len(hdr) != 64:
            raise ValueError(f"Bad DPB header: {dpb_path}")
        magic, version, dtype_code, n_atoms, n_models, stride_bytes = struct.unpack("<4sIIIII", hdr[:24])
        if magic != b"DPB1":
            raise ValueError(f"Not DPB1: {dpb_path}")
        if dtype_code != 1:
            raise ValueError(f"Unsupported dtype_code={dtype_code} (expect 1=float32) in {dpb_path}")

        n_atoms = int(n_atoms)
        n_models = int(n_models)
        stride_bytes = int(stride_bytes)

        rec_size = 4 + n_atoms * 3 * 4
        if stride_bytes != rec_size:
            raise ValueError(f"Stride mismatch: stride={stride_bytes} rec={rec_size} in {dpb_path}")

        return n_atoms, n_models

def dpb_iter_coords(dpb_path: str, n_atoms: int):
    with open(dpb_path, "rb") as f:
        _ = f.read(64)  # header
        while True:
            mid_b = f.read(4)
            if not mid_b:
                break
            if len(mid_b) != 4:
                break
            _model_id = struct.unpack("<I", mid_b)[0]

            buf = f.read(n_atoms * 3 * 4)
            if len(buf) != n_atoms * 3 * 4:
                break

            coords = np.frombuffer(buf, dtype=np.float32).reshape(n_atoms, 3)
            yield _model_id, coords

def fill_template_lines(template_lines, atom_lines_idx, coords):
    out = list(template_lines)
    if coords.shape[0] != len(atom_lines_idx):
        raise ValueError(f"Atom count mismatch: dpb_atoms={coords.shape[0]} template_atoms={len(atom_lines_idx)}")

    for k, line_i in enumerate(atom_lines_idx):
        line = out[line_i]
        x, y, z = map(float, coords[k])
        # 保留原本 PDBQT 行的非座標欄位（原子名、電荷、AutoDock type…）
        out[line_i] = f"{line[:30]}{x:8.3f}{y:8.3f}{z:8.3f}{line[54:]}"
    return out

def write_multi_model_pdbqt(out_path, template_lines, atom_lines_idx, dpb_path):
    n_atoms, n_models = dpb_read_header(dpb_path)
    # 以 template 的 ATOM 數為準做檢查
    if len(atom_lines_idx) != n_atoms:
        raise ValueError(f"Atom mismatch: DPB says {n_atoms}, template has {len(atom_lines_idx)}")

    with open(out_path, "w", encoding="utf-8") as fw:
        pose_idx = 0
        for _mid, coords in dpb_iter_coords(dpb_path, n_atoms):
            pose_idx += 1
            fw.write(f"MODEL {pose_idx}\n")
            filled = fill_template_lines(template_lines, atom_lines_idx, coords)
            fw.writelines(filled)
            if not (filled and filled[-1].endswith("\n")):
                fw.write("\n")
            fw.write("ENDMDL\n")

def write_split_models(out_dir, base_name, template_lines, atom_lines_idx, dpb_path):
    os.makedirs(out_dir, exist_ok=True)
    n_atoms, _ = dpb_read_header(dpb_path)
    if len(atom_lines_idx) != n_atoms:
        raise ValueError(f"Atom mismatch: DPB says {n_atoms}, template has {len(atom_lines_idx)}")

    pose_idx = 0
    for _mid, coords in dpb_iter_coords(dpb_path, n_atoms):
        pose_idx += 1
        out_path = os.path.join(out_dir, f"{base_name}_pose{pose_idx:04d}.pdbqt")
        filled = fill_template_lines(template_lines, atom_lines_idx, coords)
        with open(out_path, "w", encoding="utf-8") as fw:
            fw.writelines(filled)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dpb", required=True, help="input .dpb (multi-poses)")
    ap.add_argument("--template", required=True, help="reference ligand .pdbqt as template (format/records)")
    ap.add_argument("--out", required=True, help="output .pdbqt (multi-model) OR output directory if --split")
    ap.add_argument("--split", action="store_true", help="if set, write one pdbqt per pose into --out directory")
    args = ap.parse_args()

    template_lines, atom_lines_idx = ref_lines_and_atom_count(args.template)

    dpb_base = os.path.splitext(os.path.basename(args.dpb))[0]

    if args.split:
        write_split_models(args.out, dpb_base, template_lines, atom_lines_idx, args.dpb)
        print(f"[OK] wrote split poses into: {args.out}")
    else:
        write_multi_model_pdbqt(args.out, template_lines, atom_lines_idx, args.dpb)
        print(f"[OK] wrote multi-model pdbqt: {args.out}")

if __name__ == "__main__":
    main()

