"""Tests for Factor Graph Optimizer (F10)."""

import numpy as np
import pytest

from gps_denied.core.graph import FactorGraphOptimizer
from gps_denied.schemas import GPSPoint
from gps_denied.schemas.graph import FactorGraphConfig
from gps_denied.schemas.metric import Sim3Transform
from gps_denied.schemas.vo import RelativePose


@pytest.fixture
def optimizer():
    config = FactorGraphConfig()
    return FactorGraphOptimizer(config)


def test_add_relative_factor(optimizer):
    flight_id = "test_f_1"

    # Needs initial node, using chunk creation wrapper to mock initial
    optimizer.create_chunk_subgraph(flight_id, "ch_1", start_frame_id=0)

    # Mock relative pose logic: in mock graph logic it's added via "add_relative_factor_to_chunk"
    # OR we need an initial root via add_absolute/relative. Our mock for `add_relative` is simple.

    rel1 = RelativePose(
        translation=np.array([1.0, 0.0, 0.0]),
        rotation=np.eye(3),
        covariance=np.eye(6),
        confidence=0.9,
        inlier_count=100,
        total_matches=120,
        tracking_good=True
    )

    # Test global logic requires frame exist, let's inject a zero frame
    optimizer._init_flight(flight_id)
    optimizer._flights_state[flight_id]["poses"][0] = optimizer._chunks_state["ch_1"]["poses"][0]

    res = optimizer.add_relative_factor(flight_id, 0, 1, rel1, np.eye(6))
    assert res is True

    traj = optimizer.get_trajectory(flight_id)
    assert 1 in traj
    assert np.allclose(traj[1].position, [1.0, 0.0, 0.0])

def test_add_absolute_factor(optimizer):
    flight_id = "test_f_2"
    optimizer._init_flight(flight_id)

    # Inject frame 0 and frame 1
    from datetime import datetime

    from gps_denied.schemas.graph import Pose
    optimizer._flights_state[flight_id]["poses"][0] = Pose(
        frame_id=0, position=np.zeros(3), orientation=np.eye(3), timestamp=datetime.now()
    )
    optimizer._flights_state[flight_id]["poses"][1] = Pose(
        frame_id=1, position=np.zeros(3), orientation=np.eye(3), timestamp=datetime.now()
    )

    # First GPS sets the ENU origin → position becomes [0,0,0]
    origin_gps = GPSPoint(lat=49.0, lon=30.0)
    optimizer.add_absolute_factor(flight_id, 0, origin_gps, np.eye(2), is_user_anchor=True)
    assert np.allclose(optimizer.get_trajectory(flight_id)[0].position, [0, 0, 0])

    # Second GPS at different coords → should produce non-zero ENU displacement
    gps2 = GPSPoint(lat=49.1, lon=30.1)
    res = optimizer.add_absolute_factor(flight_id, 1, gps2, np.eye(2), is_user_anchor=True)
    assert res is True

    traj = optimizer.get_trajectory(flight_id)
    assert traj[1].position[1] > 1000  # ~11 km north

def test_optimize_and_retrieve(optimizer):
    flight_id = "test_f_3"
    optimizer._init_flight(flight_id)

    res = optimizer.optimize(flight_id, 5)
    assert res.converged is True
    assert res.iterations_used == 5

    traj = optimizer.get_trajectory(flight_id)
    assert isinstance(traj, dict)

def test_chunk_merging(optimizer):
    flight_id = "test_f_4"
    c1 = "chunk1"
    c2 = "chunk2"

    assert optimizer.create_chunk_subgraph(flight_id, c1, 0) is True
    assert optimizer.create_chunk_subgraph(flight_id, c2, 10) is True

    rel = RelativePose(
        translation=np.array([5.0, 0, 0]),
        rotation=np.eye(3),
        covariance=np.eye(6),
        confidence=0.9,
        inlier_count=100,
        total_matches=120,
        tracking_good=True
    )
    optimizer.add_relative_factor_to_chunk(flight_id, c2, 10, 11, rel, np.eye(6))

    sim3 = Sim3Transform(translation=np.array([100.0, 0, 0]), rotation=np.eye(3), scale=1.0)
    res = optimizer.merge_chunk_subgraphs(flight_id, c2, c1, sim3)

    assert res is True

    c1_poses = optimizer.get_chunk_trajectory(flight_id, c1)
    # 10 and 11 should now be in c1, transformed!
    assert 10 in c1_poses
    assert 11 in c1_poses

    # 100 translation
    assert c1_poses[10].position[0] == 100.0
    assert c1_poses[11].position[0] == 105.0