import pytest

import os
import pkg_resources
import numpy as np
from psoap import orbit_astrometry
from psoap import constants as C
from astropy.time import Time
import matplotlib.pyplot as plt
import matplotlib


# Create plots of all of the orbits
from astropy.io import ascii


# Create plots of all of the orbits
# If it doesn't already exist, create a "plots/basic/" directory
outdir = "tests/plots/HD10009/"
if not os.path.exists(outdir):
    print("Creating ", outdir)
    os.makedirs(outdir)

# Load the data sets for radial velocity and astrometry
astro_fname = pkg_resources.resource_filename("psoap", "data/HD10009/astro.txt")
astro_data = ascii.read(astro_fname, format="csv", fill_values=[(".", '0')])

# print(astro_data["date"].data)

# convert UT date to JD
astro_dates = Time(astro_data["date"].data, format="decimalyear")
astro_dates.format = 'jd'
astro_jds = astro_dates.value

rho_data = astro_data["rho"]
rho_err = astro_data["rho_err"]

theta_data = astro_data["PA"]
theta_err = astro_data["PA_err"]

def test_data():
    # Make a plot of the astrometric data on the sky
    fig, ax = plt.subplots(nrows=1)

    xs = rho_data * np.cos(theta_data * np.pi/180)
    ys = rho_data * np.sin(theta_data * np.pi/180)
    ax.plot(xs, ys, ".")
    ax.set_xlabel("North")
    ax.set_ylabel("East")
    ax.plot(0,0, "k*")
    ax.set_aspect("equal", "datalim")
    fig.savefig(outdir + "data_astro.png")

rv1_fname = pkg_resources.resource_filename("psoap", "data/HD10009/rv1.txt")
rv_data_1 = ascii.read(rv1_fname, format="csv")
rv2_fname = pkg_resources.resource_filename("psoap", "data/HD10009/rv2.txt")
rv_data_2 = ascii.read(rv2_fname, format="csv")

rv_jds_A = rv_data_1["date"] + 2400000
vAs_data = rv_data_1["rv"]
vAs_err = rv_data_1["err"]

rv_jds_B = rv_data_2["date"] + 2400000
vBs_data = rv_data_2["rv"]
vBs_err = rv_data_2["err"]

dpc = 37.03 # pc

# Orbital elements
a = 0.324 * dpc # [AU]
e = 0.798
i = 96.0 # [deg]
omega_2 = 251.6 # omega_1
omega = omega_2 + 180.0
Omega = 159.6
T0 = Time(1989.92, format="decimalyear")
T0.format = "jd"
T0 = T0.value # [Julian Date]

M_2 = 1.0 # [M_sun]
M_tot = 1.2 + M_2 # [M_sun]
gamma = 47.8 # [km/s]

P = np.sqrt(4 * np.pi**2 / (C.G * M_tot * C.M_sun) * (a * C.AU)**3) / (24 * 3600) # [day]
print(P/365)