1# This file is part of the Astrometry.net suite.
2# Licensed under a 3-clause BSD style license - see LICENSE
3from __future__ import print_function
4import re
5import datetime
6from astrometry.util.starutil_numpy import *
7import numpy as np
8
9## FIXME -- requires at least one digit before the decimal place.
10floatre = r'[+-]?[\d]+(.[\d]*)?([eE][+-]?[\d]*)?'
11
12def floatvar(x):
13    return '(?P<%s>%s)' % (x, floatre)
14
15# For parsing orbital elements type output.
16elemrexstr = (r'^' + floatvar('jd')
17              + '.*' + ' EC= ' + floatvar('e')
18              + '.*' + ' IN= ' + floatvar('i')
19              + '.*' + ' OM= ' + floatvar('Omega')
20              + ' W = ' + floatvar('pomega')
21              + '.*' + 'MA= ' + floatvar('M')
22              + '.*' + 'A = ' + floatvar('a')
23              )
24elemrex = re.compile(elemrexstr, re.MULTILINE | re.DOTALL)
25
26# For finding "System GM" -- only in orbital elements type
27sysgmrexstr = '^System GM *: ' + floatvar('gm') + r' '#AU\^3/d\^2$'
28sysgmrex = re.compile(sysgmrexstr)
29
30# For parsing X,V type output
31xvrexstr = ('^' + floatvar('jd') + r' = (?P<ad>A\.D\. .*?)'
32            + r'^\s+' + floatvar('x0')
33            + ' +' + floatvar('x1')
34            + ' +' + floatvar('x2')
35            + '.*?'
36            + r'^\s+' + floatvar('v0')
37            + ' +' + floatvar('v1')
38            + ' +' + floatvar('v2')
39            )
40xvrex = re.compile(xvrexstr, re.MULTILINE | re.DOTALL)
41
42# For parsing "observer" type output (RA,Dec)
43# With QUANTITIES=1; angle format=DEG
44radecrexstr = (
45    '^ ' + '(?P<datetime>[\d]{4}-[\w]{3}-[\d]{2} [\d]{2}:[\d]{2}) .*?'
46               + floatvar('ra') + ' *?' + floatvar('dec'))
47radecrex = re.compile(radecrexstr, re.MULTILINE | re.DOTALL)
48
49'''
50For output like:
51
52#Date       UT      R.A. (J2000) Decl.    Delta     r     El.    Ph.   m1     Sky Motion
53#            h m s                                                            "/min    P.A.
54# geocentric
552007 04 01 000000   23.1983  -02.463     2.953   2.070   22.9  10.8  17.1    1.45    057.9
56'''
57#radecrexstr2 = (
58#    '^' + '(?P<datetime>[\d]{4} [\d]{2} [\d]{2} [\d]{2}:[\d]{2}) .*?'
59#               + floatvar('ra') + ' *?' + floatvar('dec'))
60#radecrex = re.compile(radecrexstr, re.MULTILINE | re.DOTALL)
61
62
63
64# Returns a list of lists of elements, plus a list of the JDs.
65#   ([jd1, jd2, ...], [   [a1, e1, i1, Omega1, pomega1, M1, GM1], ... ])
66# Where  i, Omega, pomega, M   are in radians
67def parse_orbital_elements(s, needSystemGM=True):
68    if needSystemGM:
69        m = sysgmrex.search(s)
70        if not m:
71            print('Did not find "System GM" entry')
72            return None
73        gm = float(m.group('gm'))
74    else:
75        gm = 1.
76    allE = []
77    alljd = []
78    for m in elemrex.finditer(s):
79        d = m.groupdict()
80        E = [np.deg2rad(x) if rad else x
81             for (x,rad) in zip([float(d[x]) for x in
82                                 ['a',  'e',   'i',  'Omega', 'pomega','M' ]],
83                                [False, False, True, True,    True,    True])]
84        E.append(gm)
85        allE.append(E)
86        alljd.append(float(d['jd']))
87    return alljd,allE
88
89# Returns (x, v, jd), each as numpy arrays.
90#     x in AU
91#     v in AU/day
92def parse_phase_space(s):
93    all_x = []
94    all_v = []
95    all_jd = []
96    for m in xvrex.finditer(s):
97        d = m.groupdict()
98        x = np.array([float(d[k]) for k in ['x0','x1','x2']])
99        v = np.array([float(d[k]) for k in ['v0','v1','v2']])
100        all_x.append(x)
101        all_v.append(v)
102        all_jd.append(float(d['jd']))
103    return (np.array(all_x), np.array(all_v), np.array(all_jd))
104
105# Returns (ra,dec,jd), each as numpy arrays.
106#   RA,Dec in J2000 deg
107def parse_radec(s):
108    all_ra = []
109    all_dec = []
110    all_jd = []
111    for m in radecrex.finditer(s):
112        d = m.groupdict()
113        all_ra.append(float(d['ra']))
114        all_dec.append(float(d['dec']))
115        t = datetime.datetime.strptime(d['datetime'], '%Y-%b-%d %H:%M')
116        # 2000-Jan-01 12:00
117        all_jd.append(datetojd(t))
118    return (np.array(all_ra), np.array(all_dec), np.array(all_jd))
119