"""
3D metrics calculation module.
"""
import numpy as np
from collections import defaultdict

from ..class_config import CLASSES_3D as _CLASSES_3D


def normalize_angle(angle):
    """
    Normalize angle to [-π, π].
    
    Args:
        angle: float, angle in radians
        
    Returns:
        float, normalized angle
    """
    while angle > np.pi:
        angle -= 2 * np.pi
    while angle < -np.pi:
        angle += 2 * np.pi
    return angle


class Metrics3D:
    """Calculate 3D detection metrics (lateral/longitudinal errors, heading error)."""
    
    # 3D class IDs — sourced from eval_tools/class_config.py
    CLASSES_3D = _CLASSES_3D  # vehicle, bus, truck, tanker, unknown, pedestrian, bicycle, motorcyclist, tricycle
    VEHICLE_LARGE_CLASS_NAME = 'vehicle_large'
    VEHICLE_SMALL_CLASS_NAME = 'vehicle_small'
    VEHICLE_SIZE_CLASS_NAMES = [VEHICLE_LARGE_CLASS_NAME, VEHICLE_SMALL_CLASS_NAME]
    
    # Face mapping for vehicle class
    FACE_MAPPING = {
        'front': [18, 19, 20],  # Indices in GT values for front face center
        'back': [26, 27, 28],
        'rear': [26, 27, 28],   # Alias for back (some models use 'rear')
        'tail': [26, 27, 28],   # Alias for back (some models use 'tail')
        'left': [34, 35, 36],
        'right': [42, 43, 44]
    }
    
    def __init__(self, distance_ranges=None, lateral_distance_ranges=None, heading_tolerance='strict',
                 coord_system='camera', vehicle_size_split=None):
        """Initialize 3D metrics calculator.
        
        Args:
            distance_ranges: list of [min, max] longitudinal distance ranges in meters (z-axis).
                           e.g., [[0, 30], [30, 60], [60, 100]]
                           If None, only overall statistics are computed.
            lateral_distance_ranges: list of [min, max] lateral distance ranges in meters (x-axis).
                           e.g., [[-10, -5], [-5, 0], [0, 5], [5, 10]]
                           If None, no lateral distance grouping is used.
            heading_tolerance: str, heading error calculation mode
                           'strict': Standard calculation (default)
                           'relaxed': Consider 180° symmetry (for symmetric objects)
                           'both': Calculate both strict and relaxed metrics
            vehicle_size_split: dict, optional vehicle GT size split config:
                           {
                               'enabled': True,
                               'large_length_threshold': 6.0,
                               'large_height_threshold': 3.0,
                           }
        """
        if coord_system not in ('camera', 'ego'):
            raise ValueError(f"Unsupported coord_system: {coord_system}")
        self.coord_system = coord_system
        vehicle_size_split = vehicle_size_split or {}
        self.vehicle_size_split_enabled = vehicle_size_split.get('enabled', True)
        self.vehicle_large_length_threshold = float(
            vehicle_size_split.get('large_length_threshold', 6.0)
        )
        self.vehicle_large_height_threshold = float(
            vehicle_size_split.get('large_height_threshold', 3.0)
        )

        # Store errors for each class and distance range
        # Format: {class_id: {range_key: {'lateral': [...], 'longitudinal': [...], 'heading': [...], 'heading_relaxed': [...], 'heading_reversal': [...]}}}
        self.errors = defaultdict(lambda: defaultdict(lambda: {
            'lateral': [],
            'longitudinal': [],
            'longitudinal_relative': [],
            'heading': [],
            'heading_relaxed': [],
            'heading_reversal': []
        }))
        
        # Distance ranges configuration
        self.distance_ranges = distance_ranges
        self.lateral_distance_ranges = lateral_distance_ranges
        
        # Heading tolerance mode
        self.heading_tolerance = heading_tolerance
        
        # Build range keys
        self.range_keys = ['overall']
        if self.distance_ranges:
            self.range_keys.extend([f"long_{r[0]}-{r[1]}m" for r in self.distance_ranges])
        if self.lateral_distance_ranges:
            self.range_keys.extend([f"lat_{r[0]}-{r[1]}m" for r in self.lateral_distance_ranges])

    def _get_lateral_axis(self):
        return 0 if self.coord_system == 'camera' else 1

    def _get_longitudinal_axis(self):
        return 2 if self.coord_system == 'camera' else 0

    def _get_vehicle_size_class_name(self, gt):
        """Classify vehicle GT as large/small by GT 3D box dimensions."""
        dimensions = (gt.get('3d_info') or {}).get('dimensions')
        if dimensions is None or len(dimensions) < 2:
            return None

        length = float(dimensions[0])
        height = float(dimensions[1])
        if length > self.vehicle_large_length_threshold and height > self.vehicle_large_height_threshold:
            return self.VEHICLE_LARGE_CLASS_NAME
        return self.VEHICLE_SMALL_CLASS_NAME

    def _get_storage_class_keys(self, gt, class_id):
        """Return all metric buckets that should receive this sample."""
        class_keys = [class_id]
        if class_id == 0 and self.vehicle_size_split_enabled:
            vehicle_size_class = self._get_vehicle_size_class_name(gt)
            if vehicle_size_class is not None:
                class_keys.append(vehicle_size_class)
        return class_keys

    def _append_error_sample(self, class_key, range_key, lateral_error, longitudinal_error,
                             longitudinal_relative_error, heading_error_strict,
                             heading_error_relaxed, is_reversal):
        self.errors[class_key][range_key]['lateral'].append(lateral_error)
        self.errors[class_key][range_key]['longitudinal'].append(longitudinal_error)
        self.errors[class_key][range_key]['longitudinal_relative'].append(longitudinal_relative_error)
        self.errors[class_key][range_key]['heading'].append(heading_error_strict)
        self.errors[class_key][range_key]['heading_relaxed'].append(heading_error_relaxed)
        self.errors[class_key][range_key]['heading_reversal'].append(1 if is_reversal else 0)
    
    def add_sample(self, gt, det, class_id):
        """
        Add a matched GT-detection pair for 3D metric calculation.
        
        Args:
            gt: dict, ground truth with 3d_info
            det: dict, detection with 3d_info
            class_id: int, class ID
        """
        if class_id not in self.CLASSES_3D:
            return
        
        if gt['3d_info'] is None or det['3d_info'] is None:
            return
        
        # Get GT 3D center based on class type
        if class_id == 0:  # Vehicle class
            # Use face-specific center
            gt_center = self._get_vehicle_face_center(gt['3d_info'], det['3d_info']['face_type'])
        else:
            # Non-vehicle: use overall 3D box center
            gt_center = gt['3d_info']['center']
        
        # Get detection 3D center
        det_center = det['3d_info']['center']
        
        lateral_axis = self._get_lateral_axis()
        longitudinal_axis = self._get_longitudinal_axis()

        # Calculate lateral / longitudinal error in the configured coordinate system.
        lateral_error = abs(det_center[lateral_axis] - gt_center[lateral_axis])
        longitudinal_error = abs(det_center[longitudinal_axis] - gt_center[longitudinal_axis])
        
        # Calculate longitudinal relative error (relative to GT depth).
        # Use whole-object center z as the reference depth so that face centers
        # with near-zero or negative z (e.g. side faces at close range) don't
        # produce astronomically large relative errors.
        gt_whole_longitudinal = gt['3d_info']['center'][longitudinal_axis]
        gt_depth = max(abs(gt_whole_longitudinal), 1e-6)
        longitudinal_relative_error = longitudinal_error / gt_depth
        
        # Calculate heading error
        gt_rot = gt['3d_info']['rotation']
        det_rot = det['3d_info']['rotation']
        heading_error_strict = abs(normalize_angle(det_rot - gt_rot))
        
        # Calculate relaxed heading error (considering 180° symmetry)
        heading_error_relaxed = heading_error_strict
        is_reversal = False
        if heading_error_strict > np.pi - 0.2:  # Close to 180° (within ~170°)
            heading_error_relaxed = np.pi - heading_error_strict
            is_reversal = True
        
        # Get distance for range classification.
        # Always use whole-object center z so that samples with negative/near-zero
        # face center z (side faces at close range) still fall into the correct
        # longitudinal bucket instead of being silently dropped from all segments
        # while remaining in 'overall' with enormous relative errors.
        longitudinal_distance = gt['3d_info']['center'][longitudinal_axis]
        lateral_distance = gt_center[lateral_axis]
        
        class_keys = self._get_storage_class_keys(gt, class_id)

        # Store errors in overall
        for class_key in class_keys:
            self._append_error_sample(
                class_key,
                'overall',
                lateral_error,
                longitudinal_error,
                longitudinal_relative_error,
                heading_error_strict,
                heading_error_relaxed,
                is_reversal,
            )
        
        # Store errors by longitudinal distance range if configured
        if self.distance_ranges:
            range_key = self._get_longitudinal_distance_range(longitudinal_distance)
            if range_key:
                for class_key in class_keys:
                    self._append_error_sample(
                        class_key,
                        range_key,
                        lateral_error,
                        longitudinal_error,
                        longitudinal_relative_error,
                        heading_error_strict,
                        heading_error_relaxed,
                        is_reversal,
                    )
        
        # Store errors by lateral distance range if configured
        if self.lateral_distance_ranges:
            range_key = self._get_lateral_distance_range(lateral_distance)
            if range_key:
                for class_key in class_keys:
                    self._append_error_sample(
                        class_key,
                        range_key,
                        lateral_error,
                        longitudinal_error,
                        longitudinal_relative_error,
                        heading_error_strict,
                        heading_error_relaxed,
                        is_reversal,
                    )
    
    def _get_vehicle_face_center(self, gt_3d_info, face_type):
        """
        Get vehicle face center from GT based on predicted face type.
        
        Args:
            gt_3d_info: dict, GT 3D info with faces
            face_type: str, predicted face type (front/back/rear/tail/left/right)
            
        Returns:
            list, [x3d, y3d, z3d] of the specified face center
        """
        if gt_3d_info['faces'] is None:
            # Fallback to overall center if no face info
            return gt_3d_info['center']
        
        # Normalize face_type: rear/tail -> back
        normalized_face_type = face_type.lower()
        if normalized_face_type in ['rear', 'tail']:
            normalized_face_type = 'back'
        
        face_data = gt_3d_info['faces'].get(normalized_face_type)
        # face_data[7] is is_visible_from_camera; when 0 the x3d/y3d/z3d are set to
        # the sentinel -1.0, which must not be used as a real coordinate.
        if face_data is None or (len(face_data) >= 8 and face_data[7] == 0) or face_data[2] <= 0:
            # Fallback to overall center if face not found or not visible
            return gt_3d_info['center']
        
        # Extract x3d, y3d, z3d from face data (first 3 elements)
        return face_data[:3]
    
    def _get_longitudinal_distance_range(self, distance):
        """Get longitudinal distance range key for a given distance.
        
        Args:
            distance: float, longitudinal distance in meters (z-axis)
            
        Returns:
            str, range key (e.g., 'long_0-30m') or None if not in any range
        """
        if not self.distance_ranges:
            return None
        
        for range_min, range_max in self.distance_ranges:
            if range_min <= distance < range_max:
                return f"long_{range_min}-{range_max}m"
        
        return None
    
    def _get_lateral_distance_range(self, distance):
        """Get lateral distance range key for a given distance.
        
        Args:
            distance: float, lateral distance in meters (x-axis)
            
        Returns:
            str, range key (e.g., 'lat_-10--5m') or None if not in any range
        """
        if not self.lateral_distance_ranges:
            return None
        
        for range_min, range_max in self.lateral_distance_ranges:
            if range_min <= distance < range_max:
                return f"lat_{range_min}-{range_max}m"
        
        return None
    
    def compute_statistics(self, class_id, range_key='overall'):
        """
        Compute statistical metrics for a class and distance range.
        
        Args:
            class_id: int | str, class ID or synthetic class key
            range_key: str, distance range key (e.g., 'overall', '0-30m')
            
        Returns:
            dict with mean, median, std, 90th percentile for each error type
        """
        if class_id not in self.errors or range_key not in self.errors[class_id] or \
           len(self.errors[class_id][range_key]['lateral']) == 0:
            result = {
                'lateral_error': {'mean': 0, 'median': 0, 'std': 0, 'percentile_90': 0},
                'longitudinal_error': {'mean': 0, 'median': 0, 'std': 0, 'percentile_90': 0},
                'longitudinal_relative_error': {'mean': 0, 'median': 0, 'std': 0, 'percentile_90': 0},
                'heading_error': {'mean': 0, 'median': 0, 'std': 0, 'percentile_90': 0},
                'num_samples': 0
            }
            if self.heading_tolerance in ['relaxed', 'both']:
                result['heading_error_relaxed'] = {'mean': 0, 'median': 0, 'std': 0, 'percentile_90': 0}
                result['reversal_count'] = 0
                result['reversal_percentage'] = 0.0
            return result
        
        lateral = np.array(self.errors[class_id][range_key]['lateral'])
        longitudinal = np.array(self.errors[class_id][range_key]['longitudinal'])
        longitudinal_relative = np.array(self.errors[class_id][range_key]['longitudinal_relative'])
        heading = np.array(self.errors[class_id][range_key]['heading'])
        heading_relaxed = np.array(self.errors[class_id][range_key]['heading_relaxed'])
        heading_reversal = np.array(self.errors[class_id][range_key]['heading_reversal'])
        
        def stats(arr):
            return {
                'mean': float(np.mean(arr)),
                'median': float(np.median(arr)),
                'std': float(np.std(arr)),
                'percentile_90': float(np.percentile(arr, 90))
            }
        
        result = {
            'lateral_error': stats(lateral),
            'longitudinal_error': stats(longitudinal),
            'longitudinal_relative_error': stats(longitudinal_relative),
            'num_samples': len(lateral)
        }
        
        # Add heading error based on tolerance mode
        if self.heading_tolerance == 'strict':
            result['heading_error'] = stats(heading)
        elif self.heading_tolerance == 'relaxed':
            result['heading_error'] = stats(heading_relaxed)
            result['reversal_count'] = int(np.sum(heading_reversal))
            result['reversal_percentage'] = float(np.sum(heading_reversal) / len(heading_reversal) * 100)
        elif self.heading_tolerance == 'both':
            result['heading_error'] = stats(heading)
            result['heading_error_relaxed'] = stats(heading_relaxed)
            result['reversal_count'] = int(np.sum(heading_reversal))
            result['reversal_percentage'] = float(np.sum(heading_reversal) / len(heading_reversal) * 100)
        
        return result
    
    def get_summary(self):
        """
        Get complete 3D evaluation summary.
        
        Returns:
            dict with per-class 3D metrics, including distance ranges if configured
        """
        from .parser import GroundTruthParser
        
        summary = {}

        summary_class_keys = list(self.CLASSES_3D)
        if self.vehicle_size_split_enabled:
            summary_class_keys.extend(self.VEHICLE_SIZE_CLASS_NAMES)

        for class_id in summary_class_keys:
            if isinstance(class_id, str):
                class_name = class_id
            else:
                class_name = GroundTruthParser.CLASS_NAMES.get(class_id, f"class_{class_id}")
            
            # Compute statistics for each distance range
            class_summary = {}
            for range_key in self.range_keys:
                stats = self.compute_statistics(class_id, range_key)
                if stats['num_samples'] > 0 or range_key == 'overall':
                    class_summary[range_key] = stats
            
            if class_summary:
                summary[class_name] = class_summary
        
        return summary