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

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

def sdirs(p):
    return sorted([x for x in os.listdir(p) if os.path.isdir(os.path.join(p, x))])

def read_dock_table(path):
    rows = []
    with open(path, "r", encoding="utf-8", errors="ignore") as fr:
        for line in fr:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            receptor = parts[0]
            ligands = parts[1:]
            if ligands:
                rows.append((receptor, ligands))
    return rows

def ref_coords_and_lines(pdbqt_path):
    coords = []
    lines = []
    with open(pdbqt_path, "r", encoding="utf-8", errors="ignore") as fr:
        for line in fr:
            lines.append(line)
            if ATOM_RE.match(line):
                coords.append((float(line[30:38]), float(line[38:46]), float(line[46:54])))
    if not coords:
        raise ValueError(pdbqt_path)
    return np.asarray(coords, dtype=np.float32), lines

def rmsd(a, b):
    d = a - b
    return float(np.sqrt(np.mean(np.sum(d * d, axis=1), dtype=np.float32), dtype=np.float32))

def dpb_iter_coords(dpb_path):
    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} in {dpb_path}")
        n_atoms = int(n_atoms)
        stride_bytes = int(stride_bytes)

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

        while True:
            mid_b = f.read(4)
            if not mid_b:
                break
            if len(mid_b) != 4:
                break
            _ = 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 coords

def best_pose_from_dpb(dpb_path, ref):
    best_r = 1e30
    best_coords = None
    for coords in dpb_iter_coords(dpb_path):
        if coords.shape != ref.shape:
            continue
        r = rmsd(coords, ref)
        if r < best_r:
            best_r = r
            best_coords = coords
    return best_r, best_coords

def write_pose_from_template(out_path, template_lines, new_coords):
    out_lines = []
    k = 0
    for line in template_lines:
        if ATOM_RE.match(line):
            x, y, z = map(float, new_coords[k])
            k += 1
            out_lines.append(f"{line[:30]}{x:8.3f}{y:8.3f}{z:8.3f}{line[54:]}")
        else:
            out_lines.append(line)
    with open(out_path, "w", encoding="utf-8") as fw:
        fw.writelines(out_lines)

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--score-root", required=True, help="e.g. ../3_dock/0_Dis_0.7_1.3/lig_multi_pos")
    ap.add_argument("--dock-table", required=True, help="e.g. ../1_dataset/9_summary/dock_table")
    ap.add_argument("--map-root", required=True, help="e.g. ../1_dataset/9_summary/lig_exp_pos")
    ap.add_argument("--out-root", required=True, help="e.g. ./rmsd")
    ap.add_argument("--alphas", nargs="*", default=None,
                    help="指定要跑哪些 alpha 子資料夾；不給就自動掃描 score-root 下的子資料夾")
    ap.add_argument("--suffix", default=SUFFIX_DEFAULT, help="dock 檔案副檔名，預設 .dpb")
    return ap.parse_args()

def main():
    args = parse_args()

    os.makedirs(args.out_root, exist_ok=True)

    alphas = args.alphas if args.alphas is not None and len(args.alphas) > 0 else sdirs(args.score_root)
    table_rows = read_dock_table(args.dock_table)

    for alpha in alphas:
        alpha_dir = os.path.join(args.score_root, alpha)
        summary = []

        for receptor, ligands in table_rows:
            rcp_dir = os.path.join(alpha_dir, receptor)
            if not os.path.isdir(rcp_dir):
                continue

            row = [f"{receptor:>5s}"]
            out_dir = os.path.join(args.out_root, alpha, receptor)
            os.makedirs(out_dir, exist_ok=True)

            for lig in ligands:
                dock_path = os.path.join(rcp_dir, f"{lig}{args.suffix}")
                ref_path  = os.path.join(args.map_root, lig, f"{lig}_ligand.pdbqt")

                if not os.path.isfile(dock_path) or not os.path.isfile(ref_path):
                    row += [f"{lig:>5s}", "  NA   "]
                    continue

                ref, ref_lines = ref_coords_and_lines(ref_path)
                r, best_coords = best_pose_from_dpb(dock_path, ref)

                if best_coords is None or r >= 1e29:
                    row += [f"{lig:>5s}", "  NA   "]
                    continue

                out_pose = os.path.join(out_dir, f"{lig}_rmsd.pdbqt")
                write_pose_from_template(out_pose, ref_lines, best_coords)

                row += [f"{lig:>5s}", f"{r:7.3f}"]

            summary.append(" ".join(row))

        with open(os.path.join(args.out_root, f"rmsd_{alpha}"), "w", encoding="utf-8") as fw:
            fw.write("\n".join(summary) + "\n")
        print(f"[OK] alpha={alpha}")

if __name__ == "__main__":
    main()

