import os
import argparse
import torch

DTYPE = torch.float32

# Defaults (will be overridden by CLI)
USE_GPU_TRIPLETS = True
GPU_DEVICE = "cuda"
IJ_BATCH = None

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

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,
}


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):
  lines, crd, heavy, rad = [], [], [], []
  with open(pdbqt_path) as fr:
    for line in fr:
      lines.append(line)
      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 (
    lines,
    torch.tensor(crd, dtype=DTYPE),
    torch.tensor(heavy, dtype=torch.bool),
    torch.tensor(rad, dtype=DTYPE),
  )


def write_model(lines_template, new_crd, fw, midx):
  fw.write(f"MODEL {midx}\n")
  k = 0
  for line in lines_template:
    if line.startswith(("ATOM", "HETATM")) and k < new_crd.shape[0]:
      x, y, z = new_crd[k].tolist()
      fw.write(line[:30] + f"{x:8.3f}{y:8.3f}{z:8.3f}" + line[54:])
      k += 1
    else:
      fw.write(line)
  fw.write("ENDMDL\n")


def pick_3_boundary_atoms(crd, heavy_mask):
  hidx = torch.nonzero(heavy_mask, as_tuple=False).flatten()
  if hidx.numel() < 3:
    hidx = torch.arange(crd.shape[0], dtype=torch.long)
  hcrd = crd[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_cpu, target_dis, delta, device="cuda", ij_batch=None):
  M = points_cpu.shape[0]
  if M < 3:
    return []
  pts = points_cpu.to(device=device, dtype=DTYPE)
  d01, d02, d12 = map(float, target_dis)
  delta = float(delta)

  D = torch.cdist(pts, pts)
  tri = torch.triu_indices(M, M, offset=1, device=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=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].cpu().tolist(), J[p].cpu().tolist(), k.cpu().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=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].cpu().tolist(), Jb[p].cpu().tolist(), k.cpu().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, rec_heavy, rec_rad = read_pdbqt_atoms(rec_path)
  rec_crd = rec_crd[rec_heavy]
  rec_rad = rec_rad[rec_heavy]

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

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


def pass_clash_check(lig_crd, lig_rad, lig_heavy, rec_crd, rec_rad, device):
  global CLASH_SCALE, REC_CHUNK

  L = lig_crd[lig_heavy]
  if L.numel() == 0 or rec_crd.numel() == 0:
    return True
  Lr = lig_rad[lig_heavy]

  L = L.to(device=device, dtype=DTYPE)
  Lr = Lr.to(device=device, dtype=DTYPE)
  R = rec_crd.to(device=device, dtype=DTYPE)
  Rr = rec_rad.to(device=device, dtype=DTYPE)

  for s in range(0, R.shape[0], REC_CHUNK):
    e = min(s + REC_CHUNK, R.shape[0])
    Rc = R[s:e]
    Rrc = Rr[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, rec_root, lig_root, out_root, rec_cache):
  bits = xmat_bits[receptor]
  points = bits_to_points(bits, grid_xyz)
  if points.shape[0] == 0:
    return

  if ENABLE_CLASH_CHECK:
    rec_path = os.path.join(rec_root, f"{receptor}_protein.pdbqt")
    rec_crd, rec_rad = load_receptor_cached(rec_cache, rec_path)
    dev = GPU_DEVICE if torch.cuda.is_available() else "cpu"
  else:
    rec_crd, rec_rad, dev = None, None, "cpu"

  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")
    lines, crd, heavy, rad = read_pdbqt_atoms(lig_path)
    idx3, dis3 = pick_3_boundary_atoms(crd, heavy)
    P = crd[idx3]

    triplets = match_triplets_gpu_full(points, dis3, delta_dist, GPU_DEVICE, IJ_BATCH) \
      if (USE_GPU_TRIPLETS and torch.cuda.is_available()) else []

    out_path = os.path.join(out_dir, f"{lig}_dock.pdbqt")
    with open(out_path, "w") as fw:
      midx = 0
      for (i, j, k) in triplets:
        Q = points[torch.tensor([i, j, k], dtype=torch.long)]
        for perm in [(0,1,2),(0,2,1),(1,0,2),(1,2,0),(2,0,1),(2,1,0)]:
          Qp = Q[torch.tensor(perm, dtype=torch.long)]
          R, t = kabsch_rt(P, Qp)
          new_crd = apply_rt(crd, R, t)
          if ENABLE_CLASH_CHECK and (not pass_clash_check(new_crd, rad, heavy, rec_crd, rec_rad, dev)):
            continue
          midx += 1
          write_model(lines, new_crd, fw, midx)

    if os.path.getsize(out_path) == 0:
      open(out_path, "w").close()


def Main():
  ap = argparse.ArgumentParser(description="Flat select / flat receptor / flat lig / no sphere / no group.")
  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("./score"))

  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("--use-gpu-triplets", action="store_true")
  ap.add_argument("--ij-batch", type=int, default=None)

  ap.add_argument("--enable-clash-check", action="store_true")
  ap.add_argument("--rec-chunk", type=int, default=4096)

  args = ap.parse_args()

  global ELEM_VDW, DIV, REC_FILTER_RAD, CLASH_SCALE
  global USE_GPU_TRIPLETS, 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)

  USE_GPU_TRIPLETS = bool(args.use_gpu_triplets)
  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 = read_grid_xyz(args.grid)
  xmat_bits = load_xmat_dict(args.xmat)
  tasks = read_select(args.select)

  rec_cache = {}

  print(f"[INFO] select={args.select}")
  print(f"[INFO] rec_root={args.rec_root}")
  print(f"[INFO] lig_root={args.lig_root}")
  print(f"[INFO] out_root={args.out_root}")
  print(f"[INFO] vdw_table={args.vdw_table} n_vdw={len(ELEM_VDW)}")

  print(f"[INFO] div={DIV} alpha_list={alpha_list}")
  print(f"[INFO] use_gpu_triplets={USE_GPU_TRIPLETS} 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, args.rec_root, args.lig_root, out_root, rec_cache)

def main():
  Main()

if __name__ == "__main__":
  main()


