#!/usr/bin/env python3
"""
terrain_classifier_node.py — Issue #469: Terrain classification for speed adaptation.

Analyzes IMU vibration (via FFT) and camera texture (via MobileNetV2 CNN) to classify
terrain type (pavement/grass/gravel/sand/indoor) and estimate surface roughness.

Publishes TerrainState.msg with recommended speed adjustment based on terrain.
"""

from collections import deque
from dataclasses import dataclass
import math

import cv2
from cv_bridge import CvBridge
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy

from sensor_msgs.msg import Imu, Image
from std_msgs.msg import Header

# Import the TerrainState message (will be auto-generated at build time)
try:
    from saltybot_social_msgs.msg import TerrainState
except ImportError:
    # Fallback if not built yet
    pass

try:
    import tensorflow as tf
except ImportError:
    tf = None


@dataclass
class TerrainDecision:
    """Result of terrain classification."""
    terrain_type: int
    roughness: float  # 0.0-1.0
    confidence: float  # 0.0-1.0
    recommended_speed_ratio: float  # 0.1-1.0


class IMUVibrationAnalyzer:
    """Analyzes IMU data for surface roughness via FFT."""

    def __init__(self, window_size: int = 200, sample_rate: float = 200.0):
        """
        Args:
            window_size: Number of samples for FFT window (1 sec at 200 Hz = 200 samples)
            sample_rate: IMU sample rate in Hz
        """
        self.window_size = window_size
        self.sample_rate = sample_rate
        self.accel_buffer = deque(maxlen=window_size)

    def add_sample(self, accel_x: float, accel_y: float, accel_z: float):
        """Add acceleration sample to buffer."""
        # Remove gravity (9.81 m/s²) to get vibration only
        accel_mag = math.sqrt(accel_x**2 + accel_y**2 + accel_z**2)
        self.accel_buffer.append(accel_mag)

    def compute_roughness(self) -> tuple[float, float]:
        """
        Compute roughness metric from IMU vibration.

        Returns:
            (roughness: 0.0-1.0, accel_magnitude: m/s²)
        """
        if len(self.accel_buffer) < self.window_size // 2:
            return 0.0, 0.0

        # Compute FFT of acceleration
        accel_data = np.array(self.accel_buffer, dtype=np.float32)
        fft_result = np.fft.fft(accel_data)
        psd = np.abs(fft_result) ** 2 / len(accel_data)

        # Frequency axis
        freqs = np.fft.fftfreq(len(accel_data), 1.0 / self.sample_rate)

        # Extract energy in high-frequency band (10-50 Hz) for roughness
        # Lower frequencies include locomotion/body dynamics
        high_freq_mask = (freqs >= 10) & (freqs <= 50)
        high_freq_energy = np.sum(psd[high_freq_mask])

        # Total energy (RMS)
        total_energy = np.sum(psd)

        # Roughness = ratio of high-freq energy, normalized and saturated
        if total_energy > 0:
            roughness = min(1.0, high_freq_energy / total_energy * 10.0)
        else:
            roughness = 0.0

        # Mean acceleration magnitude (remove gravity offset)
        accel_mag = np.mean(accel_data)

        return roughness, accel_mag

    def reset(self):
        """Clear buffer."""
        self.accel_buffer.clear()


class CNNTextureAnalyzer:
    """Analyzes camera images for terrain texture using MobileNetV2."""

    def __init__(self):
        """Initialize MobileNetV2 for feature extraction."""
        self.cv_bridge = CvBridge()
        self.model = None
        self.input_size = 224

        # Try to load MobileNetV2 if TensorFlow is available
        if tf is not None:
            try:
                # Use pre-trained MobileNetV2 from Keras
                self.model = tf.keras.applications.MobileNetV2(
                    input_shape=(self.input_size, self.input_size, 3),
                    include_top=False,
                    weights='imagenet'
                )
                self.model.trainable = False
            except Exception as e:
                print(f"Warning: Could not load MobileNetV2: {e}")
                self.model = None

    def extract_features(self, image_msg: Image) -> np.ndarray:
        """
        Extract texture features from image using CNN.

        Args:
            image_msg: ROS Image message

        Returns:
            Feature vector (1280-dim for MobileNetV2, or None if model unavailable)
        """
        try:
            # Convert ROS image to OpenCV
            cv_image = self.cv_bridge.imgmsg_to_cv2(image_msg, desired_encoding='rgb8')

            if self.model is None:
                # Fallback: use simple edge detection features
                return self._fallback_texture_features(cv_image)

            # Preprocess image for model
            resized = cv2.resize(cv_image, (self.input_size, self.input_size))
            # Normalize to [-1, 1] (MobileNetV2 standard)
            normalized = (resized / 127.5) - 1.0

            # Add batch dimension
            batch = np.expand_dims(normalized, axis=0).astype(np.float32)

            # Extract features
            features = self.model(batch, training=False)
            features = tf.keras.layers.GlobalAveragePooling2D()(features)

            return features.numpy().flatten()

        except Exception as e:
            print(f"Error extracting CNN features: {e}")
            return self._fallback_texture_features(cv_image if 'cv_image' in locals() else None)

    def _fallback_texture_features(self, cv_image: np.ndarray) -> np.ndarray:
        """
        Fallback: simple edge-based texture features when CNN unavailable.

        Args:
            cv_image: OpenCV image (BGR)

        Returns:
            Simple feature vector (10-dim)
        """
        if cv_image is None or cv_image.size == 0:
            return np.zeros(10)

        # Convert to grayscale
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)

        # Edge detection
        edges = cv2.Canny(gray, 100, 200)

        # Texture features
        edge_density = np.sum(edges > 0) / edges.size

        # Histogram of oriented gradients (simplified)
        gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
        gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
        magnitude = np.sqrt(gx**2 + gy**2)
        direction = np.arctan2(gy, gx)

        # Bin gradients by direction (8 bins)
        hist, _ = np.histogram(direction, bins=8, range=(-np.pi, np.pi))
        hist = hist / np.sum(hist) if np.sum(hist) > 0 else hist

        # Combine features
        features = np.concatenate([
            [edge_density],
            [np.mean(magnitude), np.std(magnitude)],
            hist[:8]  # 8 direction bins
        ])

        return features


class TerrainClassifier:
    """Decision logic for terrain classification."""

    # Terrain type constants (must match TerrainState.msg)
    TERRAIN_UNKNOWN = 0
    TERRAIN_PAVEMENT = 1
    TERRAIN_GRASS = 2
    TERRAIN_GRAVEL = 3
    TERRAIN_SAND = 4
    TERRAIN_INDOOR = 5

    # Speed recommendations for each terrain
    SPEED_RECOMMENDATIONS = {
        TERRAIN_UNKNOWN: 0.7,
        TERRAIN_PAVEMENT: 1.0,
        TERRAIN_GRASS: 0.8,
        TERRAIN_GRAVEL: 0.5,
        TERRAIN_SAND: 0.4,
        TERRAIN_INDOOR: 0.7,
    }

    def __init__(self):
        """Initialize terrain classifier."""
        self.last_terrain = self.TERRAIN_INDOOR
        self.last_roughness = 0.0

    def classify(
        self,
        imu_roughness: float,
        accel_magnitude: float,
        cnn_features: np.ndarray = None,
    ) -> TerrainDecision:
        """
        Classify terrain based on IMU and camera features.

        Args:
            imu_roughness: IMU-derived roughness (0.0-1.0)
            accel_magnitude: Mean acceleration magnitude (m/s²)
            cnn_features: CNN feature vector (optional)

        Returns:
            TerrainDecision with type, roughness, confidence, speed ratio
        """
        # Simple decision tree based on features
        if imu_roughness < 0.2 and accel_magnitude < 9.8:
            # Smooth surface with low vibration
            terrain = self.TERRAIN_PAVEMENT
            confidence = 0.9
            roughness = imu_roughness
        elif imu_roughness < 0.35 and accel_magnitude < 10.0:
            # Soft, relatively smooth surface
            terrain = self.TERRAIN_GRASS
            confidence = 0.75
            roughness = imu_roughness
        elif imu_roughness > 0.45 and imu_roughness < 0.75:
            # Rough, loose surface
            terrain = self.TERRAIN_GRAVEL
            confidence = 0.8
            roughness = imu_roughness
        elif imu_roughness > 0.4 and accel_magnitude > 10.0:
            # Granular, bouncy surface
            terrain = self.TERRAIN_SAND
            confidence = 0.7
            roughness = imu_roughness
        else:
            # Default to indoor or unknown
            terrain = self.TERRAIN_INDOOR
            confidence = 0.5
            roughness = min(imu_roughness, 0.5)

        # Apply hysteresis to avoid terrain switching jitter
        if terrain != self.last_terrain and confidence < 0.7:
            terrain = self.last_terrain
            confidence = 0.3

        self.last_terrain = terrain
        self.last_roughness = roughness

        # Recommended speed based on terrain
        speed_ratio = self.SPEED_RECOMMENDATIONS.get(terrain, 0.7)

        return TerrainDecision(
            terrain_type=terrain,
            roughness=roughness,
            confidence=confidence,
            recommended_speed_ratio=speed_ratio,
        )


class TerrainClassifierNode(Node):
    """ROS2 node for terrain classification."""

    def __init__(self):
        super().__init__('terrain_classifier')

        # Parameters
        self.declare_parameter('imu_window_size', 200)
        self.declare_parameter('publish_rate_hz', 5.0)
        self.declare_parameter('roughness_threshold', 0.3)
        self.declare_parameter('terrain_timeout_s', 5.0)

        imu_window_size = self.get_parameter('imu_window_size').value
        publish_rate_hz = self.get_parameter('publish_rate_hz').value

        # Initialize analyzers
        self.imu_analyzer = IMUVibrationAnalyzer(window_size=imu_window_size)
        self.cnn_analyzer = CNNTextureAnalyzer()
        self.classifier = TerrainClassifier()

        # QoS profile for sensor data (BEST_EFFORT, low latency)
        qos = QoSProfile(
            reliability=ReliabilityPolicy.BEST_EFFORT,
            depth=1
        )

        # Subscribers
        self.imu_sub = self.create_subscription(
            Imu,
            '/camera/imu',
            self._imu_callback,
            qos_profile=qos
        )

        self.camera_sub = self.create_subscription(
            Image,
            '/camera/color/image_raw',
            self._camera_callback,
            qos_profile=qos
        )

        # Publisher
        self.terrain_pub = self.create_publisher(
            TerrainState,
            '/saltybot/terrain_state',
            10
        )

        # Timer for periodic publishing
        dt = 1.0 / publish_rate_hz
        self.create_timer(dt, self._publish_terrain)

        self.get_logger().info(f'Terrain classifier initialized (window={imu_window_size}, rate={publish_rate_hz} Hz)')

    def _imu_callback(self, msg: Imu):
        """Process IMU data."""
        # Extract linear acceleration (subtract gravity)
        accel_x = msg.linear_acceleration.x
        accel_y = msg.linear_acceleration.y
        accel_z = msg.linear_acceleration.z - 9.81  # Remove gravity

        self.imu_analyzer.add_sample(accel_x, accel_y, accel_z)

    def _camera_callback(self, msg: Image):
        """Process camera image."""
        # Store latest camera features for next classification
        self.latest_camera_features = self.cnn_analyzer.extract_features(msg)

    def _publish_terrain(self):
        """Compute and publish terrain state."""
        # Get features from latest sensors
        imu_roughness, accel_magnitude = self.imu_analyzer.compute_roughness()
        cnn_features = getattr(self, 'latest_camera_features', None)

        # Classify terrain
        decision = self.classifier.classify(imu_roughness, accel_magnitude, cnn_features)

        # Build message
        msg = TerrainState()
        msg.header = Header()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.header.frame_id = 'base_link'
        msg.terrain_type = decision.terrain_type
        msg.roughness = decision.roughness
        msg.confidence = decision.confidence
        msg.recommended_speed_ratio = decision.recommended_speed_ratio
        msg.source = 'fused'

        self.terrain_pub.publish(msg)

        # Debug logging
        terrain_names = {
            0: 'UNKNOWN',
            1: 'PAVEMENT',
            2: 'GRASS',
            3: 'GRAVEL',
            4: 'SAND',
            5: 'INDOOR',
        }
        self.get_logger().debug(
            f'Terrain: {terrain_names.get(decision.terrain_type, "?")} '
            f'(roughness={decision.roughness:.2f}, confidence={decision.confidence:.2f}, '
            f'speed_ratio={decision.recommended_speed_ratio:.2f})'
        )


def main(args=None):
    rclpy.init(args=args)
    node = TerrainClassifierNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
