import argparse
import logging
from pathlib import Path

import h5py
import numpy as np
import yaml

from dataset_generation.gridnet_hd_batch_generation import load_las_points
from dataset_generation.gridnet_hd_manifest import load_or_build_manifest
from utils.visibility_criteria import compute_normals_open3d


def compute_normals_for_tile(tile_info, output_dir, overwrite=False):
    output_dir = Path(output_dir) / tile_info["split"]
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"{tile_info['tile']}.h5"

    if output_path.exists() and not overwrite:
        logging.info("Skipping normals for %s", tile_info["tile"])
        return output_path

    points, _, _, pts_for_normals = load_las_points(tile_info["lidar_path"])
    normals = compute_normals_open3d(pts_for_normals)
    coords = np.vstack((points["X"], points["Y"], points["Z"])).T.astype(np.float32)

    with h5py.File(output_path, "w") as handle:
        group = handle.create_group("points")
        group.create_dataset("coordinates", data=coords, dtype="float32")
        group.create_dataset("normals", data=normals.astype(np.float32), dtype="float32")

    logging.info("Wrote normals to %s", output_path)
    return output_path


def compute_normal_main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default="configs/config_gridnet_hd_dataset_generation.yaml")
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text())

    manifest = load_or_build_manifest(cfg)
    overwrite = cfg.get("runtime", {}).get("overwrite_existing", False)
    normals_dir = cfg["data"]["normals_dir"]

    for tile_info in manifest:
        compute_normals_for_tile(tile_info, normals_dir, overwrite=overwrite)