import os
import numpy as np
import argparse

def read_grid_points(path):
    pts = []
    with open(path, "r") as fr:
        for line in fr:
            lx = line.split()
            if len(lx) >= 3:
                pts.append((int(lx[0]), int(lx[1]), int(lx[2])))
    return np.array(pts, dtype=np.int32).reshape(-1, 3)

def read_rec_C_map(path, n):
    eg = []
    with open(path, "r") as fr:
        for line in fr:
            lx = line.split()
            if len(lx) == 1:
                eg.append(float(lx[0]))
    arr = np.array(eg, dtype=np.float64)
    if arr.size != n * n * n:
        raise ValueError(f"{path} size mismatch: got {arr.size}, expect {n*n*n}")
    return arr.reshape(n, n, n).transpose(2, 1, 0)

def write_grid(path, Id, cmap, n, shift):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as fw:
        for i in range(Id.shape[0]):
            xi = Id[i][0] + shift
            yi = Id[i][1] + shift
            zi = Id[i][2] + shift
            if xi < 0 or xi >= n or yi < 0 or yi >= n or zi < 0 or zi >= n:
                continue
            fw.write(
                "%3d %3d %3d %11.3f\n"
                % (Id[i][0], Id[i][1], Id[i][2], cmap[xi][yi][zi])
            )

def main():
    parser = argparse.ArgumentParser(
        description="Convert *_rec.C.map files to sparse .grid files using Grid indices"
    )
    parser.add_argument(
        "--in-base",
        default=os.path.abspath("../../1_dataset/9_summary/map"),
        help="Input base directory containing *_rec.C.map files"
    )
    parser.add_argument(
        "--out-base",
        default=os.path.abspath("Map"),
        help="Output directory for <pdb>.grid files"
    )
    parser.add_argument(
        "--grid-path",
        default=os.path.abspath("../1_grid/Grid"),
        help="Grid file path (with i j k columns)"
    )
    parser.add_argument(
        "--n",
        type=int,
        default=61,
        help="Grid cube size N (default: 61)"
    )
    parser.add_argument(
        "--shift",
        type=int,
        default=30,
        help="Index shift applied to grid indices (default: 30)"
    )

    args = parser.parse_args()

    Id = read_grid_points(args.grid_path)
    os.makedirs(args.out_base, exist_ok=True)

    for root, dirs, files in os.walk(args.in_base):
        for fn in files:
            if not fn.endswith("_rec.C.map"):
                continue

            pdb = fn.split("_rec.", 1)[0]
            in_path = os.path.join(root, fn)

            cmap = read_rec_C_map(in_path, args.n)

            out_grid = os.path.join(args.out_base, f"{pdb}.grid")
            write_grid(out_grid, Id, cmap, args.n, args.shift)

if __name__ == "__main__":
    main()

