# Copyright (c) 2015-2020 by the parties listed in the AUTHORS file.
# All rights reserved.  Use of this source code is governed by
# a BSD-style license that can be found in the LICENSE file.

import warnings

import numpy as np

from scipy.constants import arcmin

import healpy as hp

from ..mpi import MPIShared

from ..timing import function_timer

from .._libtoast import fast_scanning_float32

DTYPE = np.float32


@function_timer
def plug_holes(m, verbose=False, in_place=True, nest=False):
    """Use simple downgrading to derive estimates of the missing pixel values"""
    nbad_start = np.sum(np.isclose(m, hp.UNSEEN))

    if nbad_start == m.size:
        if verbose:
            print("plug_holes: All map pixels are empty. Cannot plug holes", flush=True)
        return

    if nbad_start == 0:
        return

    nside = hp.get_nside(m)
    npix = m.size
    if nest:
        mnest = m.copy()
    else:
        mnest = hp.reorder(m, r2n=True)

    lowres = mnest
    nside_lowres = nside
    bad = np.isclose(mnest, hp.UNSEEN)
    while np.any(bad) and nside_lowres > 1:
        nside_lowres //= 2
        lowres = hp.ud_grade(lowres, nside_lowres, order_in="NESTED")
        hires = hp.ud_grade(lowres, nside, order_in="NESTED")
        bad = np.isclose(mnest, hp.UNSEEN)
        mnest[bad] = hires[bad]

    nbad_end = np.sum(bad)

    if nbad_end != 0:
        mn = np.mean(mnest[np.logical_not(bad)])
        mnest[bad] = mn

    if not in_place:
        m = m.copy()
    if nest:
        m[:] = mnest
    else:
        m[:] = hp.reorder(mnest, n2r=True)

    if verbose and nbad_start != 0:
        print(
            "plug_holes: Filled {} missing pixels ({:.2f}%), lowest "
            "resolution was Nside={}.".format(
                nbad_start, (100.0 * nbad_start) // npix, nside_lowres
            )
        )
    return m


class MapSampler:
    """
    MapSampler objects store maps in the node shared memory and allow
    bilinear interpolation of the maps into TOD.
    """

    @function_timer
    def __init__(
        self,
        map_path,
        pol=False,
        pol_fwhm=None,
        no_temperature=False,
        dtype=None,
        verbose=False,
        nside=None,
        comm=None,
        cache=None,
        preloaded_map=None,
        buflen=1000000,
        nest=False,
    ):
        """
        Instantiate the map sampler object, load a healpix
        map in a file located at map_path

        if pol==True, reads I,Q,U maps from extensions 0, 1, 2
        """

        if not pol and no_temperature:
            raise RuntimeError("You cannot have pol=False, " "no_temperature=True")

        self.path = map_path
        self.pol = pol
        self.pol_fwhm = pol_fwhm
        self._map = None
        self._map_Q = None
        self._map_U = None
        self.nest = nest
        if nest:
            self.order = "NESTED"
        else:
            self.order = "RING"
        self.buflen = buflen
        # Output data type, internal is always DTYPE
        if dtype is not None:
            warnings.warn("MapSampler no longer supports dtype", DeprecationWarning)

        # Use healpy to load the map into memory.

        if comm is None:
            self.comm = None
            self.rank = 0
            self.ntask = 1
        else:
            self.comm = comm
            self.rank = comm.Get_rank()
            self.ntask = comm.Get_size()

        self.shmem = self.ntask > 1
        self.pol = pol

        if self.rank == 0:
            if self.pol:
                if preloaded_map is not None:
                    if no_temperature:
                        (self._map_Q, self._map_U) = np.array(
                            preloaded_map, dtype=DTYPE
                        )
                    else:
                        (self._map, self._map_Q, self._map_U) = np.array(
                            preloaded_map, dtype=DTYPE
                        )
                else:
                    if no_temperature:
                        self._map_Q, self._map_U = hp.read_map(
                            self.path,
                            field=[1, 2],
                            dtype=DTYPE,
                            verbose=verbose,
                            memmmap=True,
                            nest=self.nest,
                        )
                    else:
                        try:
                            self._map, self._map_Q, self._map_U = hp.read_map(
                                self.path,
                                field=[0, 1, 2],
                                dtype=DTYPE,
                                verbose=verbose,
                                memmap=True,
                                nest=self.nest,
                            )
                        except IndexError:
                            print(
                                "WARNING: {} is not polarized".format(self.path),
                                flush=True,
                            )
                            self.pol = False
                            self._map = hp.read_map(
                                self.path,
                                dtype=DTYPE,
                                verbose=verbose,
                                memmap=True,
                                nest=self.nest,
                            )

                if nside is not None:
                    if not no_temperature:
                        self._map = hp.ud_grade(
                            self._map,
                            nside,
                            dtype=DTYPE,
                            order_in=self.order,
                            order_out=self.order,
                        )
                    if self.pol:
                        self._map_Q = hp.ud_grade(
                            self._map_Q,
                            nside,
                            dtype=DTYPE,
                            order_in=self.order,
                            order_out=self.order,
                        )
                        self._map_U = hp.ud_grade(
                            self._map_U,
                            nside,
                            dtype=DTYPE,
                            order_in=self.order,
                            order_out=self.order,
                        )

                if self.pol_fwhm is not None:
                    if not no_temperature:
                        plug_holes(self._map, verbose=verbose, nest=self.nest)
                    if self.pol:
                        plug_holes(self._map_Q, verbose=verbose, nest=self.nest)
                        plug_holes(self._map_U, verbose=verbose, nest=self.nest)
            else:
                if preloaded_map is not None:
                    self._map = np.array(preloaded_map, dtype=DTYPE)
                else:
                    self._map = hp.read_map(
                        map_path,
                        field=[0],
                        dtype=DTYPE,
                        verbose=verbose,
                        memmap=True,
                        nest=self.nest,
                    )
                if nside is not None:
                    self._map = hp.ud_grade(
                        self._map,
                        nside,
                        dtype=DTYPE,
                        order_in=self.order,
                        order_out=self.order,
                    )
                plug_holes(self._map, verbose=verbose, nest=self.nest)

        if self.ntask > 1:
            self.pol = comm.bcast(self.pol, root=0)
            npix = 0
            if self.rank == 0:
                if self.pol:
                    npix = len(self._map_Q)
                else:
                    npix = len(self._map)
            npix = comm.bcast(npix, root=0)
            if self.shmem:
                shared = MPIShared((npix,), np.dtype(DTYPE), comm)
                if not no_temperature:
                    shared.set(self._map, (0,), fromrank=0)
                    self._map = shared
                if self.pol:
                    shared_Q = MPIShared((npix,), np.dtype(DTYPE), comm)
                    shared_Q.set(self._map_Q, (0,), fromrank=0)
                    self._map_Q = shared_Q
                    shared_U = MPIShared((npix,), np.dtype(DTYPE), comm)
                    shared_U.set(self._map_U, (0,), fromrank=0)
                    self._map_U = shared_U
            else:
                if self.rank != 0:
                    if not no_temperature:
                        self._map = np.zeros(npix, dtype=DTYPE)
                    if self.pol:
                        self._map_Q = np.zeros(npix, dtype=DTYPE)
                        self._map_U = np.zeros(npix, dtype=DTYPE)

                if not no_temperature:
                    comm.Bcast(self._map, root=0)
                if self.pol:
                    comm.Bcast(self._map_Q, root=0)
                    comm.Bcast(self._map_U, root=0)

        if self.pol:
            self.npix = len(self._map_Q[:])
        else:
            self.npix = len(self._map[:])
        self.nside = hp.npix2nside(self.npix)

        self.cache = cache
        self.instance = 0
        if self.cache is not None and not self.shmem:
            # Increase the instance counter until we find an unused
            # instance.  If the user did not want to store duplicates,
            # they would not have created two identical mapsampler
            # objects.
            while self.cache.exists(self._cachename("I")):
                self.instance += 1
            if not no_temperature:
                self._map = self.cache.put(self._cachename("I"), self._map)
            if self.pol:
                self._map_Q = self.cache.put(self._cachename("Q"), self._map_Q)
                self._map_U = self.cache.put(self._cachename("U"), self._map_U)

        if self.pol_fwhm is not None:
            self.smooth(self.pol_fwhm, pol_only=True)
        return

    @function_timer
    def smooth(self, fwhm, lmax=None, pol_only=False):
        """Smooth the map with a Gaussian kernel."""
        if self.rank == 0:
            if pol_only:
                print(
                    "Smoothing the polarization to {} arcmin".format(fwhm), flush=True
                )
            else:
                print("Smoothing the map to {} arcmin".format(fwhm), flush=True)

        if lmax is None:
            lmax = min(np.int(fwhm / 60 * 512), 2 * self.nside)

        # If the map is in node-shared memory, only the root process on each
        # node does the smoothing.
        if not self.shmem or self._map.nodecomm.rank == 0:
            if self.pol:
                m = np.vstack([self._map[:], self._map_Q[:], self._map_U[:]])
            else:
                m = self._map[:]
            if self.nest:
                m = hp.reorder(m, n2r=True)
            smap = hp.smoothing(m, fwhm=fwhm * arcmin, lmax=lmax, verbose=False)
            del m
            if self.nest:
                smap = hp.reorder(smap, r2n=True)
        else:
            # Convenience dummy variable
            smap = np.zeros([3, 12])

        if not pol_only:
            if self.shmem:
                self._map.set(smap[0].astype(DTYPE), (0,), fromrank=0)
            else:
                self._map[:] = smap[0]

        if self.pol:
            if self.shmem:
                self._map_Q.set(smap[1].astype(DTYPE), (0,), fromrank=0)
                self._map_U.set(smap[2].astype(DTYPE), (0,), fromrank=0)
            else:
                self._map_Q[:] = smap[1]
                self._map_U[:] = smap[2]

        self.pol_fwhm = fwhm
        return

    def _cachename(self, stokes):
        """
        Construct a cache name string for the selected Stokes map
        """
        return "{}_ns{:04}_{}_{:04}".format(
            self.path, self.nside, stokes, self.instance
        )

    @function_timer
    def __del__(self):
        """
        Explicitly free memory taken up in the cache.
        """
        if self.cache is not None:
            if self._map is not None:
                del self._map
                if not self.shmem:
                    self.cache.destroy(self._cachename("I"))
            if self.pol:
                del self._map_Q
                del self._map_U
                if not self.shmem:
                    self.cache.destroy(self._cachename("Q"))
                    self.cache.destroy(self._cachename("U"))
        return

    @function_timer
    def __iadd__(self, other):
        """Accumulate provided Mapsampler object with this one."""
        if self.shmem:
            # One process does the manipulation on each node
            self._map._nodecomm.Barrier()
            if self._map._noderank == 0:
                self._map._data[:] += other._map[:]
            if self.pol and other.pol:
                if self._map_Q._noderank == (1 % self._map_Q._nodeprocs):
                    self._map_Q._data[:] += other._map_Q[:]
                if self._map_U._noderank == (2 % self._map_U._nodeprocs):
                    self._map_U._data[:] += other._map_U[:]
            self._map._nodecomm.Barrier()
        else:
            self._map += other._map
            if self.pol and other.pol:
                self._map_Q += other._map_Q
                self._map_U += other._map_U
        return self

    @function_timer
    def __isub__(self, other):
        """Subtract provided Mapsampler object from this one."""
        if self.shmem:
            # One process does the manipulation on each node
            self._map._nodecomm.Barrier()
            if self._map._noderank == 0:
                self._map._data[:] -= other._map[:]
            if self.pol and other.pol:
                if self._map_Q._noderank == (1 % self._map_Q._nodeprocs):
                    self._map_Q._data[:] -= other._map_Q[:]
                if self._map_U._noderank == (2 % self._map_U._nodeprocs):
                    self._map_U._data[:] -= other._map_U[:]
            self._map._nodecomm.Barrier()
        else:
            self._map -= other._map
            if self.pol and other.pol:
                self._map_Q -= other._map_Q
                self._map_U -= other._map_U
        return self

    @function_timer
    def __imul__(self, other):
        """Scale the maps in this MapSampler object"""
        if self.shmem:
            # One process does the manipulation on each node
            self._map._nodecomm.Barrier()
            if self._map._noderank == 0:
                self._map._data[:] *= other
            if self.pol:
                if self._map_Q._noderank == (1 % self._map_Q._nodeprocs):
                    self._map_Q._data[:] *= other
                if self._map_U._noderank == (2 % self._map_U._nodeprocs):
                    self._map_U._data[:] *= other
            self._map._nodecomm.Barrier()
        else:
            self._map *= other
            if self.pol:
                self._map_Q *= other
                self._map_U *= other
        return self

    @function_timer
    def __itruediv__(self, other):
        """Divide the maps in this MapSampler object"""
        if self.shmem:
            self._map._nodecomm.Barrier()
            if self._map._noderank == 0:
                self._map._data[:] /= other
            if self.pol:
                if self._map_Q._noderank == (1 % self._map_Q._nodeprocs):
                    self._map_Q._data[:] /= other
                if self._map_U._noderank == (2 % self._map_U._nodeprocs):
                    self._map_U._data[:] /= other
            self._map._nodecomm.Barrier()
        else:
            self._map /= other
            if self.pol:
                self._map_Q /= other
                self._map_U /= other
        return self

    @function_timer
    def at(self, theta, phi, interp_pix=None, interp_weights=None):
        """
        Use healpy bilinear interpolation to interpolate the
        map.  User must make sure that coordinate system used
        for theta and phi matches the map coordinate system.
        """
        if self._map is None:
            raise RuntimeError("No temperature map to sample")

        n = len(theta)
        stepsize = self.buflen
        signal = np.zeros(n, dtype=np.float32)

        # DEBUG begin
        if np.any(theta < 0) or np.any(theta > np.pi):
            raise RuntimeError("bad theta")
        if np.any(phi < 0) or np.any(phi > 2 * np.pi):
            raise RuntimeError("bad phi")
        # DEBUG end

        for istart in range(0, n, stepsize):
            istop = min(istart + stepsize, n)
            ind = slice(istart, istop)
            if interp_pix is None or interp_weights is None:
                p, w = hp.get_interp_weights(
                    self.nside, theta[ind], phi[ind], nest=self.nest
                )
            else:
                p = np.ascontiguousarray(interp_pix[:, ind])
                w = np.ascontiguousarray(interp_weights[:, ind])
            buf = np.zeros(istop - istart, dtype=np.float64)
            fast_scanning_float32(buf, p, w, self._map[:])
            signal[ind] = buf
        return signal

    @function_timer
    def atpol(
        self,
        theta,
        phi,
        IQUweight,
        onlypol=False,
        interp_pix=None,
        interp_weights=None,
        pol=True,
        pol_deriv=False,
    ):
        """
        Use healpy bilinear interpolation to interpolate the
        map.  User must make sure that coordinate system used
        for theta and phi matches the map coordinate system.
        IQUweight is an array of shape (nsamp,3) returned by the
        pointing library that gives the weights of the I,Q, and U maps.

        Args:
            pol_deriv(bool):  Return the polarization angle derivative
                of the signal instead of the actual signal.

        """
        if onlypol and not self.pol:
            return None

        if not self.pol or not pol:
            return self.at(
                theta, phi, interp_pix=interp_pix, interp_weights=interp_weights
            )

        if np.shape(IQUweight)[1] != 3:
            raise RuntimeError(
                "Cannot sample polarized map with only " "intensity weights"
            )

        n = len(theta)
        stepsize = self.buflen
        signal = np.zeros(n, dtype=np.float32)

        for istart in range(0, n, stepsize):
            istop = min(istart + stepsize, n)
            ind = slice(istart, istop)

            if interp_pix is None or interp_weights is None:
                p, w = hp.get_interp_weights(
                    self.nside, theta[ind], phi[ind], nest=self.nest
                )
            else:
                p = np.ascontiguousarray(interp_pix[:, ind])
                w = np.ascontiguousarray(interp_weights[:, ind])

            weights = np.ascontiguousarray(IQUweight[ind].T)

            buf = np.zeros(istop - istart, dtype=np.float64)
            fast_scanning_float32(buf, p, w, self._map_Q[:])
            if pol_deriv:
                signal[ind] = -2 * weights[2] * buf
            else:
                signal[ind] = weights[1] * buf

            buf[:] = 0
            fast_scanning_float32(buf, p, w, self._map_U[:])
            if pol_deriv:
                signal[ind] += 2 * weights[1] * buf
            else:
                signal[ind] += weights[2] * buf

            if not onlypol:
                if self._map is None:
                    raise RuntimeError("No temperature map to sample")
                buf[:] = 0
                fast_scanning_float32(buf, p, w, self._map[:])
                signal[ind] += weights[0] * buf

        return signal
