# SLAM_MAP.py
from __future__ import annotations

import json
import os
from dataclasses import dataclass
from pathlib import Path
import numpy as np


def load_slam_config() -> dict:
    import json, os
    from pathlib import Path

    cfg_path = os.environ.get("SLAM_CONFIG", "").strip()
    p = Path(cfg_path) if cfg_path else (Path(__file__).resolve().parent / "SLAM_Config.json")

    if not p.exists():
        raise FileNotFoundError(f"Missing config file: {p}")

    text = p.read_text(encoding="utf-8").strip()
    if not text:
        raise RuntimeError("SLAM_Config.json is empty.")

    return json.loads(text)


@dataclass
class MapConfig:
    display_voxel: float
    save_voxel: float
    data_folder: str
    save_extension: str

    @staticmethod
    def from_config_file() -> "MapConfig":
        cfg = load_slam_config()
        maps_dir = str(cfg["app"]["maps_dir"])
        return MapConfig(
            display_voxel=float(cfg["map"]["display_voxel"]),
            save_voxel=float(cfg["map"]["save_voxel"]),
            data_folder=maps_dir,
            save_extension=str(cfg["map"]["save_extension"]),
        )


def _voxel_downsample_numpy(points: np.ndarray, voxel: float) -> np.ndarray:
    if points is None or len(points) == 0:
        return np.zeros((0, 3), dtype=np.float32)
    if voxel <= 0:
        return points.astype(np.float32, copy=False)

    keys = np.floor(points / voxel).astype(np.int32)
    uniq, inv = np.unique(keys, axis=0, return_inverse=True)

    out = np.zeros((len(uniq), 3), dtype=np.float64)
    cnt = np.zeros((len(uniq),), dtype=np.int64)
    np.add.at(out, inv, points.astype(np.float64, copy=False))
    np.add.at(cnt, inv, 1)
    out /= np.maximum(cnt[:, None], 1)
    return out.astype(np.float32)


class StableMapLayer:
    def __init__(self, cfg: MapConfig | None = None):
        self.cfg = cfg if cfg is not None else MapConfig.from_config_file()
        Path(self.cfg.data_folder).mkdir(parents=True, exist_ok=True)
        self._points = np.zeros((0, 3), dtype=np.float32)

    def reset(self):
        self._points = np.zeros((0, 3), dtype=np.float32)

    def set_points(self, pts: np.ndarray):
        self._points = pts.astype(np.float32, copy=True) if pts is not None else np.zeros((0, 3), dtype=np.float32)

    def get_points(self) -> np.ndarray:
        return self._points

    def get_display_points(self) -> np.ndarray:
        return _voxel_downsample_numpy(self._points, float(self.cfg.display_voxel))

    def get_save_points(self) -> np.ndarray:
        return _voxel_downsample_numpy(self._points, float(self.cfg.save_voxel))

    def export_map(self, base_name: str) -> str:
        """
        Save stable map into maps_dir (no subfolders).
        """
        import open3d as o3d

        pts = self.get_save_points()
        if pts is None or len(pts) < 20:
            raise RuntimeError("Not enough stable points to save.")

        base = (base_name or "").strip() or "map_robot"
        ext = self.cfg.save_extension.strip() or ".ply"
        if base.endswith(ext):
            base = base[: -len(ext)]

        filename = f"{base}{ext}"
        n = 1
        folder = Path(self.cfg.data_folder)
        while (folder / filename).exists():
            filename = f"{base}({n}){ext}"
            n += 1

        full = str(folder / filename)
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(pts.astype(np.float64, copy=False))
        o3d.io.write_point_cloud(full, pcd)
        return full

    def load_map(self, filepath: str) -> np.ndarray:
        import open3d as o3d
        pcd = o3d.io.read_point_cloud(filepath)
        pts = np.asarray(pcd.points).astype(np.float32)
        self.set_points(pts)
        return pts