from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Optional

import numpy as np


@dataclass
class LocalizationFrameState:
    odom_to_map: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
    odom_to_ref: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
    last_alignment: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
    confidence: float = 0.0
    ref_valid: bool = False

    def reset(self) -> None:
        self.odom_to_map = np.eye(4, dtype=np.float64)
        self.odom_to_ref = np.eye(4, dtype=np.float64)
        self.last_alignment = np.eye(4, dtype=np.float64)
        self.confidence = 0.0
        self.ref_valid = False

    def reset_reference(self) -> None:
        self.odom_to_ref = np.eye(4, dtype=np.float64)
        self.last_alignment = np.eye(4, dtype=np.float64)
        self.confidence = 0.0
        self.ref_valid = False

    def set_reference_alignment(self, odom_to_ref: np.ndarray, confidence: float) -> None:
        tf = np.asarray(odom_to_ref, dtype=np.float64)
        if tf.shape != (4, 4):
            return
        self.odom_to_ref = np.array(tf, dtype=np.float64, copy=True)
        self.last_alignment = np.array(tf, dtype=np.float64, copy=True)
        self.confidence = float(np.clip(confidence, 0.0, 1.0))
        self.ref_valid = True

    def update_confidence(self, confidence: float) -> None:
        self.confidence = float(np.clip(confidence, 0.0, 1.0))

    def invalidate_reference(self) -> None:
        self.ref_valid = False
        self.confidence = 0.0

    def apply_map_correction(self, correction: np.ndarray) -> None:
        corr = np.asarray(correction, dtype=np.float64)
        if corr.shape != (4, 4):
            return
        self.odom_to_map = corr @ np.asarray(self.odom_to_map, dtype=np.float64)

    def apply_loop_correction(self, correction: np.ndarray, update_reference: bool = True) -> None:
        corr = np.asarray(correction, dtype=np.float64)
        if corr.shape != (4, 4):
            return
        self.odom_to_map = corr @ np.asarray(self.odom_to_map, dtype=np.float64)
        if update_reference and self.ref_valid:
            try:
                inv_corr = np.linalg.inv(corr)
                self.odom_to_ref = np.asarray(self.odom_to_ref, dtype=np.float64) @ inv_corr
                self.last_alignment = np.array(self.odom_to_ref, dtype=np.float64, copy=True)
            except Exception:
                pass

    def map_pose(self, odom_pose: Optional[np.ndarray]) -> Optional[np.ndarray]:
        if odom_pose is None:
            return None
        pose = np.asarray(odom_pose, dtype=np.float64)
        if pose.shape != (4, 4):
            return None
        return np.asarray(self.odom_to_map, dtype=np.float64) @ pose

    def ref_pose(self, odom_pose: Optional[np.ndarray]) -> Optional[np.ndarray]:
        if odom_pose is None or not self.ref_valid:
            return None
        pose = np.asarray(odom_pose, dtype=np.float64)
        if pose.shape != (4, 4):
            return None
        return np.asarray(self.odom_to_ref, dtype=np.float64) @ pose

    def snapshot(self) -> Dict[str, Any]:
        return {
            "ref_valid": bool(self.ref_valid),
            "confidence": float(self.confidence),
            "odom_to_map_t": np.asarray(self.odom_to_map[:3, 3], dtype=np.float64).tolist(),
            "odom_to_ref_t": np.asarray(self.odom_to_ref[:3, 3], dtype=np.float64).tolist(),
        }