#!/usr/bin/env python
"""Commandline script to take a nominal focalplane model and produce a new model
with a restricted patrol radius.

This script loads an existing focalplane model and creates a new date-stamped model
in the current $DESIMODEL data directory.

"""

import sys
import datetime
import argparse

import numpy as np
import matplotlib.pyplot as plt

import json
import gzip

from astropy.table import Table, Column

from desimodel.io import load_focalplane

from desimodel.inputs.focalplane_utils import restricted_positioner_phi


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--in_date",
        type=str,
        default=None,
        required=False,
        help="Input date in ISO format (e.g. '2020-01-01T00:00:00')."
    )

    parser.add_argument(
        "--out_date",
        type=str,
        default=None,
        required=False,
        help="Output date in ISO format (e.g. '2020-01-01T00:00:00').  Default = now."
    )

    parser.add_argument(
        "--reach",
        type=float,
        default=None,
        required=True,
        help="The maximum positioner reach in millimeters."
    )

    parser.add_argument(
        "--exclude_petals",
        type=str,
        default=None,
        required=False,
        help="Comma-separated list of petal locations to exclude in restriction."
    )

    args = parser.parse_args()

    now_time = datetime.datetime.now()
    now_time_str = None
    try:
        now_time_str = now_time.isoformat(timespec="seconds")
    except TypeError:
        # This must be python < 3.6, with no timespec option.
        # Since the focalplane time is read from the file name without
        # microseconds, the microseconds should be zero and so the
        # default return string will be correct.
        now_time_str = now_time.isoformat()

    in_time = None
    if args.in_date is None:
        in_time = now_time
    else:
        in_time = datetime.datetime.strptime(args.in_date, "%Y-%m-%dT%H:%M:%S")

    if args.out_date is None:
        args.out_date = now_time_str

    petals = {x for x in range(10)}
    if args.exclude_petals is not None:
        expetals = [int(x) for x in args.exclude_petals.split(",")]
        for p in expetals:
            if p < 0 or p > 9:
                raise RuntimeError("Petal locations to exclude must be in range 0-9")
            petals.remove(p)

    # Load the starting focalplane model

    fp, excl, state, in_tmstr = load_focalplane(in_time)

    # Restrict the effective patrol radius by putting limits on the phi
    # angle range.

    nrows = len(fp)
    for row in fp:
        if row["PETAL"] not in petals:
            continue
        theta_arm = row["LENGTH_R1"]
        phi_arm = row["LENGTH_R2"]
        if theta_arm <= 0 or phi_arm <= 0:
            # this is not a positioner
            continue
        row["MIN_P"] = restricted_positioner_phi(
            args.reach, theta_arm, phi_arm, row["OFFSET_P"], row["MIN_P"], row["MAX_P"]
        )

    out_fp_file = "desi-focalplane_{}.ecsv".format(args.out_date)
    out_excl_file = "desi-exclusion_{}.json.gz".format(args.out_date)
    out_state_file = "desi-state_{}.ecsv".format(args.out_date)

    fp.write(out_fp_file, format="ascii.ecsv", overwrite=True)

    out_cols = [
        Column(name="TIME", length=nrows, dtype=np.dtype("a20"),
               description="The timestamp of the event (UTC, ISO format)"),
        Column(name="PETAL", length=nrows, dtype=np.int32,
               description="Petal location [0-9]"),
        Column(name="DEVICE", length=nrows, dtype=np.int32,
               description="Device location on the petal"),
        Column(name="LOCATION", length=nrows, dtype=np.int32,
               description="Global device location (PETAL * 1000 + DEVICE)"),
        Column(name="STATE", length=nrows, dtype=np.uint32,
               description="State bit field (good == 0)"),
        Column(name="EXCLUSION", length=nrows, dtype=np.dtype("a9"),
               description="The exclusion polygon for this device"),
    ]

    out_state = Table()
    out_state.add_columns(out_cols)
    out_state["TIME"][:] = args.out_date
    out_state["PETAL"][:] = state["PETAL"]
    out_state["DEVICE"][:] = state["DEVICE"]
    out_state["LOCATION"][:] = state["LOCATION"]
    out_state["STATE"][:] = state["STATE"]
    out_state["EXCLUSION"][:] = state["EXCLUSION"]

    out_state.write(out_state_file, format="ascii.ecsv", overwrite=True)

    with gzip.open(out_excl_file, "wt", encoding="utf8") as pf:
        json.dump(excl, pf, indent=4)

if __name__ == "__main__":
    main()
