import pytest
import numpy as np
from unittest.mock import Mock, patch
from datetime import datetime

from f10_factor_graph_optimizer import FactorGraphOptimizer, FactorGraphConfig
from f07_sequential_visual_odometry import RelativePose
from f02_1_flight_lifecycle_manager import GPSPoint
from f09_local_geospatial_anchoring import Sim3Transform

@pytest.fixture
def fgo():
    return FactorGraphOptimizer()

@pytest.fixture
def dummy_relative_pose():
    return RelativePose(
        translation=np.array([1.0, 0.0, 0.0]),
        rotation=np.eye(3),
        confidence=0.9,
        inlier_count=100,
        total_matches=120,
        tracking_good=True
    )

class TestFactorGraphOptimizer:

    # --- 10.01 Core Factor Management ---

    def test_add_relative_factor_scales_translation(self, fgo, dummy_relative_pose):
        flight_id = "flight_1"
        cov = np.eye(6)
        
        res = fgo.add_relative_factor(flight_id, 1, 2, dummy_relative_pose, cov)
        assert res is True
        
        # Verify scaling logic via internal state inspection
        state = fgo._get_or_create_flight_graph(flight_id)
        if not state.isam2: # If mocking
            T = state.global_values[2]
            assert np.isclose(T[0, 3], 100.0) # Scaled by default 100m frame spacing

    def test_add_absolute_factor_converts_gps_to_enu(self, fgo):
        flight_id = "flight_abs"
        gps_origin = GPSPoint(lat=48.0, lon=37.0)
        gps_target = GPSPoint(lat=48.001, lon=37.001)
        
        fgo.add_absolute_factor(flight_id, 1, gps_origin, np.eye(3), is_user_anchor=True)
        fgo.add_absolute_factor(flight_id, 2, gps_target, np.eye(3), is_user_anchor=False)
        
        state = fgo._get_or_create_flight_graph(flight_id)
        assert state.reference_origin.lat == 48.0

    # --- 10.02 Trajectory Optimization ---

    def test_optimize_incremental(self, fgo, dummy_relative_pose):
        flight_id = "flight_opt"
        fgo.add_relative_factor(flight_id, 1, 2, dummy_relative_pose, np.eye(6))
        
        result = fgo.optimize(flight_id, iterations=3)
        assert result.converged is True
        assert result.iterations_used == 3
        
        traj = fgo.get_trajectory(flight_id)
        assert len(traj) > 0
        assert 2 in traj
        assert isinstance(traj[2].position, np.ndarray)

    def test_get_marginal_covariance_returns_6x6(self, fgo, dummy_relative_pose):
        fgo.add_relative_factor("flight_cov", 1, 2, dummy_relative_pose, np.eye(6))
        cov = fgo.get_marginal_covariance("flight_cov", 2)
        assert cov.shape == (6, 6)

    # --- 10.03 Chunk Subgraph Operations ---

    def test_chunk_creation_and_isolation(self, fgo, dummy_relative_pose):
        fid = "flight_chunk"
        assert fgo.create_chunk_subgraph(fid, "chunk1", 10) is True
        assert fgo.create_chunk_subgraph(fid, "chunk1", 10) is False # Duplicate fails
        
        assert fgo.add_relative_factor_to_chunk(fid, "chunk1", 10, 11, dummy_relative_pose, np.eye(6)) is True
        assert fgo.add_relative_factor_to_chunk(fid, "chunk_missing", 10, 11, dummy_relative_pose, np.eye(6)) is False
        
        traj = fgo.get_chunk_trajectory(fid, "chunk1")
        assert len(traj) == 2
        assert 11 in traj

    def test_optimize_chunk_isolation(self, fgo, dummy_relative_pose):
        fid = "flight_iso"
        fgo.create_chunk_subgraph(fid, "C1", 1)
        fgo.create_chunk_subgraph(fid, "C2", 20)
        
        res = fgo.optimize_chunk(fid, "C1", 5)
        assert res.converged is True

    # --- 10.04 Chunk Merging ---

    def test_merge_applies_sim3_transform(self, fgo, dummy_relative_pose):
        fid = "flight_merge"
        fgo.create_chunk_subgraph(fid, "main", 1)
        fgo.create_chunk_subgraph(fid, "new", 10)
        fgo.add_relative_factor_to_chunk(fid, "new", 10, 11, dummy_relative_pose, np.eye(6))
        
        transform = Sim3Transform(translation=np.array([50.0, 0.0, 0.0]), rotation=np.eye(3), scale=2.0)
        
        assert fgo.merge_chunk_subgraphs(fid, "new", "main", transform) is True
        
        # Verify new chunk is absorbed
        state = fgo._get_or_create_flight_graph(fid)
        assert "new" not in state.chunk_subgraphs
        assert state.frame_to_chunk[11] == "main"
        
        traj = fgo.get_chunk_trajectory(fid, "main")
        assert 11 in traj