#!/usr/bin/env python3

import numpy as np
from datetime import datetime
import os
from functools import partial
from lmatools.grid.make_grids import write_cf_netcdf_latlon, write_cf_netcdf_noproj, write_cf_netcdf_fixedgrid
from lmatools.grid.make_grids import dlonlat_at_grid_center, grid_h5flashfiles
from glmtools.grid.make_grids import grid_GLM_flashes
from glmtools.io.glm import parse_glm_filename
from lmatools.grid.fixed import get_GOESR_grid, get_GOESR_coordsys

import logging

log = logging.getLogger(__name__)

parse_desc = """Grid GLM flash data.

The start and end times can be specified
independently, or if not provided they will be inferred from the filenames.

Grid spacing is regular in latitude and longitude with the grid box
being sized to match the requested dx, dy at the center of the grid.

Therefore, this script can be used to process multiple days and they will
be written to a standardized directory structure.
"""


def create_parser():
    import argparse
    parser = argparse.ArgumentParser(description=parse_desc)
    parser.add_argument('-v', '--verbose', dest='verbosity', action="count", default=0,
                        help='each occurrence increases verbosity 1 level through ERROR-WARNING-INFO-DEBUG (default INFO)')
    parser.add_argument('-l', '--log', dest="log_fn", default=None,
                        help="specify the log filename")
    parser.add_argument('-o', '--output-dir', metavar='directory',
                        default='.')
    parser.add_argument('--ctr-lat', metavar='latitude',
                        type=float, help='center latitude')
    parser.add_argument('--ctr-lon', metavar='longitude',
                        type=float, help='center longitude')
    parser.add_argument('--start', metavar='yyyy-mm-ddThh:mm:ss',
                        help='UTC start time, e.g., 2017-07-04T08:00:00')
    parser.add_argument('--end', metavar='yyyy-mm-ddThh:mm:ss',
                        help='UTC end time, e.g., 2017-07-04T09:00:00')
    parser.add_argument('--dx', metavar='km',
                        default=10.0, type=float,
                        help='approximate east-west grid spacing')
    parser.add_argument('--dy', metavar='km',
                        default=10.0, type=float,
                        help='approximate north-south grid spacing')
    parser.add_argument('--dt', metavar='seconds',
                        default=60.0, type=float,
                        help='frame duration')
    parser.add_argument('--width', metavar='distance in km',
                        default=400.0,
                        type=float, help='total width of the grid')
    parser.add_argument('--height', metavar='distance in km',
                        default=400.0,
                        type=float, help='total height of the grid')
    parser.add_argument('--nevents', metavar='minimum events per flash',
                        type=int, dest='min_events', default=1,
                        help='minimum number of events per flash')
    parser.add_argument('--ngroups', metavar='minimum groups per flash',
                        type=int, dest='min_groups', default=1,
                        help='minimum number of groups per flash')
    parser.add_argument('--subdivide-grid', metavar='sqrt(number of subgrids)',
                        type=int, default=1,
                        help="subdivide the grid this many times along "
                             "each dimension")
    parser.add_argument('--goes-position',
                        help="One of [east|west|test]. "
                             "Requires '--goes-sector'.")
    parser.add_argument('--goes-sector',
                        help="One of [full|conus|meso]. "
                             "Requires goes_position. If sector is "
                             "meso, ctr_lon and ctr_lat are interpreted as "
                             "the ctr_x and ctr_y of the fixed grid")
    # parser.add_argument('--split-events',
    #                     action='store_true',
    #                     help='Split GLM event polygons when gridding')
    parser.add_argument('--ellipse', dest='ellipse_rev', default=-1,
                        type=int,
                        help='Lightning ellipse revision. -1 (default)=infer'
                             ' from date in each GLM file, 0=value at launch,'
                             ' 1=late 2018 revision')
    parser.add_argument('--float-output', dest='output_scale_and_offset',
                        action='store_false',
                        help='write all output variables as floating point')
    parser.add_argument(dest='filenames', metavar='filename', nargs='+')
    return parser


def nearest_resolution(args):
    """ Uses args.dx to find the closest resolution specified by the
    GOES-R PUG. Returns something like "10.0km" that can be used as the
    resolution argument to get_GOESR_grid.
    """
    goes_resln_options = np.asarray([0.5, 1.0, 2.0, 4.0, 8.0, 10.0])
    resln_idx = np.argmin(np.abs(goes_resln_options - args.dx))
    closest_resln = goes_resln_options[resln_idx]
    resln = '{0:4.1f}km'.format(closest_resln).replace(' ', '')
    return resln


def get_start_end(filenames, start_time=None, end_time=None):
    """Compute start and end time of data based on filenames."""
    base_filenames = [os.path.basename(p) for p in filenames]
    try:
        filename_infos = [parse_glm_filename(f) for f in base_filenames]
        # opsenv, algorithm, platform, start, end, created = parse_glm_filename(f)
        filename_starts = [info[3] for info in filename_infos]
        filename_ends = [info[4] for info in filename_infos]
    except ValueError:
        filename_starts = None
        filename_ends = None

    if args.start is not None:
        start_time = datetime.strptime(args.start, '%Y-%m-%dT%H:%M:%S')
    elif filename_starts is not None:
        start_time = min(filename_starts)

    if args.end is not None:
        end_time = datetime.strptime(args.end, '%Y-%m-%dT%H:%M:%S')
    elif filename_ends is not None:
        end_time = max(filename_ends)

    if start_time is None or end_time is None:
        raise ValueError("Could not determine start/end time")

    return start_time, end_time


def grid_setup(args):
    # When passed None for the minimum event or group counts, the gridder will skip
    # the check, saving a bit of time.
    min_events = int(args.min_events)
    if min_events <= 1:
        min_events = None
    min_groups = int(args.min_groups)
    if min_groups <= 1:
        min_groups = None

    try:
        start_time, end_time = get_start_end(args.filenames, args.start, args.end)
    except ValueError:
        log.error("Non-standard filenames provided, use --start and --end to specify data times.")
        raise

    date = datetime(start_time.year, start_time.month, start_time.day)
    os.makedirs(args.output_dir, exist_ok=True)
    proj_name = 'geos'

    if args.goes_position is not None and args.goes_sector is not None:
        resln = nearest_resolution(args)
        view = get_GOESR_grid(position=args.goes_position,
                              view=args.goes_sector,
                              resolution=resln)
        nadir_lon = view['nadir_lon']
        dx = dy = view['resolution']
        nx, ny = view['pixelsEW'], view['pixelsNS']
        geofixcs, grs80lla = get_GOESR_coordsys(sat_lon_nadir=nadir_lon)

        if 'centerEW' in view:
            x_ctr, y_ctr = view['centerEW'], view['centerNS']
        elif args.goes_sector == 'meso':
            # use ctr_lon, ctr_lat to get the center of the mesoscale FOV
            x_ctr, y_ctr, z_ctr = geofixcs.fromECEF(
                *grs80lla.toECEF(args.ctr_lon, args.ctr_lat, 0.0))
    elif args.goes_position is not None and args.goes_sector is None:
        # Requires goes_position, a center, and a width. Fully flexible
        # in resolution, i.e., doesn't slave it to one of the GOES-R specs
        view = get_GOESR_grid(position=args.goes_position,
                              view='full',
                              resolution='1.0km')
        nadir_lon = view['nadir_lon']
        dx1km = dy1km = view['resolution']
        geofixcs, grs80lla = get_GOESR_coordsys(sat_lon_nadir=nadir_lon)
        x_ctr, y_ctr, z_ctr = geofixcs.fromECEF(
            *grs80lla.toECEF(args.ctr_lon, args.ctr_lat, 0.0))

        # Convert the specified resolution in km given by args.dx to
        # a delta in fixed grid coordinates using the 1 km delta from the
        # GOES-R PUG.
        dx, dy = args.dx * dx1km, args.dy * dy1km
        nx, ny = int(args.width / args.dx), int(args.height / args.dy)
    else:
        raise ValueError("Gridding on the fixed grid requires "
                         "goes_position and dx. For goes_sector='meso', also specify "
                         "ctr_lon and ctr_lat. Without goes_sector, also include width "
                         "and height.")
    # Need to use +1 here to convert to xedge, yedge expected by gridder
    # instead of the pixel centroids that will result in the final image
    nx += 1
    ny += 1
    x_bnd = (np.arange(nx, dtype='float') - (nx) / 2.0) * dx + x_ctr + 0.5 * dx
    y_bnd = (np.arange(ny, dtype='float') - (ny) / 2.0) * dy + y_ctr + 0.5 * dy
    log.debug(("initial x,y_ctr", x_ctr, y_ctr))
    log.debug(("initial x,y_bnd", x_bnd.shape, y_bnd.shape))
    x_bnd = np.asarray([x_bnd.min(), x_bnd.max()])
    y_bnd = np.asarray([y_bnd.min(), y_bnd.max()])

    geofixcs, grs80lla = get_GOESR_coordsys(sat_lon_nadir=nadir_lon)
    ctr_lon, ctr_lat, ctr_alt = grs80lla.fromECEF(
        *geofixcs.toECEF(x_ctr, y_ctr, 0.0))
    fixed_grid = geofixcs
    log.debug((x_bnd, y_bnd, dx, dy, nx, ny))

    output_writer = partial(write_cf_netcdf_fixedgrid, nadir_lon=nadir_lon)

    gridder = grid_GLM_flashes
    output_filename_prefix = 'GLM'
    grid_kwargs = dict(proj_name=proj_name,
                       base_date=date, do_3d=False,
                       dx=dx, dy=dy, frame_interval=float(args.dt),
                       x_bnd=x_bnd, y_bnd=y_bnd,
                       ctr_lat=ctr_lat, ctr_lon=ctr_lon, outpath=args.output_dir,
                       min_points_per_flash=min_events,
                       output_writer=output_writer, subdivide=args.subdivide_grid,
                       output_filename_prefix=output_filename_prefix,
                       output_kwargs={'scale_and_offset': args.output_scale_and_offset},
                       spatial_scale_factor=1.0)

    if args.fixed_grid:
        grid_kwargs['fixed_grid'] = True
        grid_kwargs['nadir_lon'] = nadir_lon
    # if args.split_events:
    grid_kwargs['clip_events'] = True
    if min_groups is not None:
        grid_kwargs['min_groups_per_flash'] = min_groups
    grid_kwargs['energy_grids'] = ('total_energy',)
    if (proj_name == 'pixel_grid') or (proj_name == 'geos'):
        grid_kwargs['pixel_coords'] = fixed_grid
    grid_kwargs['ellipse_rev'] = args.ellipse_rev
    return gridder, args.filenames, start_time, end_time, grid_kwargs


if __name__ == '__main__':
    import sys
    parser = create_parser()
    args = parser.parse_args()

    # Configure logging
    levels = [logging.ERROR, logging.WARN, logging.INFO, logging.DEBUG]
    logging.basicConfig(level=levels[min(3, args.verbosity)], filename=args.log_fn)
    if levels[min(3, args.verbosity)] > logging.DEBUG:
        import warnings
        warnings.filterwarnings("ignore")
    log.info("Starting GLM Gridding")
    log.debug("Starting script with: %s", sys.argv)

    from multiprocessing import freeze_support
    freeze_support()
    gridder, glm_filenames, start_time, end_time, grid_kwargs = grid_setup(args)
    gridder(glm_filenames, start_time, end_time, **grid_kwargs)