#!/usr/bin/env python
# encoding: utf-8
#
# @Author: José Sánchez-Gallego
# @Date: Nov 3, 2017
# @Filename: spaxel.py
# @License: BSD 3-Clause
# @Copyright: José Sánchez-Gallego


from __future__ import absolute_import, division, print_function

import abc
import inspect
import itertools
import warnings

import numpy as np
import six

import marvin
import marvin.core.exceptions
import marvin.core.marvin_pickle
import marvin.tools.cube
import marvin.tools.maps
import marvin.tools.modelcube
import marvin.utils.general.general
from marvin.core.exceptions import MarvinBreadCrumb, MarvinError, MarvinUserWarning
from marvin.utils.datamodel.dap import datamodel as dap_datamodel
from marvin.utils.datamodel.drp import datamodel as drp_datamodel
from marvin.utils.general.structs import FuzzyDict


breadcrumb = MarvinBreadCrumb()


def spaxel_factory(cls, *args, **kwargs):
    """A factory that returns the right type of spaxel depending on input.

    Based on the input values, determines if the resulting spaxels should be
    binned or unbinned, returning a `.Spaxel` or a `.Bin` respectively.
    This function is intended for overrding the ``__call__`` method in the
    `abc.ABCMeta` metacalass. The reason is that we want `.SpaxelBase` to have
    `abstract methods <abc.abstractmethod>` while also being a factory.
    See `this stack overflow <https://stackoverflow.com/a/5961102>`_ for
    details in the implementation of the ``__call__`` factory pattern.

    It can be used as::

        SpaxelABC = abc.ABCMeta
        SpaxelABC.__call__ = region_factory


        class SpaxelBase(object, metaclass=RegionABC):
            ...

    Note that this will override ``__call__`` everywhere else where
    `abc.ABCMeta` is used, but since it only overrides the default behaviour
    when the input class is `.SpaxelBase`, that should not be a problem.

    """

    if cls is not SpaxelBase:
        return type.__call__(cls, *args, **kwargs)

    cube = kwargs.pop('cube', None)
    maps = kwargs.pop('maps', None)
    modelcube = kwargs.pop('modelcube', None)
    lazy = kwargs.pop('lazy', False)

    plateifu = kwargs.pop('plateifu', (getattr(cube, 'plateifu', None) or
                                       getattr(maps, 'plateifu', None) or
                                       getattr(modelcube, 'plateifu', None)))

    mangaid = kwargs.pop('mangaid', (getattr(cube, 'mangaid', None) or
                                     getattr(maps, 'mangaid', None) or
                                     getattr(modelcube, 'mangaid', None)))

    release = kwargs.pop('release', (getattr(cube, 'release', None) or
                                     getattr(maps, 'release', None) or
                                     getattr(modelcube, 'release', None)))

    spaxel_kwargs = kwargs.copy()
    spaxel_kwargs.update(cube=cube, maps=maps, modelcube=modelcube,
                         plateifu=plateifu, mangaid=mangaid, release=release,
                         lazy=lazy)

    if not cube and not maps and not modelcube:
        raise MarvinError('no inputs defined.')

    if not maps and not modelcube:
        return Spaxel(*args, **spaxel_kwargs)

    if (isinstance(maps, marvin.tools.maps.Maps) or
            isinstance(modelcube, marvin.tools.modelcube.ModelCube)):
        bintype = getattr(maps, 'bintype', None) or getattr(modelcube, 'bintype', None)
        spaxel_kwargs.update(bintype=bintype)
        if bintype.binned:
            return Bin(*args, **spaxel_kwargs)
        else:
            return Spaxel(*args, **spaxel_kwargs)

    if maps:
        maps = marvin.tools.maps.Maps((maps if maps is not True else None) or plateifu or mangaid,
                                      release=release, **kwargs)
        spaxel_kwargs.update(bintype=maps.bintype)
        if maps.bintype.binned:
            return Bin(*args, **spaxel_kwargs)

    if modelcube:
        modelcube = marvin.tools.modelcube.ModelCube(
            (modelcube if modelcube is not True else None) or
            plateifu or mangaid, release=release, **kwargs)
        spaxel_kwargs.update(bintype=modelcube.bintype)
        if modelcube.bintype.binned:
            return Bin(*args, **spaxel_kwargs)

    return Spaxel(*args, **spaxel_kwargs)

    raise MarvinError('you have reached the end of the SpaxelBase logic. '
                      'This should never happen!')


# Overrides the __call__ method in abc.ABC.
SpaxelABC = abc.ABCMeta
SpaxelABC.__call__ = spaxel_factory


class DataModel(object):
    """A single object that holds the DRP and DAP datamodel."""

    def __init__(self, release):

        self.drp = drp_datamodel[release]
        self.dap = dap_datamodel[release]


class SpaxelBase(six.with_metaclass(SpaxelABC, object)):
    """A base class that contains information about a spaxel.

    This class represents an spaxel with information from the reduced DRP
    spectrum, the DAP maps properties, and the model spectrum from the DAP
    logcube. A `.SpaxelBase` can be initialised with all or only part of that
    information, and either from a file, a database, or remotely via the
    Marvin API.

    The `~marvin.tools.cube.Cube`, `~marvin.tools.maps.Maps` , and
    `~marvin.tools.modelcube.ModelCube` quantities for the spaxel are available
    in ``cube_quantities``, ``maps_quantities``, and ``modelcube_quantities``,
    respectively. For convenience, the quantities can also be accessed directly
    from the `.SpaxelBase` itself (e.g., ``spaxel.emline_gflux_ha_6465``).

    Parameters:
        x,y (int):
            The `x` and `y` coordinates of the spaxel in the cube (0-indexed).
        mangaid (str):
            The mangaid of the cube/maps/modelcube of the spaxel to load.
        plateifu (str):
            The plate-ifu of the cube/maps/modelcube of the spaxel to load
            (either ``mangaid`` or ``plateifu`` can be used, but not both).
        cube (`~marvin.tools.cube.Cube` object or path or bool):
            If ``cube`` is a `~marvin.tools.cube.Cube` object, that
            cube will be used for the `.SpaxelBase` instantiation. This mode
            is mostly intended for `~marvin.utils.general.general.getSpaxel`
            as it significantly improves loading time. Otherwise, ``cube`` can
            be ``True`` (default), in which case a cube will be instantiated
            using the input ``filename``, ``mangaid``, or ``plateifu``. If
            ``cube=False``, no cube will be used and the cube associated
            quantities will not be available. ``cube`` can also be the
            path to the DRP cube to use.
        maps (`~marvin.tools.maps.Maps` object or path or bool)
            As ``cube`` but for the DAP measurements corresponding to the
            spaxel in the `.Maps`.
        modelcube (`marvin.tools.modelcube.ModelCube` object or path or bool)
            As ``maps`` but for the DAP measurements corresponding to the
            spaxel in the `.ModelCube`.
        lazy (bool):
            If ``False``, the spaxel data is loaded on instantiation.
            Otherwise, only the metadata is created. The associated quantities
            can be then loaded by calling `.SpaxelBase.load()`.
        kwargs (dict):
            Arguments to be passed to `.Cube`, `.Maps`, and `.ModelCube`
            when (and if) they are initialised.

    Attributes:
        cube_quantities (`~marvin.utils.general.structs.FuzzyDict`):
            A querable dictionary with the `.Spectrum` quantities
            derived from `.Cube` and matching ``x, y``.
        datamodel (object):
            An object contianing the DRP and DAP datamodels.
        maps_quantities (`~marvin.utils.general.structs.FuzzyDict`):
            A querable dictionary with the `.AnalysisProperty` quantities
            derived from `.Maps` and matching ``x, y``.
        model_quantities (`~marvin.utils.general.structs.FuzzyDict`):
            A querable dictionary with the `.Spectrum` quantities
            derived from `.ModelCube` and matching ``x, y``.
        ra,dec (float):
            Right ascension and declination of the spaxel. Not available until
            the spaxel has been `loaded <.SpaxelBase.load>`.

    """

    def __init__(self, x, y, mangaid=None, plateifu=None,
                 cube=True, maps=True, modelcube=True, lazy=False, **kwargs):

        self.cube_quantities = FuzzyDict({})
        self.maps_quantities = FuzzyDict({})
        self.modelcube_quantities = FuzzyDict({})

        self._cube = cube
        self._maps = maps
        self._modelcube = modelcube

        self._plateifu = plateifu
        self._mangaid = mangaid

        self.kwargs = kwargs

        if not self._cube and not self._maps and not self._modelcube:
            raise MarvinError('either cube, maps, or modelcube must be True or '
                              'a Marvin Cube, Maps, or ModelCube object must be specified.')

        self._parent_shape = None

        # drop breadcrumb
        breadcrumb.drop(message='Initializing MarvinSpaxel {0}'.format(self.__class__),
                        category=self.__class__)

        self.x = int(x)
        self.y = int(y)

        self.loaded = False
        self.datamodel = None

        if lazy is False:
            self.load()

    def __dir__(self):

        class_members = list(list(zip(*inspect.getmembers(self.__class__)))[0])
        instance_attr = list(self.__dict__.keys())

        items = self.cube_quantities.__dir__()
        items += self.maps_quantities.__dir__()
        items += self.modelcube_quantities.__dir__()
        items += class_members + instance_attr

        return sorted(items)

    def __getattr__(self, value):

        _getattr = super(SpaxelBase, self).__getattribute__

        if value in _getattr('cube_quantities'):
            return _getattr('cube_quantities')[value]
        if value in _getattr('maps_quantities'):
            return _getattr('maps_quantities')[value]
        if value in _getattr('modelcube_quantities'):
            return _getattr('modelcube_quantities')[value]

        return super(SpaxelBase, self).__getattribute__(value)

    @abc.abstractmethod
    def __repr__(self):

        return '<SpaxelBase>'

    def _check_versions(self, attr):
        """Checks that all input object have the same versions.

        Runs sanity checks to make sure that ``attr`` has the same value
        in the input `.Cube`, `.Maps`, and `.ModelCube`. Returns the value
        for the attribute or ``None`` if the attribute does not exist.

        """

        out_value = None

        inputs = []
        for obj in [self._cube, self._maps, self._modelcube]:
            if obj is not None and not isinstance(obj, bool):
                inputs.append(obj)

        if len(inputs) == 1:
            return getattr(inputs[0], attr, None)

        for obj_a, obj_b in itertools.combinations(inputs, 2):
            if hasattr(obj_a, attr) and hasattr(obj_b, attr):
                assert getattr(obj_a, attr) == getattr(obj_b, attr), \
                    'inconsistent {!r} between {!r} and {!r}'.format(attr, obj_a, obj_b)
            out_value = getattr(obj_a, attr, None) or getattr(obj_b, attr, None)

        return out_value

    def _set_radec(self):
        """Calculates ra and dec for this spaxel."""

        self.ra = None
        self.dec = None

        for obj in [self._cube, self._maps, self._modelcube]:
            if hasattr(obj, 'wcs'):
                if obj.wcs.naxis == 2:
                    self.ra, self.dec = obj.wcs.wcs_pix2world([[self.x, self.y]], 0)[0]
                elif obj.wcs.naxis == 3:
                    self.ra, self.dec, __ = obj.wcs.wcs_pix2world([[self.x, self.y, 0]], 0)[0]

    def load(self, force=None):
        """Loads the spaxel data.

        Loads the spaxel data for cubes/maps/modelcubes.  By default attempts
        to load whatever is specified when spaxels are instantianted from other
        Marvin Tools.  Can manually force load a data type with the force
        keyword.

        Parameters:
            force ({cube|maps|models}):
                Str datatype to force load
        """

        if self.loaded and force is None:
            warnings.warn('already loaded', MarvinUserWarning)
            return

        # setup what will be loaded
        self._setup_load(force=force)

        self._load_cube()
        self._load_maps()
        self._load_modelcube()

        self._set_radec()
        self.loaded = True

        for attr in ['mangaid', 'plateifu', 'release', 'bintype', 'template']:
            self._check_versions(attr)

        self._plateifu = self.plateifu
        self._mangaid = self.mangaid

        self.datamodel = DataModel(self.release)

    def save(self, path, overwrite=False):
        """Pickles the spaxel to a file.

        Parameters:
            path (str):
                The path of the file to which the `.Spaxel` will be saved.
                Unlike for other Marvin Tools that derive from
                `~marvin.tools.core.MarvinToolsClass`, ``path`` is
                mandatory for `.Spaxel.save` as there is no default path for a
                given spaxel.
            overwrite (bool):
                If True, and the ``path`` already exists, overwrites it.
                Otherwise it will fail.

        Returns:
            path (str):
                The realpath to which the file has been saved.

        """

        return marvin.core.marvin_pickle.save(self, path=path, overwrite=overwrite)

    @classmethod
    def restore(cls, path, delete=False):
        """Restores a Spaxel object from a pickled file.

        If ``delete=True``, the pickled file will be removed after it has been
        unplickled. Note that, for objects with ``data_origin='file'``, the
        original file must exists and be in the same path as when the object
        was first created.

        """

        return marvin.core.marvin_pickle.restore(path, delete=delete)

    def _load_cube(self):
        """Loads the cube and the associated quantities."""

        if self._cube is False or self._cube is None:
            self._cube = None
            return
        elif not isinstance(self._cube, marvin.tools.cube.Cube):
            self._cube = self.getCube()

        self._parent_shape = self._cube._shape

        self.cube_quantities = self._cube._get_spaxel_quantities(self.x, self.y)

    def _load_maps(self):
        """Loads the cube and the properties."""

        if self._maps is False or self._maps is None:
            self._maps = None
            return
        elif not isinstance(self._maps, marvin.tools.maps.Maps):
            self._maps = self.getMaps()

        self._parent_shape = self._maps._shape

        self.maps_quantities = self._maps._get_spaxel_quantities(self.x, self.y)

    def _load_modelcube(self):
        """Loads the modelcube and associated arrays."""

        if self._modelcube is False or self._modelcube is None:
            self._modelcube = None
            return
        elif not isinstance(self._modelcube, marvin.tools.modelcube.ModelCube):

            if self.release == 'MPL-4':
                warnings.warn('ModelCube cannot be instantiated for MPL-4.', MarvinUserWarning)
                self._modelcube = None
                return

            self._modelcube = self.getModelCube()

        self._parent_shape = self._modelcube._shape

        self.modelcube_quantities = self._modelcube._get_spaxel_quantities(self.x, self.y)

    def _setup_load(self, force=None):
        ''' Setup the loading for cube/maps/models

        Can manually force load a data type.

        Parameters:
            force ({cube|maps|models}):
                Str datatype to force load
        '''

        # get what is currently loaded
        flag = self._get_loaded()

        # only do something if force to set.  Otherwise do the original nothing.
        if force:
            assert force in ['cube', 'maps', 'models'], 'force can only be cube, maps, or models'

            if force in flag:
                warnings.warn('{0} is already loaded'.format(force), MarvinUserWarning)
            else:
                if force == 'cube' and not self._cube:
                    self._cube = True
                elif force == 'maps' and not self._maps:
                    self._maps = True
                elif force == 'models' and not self._modelcube:
                    self._modelcube = True
                else:
                    # do nothing
                    pass

    def _get_loaded(self):
        ''' Get a flag to indicate what is or is not loaded '''

        dataloads = [self._cube, self._maps, self._modelcube]
        flags = np.array(['cube', 'maps', 'models'])
        load_idx = np.where(dataloads)[0]
        flag = '/'.join(flags[load_idx])

        return flag

    def getCube(self):
        """Returns the associated `~marvin.tools.cube.Cube`"""

        if isinstance(self._cube, marvin.tools.cube.Cube):
            return self._cube

        cube_kwargs = self.kwargs.copy()
        cube_kwargs.pop('bintype', None)
        cube_kwargs.pop('template', None)

        return marvin.tools.cube.Cube(
            (self._cube if self._cube is not True else None) or self._plateifu or self._mangaid,
            **cube_kwargs)

    def getMaps(self):
        """Returns the associated `~marvin.tools.maps.Maps`"""

        if isinstance(self._maps, marvin.tools.maps.Maps):
            return self._maps

        maps_kwargs = self.kwargs.copy()

        bintype = maps_kwargs.pop('bintype', None) or self.bintype
        template = maps_kwargs.pop('template', None) or self.template
        release = maps_kwargs.pop('release', None) or self.release

        return marvin.tools.maps.Maps(
            (self._maps if self._maps is not True else None) or self._plateifu or self._mangaid,
            bintype=bintype, template=template, release=release, **maps_kwargs)

    def getModelCube(self):
        """Returns the associated `~marvin.tools.modelcube.ModelCube`"""

        if isinstance(self._modelcube, marvin.tools.modelcube.ModelCube):
            return self._modelcube

        modelcube_kwargs = self.kwargs.copy()

        bintype = modelcube_kwargs.pop('bintype', None) or self.bintype
        template = modelcube_kwargs.pop('template', None) or self.template
        release = modelcube_kwargs.pop('release', None) or self.release

        return marvin.tools.modelcube.ModelCube(
            ((self._modelcube if self._modelcube is not True else None) or
             self._plateifu or self._mangaid),
            bintype=bintype, template=template, release=release, **modelcube_kwargs)

    @property
    def plateifu(self):
        """Returns the plateifu."""

        return self._check_versions('plateifu')

    @property
    def mangaid(self):
        """Returns the mangaid."""

        return self._check_versions('mangaid')

    @property
    def release(self):
        """Returns the release."""

        return self._check_versions('release')

    @property
    def bintype(self):
        """Returns the bintype."""

        return self._check_versions('bintype')

    @property
    def template(self):
        """Returns the template."""

        return self._check_versions('template')

    @property
    def quality_flags(self):
        """Bundle Cube DRP3QUAL and Maps DAPQUAL flags."""

        drp3qual = self.datamodel.drp.bitmasks['MANGA_DRP3QUAL']
        cube = self.getCube()
        drp3qual.mask = int(cube.header['DRP3QUAL'])
        qual_flags = [drp3qual]

        if self.release != 'MPL-4':
            qual_flags.append(self.datamodel.dap.bitmasks['MANGA_DAPQUAL'])

        return qual_flags


class Spaxel(SpaxelBase):
    """A class representing an unbinned spaxel.

    This subclass of `.SpaxelBase` represents an spaxel belonging to an
    unbinned `~marvin.tools.maps.Maps` and `~marvin.tools.modelcube.ModelCube`.
    If initialised directly, a `.Spaxel` will fail if the input data correspond
    to a binned maps or modelcube object.

    Refer to the documentation of `.SpaxelBase` for information about the valid
    parameters and methods.

    """

    def __init__(self, *args, **kwargs):

        super(Spaxel, self).__init__(*args, **kwargs)

        if isinstance(self._maps, marvin.tools.maps.Maps):
            assert self._maps.is_binned() is False, 'a Spaxel cannot be binned.'

        if isinstance(self._modelcube, marvin.tools.modelcube.ModelCube):
            assert self._modelcube.is_binned() is False, 'a Spaxel cannot be binned.'

    def __repr__(self):
        """Spaxel representation."""

        if not self.loaded:
            return '<Marvin Spaxel (x={0.x:d}, y={0.y:d}, loaded=False)'.format(self)

        # Gets the coordinates relative to the centre of the cube/maps.
        y_mid, x_mid = np.array(self._parent_shape) / 2.
        x_centre = int(self.x - x_mid)
        y_centre = int(self.y - y_mid)

        return ('<Marvin Spaxel (plateifu={0.plateifu}, x={0.x:d}, y={0.y:d}; '
                'x_cen={1:d}, y_cen={2:d}, loaded={3})>'.format(self, x_centre, y_centre,
                                                                self._get_loaded()))


class Bin(SpaxelBase):
    """A class that represents a bin.

    Attributes
    ----------
    spaxels : list
        A list of `.Spaxel` instances that provides access to the unbinned
        spaxels that are associated with the bin.

    """

    def __init__(self, *args, **kwargs):

        self.binid = None
        self.spaxels = None

        super(Bin, self).__init__(*args, **kwargs)

        if isinstance(self._maps, marvin.tools.maps.Maps):
            assert self._maps.is_binned() is True, 'a Spaxel cannot be unbinned.'

        if isinstance(self._modelcube, marvin.tools.modelcube.ModelCube):
            assert self._modelcube.is_binned() is True, 'a Spaxel cannot be unbinned.'

        assert (isinstance(self._maps, marvin.tools.maps.Maps) or
                isinstance(self._modelcube, marvin.tools.modelcube.ModelCube)), \
            'a Bin must have a Maps or a ModelCube.'

    def load(self):
        """Loads quantities and spaxels."""

        super(Bin, self).load()
        self._create_spaxels()

    def __repr__(self):
        """Spaxel representation."""

        if not self.loaded:
            return '<Marvin Bin (x={0.x:d}, y={0.y:d}, loaded=False)'.format(self)

        # Gets the coordinates relative to the centre of the cube/maps.
        y_mid, x_mid = np.array(self._parent_shape) / 2.
        x_centre = int(self.x - x_mid)
        y_centre = int(self.y - y_mid)

        return ('<Marvin Bin (plateifu={0.plateifu}, x={0.x:d}, y={0.y:d}; '
                'x_cen={1:d}, y_cen={2:d}, n_spaxels={3}, loaded={4})>'.format(self, x_centre,
                                                                               y_centre, len(self.spaxels),
                                                                               self._get_loaded()))

    def _create_spaxels(self):
        """Returns a list of the unbinned spaxels associated with this bin."""

        if self._maps:
            binid_map = self._maps.get_binid()
        elif self._modelcube:
            binid_map = self._modelcube.get_binid()

        # TODO: this may be an overkill and an extra API call. Remove it?
        if self._maps and self._modelcube:
            assert np.all(self._maps.get_binid() == self._modelcube.get_binid()), \
                'inconsistent binid arrays between Maps and ModelCube.'

        self.binid = binid_map[self.y, self.x]

        if self.binid < 0:
            raise MarvinError('coordinates ({}, {}) do not correspond to a valid binid.'
                              .format(self.x, self.y))

        spaxel_coords = zip(*np.where(binid_map == self.binid))
        self.spaxels = []

        for ii, jj in spaxel_coords:
            self.spaxels.append(Spaxel(x=jj, y=ii, plateifu=self.plateifu, release=self.release,
                                       cube=self._cube, maps=True if self._maps else False,
                                       modelcube=True if self._modelcube else False,
                                       bintype=None, template=self.template, lazy=True))

    def load_all(self):
        """Loads all the spaxels."""

        for spaxel in self.spaxels:
            spaxel.load()
