import numpy as np
import pyarrow.compute as pc
import pytest
from adam_assist import ASSISTPropagator

from ...coordinates import CartesianCoordinates, CoordinateCovariances
from ...coordinates.origin import Origin
from ...coordinates.residuals import Residuals, calculate_reduced_chi2
from ...observers import Observers
from ...orbit_determination.evaluate import OrbitDeterminationObservations
from ...orbits import Orbits
from ...time import Timestamp
from ..least_squares import LeastSquares


@pytest.fixture
def real_data():
    # Subset of actual observers for "2023 CL3"
    obstimes = Timestamp.from_kwargs(
        days=[59989, 59989, 59989, 59989, 59990, 59990],
        nanos=[
            78339000000000,
            80234800000000,
            82933632000000,
            85858272000000,
            68911862000000,
            70200173000000,
        ],
        scale="utc",
    )
    obscodes = ["K63", "K63", "204", "204", "I93", "I93"]
    observers = Observers.from_codes(codes=obscodes, times=obstimes)

    # MPC orbit for the same object
    orbit = Orbits.from_kwargs(
        orbit_id=["test1"],
        object_id=["2023 CL3"],
        coordinates=CartesianCoordinates.from_kwargs(
            x=[-0.875655],
            y=[1.066652],
            z=[0.085413],
            vx=[-0.007748],
            vy=[-0.014017],
            vz=[0.00026],
            time=Timestamp.from_kwargs(days=[60800], nanos=[0]),
            origin=Origin.from_kwargs(code=["SUN"]),
            frame="ecliptic",
            covariance=CoordinateCovariances.from_sigmas(
                np.array(
                    [
                        [
                            1.63293252e-04,
                            2.95269724e-04,
                            5.70764235e-06,
                            2.06775246e-06,
                            2.51608186e-06,
                            2.01679472e-07,
                        ]
                    ]
                )
            ),
        ),
    )

    return observers, orbit


@pytest.mark.parametrize("use_central_difference", [True, False])
def test_least_squares_adjusts(real_data, use_central_difference) -> None:
    observers, orbit = real_data

    # Make observations that fit the orbit perfectly
    propagator = ASSISTPropagator()
    perfect_ephem = propagator.generate_ephemeris(
        orbit, observers, covariance=True, chunk_size=1, max_processes=1
    )
    perfect_observations = OrbitDeterminationObservations.from_kwargs(
        id=[f"obs{i}" for i in range(len(observers))],
        coordinates=perfect_ephem.coordinates,
        observers=observers,
    )

    fitter = LeastSquares(use_central_difference)

    # If we feed perfect observations to perfect orbit, there should be no change
    debug_info = {}
    unchanged_orbit = fitter.least_squares(
        orbit, perfect_observations, propagator, debug_info=debug_info
    )
    assert len(debug_info["iterations"]) == 1
    assert debug_info["exit_message"].startswith("RMS is zero")
    assert unchanged_orbit == orbit

    # Change the orbit a bit so that it still works as an initial guess, but doesn't fit observations perfectly.
    # This will take several iterations, with improvements
    original_cometary = orbit.coordinates.to_cometary()
    updated_cometary = original_cometary.set_column(
        "i", pc.multiply(original_cometary.i, 0.99)
    )
    guess_orbit = orbit.set_column("coordinates", updated_cometary.to_cartesian())
    debug_info = {}
    improved_orbit = fitter.least_squares(
        guess_orbit, perfect_observations, propagator, debug_info=debug_info
    )
    assert improved_orbit is not None

    # We should make several iterations. Even if it doesn't converge in the sense of getting tiny delta RMS
    # (it may overshoot on RMS), it should make several iterations and improve the fit
    assert len(debug_info["iterations"]) > 2
    assert debug_info["iterations"][-1]["rms"] < 1e-5
    guessed_ephem = propagator.generate_ephemeris(
        guess_orbit, observers, chunk_size=1, max_processes=1
    )
    guessed_rchi2 = calculate_reduced_chi2(
        Residuals.calculate(
            perfect_observations.coordinates, guessed_ephem.coordinates
        ),
        6,
    )
    improved_ephem = propagator.generate_ephemeris(
        improved_orbit, observers, chunk_size=1, max_processes=1
    )
    improved_rchi2 = calculate_reduced_chi2(
        Residuals.calculate(
            perfect_observations.coordinates, improved_ephem.coordinates
        ),
        6,
    )
    assert improved_rchi2 < guessed_rchi2

    # Check things don't blow up if we don't ask for debug info
    no_debug_orbit = fitter.least_squares(guess_orbit, perfect_observations, propagator)
    assert no_debug_orbit is not None


@pytest.mark.parametrize("use_central_difference", [True, False])
def test_least_squares_zero_base(real_data, use_central_difference) -> None:
    observers, orbit = real_data

    # Set the last component of the orbit, which is already the smallest value, to 0.
    # This verifies that we don't end up with division by zero.
    zeroed_value = orbit.coordinates.set_column("vz", [0.0])
    zeroed_orbit = orbit.set_column("coordinates", zeroed_value)

    propagator = ASSISTPropagator()
    perfect_ephem = propagator.generate_ephemeris(
        orbit, observers, covariance=True, chunk_size=1, max_processes=1
    )
    perfect_observations = OrbitDeterminationObservations.from_kwargs(
        id=[f"obs{i}" for i in range(len(observers))],
        coordinates=perfect_ephem.coordinates,
        observers=observers,
    )

    fitter = LeastSquares(use_central_difference)
    debug_info = {}
    improved_orbit = fitter.least_squares(
        zeroed_orbit, perfect_observations, propagator, debug_info=debug_info
    )
    assert improved_orbit is not None

    assert len(debug_info["iterations"]) > 2
    guessed_ephem = propagator.generate_ephemeris(
        zeroed_orbit, observers, chunk_size=1, max_processes=1
    )
    guessed_rchi2 = calculate_reduced_chi2(
        Residuals.calculate(
            perfect_observations.coordinates, guessed_ephem.coordinates
        ),
        6,
    )
    improved_ephem = propagator.generate_ephemeris(
        improved_orbit, observers, chunk_size=1, max_processes=1
    )
    improved_rchi2 = calculate_reduced_chi2(
        Residuals.calculate(
            perfect_observations.coordinates, improved_ephem.coordinates
        ),
        6,
    )
    assert improved_rchi2 < guessed_rchi2
