from slalib import *
import numpy as np
import string as st
import time as tm
from control20m import *
ant = Control20m()
ant.take_control()
ant.sync_unix_time()

def take_antenna_control ( ) :
    """
    Put the 20-meter antenna under Python program control.
    """
    ant.take_control()
    
def release_antenna ( ) :
    """
    Release the 20-meter antenna from Python program control.
    """
    ant.release_control()

def hadec_to_azel ( ha_hrs, dec_deg ) :
    """
    hadec_to_azel ( ha_hrs, dec_deg )
    Convert Hour Angle - Declination coordinates to Azimuth and Elevation
    without making any pointing or refraction corrections.  HA is in decimal
    hours and Dec, Az, and El are in decimal degrees.
    The returned value is a dictionary {'az': , 'el': }
    """
    ha_rad = np.pi * ha_hrs / 12.0
    dec_rad = np.pi * dec_deg / 180.0
    vec = sla_dcs2c(ha_rad, dec_rad)
    desig, name, of_long, of_lat, elev = sla_obs(0, 'GBVA140')
    #print of_lat * 180 / np.pi
    rmat = sla_deuler('yxx', np.pi / 2.0 - of_lat, 0.0, 0.0)
    new_vec = sla_dmxv(rmat, vec)
    az, el = sla_dcc2s(new_vec)
    return {'az':180.0 * (1.0 + az / np.pi), 'el':180.0 * el / np.pi}

#print hadec_to_azel(0.0, 0.0)
#print hadec_to_azel(-6.0, 0.0)

def str2frac ( dt_str ) :
    """
    str2frac(dt_str)
    Convert a colon separated hexagesimal string to decimal.
    'xx:yy:zz.zz' -> xx.xxxx, xx can be signed
    """
    parsed = st.split(dt_str, ':')
    sign = 1.0
    if parsed[0][0] == '-' :
        sign = -1.0
    return sign * (abs(st.atof(parsed[0])) + (st.atof(parsed[1]) +
                                 st.atof(parsed[2]) / 60.0) / 60.0)

#print str2frac('11:22:33')
#print str2frac('-11:22:33')

def j2000_to_epoch ( ra_str, dec_str, mjd ) :
    """
    j2000_to_epoch('HH:MM:SS.ss', '-DD:MM:SS.s', mjd)
    Convert hexagesimal J2000 R.A.and Dec. string values to R.A.and Dec
    for the given Modified Julian Date using J2000 precession and nutation.
    The returned value is a dictionary, {'ra_hrs': , 'dec_deg': }
    """
    rmat = sla_prenut(2000, mjd)
    ra = np.pi * str2frac(ra_str) / 12.0
    dec = np.pi * str2frac(dec_str) / 180.0
    vec = sla_dcs2c(ra, dec)
    new_vec = sla_dmxv(rmat, vec)
    pra, pdec = sla_dcc2s(new_vec)
    if pra < 0.0 :
        pra += 2.0 * np.pi
    return {'ra_hrs':12.0 * pra / np.pi, 'dec_deg':180.0 * pdec / np.pi}

#print 'Precessed:', j2000_to_epoch('1:00:00', '1:00:00', 54630)

def get_current_mjd ( ) :
    """
    get_current_mjd()
    Reads the system clock and converts to fractional Modified Julian Date
    with fractional second precision to the accuracy of the system clock.
    """
    frac_sec1 = 1.0; frac_sec2 = 0.0
    while (frac_sec1 > frac_sec2) :
        ftime1 = tm.time()
        gmt = tm.gmtime()
        ftime2 = tm.time()
        frac_sec1 = ftime1 - int(ftime1)
        frac_sec2 = ftime2 - int(ftime2)
    #print gmt
    yr = gmt[0]
    mjd = 54465 + gmt[7] + (gmt[3] + (gmt[4] + gmt[5] / 60.0) / \
                             60.0) / 24.0 + (frac_sec1 + frac_sec2) / 172800.0
    for i in range(2009, 2052) :
        if i > yr :
            break
        if (i - 1) % 4 == 0 :
            mjd += 366
        else :
            mjd += 365
    return mjd

#print get_current_mjd()

_RH = 0.8; print 'RH:', _RH * 100, '%'
_REFA, _REFB = sla_refco(881.0, 290.0, 1000.0, _RH,
                         1000.0, 0.67086, 0.0065, 1e-10)

def get_refracted_el ( el ) :
    zu = np.pi * (90.0 - el) / 180.0
    zr = sla_refz(zu, _REFA, _REFB)
    return 90.0 - 180.0 * zr / np.pi

def get_corrected_azel ( az, el, with_refraction=True ) :
    if with_refraction :
        rel = get_refracted_el(el)
    else :
        rel = el
    encoder_az = az + ant.get_az_pointing_offset(az, rel)
    encoder_el = rel + ant.get_el_pointing_offset(az, rel)
    return (encoder_az, encoder_el)

def get_object_posn ( object ) :
    if object == 'CasA' :
        object = '23:23:26.7,+58:49:03'
    elif object == 'CygA' :
        object = '19:59:28.3,+40:44:02'
    pos = st.split(object, ',')
    if len(pos) != 2 :
        print 'make_track(): Object position error', object
        return False
    return j2000_to_epoch(pos[0], pos[1], get_current_mjd())

#print get_object_posn('11:22:33,-5:33:77')

def map_azel_to_azel ( map_az_deg, map_el_deg, src_az_deg, src_el_deg ) :
    """
    map_azel_to_azel ( map_az_deg, map_el_deg, src_az_deg, src_el_deg ) :
    map_az, map_el -> Az, El
    positive map_az drives the telescope to smaller azimuth
    positive map_el drives the telescope to lower elevation
    all arguments in degrees
    """
    map_az_rad = np.pi * map_az_deg / 180.0
    map_el_rad = np.pi * map_el_deg / 180.0
    src_az_rad = np.pi * src_az_deg / 180.0
    src_el_rad = np.pi * src_el_deg / 180.0
    vec = sla_dcs2c(map_az_rad, map_el_rad)
    #print vec
    elg = src_el_rad - map_el_rad
    azg = src_az_rad - map_az_rad / np.cos(elg)
    err_tol = 2.0 * np.pi / (180.0 * 3600.0)
    num_iter = 10
    for i in range(num_iter) :
        rmat = sla_deuler('yzx', elg, -azg, 0.0)
        new_vec = sla_dmxv(rmat, vec)
        #print new_vec
        az, el = sla_dcc2s(new_vec)
        #print i, '%10.5f %10.5f %10.5f %10.5f' % \
        #      (azg * 180 / np.pi, elg * 180 / np.pi,
        #       az * 180 / np.pi, el * 180 / np.pi)
        az_diff = src_az_rad - az
        if az_diff > np.pi :
            az_diff -= 2.0 * np.pi
        el_diff = src_el_rad - el
        #print az_diff, el_diff
        if abs(az_diff) < err_tol and abs(el_diff) < err_tol :
            break
        azg += az_diff
        elg += el_diff
        if i >= num_iter - 1 :
            print 'map_azel_to_azel(%4.2f, %4.2f, %4.2f, %4.2f) did not converge' % (map_az_deg, map_el_deg, src_az_deg, src_el_deg)
    if azg < 0.0 :
        azg += 2.0 * np.pi
    if azg > 2.0 * np.pi :
        azg -= 2.0 * np.pi
    return {'tel_az':180.0 * azg / np.pi, 'tel_el':180.0 * elg / np.pi}

#for az in np.arange(0.0, 361, 5.0) :
#    for el in np.arange(0.0, 79.0, 5.0) :
#        map_azel_to_azel(10.0, -10.0, az, el)

#print map_azel_to_azel(5.0, 5.0, 0.0, 5.0), '<- result'

DUT1 = -0.43  # UT1 - UTC in seconds  UT1 = UTC + DUT1

def get_last_str ( mjd, utc_string ) :
    utc_parsed = st.split(utc_string, ':')
    utc_frac = str2frac(utc_string) / 24.0
    mjd_frac = mjd + utc_frac
    #desig, name, gbt_long, lat, elev = sla_obs(0, 'GBT')
    desig, name, of_long, lat, elev = sla_obs(0, 'GBVA140')
    lon_frac =  of_long / (2.0 * np.pi)
    gast_rad = sla_gmst(mjd_frac + DUT1 / 86400.0) + sla_eqeqx(mjd_frac)
    diff140to20m = 2.26  #seconds
    return (gast_rad - of_long) / (2.0 * np.pi) + diff140to20m / 86400.0

#print (get_last_str(54630.0, '0:0:0') - 0.50511288448786651)  * 86400.0
#print (get_last_str(54630.0, '12:12:12') - 0.014977253263429011)  * 86400.0

def get_last ( mjd_frac ) :
    #desig, name, gbt_long, lat, elev = sla_obs(0, 'GBT')
    desig, name, of_long, lat, elev = sla_obs(0, 'GBVA140')
    lon_frac =  of_long / (2.0 * np.pi)
    gast_rad = sla_gmst(mjd_frac + DUT1 / 86400.0) + sla_eqeqx(mjd_frac)
    diff140to20m = 2.26  #seconds
    return (gast_rad - of_long) / (2.0 * np.pi) + diff140to20m / 86400.0

#print (get_last(54630.5084722222222) - 0.014977253263429011)  * 86400.0
#print get_last(get_current_mjd())

def make_track ( object, start_mjd, stop_mjd, cross_el_offset, el_offset ) :
    """
    make_track ( object, start_mjd, stop_mjd, cross_el_offset, el_offset )
    object = 'CasA', CygA', or 'hh:mm:ss.s,sdd:mm:ss' J2000 coordinates
    start_mjd, stop_mjd include UTC, e.g., 54630.3456
    positive cross_el_offset drives the telescope to smaller azimuth
    positive el_offset drives the telescope to lower elevation
    both arguments in degrees
    """
    posn = get_object_posn(object)
    if posn == False :
        return
    start_last = get_last(start_mjd)
    start_ha = 24.0 * start_last - posn['ra_hrs']
    if start_ha < -12.0 :
        start_ha += 24.0
    if start_ha > 12.0 :
        start_ha -= 24.0
    stop_last = get_last(stop_mjd)
    stop_ha = 24.0 * stop_last - posn['ra_hrs']
    if stop_ha < -12.0 :
        stop_ha += 24.0
    if stop_ha > 12.0 :
        stop_ha -= 24.0
    if stop_ha < start_ha :
        stop_ha += 24.0
    duration = (stop_ha - start_ha) * 3600.0
    min_steps = 1
    while duration / min_steps > 20.0 :
        min_steps *= 2
    
    if stop_last < start_last :
        stop_last += 1.0
    num_iter = 3
    max_num_pts = min_steps * (2**num_iter)
    pt_mjd = np.zeros(max_num_pts + 1, dtype=np.float64)
    pt_az = np.zeros(max_num_pts + 1, dtype=np.float64)
    pt_az_vel = np.zeros(max_num_pts, dtype=np.float64)
    pt_el = np.zeros(max_num_pts + 1, dtype=np.float64)
    pt_el_vel = np.zeros(max_num_pts, dtype=np.float64)
    posn_tol = 2.0 / 3600.0  # degrees
    for iter in range(num_iter) :
        num_pts = min_steps * (2**(iter + 1))
        ha_step = (stop_ha - start_ha) / num_pts
        for pt in range(num_pts + 1) :
            azel = hadec_to_azel(start_ha + ha_step * pt, posn['dec_deg'])
            src_az = azel['az']
            src_el = get_refracted_el(azel['el'])
            tel_azel = map_azel_to_azel(cross_el_offset, el_offset,
                                        src_az, src_el)
            pt_mjd[pt] = start_mjd + pt * ha_step * 0.997269 / 24.0
            tel_az, tel_el = get_corrected_azel(tel_azel['tel_az'],
                                                tel_azel['tel_el'], False)
            pt_az[pt] = tel_az
            pt_el[pt] = tel_el
        tol_flag = True
        for pt in range(1, num_pts + 1, 2) :
            az_err = pt_az[pt] - (pt_az[pt - 1] + pt_az[pt + 1]) / 2.0
            az_err *= np.cos(pt_el[pt] * np.pi / 180.0)
            el_err = pt_el[pt] - (pt_el[pt - 1] + pt_el[pt + 1]) / 2.0
            err = np.sqrt(az_err**2 + el_err**2)
            if err > posn_tol :
                tol_flag = False
        if tol_flag :
            break
    if not tol_flag :
        print 'WARNING: make_track() did not converge!'
    num_pts_out = num_pts / 2
    step_sec = 2.0 * ha_step * 0.997269 * 3600.0  # seconds
    for pt in range(num_pts_out) :
        pt_mjd[pt] = pt_mjd[pt * 2]
        pt_az[pt] = pt_az[pt * 2]
        pt_el[pt] = pt_el[pt * 2]
        pt_az_vel[pt] = (pt_az[pt * 2 + 2] - pt_az[pt * 2]) / step_sec
        pt_el_vel[pt] = (pt_el[pt * 2 + 2] - pt_el[pt * 2]) / step_sec
    return {'mjd':pt_mjd[:num_pts_out],
            'az':pt_az[:num_pts_out], 'az_vel':pt_az_vel[:num_pts_out],
            'el':pt_el[:num_pts_out], 'el_vel':pt_el_vel[:num_pts_out]}

#print make_track('19:22:33,65:33:77', 54642.7803896, 54642.7803896 + 30.0/86400, 5.0, 5.0)

def wait_for_slew ( cmd_az, cmd_el, tol_deg=0.02 ) :
    old_tel_stat = False
    old_mjd = 0.0
    tm.sleep(1.0)
    while True :
        new_mjd = get_current_mjd()
        #print '1111'
        tel_stat = ant.get_time_tagged_status()
        #print tel_stat['az'], cmd_az, tel_stat['el'], cmd_el
        if (abs(tel_stat['az'] - cmd_az) < tol_deg and
            abs(tel_stat['el'] - cmd_el) < tol_deg) :
            break
        if old_tel_stat :
            delta_t = (new_mjd - old_mjd) * 86400.0
            az_vel = (tel_stat['az'] - old_tel_stat['az']) / delta_t
            el_vel = (tel_stat['el'] - old_tel_stat['el']) / delta_t
            #print 'az_vel: %5.3f, el_vel: %5.3f deg/sec' % (az_vel, el_vel)
        old_mjd = new_mjd
        old_tel_stat = tel_stat
        tm.sleep(1.0)

def execute_track ( trk, stop_mjd ) :
    #trk = {'mjd':[], 'az':[], 'az_vel':[], 'el':[], 'el_vel':[]}
    for i in range(len(trk['mjd'])) :
        while (trk['mjd'][i] - get_current_mjd()) * 86400.0 > 3.0 :
            tm.sleep(0.5)
            log_ant_posn()
        utc_sec = (trk['mjd'][i] - int(trk['mjd'][i])) * 86400.0
        #time_str = mjd_to_time_str(trk['mjd'][i])
        status = ant.set_azel(trk['az'][i], trk['el'][i], tim=utc_sec,
                              azvel=trk['az_vel'][i], elvel=trk['el_vel'][i])
        if not status :
            return status
    ctr = 0
    while stop_mjd > get_current_mjd() :
        tm.sleep(0.5)
        ctr += 1
        if ctr % 2 :
            log_ant_posn()
    return True

def track ( object, duration, cross_el_offset=0.0, el_offset=0.0 ) :
    tel_stat = ant.get_time_tagged_status()
    cur_az = tel_stat['az']
    cur_el = tel_stat['el']
    mjd = get_current_mjd()  # includes day fraction
    posn = get_object_posn(object)  # {'ra_hrs':, 'dec_deg':} precessed & nut
    last = get_last(mjd)  # in fraction of a day
    obj_ha = 24.0 * last - posn['ra_hrs']
    if obj_ha < -12.0 :
        obj_ha += 24.0
    if obj_ha > 12.0 :
        obj_ha -= 24.0
    obj_azel = hadec_to_azel(obj_ha, posn['dec_deg']) # {'az':deg, 'el':deg}
    az_dist = abs(obj_azel['az'] - cur_az)
    if az_dist > 180 :
        az_dist = abs(360 - az_dist)
    el_dist = abs(obj_azel['el'] - cur_el)
    slew_time = max(az_dist / 2.0, el_dist / 2.0)  # seconds
    print 'slew time:', slew_time, 'sec to', obj_azel
    encoder_az, encoder_el = get_corrected_azel(obj_azel['az'],
                                                obj_azel['el'])
    print encoder_az, encoder_el
    status = ant.set_azel(encoder_az, encoder_el)
    if not status :
        return status

    wait_for_slew(encoder_az, encoder_el)

    start_mjd = get_current_mjd() + 0.5 / 86400.0
    stop_mjd = start_mjd + duration / 86400.0
    trk = make_track(object, start_mjd, stop_mjd,
                     cross_el_offset, el_offset)
    print trk
    status = execute_track(trk, stop_mjd)
    ant.stop()
    return status

_st = tm.gmtime()
_fn = 'log%4d_%02d_%02d_%02d:%02d:%02d' % (_st[0], _st[1], _st[2],
                                           _st[3], _st[4], _st[5])
_log_fd = file(_fn, 'wb')

def log_ant_posn ( ) :
    mjd = get_current_mjd()
    tel_stat = ant.get_time_tagged_status()
    _log_fd.write('%02d:%02d:%02d %15.9f %8.4f %8.4f\n' % \
                      (tel_stat['UTC_hr'], tel_stat['UTC_min'],
                       tel_stat['UTC_sec'], mjd, tel_stat['az'],
                       tel_stat['el']))

def mjd_to_time_str ( mjd ) :
    utc_sec = (mjd - int(mjd)) * 86400.0
    print 'mjd sec:', utc_sec
    hrs = int(utc_sec / 3600.0)
    utc_sec -= 3600.0 * hrs
    mint = int(utc_sec / 60.0)
    utc_sec -= 60.0 * mint
    return '%02d:%02d:%05.2f' % (hrs, mint, utc_sec)

def test_one_cmd ( az, el, az_vel=0.0, el_vel=0.0 ) :
    ant.set_computer_timeout(20.0)
    mjd = get_current_mjd() + 3.0 / 86400.0
    time_str = mjd_to_time_str(mjd)
    #print time_str
    ant.set_azel(az, el, tim=time_str, azvel=az_vel, elvel=el_vel)
    while (mjd + 13.0 / 86400.0) > get_current_mjd() :
        log_ant_posn()
        tm.sleep(0.5)
    ant.stop()

def test_two_cmds ( az, el, az_vel=0.0, el_vel=0.0 ) :
    ant.set_computer_timeout(30.0)
    mjd = get_current_mjd() + 3.0 / 86400.0
    utc_sec = (mjd - int(mjd)) * 86400.0
    #time_str = mjd_to_time_str(mjd)
    ant.set_azel(az, el, tim=utc_sec, azvel=az_vel, elvel=el_vel)

    mjd += 10.0 / 86400.0
    #while mjd > get_current_mjd() :
    #    tm.sleep(1.0)
    #time_str = mjd_to_time_str(mjd)
    utc_sec = (mjd - int(mjd)) * 86400.0
    ant.set_azel(az + 1.0, el + 1.0, tim=utc_sec, azvel=az_vel, elvel=el_vel)
    while (mjd + 13.0 / 86400.0) > get_current_mjd() :
        log_ant_posn()
        tm.sleep(0.5)
    ant.stop()

def test_three_cmds ( az, el, az_vel=0.0, el_vel=0.0 ) :
    ant.set_computer_timeout(30.0)
    mjd = get_current_mjd() + 3.0 / 86400.0
    time_str = mjd_to_time_str(mjd)
    ant.set_azel(az, el, tim=time_str, azvel=az_vel, elvel=el_vel)

    mjd += 10.0 / 86400.0
    #while mjd > get_current_mjd() :
    #    tm.sleep(1.0)
    time_str = mjd_to_time_str(mjd)
    ant.set_azel(az + 1.0, el + 1.0, tim=time_str, azvel=az_vel, elvel=el_vel)

    mjd += 10.0 / 86400.0
    #while mjd > get_current_mjd() :
    #    tm.sleep(1.0)
    time_str = mjd_to_time_str(mjd)
    ant.set_azel(az + 2.0, el + 2.0, tim=time_str, azvel=az_vel, elvel=el_vel)

    while (mjd + 13.0 / 86400.0) > get_current_mjd() :
        log_ant_posn()
        tm.sleep(0.5)
    ant.stop()

def test_slew ( az, el ) :
    ant.set_computer_timeout(30.0)
    ant.set_azel(az, el)
    wait_for_slew(az, el, tol_deg=0.02)
    ant.stop()

def test_execute_track ( az, el ) :
    num_pts = 3
    tmjd = np.zeros(num_pts, dtype=np.float64)
    taz = np.zeros(num_pts, dtype=np.float64)
    taz_vel = np.zeros(num_pts, dtype=np.float64)
    tel = np.zeros(num_pts, dtype=np.float64)
    tel_vel = np.zeros(num_pts, dtype=np.float64)
    delta_t = 5.0
    delta_p = 0.1
    mjd = get_current_mjd() + 3.0 / 86400.0
    for i in range(num_pts) :
        tmjd[i] = mjd + i * delta_t
        taz[i] = az + i * delta_p
        taz_vel[i] = delta_p / delta_t
        tel[i] = el + i * delta_p
        tel_vel[i] = delta_p / delta_t
    stop_mjd = tmjd[-1] + delta_t
    trk = {'mjd':tmjd, 'az':taz, 'az_vel':taz_vel, 'el':tel, 'el_vel':tel_vel}
    execute_track(trk, stop_mjd)