#!/usr/bin/env python3
#
# Copyright (c) 2016, Ole Streicher
# License: BSD-2-Clause
#

import argparse
import datetime
import os

from casacore import tables

default_input = '/usr/share/zoneinfo/leap-seconds.list'
default_database ='/var/lib/casacore/data/geodetic/TAI_UTC'

# casacore needs a HOME environment variable, so make it happy
if not os.environ.get('HOME'):
    os.environ['HOME'] = '/nonexistent'

parser = argparse.ArgumentParser(
    usage = 'casacore-update-tai_utc [options]',
    description = 'Update casacore TAI_UTC table from tzdata')
parser.add_argument('-i', '--input-file',
                    help='leap_second input data file (default: {})'
                    .format(default_input),
                    default=default_input)
parser.add_argument('-o', '--output-path',
                    help='output database (default: {})'
                    .format(default_database),
                    default=default_database)
parser.add_argument('-f', '--force',
                    help='overwrite database, even if it has another (sub)type'
                    + ' or the same NTP timestamp',
                    action='store_true')
args = parser.parse_args()

# Old data that is not in the zoneinfo database
data = [
    { 'MJD': 37300, 'dUTC': 1.42282, 'Offset': 37300, 'Multiplier': 0.001296 },
    { 'MJD': 37512, 'dUTC': 1.37282, 'Offset': 37300, 'Multiplier': 0.001296 },
    { 'MJD': 37665, 'dUTC': 1.84586, 'Offset': 37665, 'Multiplier': 0.0011232 }, 
    { 'MJD': 38334, 'dUTC': 1.94586, 'Offset': 37665, 'Multiplier': 0.0011232 }, 
    { 'MJD': 38395, 'dUTC': 3.24013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 38486, 'dUTC': 3.34013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 38639, 'dUTC': 3.44013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 38761, 'dUTC': 3.54013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 38820, 'dUTC': 3.64013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 38942, 'dUTC': 3.74013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 39004, 'dUTC': 3.84013, 'Offset': 38761, 'Multiplier': 0.001296 }, 
    { 'MJD': 39126, 'dUTC': 4.31317, 'Offset': 39126, 'Multiplier': 0.002592 }, 
    { 'MJD': 39887, 'dUTC': 4.21317, 'Offset': 39126, 'Multiplier': 0.002592 },
]

# Read the zoneinfo leap second file and return it as a list of dicts
# MJD dUTC Offset Multiplier.
# Also return the NTP timestamp of the last update.
offset = None
update_timestamp = None
with open(args.input_file) as f:
    for line in f.readlines():
        line = line.split('#', 1)
        if not line[0]:
            if line[1].strip().startswith('$'):
                try:
                    update_timestamp = int(line[1].split()[1])
                    ntp_epoch = -2208988800 # NTP starts at 1900-01-01
                    utc = datetime.datetime.fromtimestamp(ntp_epoch
                                                          + update_timestamp)
                    update_timestamp = utc.strftime('%Y/%m/%d/00:00')
                except ValueError:
                    pass
            continue
        t, dUTC = [int(c) for c in line[0].split()]
        mjd = t/86400. + 15020.0
        if offset is None:
            offset = mjd
        data.append({'MJD': mjd, 'dUTC': dUTC,
                     'Offset': offset, 'Multiplier': 0.0})

# Check if the table exists and has the correct type and subtype,
# and a different NTP timestamp
if not args.force:
    try:
        with tables.table(args.output_path) as tbl:
            db_info = tbl.info()
            if db_info['type'] != 'IERS':
                print("Existing table type '{type}' is not 'IERS'"
                      .format(**db_info))
                exit(1)
            if db_info['subType'] != 'leapSecond':
                print("Existing table type '{subType}' is not 'leapSecond'"
                      .format(**db_info))
                exit(1)
            if tbl.getkeyword('VS_DATE') == update_timestamp:
                print('NTP timestamp unchanged. Not updating')
                exit(0)
        print('Overwrite old table {}'.format(args.output_path))
    except RuntimeError as e:
        print('Creating new table {}'.format(args.output_path))

if not os.path.exists(args.output_path):
    os.makedirs(args.output_path)

# Create data table
columns = [
    tables.makescacoldesc('MJD', 0., valuetype='double', keywords={'UNIT':'d'}),
    tables.makescacoldesc('dUTC', 0., valuetype='double', keywords={'UNIT':'s'}),
    tables.makescacoldesc('Offset', 0., valuetype='double', keywords={'UNIT':'d'}),
    tables.makescacoldesc('Multiplier', 0., valuetype='double', keywords={'UNIT':'s'}),
]

keywords = {
    'MJD0': 37300.0,
    'dMJD': 0.0,
    'TAB_VERSION': '0001.0000',
    'VS_TYPE': 'TAI_UTC difference obtained from tzone data',
    'VS_VERSION': '{:04d}.0000'.format(len(data)),
    'VS_CREATE': update_timestamp,
    'VS_DATE': update_timestamp,
}

info = {
    'type': 'IERS',
    'subType': 'leapSecond',
    'readme': 'Created from zoneinfo leap second information.',
}
with tables.table(args.output_path,
                  tables.tablecreatedesc(columns),
                  len(data)) as tbl:
    tbl.putinfo(info)
    tbl.putkeywords(keywords)
    tablerows = tbl.row()
    for i, row in enumerate(data):
        tablerows.put(i, row)
