# -*- coding: utf-8 -*-

from collections import namedtuple
import logging

from africanus.model.wsclean.file_model import load
from astropy.coordinates import SkyCoord
import numpy as np


log = logging.getLogger(__name__)

WSCleanModel = namedtuple(
    "WSCleanModel",
    ["source_type", "radec", "flux", "spi", "ref_freq", "log_poly", "gauss_shape"],
)


def import_from_wsclean(
    wsclean_comp_list, include_regions=[], point_only=False, num=None
):
    """
    Imports sources from wsclean, sorted from brightest to faintest.
    If ``include_regions`` is specified only those sources within
    are returned. Similarly if ``num`` is specified, ``num`` brightest
    sources are returned.

    Parameters
    ----------
    wsclean_comp_list : str
        wsclean component list file.
    include_regions : List[:class:`regions.CircleSkyRegion`]
        List of valid region's. Only sources within these regions
        will be loaded.
        Defaults to ``[]`` in which case, all sources are loaded.
    point_only : bool, optional
        Only include point sources. Defaults to False
    num :  integer, optional
        Number of sources to include.
        If ``None`` all sources are returned.

    Returns
    -------
    source_type : :class:`numpy.ndarray`
        Source types of shape :code:`(source,)`
    radec : :class:`numpy.ndarray`
        Source radec coordinates of shape :code:`(source, 2)`.
    stokes : :class:`numpy.ndarray`
        Source stokes paramters of shape :code:`(source, 4)`.
    spectral_index : :class:`numpy.ndarray`
        Spectral index of shape :code:`(source, spi)`
    ref_freq : :class:`numpy.ndarray`
        Reference frequency of shape :code:`(source,)`
    log_si : bool
        Boolean indicated whether to use the logarithmic spectral index.
    gauss_shape : :class:`numpy.ndarray`
        Gaussian shape parameters of shape :code:`(source, 3)`.
        The three components are MajorAxis, MinorAxis and Orientation,
        respectively.
    """

    wsclean_comps = {
        column: np.asarray(values) for column, values in load(wsclean_comp_list)
    }

    if np.unique(wsclean_comps["LogarithmicSI"]).shape[0] > 1:
        raise ValueError(
            "Can't handle Mixed log and ordinary polynomial "
            "coefficients in '%s'" % wsclean_comp_list
        )

    # Sort components by descending flux
    sort_index = np.ascontiguousarray(np.argsort(wsclean_comps["I"])[::-1])

    # re-sort all entries using this
    wsclean_comps = {key: value[sort_index] for key, value in wsclean_comps.items()}

    log.info("%s contains %d components", wsclean_comp_list, len(wsclean_comps["Type"]))

    # make mask of sources to include
    include = np.ones_like(wsclean_comps["Type"], bool)

    if include_regions:
        include[:] = False
        # NB: regions is *supposed* to have a sensible "contains" interface,
        # but it doesn't work as of Mar 2019. So hacking
        # a kludge for circular regions for now
        from regions import CircleSkyRegion

        if not all([type(reg) is CircleSkyRegion for reg in include_regions]):
            raise ValueError("Only circular DS( regions supported for now")

        coord = SkyCoord(
            wsclean_comps["Ra"],
            wsclean_comps["Dec"],
            unit="rad",
            frame=include_regions[0].center.frame,
        )
        include = (
            coord.separation(include_regions[0].center) <= include_regions[0].radius
        )

        for reg in include_regions[1:]:
            include |= coord.separation(reg.center) <= reg.radius

        log.info(
            "%d of which fall within the %d inclusive regions",
            include.sum(),
            len(include_regions),
        )

    # select points
    if point_only:
        include &= wsclean_comps["Type"] == "POINT"
        log.info("%d of which are point sources", include.sum())

    # apply filters to component list
    wsclean_comps = {key: value[include] for key, value in wsclean_comps.items()}

    # select limited number if asked
    if num is not None:
        wsclean_comps = {key: value[:num] for key, value in wsclean_comps.items()}
        log.info("Selecting up %d brightest sources", num)

    # print if small subset
    if (num is not None and num < 100) or include_regions:
        flux_types = zip(wsclean_comps["I"], wsclean_comps["Type"])
        it = enumerate(sorted(flux_types, reverse=True))
        for i, (flux, src_type) in it:
            log.info("%d: %s %s Jy", i, src_type, flux)

    log.info(
        "Total flux of %d selected components is %f Jy",
        len(wsclean_comps["I"]),
        wsclean_comps["I"].sum(),
    )

    # Create radec array
    radec = np.concatenate(
        (wsclean_comps["Ra"][:, None], wsclean_comps["Dec"][:, None]), axis=1
    )

    # Create gaussian shapes
    gauss_shape = np.stack(
        (
            wsclean_comps["MajorAxis"],
            wsclean_comps["MinorAxis"],
            wsclean_comps["Orientation"],
        ),
        axis=-1,
    )

    return WSCleanModel(
        wsclean_comps["Type"],
        radec,
        wsclean_comps["I"],
        wsclean_comps["SpectralIndex"],
        wsclean_comps["ReferenceFrequency"],
        wsclean_comps["LogarithmicSI"],
        gauss_shape,
    )
