from anise import (
    Almanac,
    Aberration,
    LocationDhallSet,
    LocationDhallSetEntry,
)
from anise.astro import Location, TerrainMask, FrameUid
from anise.analysis import Event, Condition
import anise.analysis as analysis
from anise.time import Duration, Epoch, TimeSeries, Unit
from anise.constants import Frames
from anise.astro import Frame
from pathlib import Path
from sys import platform


def test_analysis_gen_report():
    """
    Tests the generation of a scalar report with complex, nested expressions.
    """

    if platform == "win32":
        return True
    data_path = Path(__file__).parent.joinpath("..", "..", "data")
    almanac = (
        Almanac(str(data_path.joinpath("de440s.bsp")))
        .load(str(data_path.joinpath("pck08.pca")))
        .load(str(data_path.joinpath("lro.bsp")))
    )
    almanac.describe(spk=True)

    target_frame = analysis.FrameSpec.Loaded(Frames.EME2000)
    observer_frame = analysis.FrameSpec.Loaded(Frames.MOON_J2000)

    state = analysis.StateSpec(
        target_frame=target_frame,
        observer_frame=observer_frame,
        ab_corr=None,
    )

    # Build the orthogonal VNC frame of the Earth...Moon orbit
    vnc = analysis.OrthogonalFrame.XY(
        x=analysis.VectorExpr.Unit(analysis.VectorExpr.Velocity(state)),
        y=analysis.VectorExpr.Unit(analysis.VectorExpr.OrbitalMomentum(state)),
    )

    sun_state = analysis.StateSpec(
        target_frame=target_frame,
        observer_frame=analysis.FrameSpec.Loaded(Frames.SUN_J2000),
        ab_corr=Aberration("LT"),
    )

    # Project the Earth->Sun vector onto the VNC frame
    proj = analysis.VectorExpr.Project(
        v=analysis.VectorExpr.Negate(
            analysis.VectorExpr.Unit(analysis.VectorExpr.Radius(sun_state))
        ),
        frame=vnc,
        plane=analysis.Plane.XY,
    )

    # Rebuild the Local Solar Time calculation from fundamental expressions
    earth_sun = analysis.StateSpec(
        target_frame=analysis.FrameSpec.Loaded(Frames.SUN_J2000),
        observer_frame=observer_frame,
        ab_corr=Aberration("LT"),
    )
    u = analysis.VectorExpr.Unit(
        analysis.VectorExpr.CrossProduct(
            a=analysis.VectorExpr.Unit(analysis.VectorExpr.Radius(earth_sun)),
            b=analysis.VectorExpr.Unit(analysis.VectorExpr.OrbitalMomentum(state)),
        )
    )
    v = analysis.VectorExpr.CrossProduct(
        a=analysis.VectorExpr.Unit(analysis.VectorExpr.OrbitalMomentum(state)), b=u
    )
    r = analysis.VectorExpr.Radius(state)
    sin_theta = analysis.ScalarExpr.DotProduct(a=v, b=r)
    cos_theta = analysis.ScalarExpr.DotProduct(a=u, b=r)
    theta = analysis.ScalarExpr.Atan2(y=sin_theta, x=cos_theta)
    lst_prod = analysis.ScalarExpr.Mul(
        a=analysis.ScalarExpr.Mul(a=theta, b=analysis.ScalarExpr.Constant(1.0 / 180.0)),
        b=analysis.ScalarExpr.Constant(12.0),
    )
    lst_add = analysis.ScalarExpr.Add(a=lst_prod, b=analysis.ScalarExpr.Constant(6.0))
    lst = analysis.ScalarExpr.Modulo(v=lst_add, m=analysis.ScalarExpr.Constant(24.0))

    # Define all scalars to be calculated
    scalars = [
        analysis.ScalarExpr.Element(analysis.OrbitalElement.SemiMajorAxis),
        analysis.ScalarExpr.Element(analysis.OrbitalElement.Eccentricity),
        analysis.ScalarExpr.Element(analysis.OrbitalElement.Rmag),
        # IMPORTANT: Enums that do not have arguments MUST be called with ()
        # Otherwise PyO3 (the bindings) initializes them differently.
        # This causes a 'type' cannot be converted to ScalarExpr error!
        analysis.ScalarExpr.BetaAngle(),
        analysis.ScalarExpr.SolarEclipsePercentage(eclipsing_frame=Frames.VENUS_J2000),
        analysis.ScalarExpr.Norm(analysis.VectorExpr.Radius(state)),
        analysis.ScalarExpr.DotProduct(
            a=analysis.VectorExpr.EccentricityVector(state),
            b=analysis.VectorExpr.Fixed(x=1.0, y=0.0, z=0.0),
        ),
        analysis.ScalarExpr.VectorX(analysis.VectorExpr.EccentricityVector(state)),
        analysis.ScalarExpr.VectorY(analysis.VectorExpr.EccentricityVector(state)),
        analysis.ScalarExpr.VectorZ(analysis.VectorExpr.EccentricityVector(state)),
        analysis.ScalarExpr.Norm(
            analysis.VectorExpr.CrossProduct(
                a=analysis.VectorExpr.Radius(state),
                b=analysis.VectorExpr.Velocity(state),
            )
        ),
        analysis.ScalarExpr.Element(analysis.OrbitalElement.Hmag),
        analysis.ScalarExpr.AngleBetween(
            a=analysis.VectorExpr.Radius(state), b=analysis.VectorExpr.Velocity(state)
        ),
        analysis.ScalarExpr.AzimuthFromLocation(location_id=123, obstructing_body=None),
        analysis.ScalarExpr.ElevationFromLocation(
            location_id=123, obstructing_body=None
        ),
        analysis.ScalarExpr.RangeFromLocation(location_id=123, obstructing_body=None),
        analysis.ScalarExpr.RangeRateFromLocation(
            location_id=123, obstructing_body=None
        ),
        analysis.ScalarExpr.LocalTimeAscNode(),
        analysis.ScalarExpr.LocalTimeDescNode(),
        analysis.ScalarExpr.VectorX(proj),
        analysis.ScalarExpr.VectorY(proj),
        analysis.ScalarExpr.VectorZ(proj),
        analysis.ScalarExpr.LocalSolarTime(),
        lst,  # Our custom LST calculation
    ]

    # Test S-Expression serialization/deserialization
    proj_s_expr = scalars[-1].to_s_expr()
    assert len(proj_s_expr) > 0

    # Add aliases to some of the scalars for the report header
    scalars_with_aliases = [(s, None) for s in scalars]
    scalars_with_aliases[-5] = (scalars_with_aliases[-5][0], "proj VNC X")
    scalars_with_aliases[-4] = (scalars_with_aliases[-4][0], "proj VNC Y")
    scalars_with_aliases[-3] = (scalars_with_aliases[-3][0], "proj VNC Z")
    scalars_with_aliases[-1] = (scalars_with_aliases[-1][0], "LST (h)")

    # Build the final report object
    report = analysis.ReportScalars(scalars_with_aliases, state)

    # Test report serialization
    report_s_expr = report.to_s_expr()
    report_reloaded = analysis.ReportScalars.from_s_expr(report_s_expr)
    assert report_reloaded.to_s_expr() == report_s_expr

    # Generate the report data
    series = TimeSeries(
        Epoch("2025-01-01 00:00:00 UTC"),
        Epoch("2025-01-02 12:00:00 UTC"),
        Unit.Day * 0.5,
        inclusive=True,
    )
    data = almanac.report_scalars(report, series)

    assert len(data) == 4
    last_row = data["2025-01-02T12:00:00.000000"]
    assert len(last_row) == len(scalars_with_aliases)

    # Validate some computed values
    assert (
        last_row["Hmag (km)"]
        == last_row[
            "|Radius(Earth J2000 -> Moon J2000) ⨯ Velocity(Earth J2000 -> Moon J2000)|"
        ]
    )

    # Check that our manual LST is close to the built-in one
    time_diff = abs(last_row["LST (h)"] - last_row["local solar time (h)"])
    assert time_diff < (Unit.Second * 1).to_unit(Unit.Hour)

    # Check that the projections are valid (components of a unit vector)
    for col_name, value in last_row.items():
        if "proj" in col_name:
            assert abs(value) <= 1.0


def test_analysis_event():
    """
    Tests event finding for apoapsis, periapsis, eclipses, and other criteria.
    """
    if platform == "win32":
        return True
    data_path = Path(__file__).parent.joinpath("..", "..", "data")
    almanac = (
        Almanac(str(data_path.joinpath("de440s.bsp")))
        .load(str(data_path.joinpath("pck08.pca")))
        .load(str(data_path.joinpath("lro.bsp")))
    )

    lro_frame = Frame(-85, 1)  # LRO NAIF ID

    lro_state_spec = analysis.StateSpec(
        target_frame=analysis.FrameSpec.Loaded(lro_frame),
        observer_frame=analysis.FrameSpec.Loaded(Frames.MOON_J2000),
        ab_corr=None,
    )

    # Define several event criteria
    sun_has_set = analysis.Event(
        analysis.ScalarExpr.SunAngle(observer_id=-85),
        Condition.LessThan(90.0),
        Unit.Second * 0.5,
        ab_corr=None,
    )
    apolune = Event.apoapsis()
    perilune = Event.periapsis()
    eclipse = Event.total_eclipse(Frames.MOON_J2000)
    eclipse_boundary = Event(
        eclipse.scalar, Condition.Equals(99.0), eclipse.epoch_precision, None
    )

    # Test event serialization
    eclipse_s_expr = eclipse.to_s_expr()
    deserialized_eclipse = analysis.Event.from_s_expr(eclipse_s_expr)
    assert deserialized_eclipse.to_s_expr() == eclipse_s_expr

    # Get the time domain for LRO from the loaded ephemeris
    start_epoch, end_epoch = almanac.spk_domain(-85)
    start_orbit = almanac.transform(lro_frame, Frames.MOON_J2000, start_epoch, None)
    period = start_orbit.period()

    # Find apoapsis events
    apo_events = almanac.report_events(lro_state_spec, apolune, start_epoch, end_epoch)
    for event in apo_events:
        ta_deg = event.orbit.ta_deg()
        assert abs(ta_deg - 180.0) < 1e-2

    # Find periapsis events
    peri_events = almanac.report_events(
        lro_state_spec, perilune, start_epoch, end_epoch
    )
    for event in peri_events:
        ta_deg = event.orbit.ta_deg()
        assert ta_deg < 1e-2 or abs(ta_deg - 360.0) < 1e-2

    # Check that we found one apoapsis/periapsis per orbit
    dts_apo = [
        s.orbit.epoch.timedelta(f.orbit.epoch)
        for f, s in zip(apo_events[:-1], apo_events[1:])
    ]
    assert all((dt - period).abs() < Unit.Minute * 5 for dt in dts_apo)

    dts_peri = [
        s.orbit.epoch.timedelta(f.orbit.epoch)
        for f, s in zip(peri_events[:-1], peri_events[1:])
    ]
    assert all((dt - period).abs() < Unit.Minute * 5 for dt in dts_peri)

    tick = Epoch.system_now()
    # Find sunset/sunrise arcs
    sunset_arcs = almanac.report_event_arcs(
        lro_state_spec, sun_has_set, start_epoch, end_epoch
    )
    print(
        f"{len(sunset_arcs)} sunset arcs found in {Epoch.system_now().timedelta(tick)}"
    )
    assert sunset_arcs[1].rise.edge == analysis.EventEdge.Falling
    assert sunset_arcs[1].fall.edge == analysis.EventEdge.Rising
    assert len(sunset_arcs) == 452

    # Find eclipse arcs in a short time span
    tick = Epoch.system_now()
    eclipse_arcs = almanac.report_event_arcs(
        lro_state_spec, eclipse, start_epoch, start_epoch + period * 3
    )
    assert len(eclipse_arcs) == 3
    print("Eclipse arcs found in ", Epoch.system_now().timedelta(tick))

    # Validate the eclipse periods
    for arc in eclipse_arcs:
        assert Unit.Minute * 24 < arc.duration() < Unit.Minute * 40

        # Check points in and around the arc to confirm the event state
        series = TimeSeries(
            arc.start_epoch() - eclipse.epoch_precision,
            arc.end_epoch() + eclipse.epoch_precision,
            Unit.Minute * 0.5,
            inclusive=True,
        )
        for epoch in series:
            orbit = lro_state_spec.evaluate(epoch, almanac)
            # IMPORTANT: Eclipse is not defined as a Condition.Equal
            # so it cannot be used as-is for an `eval` check.
            # Instead we use the eclipse boundary. That's how's the internals work.
            eclipse_val = eclipse_boundary.eval(orbit, almanac)
            is_in_eclipse = eclipse_val >= 0.0

            if arc.start_epoch() <= epoch < arc.end_epoch():
                assert is_in_eclipse, (
                    f"Epoch {epoch} should be in eclipse: {eclipse_val}: {arc}"
                )
            else:
                # Outside the arc, it should not be in eclipse, or it's a falling value
                assert not is_in_eclipse or eclipse_val < 0.0


def test_location_accesses():
    """
    Demonstrate building a Location Dhall file, loading it into the Almanac, reporting the access times.
    """
    if platform == "win32":
        return True

    mask = [TerrainMask(0.0, 5.0), TerrainMask(35.0, 10.0), TerrainMask(270.0, 3.0)]
    dss65 = Location(
        40.427_222,
        4.250_556,
        0.834_939,
        FrameUid(399, 399),
        mask,
        terrain_mask_ignored=True,
    )

    # To build a location data set kernel, we must first build a location dhall set entry
    entry = LocationDhallSetEntry(dss65, id=1, alias="My Alias")
    # Then we append it to a LocationDhallSet
    dhallset = LocationDhallSet([entry])
    assert "data" in dir(dhallset), "missing getting on dhall set"
    # Now, we can build the kernel
    dataset = dhallset.to_dataset()
    # Save it as a Location Kernel ANISE (LKA) file, overwritting it if it exists
    dataset.save_as("pytest_loc_kernel_report.lka", True)
    data_path = Path(__file__).parent.joinpath("..", "..", "data")
    almanac = (
        Almanac(str(data_path.joinpath("de440s.bsp")))
        .load(str(data_path.joinpath("pck08.pca")))
        .load(str(data_path.joinpath("lro.bsp")))
        .load("pytest_loc_kernel_report.lka")
    )
    # Only print the loaded location info
    almanac.describe(locations=True)

    # Build the horizon event
    horizon = Event.visible_from_location_id(1)

    lro_frame = Frame(-85, 1)  # LRO NAIF ID
    start_epoch, _end_epoch = almanac.spk_domain(-85)

    lro_state_spec = analysis.StateSpec(
        target_frame=analysis.FrameSpec.Loaded(lro_frame),
        observer_frame=analysis.FrameSpec.Loaded(Frames.MOON_J2000),
        ab_corr=None,
    )

    tick = Epoch.system_now()
    comm_arcs = almanac.report_event_arcs(
        lro_state_spec, horizon, start_epoch, start_epoch + Unit.Day * 3
    )
    print(f"{len(comm_arcs)} Comm arcs found in ", Epoch.system_now().timedelta(tick))

    series = TimeSeries(
        start_epoch,
        start_epoch + Unit.Day * 3,
        Unit.Minute * 0.5,
        inclusive=True,
    )

    for arc in comm_arcs:
        # Check points in and around the arc to confirm the event state
        series = TimeSeries(
            arc.start_epoch() - horizon.epoch_precision,
            arc.end_epoch() + horizon.epoch_precision,
            Unit.Minute * 0.5,
            inclusive=True,
        )
        for epoch in series:
            orbit = lro_state_spec.evaluate(epoch, almanac)
            aer = almanac.azimuth_elevation_range_sez_from_location_name(
                orbit, "My Alias"
            )
            is_visible = aer.elevation_deg > 0.0

            if arc.start_epoch() <= epoch < arc.end_epoch():
                assert is_visible, f"Epoch {epoch} should be visible"
            else:
                assert not is_visible

    # Check visibility arc function
    visibility_arcs = almanac.report_visibility_arcs(
        lro_state_spec,
        1,
        start_epoch,
        start_epoch + Unit.Day * 3,
        Unit.Minute * 10,
        None,
    )
    first_pass = visibility_arcs[0]
    assert first_pass.duration().round(Unit.Second * 1) == Duration("8 h 56 min 55 s")
    # Visibility Arc contains the location reference info for pretty printing
    print(first_pass.location_ref)
    # And the location data itself
    print(first_pass.location, first_pass.location.height_km)
    # AER data is a list of AzElRange objects
    assert len(first_pass.aer_data) == 54
    assert abs(first_pass.aer_data[0].elevation_deg - 1e-11) < 1e-11, (
        "elevation should be nearly zero"
    )


# TODO: Add FOV Margin check with LRO data
# TODO: Add minimum RIC diff analysis
