import os
import argparse
import struct
import torch

DTYPE = torch.float32
GPU_DEVICE = "cuda"

# Defaults (overridden by CLI)
IJ_BATCH = None

DIV = 0.375
CLASH_SCALE = 0.50
REC_FILTER_RAD = 13.0
REC_CHUNK = 4096
ENABLE_CLASH_CHECK = True

ELEM_VDW = {
  "H": 1.10, "C": 1.90, "N": 1.75, "O": 1.60, "S": 2.00, "P": 2.10,
  "F": 1.47, "CL": 1.75, "BR": 1.85, "I": 1.98, "B": 1.92, "SI": 2.10,
  "FE": 1.60, "ZN": 1.39, "MG": 1.73, "CA": 2.31, "MN": 1.61, "CU": 1.40,
  "NA": 2.27, "K": 2.75,
}


# -----------------------------
# DPB (Dock Pose Binary) Writer
# -----------------------------
# Header (64 bytes, little-endian):
#   magic        4s   = b"DPB1"
#   version      u32  = 1
#   dtype        u32  = 1 (float32)
#   n_atoms      u32
#   n_models     u32  (placeholder, backfilled on close)
#   stride_bytes u32  = 4 + n_atoms*3*4
#   reserved     bytes to 64
#
# Record (repeated n_models times):
#   model_id     u32
#   coords       float32[n_atoms*3]  (x,y,z,x,y,z,...)
class DPBWriter:
  def __init__(self, path, n_atoms):
    self.path = path
    self.n_atoms = int(n_atoms)
    self.n_models = 0
    self.dtype_code = 1  # float32
    self.stride_bytes = 4 + self.n_atoms * 3 * 4

    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    self.fw = open(path, "wb")

    header = struct.pack(
      "<4sIIIII",
      b"DPB1",
      1,
      self.dtype_code,
      self.n_atoms,
      0,
      self.stride_bytes
    )
    if len(header) > 64:
      raise RuntimeError("Header packing exceeded 64 bytes")
    self.fw.write(header)
    self.fw.write(b"\x00" * (64 - len(header)))

  def write(self, model_id, coords_cpu):
    # coords_cpu: torch.Tensor on CPU OR numpy-like with shape (n_atoms,3)
    # We write float32 contiguous in (n_atoms*3)
    self.n_models += 1
    self.fw.write(struct.pack("<I", int(model_id)))

    if isinstance(coords_cpu, torch.Tensor):
      arr = coords_cpu.to(dtype=torch.float32).contiguous().view(-1).cpu().numpy()
      self.fw.write(arr.tobytes(order="C"))
    else:
      # assume numpy array
      import numpy as np
      arr = np.asarray(coords_cpu, dtype=np.float32).reshape(-1)
      self.fw.write(arr.tobytes(order="C"))

  def close(self):
    # backfill n_models at header offset:
    # layout: magic(4) + version(4) + dtype(4) + n_atoms(4) = 16
    # n_models is next u32 at offset 16
    self.fw.flush()
    self.fw.seek(16)
    self.fw.write(struct.pack("<I", int(self.n_models)))
    self.fw.flush()
    self.fw.close()


def read_select(path):
  tasks = []
  with open(path) as fr:
    for line in fr:
      lx = line.split()
      if len(lx) >= 2:
        tasks.append((lx[0], lx[1:]))
  return tasks


def read_grid_xyz(path):
  xyz = []
  with open(path) as fr:
    for line in fr:
      lx = line.split()
      xyz.append((float(lx[3]), float(lx[4]), float(lx[5])))
  return torch.tensor(xyz, dtype=DTYPE)


def load_xmat_dict(path):
  d = {}
  with open(path) as fr:
    for line in fr:
      lx = line.split()
      if len(lx) >= 2:
        d[lx[0]] = lx[1].strip()
  return d


def load_vdw_table(path):
  d = {}
  with open(path) as fr:
    for line in fr:
      s = line.strip()
      if not s or s.startswith("#"):
        continue
      parts = s.split()
      if len(parts) < 2:
        continue
      k = parts[0].upper()
      v = float(parts[1])
      d[k] = v
  if not d:
    raise ValueError(f"Empty VDW table: {path}")
  return d


def parse_alpha_list(s):
  if s is None:
    return []
  s = s.strip()
  if not s:
    return []
  out = []
  for x in s.split(","):
    x = x.strip()
    if x:
      out.append(float(x))
  return out


def bits_to_points(bits, grid_xyz):
  idx = [i for i, ch in enumerate(bits) if ch == "1"]
  if not idx:
    return torch.empty((0, 3), dtype=DTYPE)
  return grid_xyz[idx]


def adtype_to_radius(ad_type, atom_name=""):
  global ELEM_VDW

  if not ad_type:
    s = atom_name.strip().upper()
    if len(s) >= 2 and s[:2] in ELEM_VDW:
      return ELEM_VDW[s[:2]]
    return ELEM_VDW.get(s[:1], 1.80)

  t = ad_type.strip().upper()
  if t.startswith("H"):
    return ELEM_VDW["H"]
  if t.startswith("CL"):
    return ELEM_VDW["CL"]
  if t.startswith("BR"):
    return ELEM_VDW["BR"]
  if t == "I":
    return ELEM_VDW["I"]
  if t == "F":
    return ELEM_VDW["F"]
  if len(t) >= 2 and t[:2] in ELEM_VDW:
    return ELEM_VDW[t[:2]]
  if t[:1] in ELEM_VDW:
    return ELEM_VDW[t[:1]]
  return 1.80


def read_pdbqt_atoms(pdbqt_path):
  crd, heavy, rad = [], [], []
  with open(pdbqt_path) as fr:
    for line in fr:
      if line.startswith(("ATOM", "HETATM")):
        x = float(line[30:38]); y = float(line[38:46]); z = float(line[46:54])
        lx = line.split()
        ad_type = lx[-1].upper() if lx else ""
        is_h = ad_type.startswith("H")
        r = adtype_to_radius(ad_type, atom_name=line[12:16])
        crd.append((x, y, z))
        heavy.append(not is_h)
        rad.append(r)
  if not crd:
    raise ValueError(pdbqt_path)
  return (
    torch.tensor(crd, dtype=DTYPE),
    torch.tensor(heavy, dtype=torch.bool),
    torch.tensor(rad, dtype=DTYPE),
  )


def pick_3_boundary_atoms(crd_cpu, heavy_mask_cpu):
  hidx = torch.nonzero(heavy_mask_cpu, as_tuple=False).flatten()
  if hidx.numel() < 3:
    hidx = torch.arange(crd_cpu.shape[0], dtype=torch.long)

  hcrd = crd_cpu[hidx]
  D = torch.cdist(hcrd, hcrd, p=2.0, compute_mode="donot_use_mm_for_euclid_dist")

  dmax_row, arg_row = torch.max(D, dim=1)
  _, xi = torch.max(dmax_row, dim=0)
  yi = arg_row[xi].item()
  xi = xi.item()

  best_s, zi = -1.0, -1
  for k in range(hcrd.shape[0]):
    if k == xi or k == yi:
      continue
    s = D[xi, k].item() + D[yi, k].item()
    if s > best_s:
      best_s, zi = s, k

  z0 = D[xi, zi].item()
  z1 = D[zi, yi].item()
  if z0 > z1:
    idx_local = [xi, yi, zi]
    dis = [D[xi, yi].item(), D[xi, zi].item(), D[yi, zi].item()]
  else:
    idx_local = [yi, xi, zi]
    dis = [D[xi, yi].item(), D[yi, zi].item(), D[xi, zi].item()]

  idx_global = [hidx[idx_local[0]].item(), hidx[idx_local[1]].item(), hidx[idx_local[2]].item()]
  return idx_global, dis


def kabsch_rt(P, Q):
  Pc = P.mean(dim=0, keepdim=True)
  Qc = Q.mean(dim=0, keepdim=True)
  X = P - Pc
  Y = Q - Qc
  C = X.t().mm(Y)
  U, S, Vt = torch.linalg.svd(C)
  V = Vt.t()
  I = torch.eye(3, dtype=DTYPE, device=P.device)
  if torch.det(V.mm(U.t())) < 0:
    I[2, 2] = -1.0
  R = V.mm(I).mm(U.t())
  t = Qc.squeeze(0) - R.mv(Pc.squeeze(0))
  return R, t


def apply_rt(crd, R, t):
  return crd.mm(R.t()) + t


def match_triplets_gpu_full(points_dev, target_dis, delta, ij_batch=None):
  M = points_dev.shape[0]
  if M < 3:
    return []

  d01, d02, d12 = map(float, target_dis)
  delta = float(delta)

  D = torch.cdist(points_dev, points_dev)
  tri = torch.triu_indices(M, M, offset=1, device=points_dev.device)
  I_all, J_all = tri[0], tri[1]

  keep = torch.abs(D[I_all, J_all] - d01) <= delta
  if not torch.any(keep):
    return []

  I = I_all[keep]
  J = J_all[keep]
  n_pairs = I.numel()
  triplets = []

  if ij_batch is None:
    mask = (torch.abs(D[I, :] - d02) <= delta) & (torch.abs(D[J, :] - d12) <= delta)
    rows = torch.arange(n_pairs, device=points_dev.device)
    mask[rows, I] = False
    mask[rows, J] = False
    bk = torch.nonzero(mask, as_tuple=False)
    if bk.numel() > 0:
      p = bk[:, 0]
      k = bk[:, 1]
      triplets = list(zip(I[p].tolist(), J[p].tolist(), k.tolist()))
  else:
    B = int(ij_batch)
    for s in range(0, n_pairs, B):
      e = min(s + B, n_pairs)
      Ib, Jb = I[s:e], J[s:e]
      bsz = Ib.numel()
      mask = (torch.abs(D[Ib, :] - d02) <= delta) & (torch.abs(D[Jb, :] - d12) <= delta)
      rows = torch.arange(bsz, device=points_dev.device)
      mask[rows, Ib] = False
      mask[rows, Jb] = False
      bk = torch.nonzero(mask, as_tuple=False)
      if bk.numel() > 0:
        p = bk[:, 0]
        k = bk[:, 1]
        triplets.extend(list(zip(Ib[p].tolist(), Jb[p].tolist(), k.tolist())))

  return triplets


def load_receptor_cached(cache, rec_path):
  global REC_FILTER_RAD

  if rec_path in cache:
    return cache[rec_path]

  rec_crd_cpu, rec_heavy, rec_rad_cpu = read_pdbqt_atoms(rec_path)
  rec_crd_cpu = rec_crd_cpu[rec_heavy]
  rec_rad_cpu = rec_rad_cpu[rec_heavy]

  if REC_FILTER_RAD is not None:
    keep = torch.linalg.norm(rec_crd_cpu, dim=1) <= float(REC_FILTER_RAD)
    rec_crd_cpu = rec_crd_cpu[keep]
    rec_rad_cpu = rec_rad_cpu[keep]

  rec_crd = rec_crd_cpu.to(device=GPU_DEVICE, dtype=DTYPE)
  rec_rad = rec_rad_cpu.to(device=GPU_DEVICE, dtype=DTYPE)

  cache[rec_path] = (rec_crd, rec_rad)
  return rec_crd, rec_rad


def pass_clash_check(lig_crd_dev, lig_rad_cpu, lig_heavy_cpu, rec_crd_dev, rec_rad_dev):
  global CLASH_SCALE, REC_CHUNK

  lig_heavy_dev = lig_heavy_cpu.to(device=lig_crd_dev.device)
  L = lig_crd_dev[lig_heavy_dev]
  if L.numel() == 0 or rec_crd_dev.numel() == 0:
    return True

  Lr = lig_rad_cpu[lig_heavy_cpu].to(device=lig_crd_dev.device, dtype=DTYPE)

  for s in range(0, rec_crd_dev.shape[0], REC_CHUNK):
    e = min(s + REC_CHUNK, rec_crd_dev.shape[0])
    Rc = rec_crd_dev[s:e]
    Rrc = rec_rad_dev[s:e]
    diff = L[:, None, :] - Rc[None, :, :]
    d2 = (diff * diff).sum(dim=2)
    thr = (CLASH_SCALE * (Lr[:, None] + Rrc[None, :])) ** 2
    if torch.any(d2 < thr):
      return False
  return True


def run_one(receptor, ligands, xmat_bits, grid_xyz_cpu, rec_root, lig_root, out_root, rec_cache):
  global delta_dist

  bits = xmat_bits[receptor]
  points_cpu = bits_to_points(bits, grid_xyz_cpu)
  if points_cpu.shape[0] == 0:
    return

  points_dev = points_cpu.to(device=GPU_DEVICE, dtype=DTYPE)

  if ENABLE_CLASH_CHECK:
    rec_path = os.path.join(rec_root, f"{receptor}_protein.pdbqt")
    rec_crd_dev, rec_rad_dev = load_receptor_cached(rec_cache, rec_path)
  else:
    rec_crd_dev, rec_rad_dev = None, None

  out_dir = os.path.join(out_root, receptor)
  os.makedirs(out_dir, exist_ok=True)

  for lig in ligands:
    lig_path = os.path.join(lig_root, f"{lig}_ligand.pdbqt")
    crd_cpu, heavy_cpu, rad_cpu = read_pdbqt_atoms(lig_path)

    n_atoms = crd_cpu.shape[0]
    crd_dev = crd_cpu.to(device=GPU_DEVICE, dtype=DTYPE)

    idx3, dis3 = pick_3_boundary_atoms(crd_cpu, heavy_cpu)
    idx3_t = torch.tensor(idx3, dtype=torch.long, device=GPU_DEVICE)
    P = crd_dev[idx3_t]

    triplets = match_triplets_gpu_full(points_dev, dis3, delta_dist, IJ_BATCH)

    out_path = os.path.join(out_dir, f"{lig}.dpb")
    writer = DPBWriter(out_path, n_atoms)

    try:
      midx = 0
      for (i, j, k) in triplets:
        ijk = torch.tensor([i, j, k], dtype=torch.long, device=GPU_DEVICE)
        Q = points_dev[ijk]

        for perm in [(0,1,2),(0,2,1),(1,0,2),(1,2,0),(2,0,1),(2,1,0)]:
          perm_t = torch.tensor(perm, dtype=torch.long, device=GPU_DEVICE)
          Qp = Q[perm_t]

          R, t = kabsch_rt(P, Qp)
          new_crd_dev = apply_rt(crd_dev, R, t)

          if ENABLE_CLASH_CHECK:
            if not pass_clash_check(new_crd_dev, rad_cpu, heavy_cpu, rec_crd_dev, rec_rad_dev):
              continue

          midx += 1
          writer.write(midx, new_crd_dev.detach().cpu().contiguous())
    finally:
      writer.close()


def Main():
  if not torch.cuda.is_available():
    raise RuntimeError("CUDA not available. This script is GPU-only.")

  ap = argparse.ArgumentParser(description="GPU-only docking; output DPB (binary float32) per ligand.")
  ap.add_argument("--select", required=True)
  ap.add_argument("--xmat", required=True)
  ap.add_argument("--grid", required=True)
  ap.add_argument("--rec-root", required=True)
  ap.add_argument("--lig-root", required=True)
  ap.add_argument("--out-root", default=os.path.abspath("./lig_multi_pos"))

  ap.add_argument("--vdw-table", required=True)

  ap.add_argument("--div", type=float, default=0.375)
  ap.add_argument("--alpha-list", type=str, default="0.1")

  ap.add_argument("--rec-filter-rad", type=float, default=13.0)
  ap.add_argument("--clash-scale", type=float, default=0.50)

  ap.add_argument("--ij-batch", type=int, default=None)

  ap.add_argument("--no-clash-check", dest="enable_clash_check", action="store_false")
  ap.set_defaults(enable_clash_check=True)

  ap.add_argument("--rec-chunk", type=int, default=4096)

  args = ap.parse_args()

  global ELEM_VDW, DIV, REC_FILTER_RAD, CLASH_SCALE
  global IJ_BATCH, ENABLE_CLASH_CHECK, REC_CHUNK

  ELEM_VDW = load_vdw_table(args.vdw_table)

  DIV = float(args.div)

  REC_FILTER_RAD = float(args.rec_filter_rad)
  if REC_FILTER_RAD <= 0:
    REC_FILTER_RAD = None

  CLASH_SCALE = float(args.clash_scale)

  IJ_BATCH = args.ij_batch

  ENABLE_CLASH_CHECK = bool(args.enable_clash_check)
  REC_CHUNK = int(args.rec_chunk)

  alpha_list = parse_alpha_list(args.alpha_list)
  if not alpha_list:
    raise ValueError("--alpha-list is empty")

  grid_xyz_cpu = read_grid_xyz(args.grid)
  xmat_bits = load_xmat_dict(args.xmat)
  tasks = read_select(args.select)

  rec_cache = {}

  print(f"[INFO] GPU-only device={GPU_DEVICE}")
  print(f"[INFO] out_root={args.out_root} (DPB output per ligand)")
  print(f"[INFO] div={DIV} alpha_list={alpha_list}")
  print(f"[INFO] ij_batch={IJ_BATCH}")
  print(f"[INFO] clash={ENABLE_CLASH_CHECK} clash_scale={CLASH_SCALE} rec_filter={REC_FILTER_RAD} rec_chunk={REC_CHUNK}")

  for alpha in alpha_list:
    global delta_dist
    delta_dist = DIV * float(alpha)
    out_root = os.path.join(args.out_root, f"{alpha:.2f}")
    os.makedirs(out_root, exist_ok=True)
    print(f"[INFO] alpha={alpha:.6f} delta={delta_dist:.6f}")

    for receptor, ligands in tasks:
      if receptor not in xmat_bits:
        raise KeyError(f"Xmat missing receptor: {receptor}")
      run_one(receptor, ligands, xmat_bits, grid_xyz_cpu, args.rec_root, args.lig_root, out_root, rec_cache)


def main():
  Main()


if __name__ == "__main__":
  main()

