import os
import sys
import argparse
import numpy as np
import torch
from torch import optim


def _require_file(path: str, what: str):
    if not os.path.isfile(path):
        raise FileNotFoundError(f"Missing {what}: {path}")


def _require_dir(path: str, what: str):
    if not os.path.isdir(path):
        raise FileNotFoundError(f"Missing {what}: {path}")


def read_grid_xyz(grid_path: str) -> np.ndarray:
    Id = []
    with open(grid_path, "r") as fr:
        for line in fr:
            lx = line.split()
            if len(lx) < 3:
                continue
            Id.append(int(lx[0]))
            Id.append(int(lx[1]))
            Id.append(int(lx[2]))
    arr = np.array(Id, dtype=np.int32).reshape(-1, 3)
    if arr.shape[0] == 0:
        raise ValueError(f"Grid seems empty or invalid: {grid_path}")
    return arr


def read_j_dense(Q: np.ndarray, j_path: str, E_vdw: float, E_bond: float) -> np.ndarray:
    with open(j_path, "r") as fr:
        for line in fr:
            lx = line.split()
            if len(lx) < 3:
                continue
            ij = int(lx[2])
            if ij == 1:
                s = E_vdw
            elif ij == -1:
                s = E_bond
            else:
                raise ValueError(f"Bad J value: {line.strip()}")
            a = int(lx[0])
            b = int(lx[1])
            Q[a][b] = s
    return Q


def read_h_to_diag(grid_path: str, Q: np.ndarray) -> np.ndarray:
    Vdw = []
    with open(grid_path, "r") as fr:
        for line in fr:
            lx = line.split()
            if len(lx) >= 4:
                Vdw.append(float(lx[3]))
    if len(Vdw) != Q.shape[0]:
        raise ValueError(f"Grid length mismatch: {len(Vdw)} vs {Q.shape[0]} for {grid_path}")
    for i, v in enumerate(Vdw):
        Q[i][i] = v
    return Q


def list_targets(h_base: str):
    targets = []
    for fn in sorted(os.listdir(h_base)):
        if not fn.endswith(".grid"):
            continue
        pdb = fn[:-5]
        grid_path = os.path.join(h_base, fn)
        if os.path.isfile(grid_path):
            targets.append((pdb, grid_path))
    return targets


def convert_x(X: torch.Tensor, half: torch.Tensor, zero: torch.Tensor) -> torch.Tensor:
    X2 = torch.sigmoid(X.detach())
    X3 = torch.heaviside(X2 - half, zero)
    return X3


def train_x(opt: optim.Optimizer, Q: torch.Tensor, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    opt.zero_grad()
    X2 = torch.sigmoid(X).reshape(1, -1)
    X3 = X2.reshape(-1, 1)
    y1 = torch.mm(X2, Q)
    loss = torch.mm(y1, X3)
    loss.backward()
    opt.step()
    return X, loss


def est_score_discrete(X: torch.Tensor, Q: torch.Tensor, half: torch.Tensor, zero: torch.Tensor) -> torch.Tensor:
    X1 = convert_x(X, half, zero).reshape(1, -1)
    X2 = X1.reshape(-1, 1)
    y1 = torch.mm(X1, Q)
    loss = torch.mm(y1, X2)
    return loss


def write_xmat_line(fw, tag: str, X: torch.Tensor, half: torch.Tensor, zero: torch.Tensor):
    # One target per line, no blank line between targets
    fw.write(f"{tag} ")
    XB = convert_x(X, half, zero).to("cpu").numpy()
    for i in range(XB.shape[0]):
        fw.write(f"{int(XB[i])}")
    fw.write("\n")
    fw.flush()


def write_scor_block(fw, tag: str, loss1: list[float], loss2: list[float]):
    # One target block, no extra blank line between targets
    fw.write(f"{tag} ")
    for i in range(len(loss1)):
        fw.write("%10.4f/%10.4f  " % (loss1[i], loss2[i]))
        if i % 5 == 4:
            fw.write("\n")
    if len(loss1) > 0 and (len(loss1) - 1) % 5 != 4:
        fw.write("\n")
    fw.flush()


def main():
    ap = argparse.ArgumentParser()

    # You said you WILL input these (no guessing, no auto defaults)
    ap.add_argument("--device", choices=["cpu", "cuda"], required=True)
    ap.add_argument("--learning-rate", type=float, required=True)
    ap.add_argument("--E-bond", type=float, required=True)
    ap.add_argument("--E-vdw", type=float, required=True)
    ap.add_argument("--num-step", type=int, required=True)

    ap.add_argument("--grid-path", required=True)
    ap.add_argument("--h-base", required=True)
    ap.add_argument("--j-path", required=True)

    ap.add_argument("--out-xmat", required=True)
    ap.add_argument("--out-scor", required=True)

    ap.add_argument("--dtype", choices=["float32", "float64"], default="float64")

    args = ap.parse_args()

    _require_file(args.grid_path, "Grid file (--grid-path)")
    _require_dir(args.h_base, "H matrix directory (--h-base)")
    _require_file(args.j_path, "J matrix file (--j-path)")

    device = torch.device(args.device)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA not available. Use --device cpu, or install CUDA-enabled PyTorch.")

    dtype = torch.float64 if args.dtype == "float64" else torch.float32

    Idn_G = read_grid_xyz(args.grid_path)
    Num_X = Idn_G.shape[0]

    targets = list_targets(args.h_base)
    if not targets:
        raise FileNotFoundError(f"No *.grid files found in: {args.h_base}")

    print(f"[info] grid={args.grid_path} Num_X={Num_X}", file=sys.stderr)
    print(f"[info] h_base={args.h_base} targets={len(targets)}", file=sys.stderr)
    print(f"[info] j_path={args.j_path}", file=sys.stderr)
    print(f"[info] device={device} dtype={args.dtype}", file=sys.stderr)
    print(
        f"[info] learning_rate={args.learning_rate} E_bond={args.E_bond} E_vdw={args.E_vdw} num_step={args.num_step}",
        file=sys.stderr,
    )
    print(f"[info] out_xmat={args.out_xmat} out_scor={args.out_scor}", file=sys.stderr)

    half = torch.tensor([0.5], dtype=dtype, device=device)
    zero = torch.tensor([0.0], dtype=dtype, device=device)

    # Build base Q on CPU (dense), then copy per target
    Q_base = np.zeros((Num_X, Num_X), dtype=np.float64)
    Q_base = read_j_dense(Q_base, args.j_path, args.E_vdw, args.E_bond)

    fw1 = open(args.out_xmat, "w")
    fw2 = open(args.out_scor, "w")

    for pdb, grid_path in targets:
        tag = pdb

        X = torch.randn((Num_X,), device=device, dtype=dtype, requires_grad=True)
        Opt = optim.SGD([X], lr=args.learning_rate, weight_decay=0.0, momentum=0.9)

        Q_np = Q_base.copy()
        Q_np = read_h_to_diag(grid_path, Q_np)
        Q2 = torch.from_numpy(Q_np).to(device=device, dtype=dtype)

        Loss1 = []
        Loss2 = []
        for j in range(args.num_step):
            X, loss1 = train_x(Opt, Q2, X)
            if j % 100 == 99:
                Loss1.append(float(loss1.detach().to("cpu").item()))
                loss2 = est_score_discrete(X, Q2, half, zero)
                Loss2.append(float(loss2.detach().to("cpu").item()))

        write_xmat_line(fw1, tag, X, half, zero)
        write_scor_block(fw2, tag, Loss1, Loss2)

        del Q2

    fw1.close()
    fw2.close()


if __name__ == "__main__":
    main()

