#!/usr/bin/env python
"""Generate a DCT/LMI dither macro.

2015-10-24: Initial version. Requires Python 2.6 or later.

2015-01-12: Requires Python 2.7 or later.

2016-02-12: Added context r' at start, configurable context
            integration time, tweaked default itimes.

2016-04-18: New --offset, --no-filter, and --no-alternate options.

2016-05-26: New title argument, new -g (group) option.

2017-04-05: New --subframe option.

2021-03-06: Dither context images to mitigate LMI's bad column.

- Michael S. P. Kelley, UMD

"""

from __future__ import print_function
import sys
import argparse


patterns = {
    'hex': [(0, -1), (-0.866, 0.5), (0.866, 0.5),
            (-0.866, -0.5), (0, 1), (0.866, -0.5)],
    'cross': [(-1, 0), (1, 0), (0, -1), (0, 1)]
}

sets = {
    'OH': [('RC', 50), ('OH', 150), ('BC', 50)],
    'OH-long': [('RC', 100), ('OH', 300), ('BC', 100)],
    'OH-CN': [('RC', 50), ('OH', 200), ('CN', 100), ('BC', 50)],
    'RC-BC': [('RC', 120), ('BC', 120)]
}


def to_list(_type=str):
    if _type == str:
        return lambda s: [x.strip() for x in s.split(',')]
    else:
        return lambda s: [_type(x) for x in s.split(',')]


def exposures(n, group, pattern, scale, offset, filters, itimes, subframe):
    from itertools import cycle

    # empty filters is allowed, but we still need the cycling
    if len(filters) == 0:
        filters = ['']

    filter = cycle(filters + filters[::-1])
    itime = cycle(itimes + itimes[::-1])
    position = cycle(pattern)

    for i in range(0, n, group):
        subpattern = []
        for j in range(group):
            subpattern.append(next(position))

        p = cycle(subpattern)
        for j in range(len(filters)):
            f = next(filter)
            t = next(itime)
            for k in range(i, min(n, i + group)):
                px, py = next(p)
                x = scale * px + offset[0]
                y = scale * py + offset[1]

                yield dict(filter=f, itime=t, x=x, y=y, subframe=subframe,
                           comment='Position {}'.format(k))


def dither(n, scale, offset, pattern, filters, itimes, group,
           context_itime, context, subframe):
    from itertools import product

    line = ['{itime:6.1f}']
    if len(filters) > 0:
        line.append('{filter:6s}')

    if subframe is not None:
        line.append('{subframe:1d}')

    line.extend(['{x:4.0f}', '{y:4.0f}', '"{comment:4}"'])
    line = '  '.join(line)

    if context:
        yield line.format(itime=context_itime, filter='SL-r', subframe=subframe,
                          x=offset[0], y=offset[1] + 2, comment="Context")

    for exposure in exposures(n, group, pattern, scale, offset, filters,
                              itimes, subframe):
        yield line.format(**exposure)

    if context:
        yield line.format(itime=context_itime, filter='SL-r', subframe=subframe,
                          x=offset[0] - 2, y=offset[1], comment="Context")


parser = argparse.ArgumentParser(
    description='Generate a DCT/LMI dither macro.')

parser.add_argument('title', type=str, default=None, nargs='*',
                    help='Image title.')
parser.add_argument('-n', type=int, default=0,
                    help='Number of dithers; even is best for airmass matching')
parser.add_argument('-g', type=int, default=0,
                    help='Number of dithers to group before changing filters.  Set to 1 to disable.')
parser.add_argument('--scale', type=int, default=60,
                    help='Scale factor for offsets.  [arcsec]')
parser.add_argument('--offset', type=to_list(float), default=[0, 0],
                    help='Absolute dither pattern offset (xi, eta).')
parser.add_argument('--pattern', choices=sorted(patterns.keys()), default='hex',
                    help='Dither pattern.')
parser.add_argument('--subframe', type=int, default=None,
                    help='Use this subframe number.')
parser.add_argument('--no-filter', action='store_true',
                    help='Set to omit filter column.')
parser.add_argument('--filters', type=to_list(), default=['SL-r'],
                    help='Comma-separated list of filter names')
parser.add_argument('--itimes', type=to_list(float), default=[60],
                    help='Comma-separated list of filter integration times, or one time for all filters.  [s]')
parser.add_argument('--no-alternate', dest='alternate', action='store_false',
                    help='Cycle through each step before each filter.')
parser.add_argument('--set', choices=sorted(sets.keys()),
                    help='Use a predefined filter/itime sequence.')
parser.add_argument('--no-context', action='store_false', dest='context',
                    help='Do not bracket with a context images.')
parser.add_argument('--context-itime', type=float, default=30.0,
                    help='Integration time to use for context images.  [s]')

args = parser.parse_args()

# sort out filters and integration times
filters = []
itimes = []
if args.set is not None:
    for f, t in sets[args.set]:
        filters.append(f)
        itimes.append(t)
else:
    filters = args.filters
    itimes = args.itimes

if args.no_filter:
    filter_col = 'false'
    filters = []
else:
    filter_col = 'true'

# set the subframe option
subframe_col = 'true' if args.subframe is not None else 'false'

assert len(itimes) > 0, "No itimes given."

if len(filters) > len(itimes):
    assert len(itimes) == 1, "Either 1 or {0} itimes expected, but found {1}.".format(
        len(filters), len(itimes))
    itimes = [itimes[0]] * len(filters)

if len(filters) > 0:
    assert len(itimes) == len(filters), "{0} itimes expected, but found {1}.".format(
        len(filters), len(itimes))

# handle --offset
assert len(args.offset) == 2, "Offset must have length 2: xi, eta."

pattern = patterns[args.pattern]

n = args.n if args.n > 0 else len(pattern)

# consider grouping
if not args.alternate:
    assert args.g == 0, "Only one of group (-g) or --no-alternate can be specified."
    group = n
elif args.g > 0:
    group = args.g
else:
    if n % 3 == 0 and n != 3:
        group = 3
    elif n % 2 == 0 and n != 2:
        group = 2
    else:
        group = 1

# set the title column
if args.title == []:
    title_col = 'false'
    title = ''
else:
    title_col = 'true'
    title = '"{0:s}" '.format(' '.join(args.title))

print("""#title={0:s} ra=false dec=false exposureTime=true numExposures=false filter={1:s} subframe={2:s} muRA=false muDec=false epoch=false dRA=false dDec=false rotatorPA=false rotatorFrame=false xi=true eta=true comment=true commandOption=false
#""".format(title_col, filter_col, subframe_col))
it = 0
wt = 0
nf = 0
for line in dither(n, args.scale, args.offset, pattern, filters, itimes,
                   group, args.context_itime, args.context, args.subframe):
    print(title + line)
    nf += 1
    it += float(line.split()[0])
    wt += float(line.split()[0]) + 24

print("""#
# {0} frames
# {1} s = {2:.0f} min integration time
# {3} s = {4:.0f} min estimated wall time""".format(
    nf, it, it / 60., wt, wt / 60.))