import os
import numpy as np
import astropy.table as table
import astropy.units as u
import orbitize
import orbitize.read_input as read_input
import orbitize.kepler as kepler
import orbitize.system as system
import orbitize.basis as basis

# Skip this test on Windows since REBOUND doesn't work on Windows
import sys
import pytest

if sys.platform.startswith("win"):
    pytest.skip("Skipping REBOUND tests on Windows", allow_module_level=True)

try:
    import rebound
except ImportError:
    pytest.skip(
        "Skipping REBOUND tests because REBOUND is not installed",
        allow_module_level=True,
    )


def test_1planet():
    """
    Sanity check that things agree for 1 planet case
    """
    # generate a planet orbit
    sma = 1
    ecc = 0.1
    inc = np.radians(45)
    aop = np.radians(45)
    pan = np.radians(45)
    tau = 0.5
    plx = 1
    mtot = 1
    tau_ref_epoch = 0
    mjup = u.Mjup.to(u.Msun)
    mass_b = 12 * mjup

    epochs = np.linspace(0, 300, 100) + tau_ref_epoch  # nearly the full period, MJD

    ra_model, dec_model, vz_model = kepler.calc_orbit(
        epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch
    )

    # generate some fake measurements just to feed into system.py to test bookkeeping
    t = table.Table(
        [
            epochs,
            np.ones(epochs.shape, dtype=int),
            ra_model,
            np.zeros(ra_model.shape),
            dec_model,
            np.zeros(dec_model.shape),
        ],
        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
    )
    filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv")
    t.write(filename, overwrite=True)

    # create the orbitize system and generate model predictions using the ground truth
    astrom_dat = read_input.read_file(filename)

    sys = system.System(1, astrom_dat, mtot, plx, tau_ref_epoch=tau_ref_epoch)

    params = np.array([sma, ecc, inc, aop, pan, tau, plx, mtot])
    radec_orbitize, _ = sys.compute_model(params)
    ra_orb = radec_orbitize[:, 0]
    dec_orb = radec_orbitize[:, 1]

    # now project the orbit with rebound
    manom = basis.tau_to_manom(epochs[0], sma, mtot, tau, tau_ref_epoch)

    sim = rebound.Simulation()
    sim.units = ("yr", "AU", "Msun")

    # add star
    sim.add(m=mtot - mass_b)

    # add one planet
    sim.add(m=mass_b, a=sma, e=ecc, M=manom, omega=aop, Omega=pan + np.pi / 2, inc=inc)
    ps = sim.particles

    sim.move_to_com()

    # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner
    # planet's orbital period
    sim.integrator = "ias15"
    sim.dt = ps[1].P / 1000.0

    # integrate and measure star/planet separation
    ra_reb = []
    dec_reb = []

    for t in epochs:
        sim.integrate(t / 365.25)

        ra_reb.append(-(ps[1].x - ps[0].x))  # ra is negative x
        dec_reb.append(ps[1].y - ps[0].y)

    ra_reb = np.array(ra_reb)
    dec_reb = np.array(dec_reb)

    diff_ra = ra_reb - ra_orb / plx
    diff_dec = dec_reb - dec_orb / plx

    assert np.all(np.abs(diff_ra) < 1e-9)
    assert np.all(np.abs(diff_dec) < 1e-9)

    # clean up
    # os.system("rm {}".format(filename))


def test_2planet_massive():
    """
    Compare multiplanet to rebound for planets with mass.
    """
    # generate a planet orbit
    mjup = u.Mjup.to(u.Msun)
    mass_b = 12 * mjup
    mass_c = 9 * mjup

    params = np.array(
        [
            10,
            0.1,
            np.radians(45),
            np.radians(45),
            np.radians(45),
            0.5,
            3,
            0.1,
            np.radians(45),
            np.radians(190),
            np.radians(45),
            0.2,
            50,
            mass_b,
            mass_c,
            1.5 - mass_b - mass_c,
        ]
    )
    params_noc = np.array(
        [
            10,
            0.1,
            np.radians(45),
            np.radians(45),
            np.radians(45),
            0.5,
            3,
            0.1,
            np.radians(45),
            np.radians(190),
            np.radians(45),
            0.2,
            50,
            mass_b,
            0,
            1.5 - mass_b,
        ]
    )
    tau_ref_epoch = 0

    epochs = (
        np.linspace(0, 365.25 * 10, 100) + tau_ref_epoch
    )  # nearly the full period, MJD

    # doesn't matter that this is right, just needs to be the same size. below doesn't
    # include effect of c
    # just want to generate some measurements of plaent b to test compute model
    b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit(
        epochs,
        params[0],
        params[1],
        params[2],
        params[3],
        params[4],
        params[5],
        params[-2],
        params[-1],
        tau_ref_epoch=tau_ref_epoch,
    )

    # generate some fake measurements of planet b, just to feed into system.py
    # to test bookkeeping
    t = table.Table(
        [
            epochs,
            np.ones(epochs.shape, dtype=int),
            b_ra_model,
            np.zeros(b_ra_model.shape),
            b_dec_model,
            np.zeros(b_dec_model.shape),
        ],
        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
    )
    filename = os.path.join(orbitize.DATADIR, "rebound_2planet_outer.csv")
    t.write(filename, overwrite=True)

    # TEST THE OUTER PLANET

    # create the orbitize system and generate model predictions using the ground truth
    astrom_dat = read_input.read_file(filename)

    sys = system.System(
        2,
        astrom_dat,
        params[-1],
        params[-4],
        tau_ref_epoch=tau_ref_epoch,
        fit_secondary_mass=True,
    )

    # generate measurement
    radec_orbitize, _ = sys.compute_model(params)
    b_ra_orb = radec_orbitize[:, 0]
    b_dec_orb = radec_orbitize[:, 1]
    # debug, generate measurement without c having any mass
    radec_orb_noc, _ = sys.compute_model(params_noc)
    b_ra_orb_noc = radec_orb_noc[:, 0]
    _ = radec_orb_noc[:, 1]

    # check that planet c's perturbation is imprinted (nonzero))
    assert np.all(b_ra_orb_noc != b_ra_orb)

    # now project the orbit with rebound
    b_manom = basis.tau_to_manom(
        epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch
    )
    c_manom = basis.tau_to_manom(
        epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch
    )

    sim = rebound.Simulation()
    sim.units = ("yr", "AU", "Msun")

    # add star
    sim.add(m=params[-1])

    # add two planets
    sim.add(
        m=mass_c,
        a=params[0 + 6],
        e=params[1 + 6],
        M=c_manom,
        omega=params[3 + 6],
        Omega=params[4 + 6] + np.pi / 2,
        inc=params[2 + 6],
    )
    sim.add(
        m=mass_b,
        a=params[0],
        e=params[1],
        M=b_manom,
        omega=params[3],
        Omega=params[4] + np.pi / 2,
        inc=params[2],
    )
    ps = sim.particles

    sim.move_to_com()

    # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner
    # planet's orbital period
    sim.integrator = "ias15"
    sim.dt = ps[1].P / 1000.0

    # integrate and measure star/planet separation
    b_ra_reb = []
    b_dec_reb = []

    for t in epochs:
        sim.integrate(t / 365.25)

        b_ra_reb.append(-(ps[2].x - ps[0].x))  # ra is negative x
        b_dec_reb.append(ps[2].y - ps[0].y)

    b_ra_reb = np.array(b_ra_reb)
    b_dec_reb = np.array(b_dec_reb)

    diff_ra = b_ra_reb - b_ra_orb / params[6 * 2]
    diff_dec = b_dec_reb - b_dec_orb / params[6 * 2]

    # we placed the planets far apart to minimize secular interactions but there are
    # still some, so relax precision
    assert np.all(np.abs(diff_ra) / (params[0]) < 1e-3)
    assert np.all(np.abs(diff_dec) / (params[0]) < 1e-3)

    # clean up
    # os.system("rm {}".format(filename))

    # NOW TEST THE INNER PLANET

    # generate some fake measurements of planet c, just to feed into system.py to
    # test bookkeeping
    t = table.Table(
        [
            epochs,
            np.ones(epochs.shape, dtype=int) * 2,
            b_ra_model,
            np.zeros(b_ra_model.shape),
            b_dec_model,
            np.zeros(b_dec_model.shape),
        ],
        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
    )
    filename = os.path.join(orbitize.DATADIR, "rebound_2planet_inner.csv")
    t.write(filename, overwrite=True)

    # create the orbitize system and generate model predictions using the ground truth
    astrom_dat = read_input.read_file(filename)

    sys = system.System(
        2,
        astrom_dat,
        params[-1],
        params[-2],
        tau_ref_epoch=tau_ref_epoch,
        fit_secondary_mass=True,
    )

    # generate measurement
    radec_orbitize, _ = sys.compute_model(params)
    c_ra_orb = radec_orbitize[:, 0]
    c_dec_orb = radec_orbitize[:, 1]

    # start the REBOUND sim again
    sim = rebound.Simulation()
    sim.units = ("yr", "AU", "Msun")

    # add star
    sim.add(m=params[-1])

    # add two planets
    sim.add(
        m=mass_c,
        a=params[0 + 6],
        e=params[1 + 6],
        M=c_manom,
        omega=params[3 + 6],
        Omega=params[4 + 6] + np.pi / 2,
        inc=params[2 + 6],
    )
    sim.add(
        m=mass_b,
        a=params[0],
        e=params[1],
        M=b_manom,
        omega=params[3],
        Omega=params[4] + np.pi / 2,
        inc=params[2],
    )
    ps = sim.particles

    sim.move_to_com()

    # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner
    # planet's orbital period
    sim.integrator = "ias15"
    sim.dt = ps[1].P / 1000.0

    # integrate and measure star/planet separation
    c_ra_reb = []
    c_dec_reb = []

    for t in epochs:
        sim.integrate(t / 365.25)

        c_ra_reb.append(-(ps[1].x - ps[0].x))  # ra is negative x
        c_dec_reb.append(ps[1].y - ps[0].y)

    c_ra_reb = np.array(c_ra_reb)
    c_dec_reb = np.array(c_dec_reb)

    diff_ra = c_ra_reb - c_ra_orb / params[6 * 2]
    diff_dec = c_dec_reb - c_dec_orb / params[6 * 2]

    # planet is 3 times closer, so roughly 3 times larger secular errors.
    assert np.all(np.abs(diff_ra) / (params[0]) < 3e-3)
    assert np.all(np.abs(diff_dec) / (params[0]) < 3e-3)

    # clean up
    # os.system("rm {}".format(filename))


def test_2planet_massive_reverse_order():
    """
    Compare multiplanet to rebound for planets with mass. reverse planet indicies to
    check bookkeeping.
    """
    # generate a planet orbit
    mjup = u.Mjup.to(u.Msun)
    mass_b = 12 * mjup
    mass_c = 9 * mjup

    params = np.array(
        [
            3,
            0.1,
            np.radians(45),
            np.radians(190),
            np.radians(45),
            0.2,
            10,
            0.1,
            np.radians(45),
            np.radians(45),
            np.radians(45),
            0.5,
            50,
            mass_c,
            mass_b,
            1.5 - mass_b - mass_c,
        ]
    )
    tau_ref_epoch = 0

    epochs = (
        np.linspace(0, 365.25 * 10, 100) + tau_ref_epoch
    )  # nearly the full period, MJD

    # doesn't matter that this is right, just needs to be the same size. below doesn't
    # include effect of c
    # just want to generate some measurements of plaent b to test compute model
    b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit(
        epochs,
        params[6 + 0],
        params[6 + 1],
        params[6 + 2],
        params[6 + 3],
        params[6 + 4],
        params[6 + 5],
        params[-2],
        params[-1],
        tau_ref_epoch=tau_ref_epoch,
    )

    # generate some fake measurements of planet b, just to feed into system.py to test
    # bookkeeping
    t = table.Table(
        [
            epochs,
            np.ones(epochs.shape, dtype=int) * 2,
            b_ra_model,
            np.zeros(b_ra_model.shape),
            b_dec_model,
            np.zeros(b_dec_model.shape),
        ],
        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
    )
    filename = os.path.join(orbitize.DATADIR, "rebound_2planet_swapped_outer.csv")
    t.write(filename, overwrite=True)

    # TEST THE OUTER PLANET ####

    # create the orbitize system and generate model predictions using the ground truth
    astrom_dat = read_input.read_file(filename)

    sys = system.System(
        2,
        astrom_dat,
        params[-1],
        params[-4],
        tau_ref_epoch=tau_ref_epoch,
        fit_secondary_mass=True,
    )

    # generate measurement
    radec_orbitize, _ = sys.compute_model(params)
    b_ra_orb = radec_orbitize[:, 0]
    b_dec_orb = radec_orbitize[:, 1]

    # now project the orbit with rebound
    c_manom = basis.tau_to_manom(
        epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch
    )
    b_manom = basis.tau_to_manom(
        epochs[0],
        params[0 + 6],
        params[-1] + params[-2] + params[-3],
        params[5 + 6],
        tau_ref_epoch,
    )

    sim = rebound.Simulation()
    sim.units = ("yr", "AU", "Msun")

    # add star
    sim.add(m=params[-1])

    # add two planets
    sim.add(
        m=mass_c,
        a=params[0],
        e=params[1],
        M=c_manom,
        omega=params[3],
        Omega=params[4] + np.pi / 2,
        inc=params[2],
    )
    sim.add(
        m=mass_b,
        a=params[0 + 6],
        e=params[1 + 6],
        M=b_manom,
        omega=params[3 + 6],
        Omega=params[4 + 6] + np.pi / 2,
        inc=params[2 + 6],
    )
    ps = sim.particles

    sim.move_to_com()

    # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner
    # planet's orbital period
    sim.integrator = "ias15"
    sim.dt = ps[1].P / 1000.0

    # integrate and measure star/planet separation
    b_ra_reb = []
    b_dec_reb = []

    for t in epochs:
        sim.integrate(t / 365.25)

        b_ra_reb.append(-(ps[2].x - ps[0].x))  # ra is negative x
        b_dec_reb.append(ps[2].y - ps[0].y)

    b_ra_reb = np.array(b_ra_reb)
    b_dec_reb = np.array(b_dec_reb)

    diff_ra = b_ra_reb - b_ra_orb / params[6 * 2]
    diff_dec = b_dec_reb - b_dec_orb / params[6 * 2]

    # we placed the planets far apart to minimize secular interactions but there are
    # still some, so relax precision
    assert np.all(np.abs(diff_ra) / (params[6]) < 1e-3)
    assert np.all(np.abs(diff_dec) / (params[6]) < 1e-3)

    # clean up
    # os.system("rm {}".format(filename))


def test_2planet_nomass():
    """
    Compare multiplanet to rebound for planets with mass.
    """
    # generate a planet orbit
    mjup = u.Mjup.to(u.Msun)
    mass_b = 0 * mjup
    mass_c = 0 * mjup

    params = np.array(
        [
            10,
            0.1,
            np.radians(45),
            np.radians(45),
            np.radians(45),
            0.5,
            3,
            0.1,
            np.radians(45),
            np.radians(190),
            np.radians(45),
            0.2,
            1,
            mass_b,
            mass_c,
            1.5 - mass_b - mass_c,
        ]
    )
    tau_ref_epoch = 0

    epochs = (
        np.linspace(0, 365.25 * 4, 100) + tau_ref_epoch
    )  # nearly the full period, MJD

    # doesn't matter that this is right, just needs to be the same size. below
    # doesn't include effect of c
    # just want to generate some measurements of plaent b to test compute model
    b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit(
        epochs,
        params[0],
        params[1],
        params[2],
        params[3],
        params[4],
        params[5],
        params[-2],
        params[-1],
        tau_ref_epoch=tau_ref_epoch,
    )

    # generate some fake measurements of planet b, just to feed into system.py to
    # test bookkeeping
    t = table.Table(
        [
            epochs,
            np.ones(epochs.shape, dtype=int),
            b_ra_model,
            np.zeros(b_ra_model.shape),
            b_dec_model,
            np.zeros(b_dec_model.shape),
        ],
        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
    )
    filename = os.path.join(orbitize.DATADIR, "rebound_2planet.csv")
    t.write(filename, overwrite=True)

    # create the orbitize system and generate model predictions using the ground truth
    astrom_dat = read_input.read_file(filename)

    sys = system.System(
        2,
        astrom_dat,
        params[-1],
        params[-2],
        tau_ref_epoch=tau_ref_epoch,
        fit_secondary_mass=True,
    )

    # generate measurement
    radec_orbitize, _ = sys.compute_model(params)
    b_ra_orb = radec_orbitize[:, 0]
    b_dec_orb = radec_orbitize[:, 1]

    # now project the orbit with rebound
    b_manom = basis.tau_to_manom(
        epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch
    )
    c_manom = basis.tau_to_manom(
        epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch
    )

    sim = rebound.Simulation()
    sim.units = ("yr", "AU", "Msun")

    # add star
    sim.add(m=params[-1])

    # add two planets
    sim.add(
        m=mass_c,
        a=params[0 + 6],
        e=params[1 + 6],
        M=c_manom,
        omega=params[3 + 6],
        Omega=params[4 + 6] + np.pi / 2,
        inc=params[2 + 6],
    )
    sim.add(
        m=mass_b,
        a=params[0],
        e=params[1],
        M=b_manom,
        omega=params[3],
        Omega=params[4] + np.pi / 2,
        inc=params[2],
    )
    ps = sim.particles

    sim.move_to_com()

    # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner
    # planet's orbital period
    sim.integrator = "ias15"
    sim.dt = ps[1].P / 1000.0

    # integrate and measure star/planet separation
    b_ra_reb = []
    b_dec_reb = []

    for t in epochs:
        sim.integrate(t / 365.25)

        b_ra_reb.append(-(ps[2].x - ps[0].x))  # ra is negative x
        b_dec_reb.append(ps[2].y - ps[0].y)

    b_ra_reb = np.array(b_ra_reb)
    b_dec_reb = np.array(b_dec_reb)

    diff_ra = b_ra_reb - b_ra_orb / params[6 * 2]
    diff_dec = b_dec_reb - b_dec_orb / params[6 * 2]

    # should be as good as the one planet case
    assert np.all(np.abs(diff_ra) / (params[0] * params[6 * 2]) < 1e-9)
    assert np.all(np.abs(diff_dec) / (params[0] * params[6 * 2]) < 1e-9)

    # clean up
    # os.system("rm {}".format(filename))


if __name__ == "__main__":
    test_1planet()
    test_2planet_massive()
    test_2planet_massive_reverse_order()
    test_2planet_nomass()
