import os
import torch
import argparse

E_rep1 = torch.tensor(1, dtype=torch.int8)
E_bond = torch.tensor(-1, dtype=torch.int8)

def Read_Grid(grid_path):
    crd = []
    with open(grid_path, "r") as fr:
        for line in fr:
            lx = line.split()
            if not lx:
                continue
            crd.append(float(lx[3]))
            crd.append(float(lx[4]))
            crd.append(float(lx[5]))
    return torch.tensor(crd).reshape(-1, 3)


def Read_Type(dis_path):
    dis = []
    with open(dis_path, "r") as fr:
        for line in fr:
            s = line.strip()
            if not s:
                continue
            if s.startswith("#"):
                continue
            lx = s.split()
            if len(lx) < 2:
                continue
            dis.append((float(lx[0]), float(lx[1])))
    return torch.tensor(dis).reshape(-1, 2)


def Write_Score(fw, D):
    DN = D.to("cpu").numpy()
    Idx = torch.nonzero(D, as_tuple=False).to("cpu").numpy()
    for i in range(Idx.shape[0]):
        m = int(Idx[i][0])
        n = int(Idx[i][1])
        fw.write("%d %d %2d\n" % (m, n, DN[m][n]))
    fw.flush()


def Score_Grid(crd, disD, out_dir):
    ln = crd.shape[0]
    D = torch.cdist(
        crd, crd,
        p=2.0,
        compute_mode="donot_use_mm_for_euclid_dist"
    )

    os.makedirs(out_dir, exist_ok=True)

    for i in range(disD.shape[0]):
        d0 = float(disD[i, 0])
        d1 = float(disD[i, 1])

        out_path = os.path.join(
            out_dir, f"Score_Bond_{d0:.1f}_{d1:.1f}"
        )

        with open(out_path, "w") as fw:
            Bond_Score = torch.zeros((ln, ln), dtype=torch.int8)

            D1 = D < d0
            D4 = (D >= d0) & (D < d1)

            Bond_Score = Bond_Score + E_rep1 * D1 + E_bond * D4
            Bond_Score = torch.triu(Bond_Score, diagonal=1)

            Write_Score(fw, Bond_Score)


def main():
    parser = argparse.ArgumentParser(
        description="Generate J matrix (bond score) from grid and distance types"
    )
    parser.add_argument(
        "--div", type=float, default=0.375,
        help="Grid spacing (default: 0.375)"
    )
    parser.add_argument(
        "--grid-path", default="../1_grid/Grid",
        help="Path to Grid file"
    )
    parser.add_argument(
        "--dis-path", default="Dis",
        help="Distance type definition file"
    )
    parser.add_argument(
        "--out-dir", default="J_Matrix",
        help="Output directory"
    )

    args = parser.parse_args()

    Crd = Read_Grid(args.grid_path)
    DisD = Read_Type(args.dis_path)
    Score_Grid(Crd, DisD, args.out_dir)


if __name__ == "__main__":
    main()

