#
# Copyright 2017-2018 Universidad Complutense de Madrid
#
# This file is part of Megara DRP
#
# SPDX-License-Identifier: GPL-3.0+
# License-Filename: LICENSE.txt
#

from __future__ import print_function

import math

import matplotlib.cbook as cbook
import matplotlib.collections as mcoll
import matplotlib.colors as mcolors
import matplotlib.transforms as mtrans
import matplotlib.transforms as mtransforms
import numpy as np


M_SQRT3 = math.sqrt(3)

PLATESCALE = 1.2120  # arcsec / mm
SCALE = 0.443  # mm from center to center, upwards


def hexplot(axis, x, y, z, scale=1.0, extent=None,
            cmap=None, norm=None, vmin=None, vmax=None,
            alpha=None, linewidths=None, edgecolors='none',
            **kwargs):
    """
    Make a hexagonal grid plot.

    Returns
    =======
    object: matplotlib.collections.PolyCollection

    """
    axis._process_unit_info(xdata=x, ydata=y, kwargs=kwargs)

    x, y, z = cbook.delete_masked_points(x, y, z)

    x = np.array(x, float)
    y = np.array(y, float)

    sx = 2 / M_SQRT3 * scale * 0.99
    sy = scale * 0.99

    if extent is not None:
        xmin, xmax, ymin, ymax = extent
    else:
        xmin, xmax = (np.amin(x - sx), np.amax(x + sx)) if len(x) else (0, 1)
        ymin, ymax = (np.amin(y - sy), np.amax(y + sy)) if len(y) else (0, 1)

        # to avoid issues with singular data, expand the min/max pairs
        xmin, xmax = mtrans.nonsingular(xmin, xmax, expander=0.1)
        ymin, ymax = mtrans.nonsingular(ymin, ymax, expander=0.1)

    padding = 1.0e-9 * (xmax - xmin)
    xmin -= padding
    xmax += padding

    n = len(x)
    polygon = np.zeros((6, 2), float)

    S = 1 / M_SQRT3
    mx = my = 0.99 * scale
    polygon[:, 0] = mx * np.array([-0.5 * S, 0.5 * S, 1.0 * S, 0.5 * S, -0.5 * S, -1.0 * S])
    polygon[:, 1] = my * np.array([0.5, 0.5, 0.0, -0.5, -0.5, 0.0])

    offsets = np.zeros((n, 2), float)
    offsets[:, 0] = x
    offsets[:, 1] = y

    collection = mcoll.PolyCollection(
        [polygon],
        edgecolors=edgecolors,
        linewidths=linewidths,
        offsets=offsets,
        transOffset=mtransforms.IdentityTransform(),
        offset_position="data"
    )

    if isinstance(norm, mcolors.LogNorm):
        if (z == 0).any():
            # make sure we have not zeros
            z += 1

    if norm is not None:
        if norm.vmin is None and norm.vmax is None:
            norm.autoscale(z)

    if norm is not None and not isinstance(norm, mcolors.Normalize):
        msg = "'norm' must be an instance of 'mcolors.Normalize'"
        raise ValueError(msg)

    collection.set_array(z)
    collection.set_cmap(cmap)
    collection.set_norm(norm)
    collection.set_alpha(alpha)
    collection.update(kwargs)

    if vmin is not None or vmax is not None:
        collection.set_clim(vmin, vmax)
    else:
        collection.autoscale_None()

    corners = ((xmin, ymin), (xmax, ymax))
    axis.update_datalim(corners)
    axis.autoscale_view(tight=True)

    # add the collection last
    axis.add_collection(collection, autolim=False)
    return collection


def transform_to_pixel(wcs, coords, ref=0):
    # We need only the WL axis
    coords_m = [coords, np.zeros_like(coords)]
    r = wcs.wcs_world2pix(np.transpose(coords_m), ref)
    first_col = r[:, 0].astype('int')
    return first_col


def is_inside(shape, *coords):

    for coord in coords:
        for shape1d, crds in zip(shape, coord):
            if not is_inside_1d(shape1d, crds):
                return False
    return True


def is_inside_1d(shape1d, coords):
    for coord in coords:
        if coord < 0:
            return False
        if coord >= shape1d:
            return False
    return True


def extract_region(wcs_wl, region, rssdata):
    cut1, cut2 = transform_to_pixel(wcs_wl, region)

    if is_inside_1d(rssdata.shape[1], [cut1, cut2]):
        z = rssdata[:, cut1:cut2].mean(axis=1)
        return z
    else:
        raise ValueError('region to average is outside image')


def extract_column(wcs_wl, col, rssdata):
    cut, = transform_to_pixel(wcs_wl, [col])

    if is_inside_1d(rssdata.shape[1], [cut]):
        # cut1, cut2 = args.average_region
        z = rssdata[:, cut]
        return z
    else:
        raise ValueError('column is outside image')


def extract_zval(rssdata, wcs_wl, coordinate_type, column=None, average_region=None, continuum_region=None):

    if continuum_region is None and average_region is None:
        raise ValueError('continuum_region and average_region are None')

    if rssdata.ndim == 1:
        # data is 1d
        print('RSS is 1D')
        zval = rssdata
    elif rssdata.ndim == 2:
        print('world coordinates in', coordinate_type)
        if column is not None:
            print('use column', column)
            zval = extract_column(wcs_wl, column, rssdata)
        else:
            print('compute average in', average_region)
            zval = extract_region(wcs_wl, average_region, rssdata)
        # Subtract continuum
        if continuum_region:
            print('compute continuum in', continuum_region)
            z0 = extract_region(wcs_wl, continuum_region, rssdata)
            zval -= z0
    else:
        raise ValueError('image must be 1D or 2D')
    return zval


def main(argv=None):
    """Process command line"""

    import argparse
    import json

    import matplotlib.pyplot as plt
    import astropy.io.fits as fits
    from astropy.wcs import WCS
    from astropy.visualization import simple_norm
    import matplotlib.transforms as mtransforms

    import megaradrp.datamodel as dm
    from megaradrp.processing.wcs import update_wcs_from_ipa, compute_pa_from_ipa

    try:
        from megaradrp.processing.cube import create_cube_from_rss
        has_contours = True
    except ImportError:
        create_cube_from_rss = None
        has_contours = False

    parser = argparse.ArgumentParser(description='Display MEGARA RSS images')
    parser.add_argument('--wcs-grid', action='store_true',
                        help='Display WCS grid')
    parser.add_argument('--wcs-pa-from-header', action='store_true',
                        help="Use PA angle from PC keys", dest='pa_from_header')
    parser.add_argument('--average-region', nargs=2, default=[1000, 3000],
                        type=int, help='Region of the RSS averaged on display')
    parser.add_argument('--extname', '-e', default='PRIMARY',
                        help='Name of the extension used')
    parser.add_argument('--column', '-c', type=int,
                        help='Column of the RSS on display')
    parser.add_argument('--continuum-region', nargs=2,
                        type=int,
                        help='Region of the RSS used for continuum subtraction')
    parser.add_argument('--coordinate-type', choices=['pixel', 'wcs'],
                        default='pixel',
                        help='Types of coordinates used')
    parser.add_argument('--colormap', type=plt.get_cmap,
                        help='Name of a valid matplotlib colormap')
    parser.add_argument('--plot-sky', action='store_true',
                        help='Plot SKY bundles')
    parser.add_argument('--plot-nominal-config', action='store_true',
                        help='Plot nominal configuration, do not use the header')
    parser.add_argument('--hide-values', action='store_true',
                        help='Do not show values out of range')
    parser.add_argument('--title', help='Title of the plot')
    parser.add_argument('--label', help='Legend of the colorbar')
    parser.add_argument('--hex-size', type=float,
                        help='Size of the hexagons (default is {})'.format(SCALE),
                        default=SCALE)
    parser.add_argument('--hex-rel-size', type=float,
                        help='Scale the size of hexagons by a factor',
                        default=1.0)
    parser.add_argument('--min-cut', type=float,
                        help='Inferior cut level')
    parser.add_argument('--max-cut', type=float,
                        help='Superior cut level')
    parser.add_argument('--percent', type=float,
                        help='Compute cuts using percentiles')
    parser.add_argument('--stretch',
                        choices=['linear', 'sqrt', 'power', 'log', 'asinh'],
                        default='linear',
                        help='Name of the strech method used for display'
                        )
    if has_contours:
        parser.set_defaults(has_contours=True)
        group_c = parser.add_argument_group('contouring')
        group_c.add_argument('--contour-pixel-size', type=float, default=0.4,
                             help="Pixel size in arc seconds for image reconstruction")
        group_c.add_argument('--contour-levels',
                             help="Contour levels")
        group_c.add_argument('--contour', action='store_true',
                             help="Draw contours")
        group_c.add_argument('--contour-image',
                             help="Image for computing contours")
        group_c.add_argument('--contour-image-column', type=int,
                             help='Column of image used for contouring')
        group_c.add_argument('--contour-image-save',
                             help='Save image used for contouring')
        group_c.add_argument('--contour-image-region', nargs=2, default=[1000, 3000],
                             type=int, help='Region of the image used for contouring')
        group_c.add_argument('--contour-is-density', action='store_true',
                             help='The data is a magnitude that does not require scaling')

    else:
        parser.set_defaults(has_contours=False)

    parser.add_argument('rss', metavar='RSS', nargs='+',
                        help='RSS images to process')

    args = parser.parse_args(args=argv)

    for fname in args.rss:
        with fits.open(fname) as img:

            datamodel = dm.MegaraDataModel()
            if args.plot_nominal_config:
                insmode = img['FIBERS'].header['INSMODE']
                fiberconf = datamodel.get_fiberconf_default(insmode)
            else:
                fiberconf = datamodel.get_fiberconf(img)
            plot_mask = np.ones((fiberconf.nfibers,), dtype=np.bool)
            if not args.plot_sky:
                skyfibers = fiberconf.sky_fibers()
                skyfibers.sort()
                skyfibers_idx = [(fibid - 1) for fibid in skyfibers]
                plot_mask[skyfibers_idx] = False

            x = np.empty((fiberconf.nfibers,))
            y = np.empty((fiberconf.nfibers,))
            # Key is fibid
            for _, fiber in sorted(fiberconf.fibers.items()):
                idx = fiber.fibid - 1
                x[idx] = fiber.x
                y[idx] = fiber.y

            extname = args.extname
            if args.coordinate_type == 'wcs':
                # read WCS from extname
                wcs_wl = WCS(img[extname].header)
            else:
                wcs_wl = WCS(naxis=len(img[extname].shape))

            rssdata = np.squeeze(img[extname].data)

            zval = extract_zval(rssdata,
                                wcs_wl,
                                args.coordinate_type,
                                args.column,
                                args.average_region,
                                args.continuum_region
                                )
            fig = plt.figure()
            
            x = x[plot_mask]
            y = y[plot_mask]
            zval = zval[plot_mask]
            projection = None

            if args.wcs_grid:
                print('display wcs grid')
                if not args.pa_from_header:
                    print('compute PA from IPA')
                    ipa = img[extname].header['IPA']
                    pa = compute_pa_from_ipa(ipa)
                    print('IPA is', ipa, ' PA is', pa)
                    hdr_fibers = update_wcs_from_ipa(img['FIBERS'].header, pa)
                else:
                    hdr_fibers = img['FIBERS'].header

                projection = WCS(hdr_fibers)

            # plt.subplots_adjust(hspace=0.5)
            ax = fig.add_axes([0.15, 0.1, 0.8, 0.8], projection=projection)

#            ax.set_xlim([-6.5, 6.5])
#            ax.set_ylim([-6, 6])

            if args.wcs_grid:
                ax.coords.grid(color='black', alpha=1.0, linestyle='solid')

            norm = simple_norm(zval, args.stretch,
                               min_cut=args.min_cut,
                               max_cut=args.max_cut,
                               percent=args.percent
                               )

            if args.hide_values:
                zdisp = zval.copy()
                if norm.vmin is not None:
                    zdisp[zval < norm.vmin] = np.nan
                if norm.vmax is not None:
                    zdisp[zval > norm.vmax] = np.nan
            else:
                zdisp = zval
            scale = args.hex_rel_size * args.hex_size
            col = hexplot(ax, x, y, zdisp, scale=scale, cmap=args.colormap, norm=norm)

            if args.title is not None:
                ax.set_title(args.title)

            cb = fig.colorbar(col)
            if args.label is not None:
                cb.set_label(args.label)

            if args.has_contours and args.contour:
                target_scale_arcsec = args.contour_pixel_size

                if args.contour_image is not None:
                    with fits.open(args.contour_image) as hdu_con:
                        extname = args.extname

                        if args.coordinate_type == 'wcs':
                            # read WCS from extname
                            wcs_wl_c = WCS(hdu_con[extname].header)
                        else:
                            wcs_wl_c = WCS(naxis=len(hdu_con[extname].shape))

                        rssdata_con = np.squeeze(hdu_con[extname].data)

                        zval_c = extract_zval(rssdata_con,
                                              wcs_wl_c,
                                              args.coordinate_type,
                                              column=args.contour_image_column,
                                              average_region=args.contour_image_region,
                                              continuum_region=None)
                else:
                    zval_c = zval

                # Build synthetic rss... for reconstruction
                primary = fits.PrimaryHDU(data=zval_c[:, np.newaxis], header=img[extname].header)
                synt = fits.HDUList([primary, img['FIBERS']])

                if args.contour_image_save is not None:
                    synt.writeto(args.contour_image_save)

                conserve_flux = not args.contour_is_density
                s_cube = create_cube_from_rss(synt, target_scale_arcsec, conserve_flux=conserve_flux)
                cube_wcs = WCS(s_cube[0].header).celestial
                px, py = cube_wcs.wcs.crpix
                interp = np.squeeze(s_cube[0].data)
                td = mtransforms.Affine2D().translate(-px, -py).scale(target_scale_arcsec, target_scale_arcsec)
                tt_d = td + ax.transData
                # im = ax.imshow(interp, alpha=0.9, cmap='jet', transform=tt_d)
                # im = ax.imshow(interp, alpha=0.9, cmap='jet', transform=ax.get_transform(cube_wcs))
                if args.contour_levels is not None:
                    levels = json.loads(args.contour_levels)
                    mm = ax.contour(interp, levels, transform=tt_d)
                else:
                    mm = ax.contour(interp, transform=tt_d)
                print('contour levels', mm.levels)

            plt.show()


if __name__ == '__main__':
    main()
