#!/usr/bin/env python3

import math
import time
import threading
import yaml
from pathlib import Path
from datetime import datetime, timedelta
from typing import Optional, Tuple, List, Dict
from collections import defaultdict

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from sensor_msgs.msg import LaserScan
from nav_msgs.msg import Odometry, OccupancyGrid
from geometry_msgs.msg import MapMetaData


class GridCell:
    """Represents a single cell in the occupancy grid."""

    STATES = {
        "free": -1,
        "unknown": -1,
        "obstacle": 100,
        "hazard": 100,
    }

    def __init__(self):
        self.state = "unknown"
        self.detection_count = 0
        self.last_update = time.time()
        self.confidence = 0.0
        self.decay_rate = 0.95

    def update_obstacle(self):
        """Mark this cell as obstacle."""
        self.detection_count += 1
        self.last_update = time.time()

        if self.detection_count >= 3:
            self.state = "hazard"
        else:
            self.state = "obstacle"

        self.confidence = min(1.0, self.detection_count / 5.0)

    def mark_free(self):
        """Mark this cell as definitely free."""
        if self.state != "hazard":
            self.state = "free"
            self.confidence = 0.5

    def apply_decay(self, dt: float) -> float:
        """Apply temporal decay to confidence."""
        if self.state == "hazard":
            days_elapsed = dt / 86400.0
            decay_factor = (0.95 ** days_elapsed)
            self.confidence = max(0.3, self.confidence * decay_factor)
        elif self.state == "obstacle":
            days_elapsed = dt / 86400.0
            decay_factor = (0.85 ** days_elapsed)
            self.confidence = max(0.0, self.confidence * decay_factor)

            if self.confidence < 0.2:
                self.state = "free"
                self.confidence = 0.0

        return self.confidence

    def get_occupancy_value(self) -> int:
        """Get OccupancyGrid value (-1 to 100)."""
        if self.state == "unknown":
            return -1
        elif self.state == "free":
            return 0
        else:
            return int(self.confidence * 100)

    def to_dict(self) -> Dict:
        """Serialize to dict for YAML persistence."""
        return {
            "state": self.state,
            "detection_count": self.detection_count,
            "confidence": float(self.confidence),
            "last_update": self.last_update,
        }

    @staticmethod
    def from_dict(data: Dict) -> "GridCell":
        """Deserialize from dict."""
        cell = GridCell()
        cell.state = data.get("state", "unknown")
        cell.detection_count = data.get("detection_count", 0)
        cell.confidence = data.get("confidence", 0.0)
        cell.last_update = data.get("last_update", time.time())
        return cell


class OccupancyGridMap:
    """2D occupancy grid with spatial persistence."""

    def __init__(
        self,
        width: int = 100,
        height: int = 100,
        resolution: float = 0.1,
        origin_x: float = 0.0,
        origin_y: float = 0.0,
    ):
        """Initialize occupancy grid."""
        self.width = width
        self.height = height
        self.resolution = resolution
        self.origin_x = -width * resolution / 2
        self.origin_y = -height * resolution / 2

        self.grid = [[GridCell() for _ in range(width)] for _ in range(height)]
        self.last_decay_time = time.time()
        self.lock = threading.Lock()

    def world_to_grid(self, x: float, y: float) -> Optional[Tuple[int, int]]:
        """Convert world coordinates to grid indices."""
        grid_x = int((x - self.origin_x) / self.resolution)
        grid_y = int((y - self.origin_y) / self.resolution)

        if 0 <= grid_x < self.width and 0 <= grid_y < self.height:
            return (grid_x, grid_y)
        return None

    def grid_to_world(self, grid_x: int, grid_y: int) -> Tuple[float, float]:
        """Convert grid indices to world coordinates."""
        x = self.origin_x + (grid_x + 0.5) * self.resolution
        y = self.origin_y + (grid_y + 0.5) * self.resolution
        return (x, y)

    def mark_obstacle(self, x: float, y: float):
        """Mark a world point as obstacle."""
        with self.lock:
            coords = self.world_to_grid(x, y)
            if coords:
                grid_x, grid_y = coords
                self.grid[grid_y][grid_x].update_obstacle()

    def mark_free(self, x: float, y: float):
        """Mark a world point as free space."""
        with self.lock:
            coords = self.world_to_grid(x, y)
            if coords:
                grid_x, grid_y = coords
                self.grid[grid_y][grid_x].mark_free()

    def apply_temporal_decay(self):
        """Apply decay to all cells."""
        with self.lock:
            current_time = time.time()
            dt = current_time - self.last_decay_time

            if dt > 3600:
                for row in self.grid:
                    for cell in row:
                        cell.apply_decay(dt)
                self.last_decay_time = current_time

    def to_occupancy_grid(self, frame_id: str = "map", timestamp=None) -> OccupancyGrid:
        """Convert to ROS2 OccupancyGrid message."""
        msg = OccupancyGrid()
        msg.header.frame_id = frame_id
        msg.header.stamp = timestamp if timestamp else rclpy.clock.Clock().now().to_msg()

        msg.info.width = self.width
        msg.info.height = self.height
        msg.info.resolution = self.resolution
        msg.info.origin.position.x = self.origin_x
        msg.info.origin.position.y = self.origin_y
        msg.info.origin.position.z = 0.0

        with self.lock:
            data = []
            for y in range(self.height):
                for x in range(self.width):
                    data.append(self.grid[y][x].get_occupancy_value())
        msg.data = data

        return msg

    def to_yaml(self) -> Dict:
        """Serialize grid to YAML-compatible dict."""
        with self.lock:
            grid_data = []
            for y in range(self.height):
                row = []
                for x in range(self.width):
                    row.append(self.grid[y][x].to_dict())
                grid_data.append(row)

        return {
            "metadata": {
                "width": self.width,
                "height": self.height,
                "resolution": self.resolution,
                "origin": {"x": self.origin_x, "y": self.origin_y},
                "timestamp": datetime.now().isoformat(),
            },
            "grid": grid_data,
        }

    def from_yaml(self, data: Dict):
        """Load grid from YAML-compatible dict."""
        with self.lock:
            meta = data.get("metadata", {})
            self.width = meta.get("width", self.width)
            self.height = meta.get("height", self.height)
            self.resolution = meta.get("resolution", self.resolution)

            origin = meta.get("origin", {})
            self.origin_x = origin.get("x", self.origin_x)
            self.origin_y = origin.get("y", self.origin_y)

            grid_data = data.get("grid", [])
            self.grid = []
            for y in range(self.height):
                row = []
                for x in range(self.width):
                    if y < len(grid_data) and x < len(grid_data[y]):
                        cell = GridCell.from_dict(grid_data[y][x])
                    else:
                        cell = GridCell()
                    row.append(cell)
                self.grid.append(row)

    def get_statistics(self) -> Dict:
        """Get grid statistics."""
        with self.lock:
            total_cells = self.width * self.height
            obstacle_count = 0
            hazard_count = 0
            free_count = 0

            for row in self.grid:
                for cell in row:
                    if cell.state == "obstacle":
                        obstacle_count += 1
                    elif cell.state == "hazard":
                        hazard_count += 1
                    elif cell.state == "free":
                        free_count += 1

            return {
                "total_cells": total_cells,
                "obstacle_cells": obstacle_count,
                "hazard_cells": hazard_count,
                "free_cells": free_count,
                "unknown_cells": total_cells - obstacle_count - hazard_count - free_count,
                "coverage_percent": (obstacle_count + hazard_count + free_count) / total_cells * 100,
            }


class ObstacleMemoryNode(Node):
    """ROS2 node for persistent obstacle memory mapping."""

    def __init__(self):
        super().__init__("saltybot_obstacle_memory")

        self.declare_parameter("grid_width", 100)
        self.declare_parameter("grid_height", 100)
        self.declare_parameter("grid_resolution", 0.1)
        self.declare_parameter("map_file", "/home/seb/saltybot-data/obstacle_map.yaml")
        self.declare_parameter("scan_max_range", 10.0)
        self.declare_parameter("scan_min_range", 0.1)
        self.declare_parameter("publish_rate_hz", 5.0)
        self.declare_parameter("decay_interval_seconds", 3600.0)
        self.declare_parameter("auto_save_interval", 300.0)

        width = self.get_parameter("grid_width").value
        height = self.get_parameter("grid_height").value
        resolution = self.get_parameter("grid_resolution").value
        map_file = self.get_parameter("map_file").value

        self.map_file = Path(map_file)
        self.map_file.parent.mkdir(parents=True, exist_ok=True)

        self.grid = OccupancyGridMap(
            width=width,
            height=height,
            resolution=resolution,
        )

        if self.map_file.exists():
            try:
                with open(self.map_file) as f:
                    data = yaml.safe_load(f)
                    self.grid.from_yaml(data)
                    self.get_logger().info(f"Loaded map from {self.map_file}")
            except Exception as e:
                self.get_logger().error(f"Failed to load map: {e}")

        self.robot_x = 0.0
        self.robot_y = 0.0
        self.robot_angle = 0.0
        self.scan_max_range = self.get_parameter("scan_max_range").value
        self.scan_min_range = self.get_parameter("scan_min_range").value

        qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            history=HistoryPolicy.KEEP_LAST,
            depth=5,
        )

        self.scan_sub = self.create_subscription(
            LaserScan,
            "/scan",
            self.scan_callback,
            qos,
        )

        self.odom_sub = self.create_subscription(
            Odometry,
            "/odom",
            self.odom_callback,
            qos,
        )

        self.map_pub = self.create_publisher(
            OccupancyGrid,
            "/saltybot/obstacle_map",
            10,
        )

        publish_rate = self.get_parameter("publish_rate_hz").value
        self.publish_timer = self.create_timer(1.0 / publish_rate, self.publish_map_callback)

        decay_interval = self.get_parameter("decay_interval_seconds").value
        self.decay_timer = self.create_timer(decay_interval, self.decay_callback)

        auto_save_interval = self.get_parameter("auto_save_interval").value
        self.save_timer = self.create_timer(auto_save_interval, self.save_map_callback)

        self.get_logger().info(
            f"Obstacle memory initialized: {width}x{height} grid @ {resolution}m/cell, "
            f"map file: {self.map_file}"
        )

    def odom_callback(self, msg: Odometry):
        """Update robot position from odometry."""
        self.robot_x = msg.pose.pose.position.x
        self.robot_y = msg.pose.pose.position.y

        q = msg.pose.pose.orientation
        siny_cosp = 2 * (q.w * q.z + q.x * q.y)
        cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
        self.robot_angle = math.atan2(siny_cosp, cosy_cosp)

    def scan_callback(self, msg: LaserScan):
        """Process LIDAR scan and update grid."""
        angle_min = msg.angle_min
        angle_increment = msg.angle_increment

        for i, distance in enumerate(msg.ranges):
            if math.isnan(distance) or math.isinf(distance):
                continue

            if distance < self.scan_min_range or distance > self.scan_max_range:
                continue

            ray_angle = angle_min + i * angle_increment + self.robot_angle

            hit_x = self.robot_x + distance * math.cos(ray_angle)
            hit_y = self.robot_y + distance * math.sin(ray_angle)

            self.grid.mark_obstacle(hit_x, hit_y)

            steps = max(1, int(distance / (self.grid.resolution * 2)))
            for step in range(steps):
                free_dist = (distance * step) / steps
                free_x = self.robot_x + free_dist * math.cos(ray_angle)
                free_y = self.robot_y + free_dist * math.sin(ray_angle)
                self.grid.mark_free(free_x, free_y)

    def publish_map_callback(self):
        """Publish current occupancy grid."""
        grid_msg = self.grid.to_occupancy_grid("map")
        self.map_pub.publish(grid_msg)

    def decay_callback(self):
        """Apply temporal decay to grid."""
        self.grid.apply_temporal_decay()
        stats = self.grid.get_statistics()
        self.get_logger().info(
            f"Grid decay: {stats['hazard_cells']} hazards, "
            f"{stats['obstacle_cells']} obstacles, {stats['free_cells']} free"
        )

    def save_map_callback(self):
        """Save grid to YAML file."""
        try:
            with open(self.map_file, "w") as f:
                yaml.dump(self.grid.to_yaml(), f, default_flow_style=False)
        except Exception as e:
            self.get_logger().error(f"Failed to save map: {e}")

    def destroy_node(self):
        """Cleanup: save map before shutdown."""
        self.save_map_callback()
        super().destroy_node()


def main(args=None):
    rclpy.init(args=args)
    node = ObstacleMemoryNode()

    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
