# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:twoface]
#     language: python
#     name: conda-env-twoface-py
# ---

# +
import os
from os import path
from astropy.time import Time
from astropy.io import fits, ascii
import astropy.units as u
from astropy.table import Table, QTable
from astropy.constants import G
import astropy.coordinates as coord
from astropy.stats import LombScargle

import numpy as np
import matplotlib.pyplot as plt
# %matplotlib inline
import h5py
from sqlalchemy import func
from sqlalchemy.orm.exc import NoResultFound
from scipy.optimize import root
from scipy.stats import scoreatpercentile
import tqdm
import schwimmbad
from scipy.optimize import root
import pandas as pd
from scipy.signal import argrelmax

from thejoker import JokerSamples
from thejoker.sampler import JokerParams, TheJoker
from thejoker.plot import plot_rv_curves

from twoface.config import TWOFACE_CACHE_PATH
from twoface.data import APOGEERVData
from twoface.plot import plot_data_orbits, plot_two_panel, plot_phase_fold
from twoface.mass import get_m2_min, mf, stellar_radius
from twoface.samples_analysis import MAP_sample
from twoface.ext import FileLoader
from apogeebh.plot import plot_diag_panel
# -

asassn = ascii.read('/Users/adrian/Downloads/ASASSN-variables.csv')

# +
# from apogeebh.db import db_connect, AllStar, AllVisit, StarResult, Status
# Session, _ = db_connect('../cache/apogeebh.sqlite')
# cache_file = '../cache/apogeebh-joker.hdf5'

from twoface.db import db_connect, AllStar, AllVisit, StarResult, Status
Session, _ = db_connect('../../twoface/cache/apogee.sqlite')
cache_file = '../../twoface/cache/apogee-jitter.hdf5'
joker_pars = JokerParams(P_min=1*u.day, P_max=32768*u.day, jitter=150*u.m/u.s)

session = Session()
# -

print(session.query(AllStar).join(StarResult).filter(StarResult.status_id > 0).count())
print(session.query(AllStar).count())

q = session.query(AllStar.apogee_id, AllStar.ra, AllStar.dec).join(StarResult).filter(StarResult.status_id > 0)
stars = np.array(q.all(), 
                 dtype=[('apogee_id', 'U20'), ('ra', 'f8'), ('dec', 'f8')])

star_c = coord.SkyCoord(ra=stars['ra']*u.deg, 
                        dec=stars['dec']*u.deg)

asas_c = coord.SkyCoord(ra=asassn['Right Ascension']*u.deg, 
                        dec=asassn['Declination']*u.deg)

idx, sep, _ = asas_c.match_to_catalog_sky(star_c)
mask = (sep < 2*u.arcsec)

cross = asassn[mask].copy()
len(cross)

cross['APOGEE_ID'] = stars[idx[mask]]['apogee_id']

os.makedirs('../plots/asas-sn-variable', exist_ok=True)

with h5py.File(cache_file) as f:
    for row in cross:
        star = AllStar.get_apogee_id(session, row['APOGEE_ID'])
        samples = JokerSamples.from_hdf5(f[star.apogee_id])
        data = star.apogeervdata()
        
        fig = plot_diag_panel(data, samples, joker_params=joker_pars, M1=1.*u.Msun)
        fig.axes[3].axvline(row['Period'])
        fig.tight_layout()
        
        fig.savefig('../plots/asas-sn-variable/{0}.png'.format(star.apogee_id), dpi=250)
        plt.close(fig)

https://asas-sn.osu.edu/uploads/database/light_curve/data_file/60696/lc113200.dat

row

import time
import urllib

fmt_url = 'https://asas-sn.osu.edu/uploads/database/light_curve/data_file/{0}/lc{1}.dat'

# +
apw_P = []
apw_P_err = []
apw_e = []
apw_e_err = []
with h5py.File(cache_file) as f:
    for row in cross:
        star = AllStar.get_apogee_id(session, row['APOGEE_ID'])
        samples = JokerSamples.from_hdf5(f[star.apogee_id])
        data = star.apogeervdata()
        
        url = fmt_url.format(path.split(row['Url'])[-1], row['LC Number'])
        urllib.request.urlretrieve(url, '../asas-sn-variable/{0}-lc.dat'.format(row['APOGEE_ID']))
        
        lnP = np.log(samples['P'].to(u.day).value)
        i = np.abs(lnP - np.log(row['Period'])).argmin()
        idx = np.abs(lnP - lnP[i]) < 0.2
        apw_P.append(np.median(samples['P'][idx]))
        apw_P_err.append(np.std(samples['P'][idx]))
        apw_e.append(np.median(samples['e'][idx]))
        apw_e_err.append(np.std(samples['e'][idx]))
        jitter = np.median(samples[idx]['jitter'])
        
        rvdata = Table()
        rvdata['t'] = data.t.tcb.mjd
        rvdata['rv'] = data.rv.to(u.km/u.s).value
        rvdata['rv_err'] = np.sqrt(data.stddev**2 + jitter**2).to(u.km/u.s).value
        rvdata.write('../asas-sn-variable/{0}-rv.csv'.format(row['APOGEE_ID']))
        
        time.sleep(0.5)
        
apw_P = u.Quantity(apw_P)
apw_P_err = u.Quantity(apw_P_err)
apw_e = u.Quantity(apw_e)
apw_e_err = u.Quantity(apw_e_err)
# -

cross['APW_P'] = apw_P.value
cross['APW_P_err'] = apw_P_err.value
cross['APW_e'] = apw_e
cross['APW_e_err'] = apw_e_err

cross.write('../asas-sn-variable/catalog.csv')



StarcoderdataJupyterScriptsDedupFiltered