# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function
import copy
import numpy as np
import healpy as hp
from scipy.interpolate import RegularGridInterpolator
from scipy.ndimage.interpolation import map_coordinates
from astropy.io import fits
from astropy.wcs import WCS
from astropy.table import Table
from astropy.coordinates import SkyCoord
from astropy.coordinates import Galactic, ICRS
import gammapy
import fermipy.utils as utils
import fermipy.wcs_utils as wcs_utils
import fermipy.hpx_utils as hpx_utils
import fermipy.fits_utils as fits_utils
from fermipy.hpx_utils import HPX, HpxToWcsMapping

def coadd_maps(geom, maps, preserve_counts=True):
    """Coadd a sequence of `~gammapy.maps.Map` objects."""

    # FIXME: This functionality should be built into the Map.coadd method
    map_out = gammapy.maps.Map.from_geom(geom)
    for m in maps:
        m_tmp = m
        if isinstance(m, gammapy.maps.HpxNDMap):
            if m.geom.order < map_out.geom.order:
                factor = map_out.geom.nside // m.geom.nside
                m_tmp = m.upsample(factor, preserve_counts=preserve_counts)
        map_out.coadd(m_tmp)

    return map_out

def make_coadd_map(maps, proj, shape, preserve_counts=True):

    if isinstance(proj, WCS):
        return make_coadd_wcs(maps, proj, shape)
    elif isinstance(proj, HPX):
        return make_coadd_hpx(maps, proj, shape, preserve_counts=preserve_counts)
    else:
        raise Exception("Can't co-add map of unknown type %s" % type(proj))

def make_coadd_wcs(maps, wcs, shape):
    data = np.zeros(shape)
    axes = wcs_utils.wcs_to_axes(wcs, shape)

    for m in maps:
        c = wcs_utils.wcs_to_coords(m.wcs, m.counts.shape)
        o = np.histogramdd(c.T, bins=axes[::-1], weights=np.ravel(m.counts))[0]
        data += o

    return Map(data, copy.deepcopy(wcs))

def make_coadd_hpx(maps, hpx, shape, preserve_counts=True):
    data = np.zeros(shape)
    axes = hpx_utils.hpx_to_axes(hpx, shape)
    for m in maps:
        if m.hpx.order != hpx.order:
            m_copy = m.ud_grade(hpx.order, preserve_counts)
        else:
            m_copy = m
        c = hpx_utils.hpx_to_coords(m_copy.hpx, m_copy.counts.shape)
        o = np.histogramdd(c.T, bins=axes, weights=np.ravel(m_copy.counts))[0]
        data += o
    return HpxMap(data, copy.deepcopy(hpx))

def read_map_from_fits(fitsfile, extname=None):
    """
    """
    proj, f, hdu = fits_utils.read_projection_from_fits(fitsfile, extname)
    if isinstance(proj, WCS):
        ebins = fits_utils.find_and_read_ebins(f)
        m = Map(hdu.data, proj, ebins=ebins)
    elif isinstance(proj, HPX):
        m = HpxMap.create_from_hdu(hdu, proj.ebins)
    else:
        raise Exception("Did not recognize projection type %s" % type(proj))
    return m

class Map_Base(object):
    """ Abstract representation of a 2D or 3D counts map."""

    def __init__(self, counts):
        self._counts = counts

    @property
    def counts(self):
        return self._counts

    @property
    def data(self):
        return self._counts

    @data.setter
    def data(self, val):
        if val.shape != self.data.shape:
            raise Exception('Wrong shape.')
        self._counts = val

    def get_pixel_skydirs(self):
        """Get a list of sky coordinates for the centers of every pixel. """
        raise NotImplementedError("MapBase.get_pixel_skydirs()")

    def get_pixel_indices(self, lats, lons):
        """Return the indices in the flat array corresponding to a set of coordinates """
        raise NotImplementedError("MapBase.get_pixel_indices()")

    def sum_over_energy(self):
        """Reduce a counts cube to a counts map by summing over the energy planes """
        raise NotImplementedError("MapBase.sum_over_energy()")

    def get_map_values(self, lons, lats, ibin=None):
        """Return the map values corresponding to a set of coordinates. """
        raise NotImplementedError("MapBase.get_map_values()")

    def interpolate(self, lon, lat, egy=None):
        """Return the interpolated map values corresponding to a set of coordinates. """
        raise NotImplementedError("MapBase.interpolate()")

class Map(Map_Base):
    """ Representation of a 2D or 3D counts map using WCS. """

    def __init__(self, counts, wcs, ebins=None):
        """
        Parameters
        ----------
        counts : `~numpy.ndarray`
            Counts array in row-wise ordering (LON is first dimension).
        """
        Map_Base.__init__(self, counts)
        self._wcs = wcs

        self._npix = counts.shape[::-1]

        if len(self._npix) == 3:
            self._xindex = 2
            self._yindex = 1
        elif len(self._npix) == 2:
            self._xindex = 1
            self._yindex = 0
        else:
            raise Exception('Wrong number of dimensions for Map object.')

        # if len(self._npix) != 3 and len(self._npix) != 2:
        #    raise Exception('Wrong number of dimensions for Map object.')

        self._width = np.array([np.abs(self.wcs.wcs.cdelt[0]) * self.npix[0],
                                np.abs(self.wcs.wcs.cdelt[1]) * self.npix[1]])
        self._pix_center = np.array([(self.npix[0] - 1.0) / 2.,
                                     (self.npix[1] - 1.0) / 2.])
        self._pix_size = np.array([np.abs(self.wcs.wcs.cdelt[0]),
                                   np.abs(self.wcs.wcs.cdelt[1])])

        self._skydir = SkyCoord.from_pixel(self._pix_center[0],
                                           self._pix_center[1],
                                           self.wcs)
        self._ebins = ebins
        if ebins is not None:
            self._ectr = np.exp(utils.edge_to_center(np.log(ebins)))
        else:
            self._ectr = None

    @property
    def wcs(self):
        return self._wcs

    @property
    def npix(self):
        return self._npix

    @property
    def skydir(self):
        """Return the sky coordinate of the image center."""
        return self._skydir

    @property
    def width(self):
        """Return the dimensions of the image."""
        return self._width

    @property
    def pix_size(self):
        """Return the pixel size along the two image dimensions."""
        return self._pix_size

    @property
    def pix_center(self):
        """Return the ROI center in pixel coordinates."""
        return self._pix_center

    @classmethod
    def create_from_hdu(cls, hdu, wcs):
        return cls(hdu.data.T, wcs)

    @classmethod
    def create_from_fits(cls, fitsfile, **kwargs):
        hdu = kwargs.get('hdu', 0)

        with fits.open(fitsfile) as hdulist:
            header = hdulist[hdu].header
            data = hdulist[hdu].data
            header = fits.Header.fromstring(header.tostring())
            wcs = WCS(header)

            ebins = None
            if 'ENERGIES' in hdulist:
                tab = Table.read(fitsfile, 'ENERGIES')
                ectr = np.array(tab.columns[0])
                ebins = np.exp(utils.center_to_edge(np.log(ectr)))
            elif 'EBOUNDS' in hdulist:
                tab = Table.read(fitsfile, 'EBOUNDS')
                emin = np.array(tab['E_MIN']) / 1E3
                emax = np.array(tab['E_MAX']) / 1E3
                ebins = np.append(emin, emax[-1])

        return cls(data, wcs, ebins)

    @classmethod
    def create(cls, skydir, cdelt, npix, coordsys='CEL', projection='AIT', ebins=None, differential=False):
        crpix = np.array([n / 2. + 0.5 for n in npix])

        if ebins is not None:
            if differential:
                nebins = len(ebins)
            else:
                nebins = len(ebins) - 1
            data = np.zeros(list(npix) + [nebins]).T
            naxis = 3
        else:
            data = np.zeros(npix).T
            naxis = 2

        wcs = wcs_utils.create_wcs(skydir, coordsys, projection,
                                   cdelt, crpix, naxis=naxis, energies=ebins)
        return cls(data, wcs, ebins=ebins)

    def create_image_hdu(self, name=None, **kwargs):
        return fits.ImageHDU(self.counts, header=self.wcs.to_header(),
                             name=name)

    def create_primary_hdu(self):
        return fits.PrimaryHDU(self.counts, header=self.wcs.to_header())

    def sum_over_energy(self):
        """ Reduce a 3D counts cube to a 2D counts map
        """
        # Note that the array is using the opposite convention from WCS
        # so we sum over axis 0 in the array, but drop axis 2 in the WCS object
        return Map(np.sum(self.counts, axis=0), self.wcs.dropaxis(2))

    def xypix_to_ipix(self, xypix, colwise=False):
        """Return the flattened pixel indices from an array multi-dimensional
        pixel indices.

        Parameters
        ----------
        xypix : list
            List of pixel indices in the order (LON,LAT,ENERGY).

        colwise : bool
            Use column-wise pixel indexing.
        """
        return np.ravel_multi_index(xypix, self.npix,
                                    order='F' if colwise else 'C',
                                    mode='raise')

    def ipix_to_xypix(self, ipix, colwise=False):
        """Return array multi-dimensional pixel indices from flattened index.

        Parameters
        ----------
        colwise : bool
            Use column-wise pixel indexing.
        """
        return np.unravel_index(ipix, self.npix,
                                order='F' if colwise else 'C')

    def ipix_swap_axes(self, ipix, colwise=False):
        """ Return the transposed pixel index from the pixel xy coordinates

        if colwise is True (False) this assumes the original index was
        in column wise scheme
        """
        xy = self.ipix_to_xypix(ipix, colwise)
        return self.xypix_to_ipix(xy, not colwise)

    def get_pixel_skydirs(self):
        """Get a list of sky coordinates for the centers of every pixel.

        """

        xpix = np.linspace(0, self.npix[0] - 1., self.npix[0])
        ypix = np.linspace(0, self.npix[1] - 1., self.npix[1])
        xypix = np.meshgrid(xpix, ypix, indexing='ij')
        return SkyCoord.from_pixel(np.ravel(xypix[0]),
                                   np.ravel(xypix[1]), self.wcs)

    def get_pixel_indices(self, lons, lats, ibin=None):
        """Return the indices in the flat array corresponding to a set of coordinates

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all energy bins.

        Returns
        ----------
        pixcrd : list
           Pixel indices along each dimension of the map.
        """
        lons = np.array(lons, ndmin=1)
        lats = np.array(lats, ndmin=1)

        if len(lats) != len(lons):
            raise RuntimeError('Map.get_pixel_indices, input lengths '
                               'do not match %i %i' % (len(lons), len(lats)))
        if len(self._npix) == 2:
            pix_x, pix_y = self._wcs.wcs_world2pix(lons, lats, 0)
            pixcrd = [np.floor(pix_x).astype(int), np.floor(pix_y).astype(int)]
        elif len(self._npix) == 3:
            all_lons = np.expand_dims(lons, -1)
            all_lats = np.expand_dims(lats, -1)
            if ibin is None:
                all_bins = (np.expand_dims(
                    np.arange(self.npix[2]), -1) * np.ones(lons.shape)).T
            else:
                all_bins = ibin

            l = self.wcs.wcs_world2pix(all_lons, all_lats, all_bins, 0)
            pix_x = l[0]
            pix_y = l[1]
            pixcrd = [np.floor(l[0]).astype(int), np.floor(l[1]).astype(int),
                      all_bins.astype(int)]

        return pixcrd

    def get_map_values(self, lons, lats, ibin=None):
        """Return the map values corresponding to a set of coordinates.

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all bins

        Returns
        ----------
        vals : numpy.ndarray((n))
           Values of pixels in the flattened map, np.nan used to flag
           coords outside of map
        """
        pix_idxs = self.get_pixel_indices(lons, lats, ibin)
        idxs = copy.copy(pix_idxs)

        m = np.empty_like(idxs[0], dtype=bool)
        m.fill(True)
        for i, p in enumerate(pix_idxs):
            m &= (pix_idxs[i] >= 0) & (pix_idxs[i] < self._npix[i])
            idxs[i][~m] = 0

        vals = self.counts.T[idxs]
        vals[~m] = np.nan
        return vals

    def interpolate(self, lon, lat, egy=None):

        if len(self.npix) == 2:
            pixcrd = self.wcs.wcs_world2pix(lon, lat, 0)
        else:
            if egy is None:
                egy = self._ectr
            pixcrd = self.wcs.wcs_world2pix(lon, lat, egy, 0)
            pixcrd[2] = np.array(utils.val_to_pix(np.log(self._ectr),
                                                  np.log(egy)), ndmin=1)

        points = []
        for npix in self.npix:
            points += [np.linspace(0, npix - 1., npix)]
        data = self.counts
        fn = RegularGridInterpolator(points, data.T,
                                     bounds_error=False,
                                     fill_value=None)
        return fn(np.column_stack(pixcrd))

    def interpolate_at_skydir(self, skydir):

        coordsys = wcs_utils.get_coordsys(self.wcs)
        if coordsys == 'CEL':
            skydir = skydir.transform_to('icrs')
            return self.interpolate(skydir.ra.deg, skydir.dec.deg)
        else:
            skydir = skydir.transform_to('galactic')
            return self.interpolate(skydir.l.deg, skydir.b.deg)

class HpxMap(Map_Base):
    """ Representation of a 2D or 3D counts map using HEALPix. """

    def __init__(self, counts, hpx):
        """ C'tor, fill with a counts vector and a HPX object """
        super(HpxMap, self).__init__(counts)
        self._hpx = hpx
        self._wcs2d = None
        self._hpx2wcs = None

    @property
    def hpx(self):
        return self._hpx

    @classmethod
    def create_from_hdu(cls, hdu, ebins):
        """ Creates and returns an HpxMap object from a FITS HDU.

        hdu    : The FITS
        ebins  : Energy bin edges [optional]
        """
        hpx = HPX.create_from_hdu(hdu, ebins)
        colnames = hdu.columns.names
        cnames = []
        if hpx.conv.convname == 'FGST_SRCMAP_SPARSE':
            pixs = hdu.data.field('PIX')
            chans = hdu.data.field('CHANNEL')
            keys = chans * hpx.npix + pixs
            vals = hdu.data.field('VALUE')
            nebin = len(ebins)
            data = np.zeros((nebin, hpx.npix))
            data.flat[keys] = vals
        else:
            for c in colnames:
                if c.find(hpx.conv.colstring) == 0:
                    cnames.append(c)
            nebin = len(cnames)
            data = np.ndarray((nebin, hpx.npix))
            for i, cname in enumerate(cnames):
                data[i, 0:] = hdu.data.field(cname)

        return cls(data, hpx)

    @classmethod
    def create_from_hdulist(cls, hdulist, **kwargs):
        """ Creates and returns an HpxMap object from a FITS HDUList

        extname : The name of the HDU with the map data
        ebounds : The name of the HDU with the energy bin data
        """
        extname = kwargs.get('hdu', hdulist[1].name)
        ebins = fits_utils.find_and_read_ebins(hdulist)
        return cls.create_from_hdu(hdulist[extname], ebins)

    @classmethod
    def create_from_fits(cls, fitsfile, **kwargs):
        hdulist = fits.open(fitsfile)
        return cls.create_from_hdulist(hdulist, **kwargs)

    def create_image_hdu(self, name=None, **kwargs):
        kwargs['extname'] = name
        return self.hpx.make_hdu(self.counts, **kwargs)

    def make_wcs_from_hpx(self, sum_ebins=False, proj='CAR', oversample=2,
                          normalize=True):
        """Make a WCS object and convert HEALPix data into WCS projection

        NOTE: this re-calculates the mapping, if you have already
        calculated the mapping it is much faster to use
        convert_to_cached_wcs() instead

        Parameters
        ----------
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting

        proj       : str
           WCS-projection

        oversample : int
           Oversampling factor for WCS map

        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)

        """
        self._wcs_proj = proj
        self._wcs_oversample = oversample
        self._wcs_2d = self.hpx.make_wcs(2, proj=proj, oversample=oversample)
        self._hpx2wcs = HpxToWcsMapping(self.hpx, self._wcs_2d)
        wcs, wcs_data = self.convert_to_cached_wcs(self.counts, sum_ebins,
                                                   normalize)
        return wcs, wcs_data

    def convert_to_cached_wcs(self, hpx_in, sum_ebins=False, normalize=True):
        """ Make a WCS object and convert HEALPix data into WCS projection

        Parameters
        ----------
        hpx_in     : `~numpy.ndarray`
           HEALPix input data
        sum_ebins  : bool
           sum energy bins over energy bins before reprojecting
        normalize  : bool
           True -> perserve integral by splitting HEALPix values between bins

        returns (WCS object, np.ndarray() with reprojected data)
        """
        if self._hpx2wcs is None:
            raise Exception('HpxMap.convert_to_cached_wcs() called '
                            'before make_wcs_from_hpx()')

        if len(hpx_in.shape) == 1:
            wcs_data = np.ndarray(self._hpx2wcs.npix)
            loop_ebins = False
            hpx_data = hpx_in
        elif len(hpx_in.shape) == 2:
            if sum_ebins:
                wcs_data = np.ndarray(self._hpx2wcs.npix)
                hpx_data = hpx_in.sum(0)
                loop_ebins = False
            else:
                wcs_data = np.ndarray((self.counts.shape[0],
                                       self._hpx2wcs.npix[0],
                                       self._hpx2wcs.npix[1]))
                hpx_data = hpx_in
                loop_ebins = True
        else:
            raise Exception('Wrong dimension for HpxMap %i' %
                            len(hpx_in.shape))

        if loop_ebins:
            for i in range(hpx_data.shape[0]):
                self._hpx2wcs.fill_wcs_map_from_hpx_data(
                    hpx_data[i], wcs_data[i], normalize)
                pass
            wcs_data.reshape((self.counts.shape[0], self._hpx2wcs.npix[
                             0], self._hpx2wcs.npix[1]))
            # replace the WCS with a 3D one
            wcs = self.hpx.make_wcs(3, proj=self._wcs_proj,
                                    energies=np.log10(self.hpx.ebins),
                                    oversample=self._wcs_oversample)
        else:
            self._hpx2wcs.fill_wcs_map_from_hpx_data(
                hpx_data, wcs_data, normalize)
            wcs_data.reshape(self._hpx2wcs.npix)
            wcs = self._wcs_2d

        return wcs, wcs_data

    def get_pixel_skydirs(self):
        """Get a list of sky coordinates for the centers of every pixel. """
        sky_coords = self._hpx.get_sky_coords()
        if self.hpx.coordsys == 'GAL':
            return SkyCoord(l=sky_coords.T[0], b=sky_coords.T[1], unit='deg', frame='galactic')
        else:
            return SkyCoord(ra=sky_coords.T[0], dec=sky_coords.T[1], unit='deg', frame='icrs')

    def get_pixel_indices(self, lats, lons):
        """Return the indices in the flat array corresponding to a set of coordinates """
        return self._hpx.get_pixel_indices(lats, lons)

    def sum_over_energy(self):
        """ Reduce a counts cube to a counts map """
        # We sum over axis 0 in the array, and drop the energy binning in the
        # hpx object
        return HpxMap(np.sum(self.counts, axis=0), self.hpx.copy_and_drop_energy())

    def get_map_values(self, lons, lats, ibin=None):
        """Return the indices in the flat array corresponding to a set of coordinates

        Parameters
        ----------
        lons  : array-like
           'Longitudes' (RA or GLON)

        lats  : array-like
           'Latitidues' (DEC or GLAT)

        ibin : int or array-like
           Extract data only for a given energy bin.  None -> extract data for all bins

        Returns
        ----------
        vals : numpy.ndarray((n))
           Values of pixels in the flattened map, np.nan used to flag
           coords outside of map
        """
        theta = np.pi / 2. - np.radians(lats)
        phi = np.radians(lons)

        pix = hp.ang2pix(self.hpx.nside, theta, phi, nest=self.hpx.nest)

        if self.data.ndim == 2:
            return self.data[:, pix] if ibin is None else self.data[ibin, pix]
        else:
            return self.data[pix]

    def interpolate(self, lon, lat, egy=None, interp_log=True):
        """Interpolate map values.

        Parameters
        ----------
        interp_log : bool
            Interpolate the z-coordinate in logspace.

        """

        if self.data.ndim == 1:
            theta = np.pi / 2. - np.radians(lat)
            phi = np.radians(lon)
            return hp.pixelfunc.get_interp_val(self.counts, theta,
                                               phi, nest=self.hpx.nest)
        else:
            return self._interpolate_cube(lon, lat, egy, interp_log)

    def _interpolate_cube(self, lon, lat, egy=None, interp_log=True):
        """Perform interpolation on a healpix cube.  If egy is None
        then interpolation will be performed on the existing energy
        planes.

        """

        shape = np.broadcast(lon, lat, egy).shape
        lon = lon * np.ones(shape)
        lat = lat * np.ones(shape)
        theta = np.pi / 2. - np.radians(lat)
        phi = np.radians(lon)
        vals = []
        for i, _ in enumerate(self.hpx.evals):
            v = hp.pixelfunc.get_interp_val(self.counts[i], theta,
                                            phi, nest=self.hpx.nest)
            vals += [np.expand_dims(np.array(v, ndmin=1), -1)]

        vals = np.concatenate(vals, axis=-1)

        if egy is None:
            return vals.T

        egy = egy * np.ones(shape)

        if interp_log:
            xvals = utils.val_to_pix(np.log(self.hpx.evals), np.log(egy))
        else:
            xvals = utils.val_to_pix(self.hpx.evals, egy)

        vals = vals.reshape((-1, vals.shape[-1]))
        xvals = np.ravel(xvals)
        v = map_coordinates(vals, [np.arange(vals.shape[0]), xvals],
                            order=1)
        return v.reshape(shape)

    def swap_scheme(self):
        """
        """
        hpx_out = self.hpx.make_swapped_hpx()
        if self.hpx.nest:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], n2r=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, n2r=True)
        else:
            if self.data.ndim == 2:
                data_out = np.vstack([hp.pixelfunc.reorder(
                    self.data[i], r2n=True) for i in range(self.data.shape[0])])
            else:
                data_out = hp.pixelfunc.reorder(self.data, r2n=True)
        return HpxMap(data_out, hpx_out)

    def expanded_counts_map(self):
        """ return the full counts map """
        if self.hpx._ipix is None:
            return self.counts

        output = np.zeros(
            (self.counts.shape[0], self.hpx._maxpix), self.counts.dtype)
        for i in range(self.counts.shape[0]):
            output[i][self.hpx._ipix] = self.counts[i]
        return output

    def explicit_counts_map(self, pixels=None):
        """ return a counts map with explicit index scheme

        Parameters
        ----------
        pixels : `np.ndarray` or None
            If set, grab only those pixels.
            If none, grab only non-zero pixels
        """
        # No pixel index, so build one
        if self.hpx._ipix is None:
            if self.data.ndim == 2:
                summed = self.counts.sum(0)
                if pixels is None:
                    nz = summed.nonzero()[0]
                else:
                    nz = pixels
                data_out = np.vstack(self.data[i].flat[nz]
                                     for i in range(self.data.shape[0]))
            else:
                if pixels is None:
                    nz = self.data.nonzero()[0]
                else:
                    nz = pixels
                data_out = self.data[nz]
            return (nz, data_out)
        else:
            if pixels is None:
                return (self.hpx._ipix, self.data)
        # FIXME, can we catch this
        raise RuntimeError(
            'HPX.explicit_counts_map called with pixels for a map that already has pixels')

    def sparse_counts_map(self):
        """ return a counts map with sparse index scheme
        """
        if self.hpx._ipix is None:
            flatarray = self.data.flattern()
        else:
            flatarray = self.expanded_counts_map()
        nz = flatarray.nonzero()[0]
        data_out = flatarray[nz]
        return (nz, data_out)

    def ud_grade(self, order, preserve_counts=False):
        """
        """
        new_hpx = self.hpx.ud_graded_hpx(order)
        if new_hpx.evals is None:
            nebins = 1
        else:
            nebins = len(new_hpx.evals)
        shape = self.counts.shape

        if preserve_counts:
            power = -2.
        else:
            power = 0

        if len(shape) == 1:
            new_data = hp.pixelfunc.ud_grade(self.counts,
                                             nside_out=new_hpx.nside,
                                             order_in=new_hpx.ordering,
                                             order_out=new_hpx.ordering,
                                             power=power)
        else:
            new_data = np.vstack([hp.pixelfunc.ud_grade(self.counts[i],
                                                        nside_out=new_hpx.nside,
                                                        order_in=new_hpx.ordering,
                                                        order_out=new_hpx.ordering,
                                                        power=power) for i in range(shape[0])])
        return HpxMap(new_data, new_hpx)


#!/usr/bin/env python
import emcee
import celerite
from celerite import terms

import numpy as np
import astropy.units as u
from astropy.modeling.blackbody import blackbody_lambda

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import seaborn as sns

from pylcurve.lcurve import Lcurve
from pylcurve.modeling import Model
import pylcurve.mcmc_utils as m
from pylcurve.utils import log_g, separation, get_limbdark_params, wd_to_bb_temp

band_wavs = dict(
    u=3543*u.AA,
    g=4770*u.AA,
    r=6231*u.AA,
    i=7625*u.AA,
    z=9134*u.AA
)

def scale_pulsation(band, temp):
    return blackbody_lambda(band_wavs[band], temp*u.K) / blackbody_lambda(band_wavs['g'], temp*u.K)

class EclipseLC(Model):
    parameter_names = ('t1', 't2', 'm1', 'm2', 'incl', 'r1', 'r2', 't0', 'per',
                       'pulse_omega', 'pulse_q', 'pulse_temp', 'pulse_amp')

    def __init__(self, model_file, lightcurves, *args, **kwargs):
        """
        A lightcurve model for an eclipsing DWD with pulsations

        Parameters
        ----------
        model_file: model containing LCURVE file with auxillary (fixed) params
        lightcurves: a dictionary of band: filename pairs

        The remaining parameters are either passed in as a list of arguments
        (in order) or specified as a dictionary:

        t1, t2 :  white dwarf temp in K
        m1, m2 : white dwarf masses in solar masses (constrained through RV prior)
        incl : inclination of system
        r1, r2 : white dwarf radii in solar radii
        t0 : mid-eclipse time of primary eclipse
        pulse_omega : frequency of pulsations
        pulse_q : q factor of pulsations
        pulse_temp : BB temperature of pulsations
        pulse_amp : amplitude of pulsations in g band
        """
        super().__init__(*args, **kwargs)
        self.lightcurves = lightcurves
        self.model_file = model_file

    def get_value(self, band):
        """
        Calculate lightcurve

        Parameters
        ----------
        band : string
            SDSS/HiPERCAM bandpass

        Returns
        -------
        ym : np.ndarray
            model values
        """
        # setup LCURVE file for this band
        lcurve_model = Lcurve(self.model_file)
        lcurve_pars = dict()

        log_g1 = log_g(self.m1, self.r1)
        log_g2 = log_g(self.m2, self.r2)
        a = separation(self.m1, self.m2, self.per)
        lcurve_pars['t1'] = wd_to_bb_temp(band, self.t1, log_g1)
        lcurve_pars['t2'] = wd_to_bb_temp(band, self.t2, log_g2)
        lcurve_pars['r1'] = self.r1/a  # scale from solar radii to separation units
        lcurve_pars['r2'] = self.r2/a  # scale from solar radii to separation units
        lcurve_pars['t0'] = self.t0
        lcurve_pars['period'] = self.per
        lcurve_pars['iangle'] = self.incl
        lcurve_pars['q'] = self.m1/self.m2
        lcurve_pars['wavelength'] = band_wavs[band].to_value(u.nm)
        lcurve_model.set(lcurve_pars)
        lcurve_model.set(get_limbdark_params(self.t1, log_g1, self.t2, log_g2, band))

        if not lcurve_model.ok():
            raise ValueError('invalid parameter combination')
        x, y, e, ym = lcurve_model(self.lightcurves[band])
        return ym

    def log_prior(self):
        """
        Prior probabilities
        """
        # first call parent class log_prior -> checks params in bounds
        val = super().log_prior()
        if np.isinf(val):
            return val

        # OK, we are within limits, check M1, M2 against RV constraints
        # K1 (brighter, hotter WD rv): 186.3 +/- 1.6 km/s
        # K2 (fainter, cooler WD rv): 213.6 +/- 4.6 km/s
        # mass ratio
        q_prior = m.Prior('gauss', 1.146, 0.027)
        q_act = self.m1/self.m2
        val += q_prior.ln_prob(q_act)
        # (m1 + m2) * sini**3 = (P/2piG) * (v1+v2)**3
        mt_prior = m.Prior('gauss', 0.66, 0.02)
        val += mt_prior.ln_prob((self.m1 + self.m2) * np.sin(self.incl)**3)

        # priors on t0, p from ephemeris
        prior = m.Prior('gauss', 57460.6510218, 0.0000010)
        val += prior.ln_prob(self.t0)

        prior = m.Prior('gauss', 0.09986526542, 0.00000000010)
        val += prior.ln_prob(self.per)

        return val

    def plot(self, ax, band, params, style='whole', dcolor='k', gpcolor='r'):
        """
        Plots data and model on axis

        style is either whole, model or residuals.
            'whole' plots the raw data and the full model (mean model + GP).

            'model' plots the mean model and the data after subtraction
             of the mean of the GP - i.e the data minus the pulsations

            'residuals' plots the data minus the mean model, together with the
            mean and range of the GP
        """
        self.set_parameter_vector(params)
        t, _, y, ye, _, _ = np.loadtxt(self.lightcurves[band]).T
        ym = self.get_value(band)

        pulsation_amp = self.pulse_amp * scale_pulsation(band, self.pulse_temp)
        gp = self.gpdict[band]
        gp.set_parameter_vector((np.log(pulsation_amp),
                                 np.log(self.pulse_q),
                                 np.log(self.pulse_omega)))

        samples = gp.sample_conditional(y-ym, t, size=300)
        mu = np.mean(samples, axis=0)
        std = np.std(samples, axis=0)

        toff = int(np.floor(np.min(t)))
        tplot = t - toff

        if style == 'whole':
            ax.errorbar(tplot, y, yerr=ye, fmt='none', color=dcolor, alpha=0.5)
            ax.plot(tplot, ym + mu, color=gpcolor, lw=2)
            ax.plot(tplot, ym, color='k', lw=2, ls=':')
            ax.fill_between(tplot, ym+mu+std, ym+mu-std, color=gpcolor, alpha=0.6)
        elif style == 'model':
            ax.errorbar(tplot, y-mu, yerr=ye, fmt='none', color=dcolor, alpha=0.5)
            ax.plot(tplot, ym, color=gpcolor, lw=2)
            ax.fill_between(tplot, ym+std, ym-std, color=gpcolor, alpha=0.6)
        elif style == 'residuals':
            ax.errorbar(tplot, y-ym, yerr=ye, fmt='none', color=dcolor, alpha=0.5)
            ax.plot(tplot, mu, color=gpcolor, lw=2)
            ax.fill_between(tplot, mu+std, mu-std, color=gpcolor, alpha=0.6)
        else:
            raise ValueError('style not recognised')

    def log_probability(self, params, band):
        """
        Calculate log of posterior probability

        Parameters
        -----------
        params : iterable
            list of parameter values
        band : string
            SDSS/HiPERCAM band
        """
        self.set_parameter_vector(params)
        t, _, y, ye, _, _ = np.loadtxt(self.lightcurves[band]).T

        # check model params are valid - checks against bounds
        lp = self.log_prior()
        if not np.isfinite(lp):
            return -np.inf

        # make a GP for this band, if it doesnt already exist
        if not hasattr(self, 'gpdict'):
            self.gpdict = dict()

        # Oscillation params
        pulsation_amp = self.pulse_amp * scale_pulsation(band, self.pulse_temp)

        if band not in self.gpdict:
            kernel = terms.SHOTerm(np.log(pulsation_amp),
                                   np.log(self.pulse_q),
                                   np.log(self.pulse_omega))
            gp = celerite.GP(kernel)
            gp.compute(t, ye)
            self.gpdict[band] = gp
        else:
            gp = self.gpdict[band]
            gp.set_parameter_vector((
                np.log(pulsation_amp),
                np.log(self.pulse_q),
                np.log(self.pulse_omega)
            ))
            gp.compute(t, ye)

        # now add prior of Gaussian process
        lp += gp.log_prior()
        if not np.isfinite(lp):
            return -np.inf

        try:
            ym = self.get_value(band)
        except ValueError as err:
            # invalid lcurve params
            print('warning: model failed ', err)
            return -np.inf
        else:
            return gp.log_likelihood(y - ym) + lp

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Fit or plot model of LC')
    parser.add_argument('--nwalkers', action='store', type=int, default=40)
    parser.add_argument('--fit', '-f', action='store_true')
    parser.add_argument('--nburn', action='store', type=int, default=100)
    parser.add_argument('--nprod', action='store', type=int, default=100)
    parser.add_argument('--nthreads', action='store', type=int, default=4)
    args = parser.parse_args()

    nameList = np.array(['t1', 't2', 'm1', 'm2', 'incl',
                         'r1', 'r2', 't0', 'per', 'pulse_omega',
                         'pulse_q', 'pulse_temp', 'pulse_amp'])
    params = np.array([25000, 9000, 0.40, 0.35, 89.4, 0.0193, 0.0186, 57460.6510218, 0.09986526542,
                       60.0, 10, 10000, 0.005])
    ndim = len(params)

    model_bounds = dict(
        t1=(20000, 30000),
        t2=(8000, 10000),
        m1=(0.3, 0.5),
        m2=(0.3, 0.5),
        incl=(89, 90),
        r1=(0.01, 0.025), r2=(0.01, 0.025),
        t0=(57460.6508313, 57460.6512313),
        per=(0.0998650, 0.099866),
        pulse_omega=(1, 200), pulse_q=(1, 100),
        pulse_temp=(5000, 40000), pulse_amp=(0.00001, 0.1)
    )
    # dictionary of Tseries objects, one for each band
    light_curves = dict(
        u='u.dat',
        g='g.dat',
        r='r.dat',
        i='i.dat',
        z='z.dat'
    )
    model = EclipseLC('lcurve_model', light_curves, *params, bounds=model_bounds)

    # wrapper to combine log probability from all bands
    def log_probability(params):
        val = 0
        for band in ('u', 'g', 'r', 'i', 'z'):
            val += model.log_probability(params, band)
        return val

    if args.fit:
        nwalkers = args.nwalkers

        def log_prior(params):
            model.set_parameter_vector(params)
            return model.log_prior()

        # amount to scatter initial ball of walkers
        scatter = 0.01*np.ones_like(params)
        # small scatter for t0 and period
        scatter[7] = 1.0e-9
        scatter[8] = 1.0e-9
        p0 = m.initialise_walkers(params, scatter, nwalkers, log_prior)
        sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, threads=args.nthreads)

        pos, prob, state = m.run_burnin(sampler, p0, args.nburn)
        sampler.reset()

        sampler = m.run_mcmc_save(sampler, pos, args.nprod, state, 'chain.txt')
        chain = m.flatchain(sampler.chain, ndim, thin=3)

        bestPars = []
        for i in range(ndim):
            par = chain[:, i]
            lolim, best, uplim = np.percentile(par, [16, 50, 84])
            print("%s = %f +%f -%f" % (nameList[i], best, uplim-best, best-lolim))
            bestPars.append(best)
        fig = m.thumbPlot(chain, nameList)
        fig.savefig('cornerPlot.pdf')
        plt.close()

    else:
        try:
            chain = m.readchain('chain.txt')
            fchain = m.flatchain(chain, ndim+1, thin=3)[:, :-1]
            bestPars = np.median(fchain, axis=0)
            lolim, uplim = np.percentile(fchain, (16, 84), axis=0)
            for name, par, lo, hi in zip(nameList, bestPars, lolim, uplim):
                print('{} = {} + {} - {}'.format(name, par, hi-par, par-lo))
        except Exception as err:
            print('no chain read, falling back to guess ' + str(err))
            bestPars = params

    print('Best fit has ln_prob of {}'.format(log_probability(bestPars)))

    gs = gridspec.GridSpec(5, 2)
    gs.update(hspace=0.0)

    shared_ax = None
    for iband, band in enumerate(('u', 'g', 'r', 'i', 'z')):
        if shared_ax:
            ax_main = plt.subplot(gs[iband, 0], sharex=shared_ax)
            ax_res = plt.subplot(gs[iband, 1], sharex=ax_main)
        else:
            ax_main = plt.subplot(gs[iband, 0])
            shared_ax = ax_main
            ax_res = plt.subplot(gs[iband, 1], sharex=ax_main)

        color = sns.color_palette('nipy_spectral', 5)[iband-1]
        model.plot(ax_main, band, bestPars, style='whole', dcolor=color)
        model.plot(ax_res, band, bestPars, style='residuals', dcolor=color)
        if band != 'z':
            plt.setp(ax_main.get_xticklabels(), visible=False)
            plt.setp(ax_res.get_xticklabels(), visible=False)

    plt.show()

import numpy,import astropy,from astropy


#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''FCI 1, 2, 3, 4-particle density matrices.

Note the 1-particle density matrix has the same convention as the mean-field
1-particle density matrix (see McWeeney's book Eq 5.4.20), which is
        dm[p,q] = < q^+ p >
The contraction between 1-particle Hamiltonian and 1-pdm is
        E = einsum('pq,qp', h1, 1pdm)
Different conventions are used in the high order density matrices:
        dm[p,q,r,s,...] = < p^+ r^+ ... s q >
'''

import ctypes
import numpy
from pyscf import lib
from pyscf.fci import cistring
from pyscf.fci.addons import _unpack_nelec

librdm = lib.load_library('libfci')

def reorder_rdm(rdm1, rdm2, inplace=False):
    nmo = rdm1.shape[0]
    if not inplace:
        rdm2 = rdm2.copy()
    for k in range(nmo):
        rdm2[:,k,k,:] -= rdm1.T
    #return rdm1, rdm2
    rdm2 = lib.transpose_sum(rdm2.reshape(nmo*nmo,-1), inplace=True) * .5
    return rdm1, rdm2.reshape(nmo,nmo,nmo,nmo)

# dm[p,q] = <|q^+ p|>
def make_rdm1_ms0(fname, cibra, ciket, norb, nelec, link_index=None):
    assert(cibra is not None and ciket is not None)
    cibra = numpy.asarray(cibra, order='C')
    ciket = numpy.asarray(ciket, order='C')
    if link_index is None:
        neleca, nelecb = _unpack_nelec(nelec)
        assert(neleca == nelecb)
        link_index = cistring.gen_linkstr_index(range(norb), neleca)
    na, nlink = link_index.shape[:2]
    assert(cibra.size == na**2)
    assert(ciket.size == na**2)
    rdm1 = numpy.empty((norb,norb))
    fn = getattr(librdm, fname)
    fn(rdm1.ctypes.data_as(ctypes.c_void_p),
       cibra.ctypes.data_as(ctypes.c_void_p),
       ciket.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(norb),
       ctypes.c_int(na), ctypes.c_int(na),
       ctypes.c_int(nlink), ctypes.c_int(nlink),
       link_index.ctypes.data_as(ctypes.c_void_p),
       link_index.ctypes.data_as(ctypes.c_void_p))
    return rdm1.T

# NOTE rdm1 in this function is calculated as rdm1[p,q] = <q^+ p>;
# rdm2 is calculated as <p^+ q r^+ s>. Call reorder_rdm to transform to the
# normal rdm2, which is  dm2[p,q,r,s] = <p^+ r^+ s q>.
# symm = 1: bra, ket symmetry
# symm = 2: particle permutation symmetry
def make_rdm12_ms0(fname, cibra, ciket, norb, nelec, link_index=None, symm=0):
    if link_index is None:
        neleca, nelecb = _unpack_nelec(nelec)
        assert(neleca == nelecb)
        link_index = cistring.gen_linkstr_index(range(norb), neleca)
    link_index = (link_index, link_index)
    return make_rdm12_spin1(fname, cibra, ciket, norb, nelec, link_index, symm)

make_rdm1 = make_rdm1_ms0
make_rdm12 = make_rdm12_ms0

###################################################
#
# nelec and link_index are tuples of (alpha,beta)
#
def make_rdm1_spin1(fname, cibra, ciket, norb, nelec, link_index=None):
    assert(cibra is not None and ciket is not None)
    cibra = numpy.asarray(cibra, order='C')
    ciket = numpy.asarray(ciket, order='C')
    if link_index is None:
        neleca, nelecb = _unpack_nelec(nelec)
        link_indexa = link_indexb = cistring.gen_linkstr_index(range(norb), neleca)
        if neleca != nelecb:
            link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index
    na,nlinka = link_indexa.shape[:2]
    nb,nlinkb = link_indexb.shape[:2]
    assert(cibra.size == na*nb)
    assert(ciket.size == na*nb)
    rdm1 = numpy.empty((norb,norb))
    fn = getattr(librdm, fname)
    fn(rdm1.ctypes.data_as(ctypes.c_void_p),
       cibra.ctypes.data_as(ctypes.c_void_p),
       ciket.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(norb),
       ctypes.c_int(na), ctypes.c_int(nb),
       ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
       link_indexa.ctypes.data_as(ctypes.c_void_p),
       link_indexb.ctypes.data_as(ctypes.c_void_p))
    return rdm1.T

# NOTE rdm1 in this function is calculated as rdm1[p,q] = <q^+ p>;
# rdm2 is calculated as <p^+ q r^+ s>. Call reorder_rdm to transform to the
# normal rdm2, which is  dm2[p,q,r,s] = <p^+ r^+ s q>.
# symm = 1: bra, ket symmetry
# symm = 2: particle permutation symmetry
def make_rdm12_spin1(fname, cibra, ciket, norb, nelec, link_index=None, symm=0):
    assert(cibra is not None and ciket is not None)
    cibra = numpy.asarray(cibra, order='C')
    ciket = numpy.asarray(ciket, order='C')
    if link_index is None:
        neleca, nelecb = _unpack_nelec(nelec)
        link_indexa = link_indexb = cistring.gen_linkstr_index(range(norb), neleca)
        if neleca != nelecb:
            link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    else:
        link_indexa, link_indexb = link_index
    na,nlinka = link_indexa.shape[:2]
    nb,nlinkb = link_indexb.shape[:2]
    assert(cibra.size == na*nb)
    assert(ciket.size == na*nb)
    rdm1 = numpy.empty((norb,norb))
    rdm2 = numpy.empty((norb,norb,norb,norb))
    librdm.FCIrdm12_drv(getattr(librdm, fname),
                        rdm1.ctypes.data_as(ctypes.c_void_p),
                        rdm2.ctypes.data_as(ctypes.c_void_p),
                        cibra.ctypes.data_as(ctypes.c_void_p),
                        ciket.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_int(norb),
                        ctypes.c_int(na), ctypes.c_int(nb),
                        ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                        link_indexa.ctypes.data_as(ctypes.c_void_p),
                        link_indexb.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_int(symm))
    return rdm1.T, rdm2

##############################
#
# 3-particle and 4-particle density matrix for RHF-FCI wfn
#
# NOTE the dm3[p,q,r,s,t,u] is calculated as <p^+ q r^+ s t^+ u>
# call reorder_dm123 to transform dm3 to regular 3-pdm
def make_dm123(fname, cibra, ciket, norb, nelec):
    r'''Spin traced 1, 2 and 3-particle density matrices.

    .. note::
        In this function, 2pdm[p,q,r,s] is :math:`\langle p^\dagger q r^\dagger s\rangle`;
        3pdm[p,q,r,s,t,u] is :math:`\langle p^\dagger q r^\dagger s t^\dagger u\rangle`.

        After calling reorder_dm123, the 2pdm and 3pdm are transformed to
        the normal density matrices:
        2pdm[p,r,q,s] = :math:`\langle p^\dagger q^\dagger s r\rangle`
        3pdm[p,s,q,t,r,u] = :math:`\langle p^\dagger q^\dagger r^\dagger u t s\rangle`.
    '''
    cibra = numpy.asarray(cibra, order='C')
    ciket = numpy.asarray(ciket, order='C')
    neleca, nelecb = _unpack_nelec(nelec)
    link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    na,nlinka = link_indexa.shape[:2]
    nb,nlinkb = link_indexb.shape[:2]
    assert(cibra.size == na*nb)
    assert(ciket.size == na*nb)
    rdm1 = numpy.empty((norb,)*2)
    rdm2 = numpy.empty((norb,)*4)
    rdm3 = numpy.empty((norb,)*6)
    librdm.FCIrdm3_drv(getattr(librdm, fname),
                       rdm1.ctypes.data_as(ctypes.c_void_p),
                       rdm2.ctypes.data_as(ctypes.c_void_p),
                       rdm3.ctypes.data_as(ctypes.c_void_p),
                       cibra.ctypes.data_as(ctypes.c_void_p),
                       ciket.ctypes.data_as(ctypes.c_void_p),
                       ctypes.c_int(norb),
                       ctypes.c_int(na), ctypes.c_int(nb),
                       ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                       link_indexa.ctypes.data_as(ctypes.c_void_p),
                       link_indexb.ctypes.data_as(ctypes.c_void_p))
    rdm3 = _complete_dm3_(rdm2, rdm3)
    return rdm1.T, rdm2, rdm3
def _complete_dm3_(dm2, dm3):
# fci_4pdm.c assumed symmetry p >= r >= t for 3-pdm <p^+ q r^+ s t^+ u>
# Using E^r_sE^p_q = E^p_qE^r_s - \delta_{qr}E^p_s + \delta_{ps}E^r_q to
# complete the full 3-pdm
    def transpose01(ijk, i, j, k):
        jik = ijk.transpose(1,0,2)
        jik[:,j] -= dm2[i,:,k,:]
        jik[i,:] += dm2[j,:,k,:]
        dm3[j,:,i,:,k,:] = jik
        return jik
    def transpose12(ijk, i, j, k):
        ikj = ijk.transpose(0,2,1)
        ikj[:,:,k] -= dm2[i,:,j,:]
        ikj[:,j,:] += dm2[i,:,k,:]
        dm3[i,:,k,:,j,:] = ikj
        return ikj

# ijk -> jik -> jki -> kji -> kij -> ikj
    norb = dm2.shape[0]
    for i in range(norb):
        for j in range(i+1):
            for k in range(j+1):
                tmp = transpose01(dm3[i,:,j,:,k,:].copy(), i, j, k)
                tmp = transpose12(tmp, j, i, k)
                tmp = transpose01(tmp, j, k, i)
                tmp = transpose12(tmp, k, j, i)
                tmp = transpose01(tmp, k, i, j)
    return dm3

def make_dm1234(fname, cibra, ciket, norb, nelec):
    r'''Spin traced 1, 2, 3 and 4-particle density matrices.

    .. note::
        In this function, 2pdm[p,q,r,s] is :math:`\langle p^\dagger q r^\dagger s\rangle`;
        3pdm[p,q,r,s,t,u] is :math:`\langle p^\dagger q r^\dagger s t^\dagger u\rangle`;
        4pdm[p,q,r,s,t,u,v,w] is :math:`\langle p^\dagger q r^\dagger s t^\dagger u v^\dagger w\rangle`.

        After calling reorder_dm123, the 2pdm and 3pdm are transformed to
        the normal density matrices:
        2pdm[p,r,q,s] = :math:`\langle p^\dagger q^\dagger s r\rangle`
        3pdm[p,s,q,t,r,u] = :math:`\langle p^\dagger q^\dagger r^\dagger u t s\rangle`.
        4pdm[p,t,q,u,r,v,s,w] = :math:`\langle p^\dagger q^\dagger r^\dagger s^dagger w v u t\rangle`.
    '''
    cibra = numpy.asarray(cibra, order='C')
    ciket = numpy.asarray(ciket, order='C')
    neleca, nelecb = _unpack_nelec(nelec)
    link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    na,nlinka = link_indexa.shape[:2]
    nb,nlinkb = link_indexb.shape[:2]
    assert(cibra.size == na*nb)
    assert(ciket.size == na*nb)
    rdm1 = numpy.empty((norb,)*2)
    rdm2 = numpy.empty((norb,)*4)
    rdm3 = numpy.empty((norb,)*6)
    rdm4 = numpy.empty((norb,)*8)
    librdm.FCIrdm4_drv(getattr(librdm, fname),
                       rdm1.ctypes.data_as(ctypes.c_void_p),
                       rdm2.ctypes.data_as(ctypes.c_void_p),
                       rdm3.ctypes.data_as(ctypes.c_void_p),
                       rdm4.ctypes.data_as(ctypes.c_void_p),
                       cibra.ctypes.data_as(ctypes.c_void_p),
                       ciket.ctypes.data_as(ctypes.c_void_p),
                       ctypes.c_int(norb),
                       ctypes.c_int(na), ctypes.c_int(nb),
                       ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                       link_indexa.ctypes.data_as(ctypes.c_void_p),
                       link_indexb.ctypes.data_as(ctypes.c_void_p))
    rdm3 = _complete_dm3_(rdm2, rdm3)
    rdm4 = _complete_dm4_(rdm3, rdm4)
    return rdm1.T, rdm2, rdm3, rdm4
def _complete_dm4_(dm3, dm4):
# fci_4pdm.c assumed symmetry p >= r >= t >= v for 4-pdm <p^+ q r^+ s t^+ u v^+ w>
# Using E^r_sE^p_q = E^p_qE^r_s - \delta_{qr}E^p_s + \delta_{ps}E^r_q to
# complete the full 4-pdm
    def transpose01(ijkl, i, j, k, l):
        jikl = ijkl.transpose(1,0,2,3)
        jikl[:,j] -= dm3[i,:,k,:,l,:]
        jikl[i,:] += dm3[j,:,k,:,l,:]
        dm4[j,:,i,:,k,:,l,:] = jikl
        return jikl
    def transpose12(ijkl, i, j, k, l):
        ikjl = ijkl.transpose(0,2,1,3)
        ikjl[:,:,k] -= dm3[i,:,j,:,l,:]
        ikjl[:,j,:] += dm3[i,:,k,:,l,:]
        dm4[i,:,k,:,j,:,l,:] = ikjl
        return ikjl
    def transpose23(ijkl, i, j, k, l):
        ijlk = ijkl.transpose(0,1,3,2)
        ijlk[:,:,:,l] -= dm3[i,:,j,:,k,:]
        ijlk[:,:,k,:] += dm3[i,:,j,:,l,:]
        dm4[i,:,j,:,l,:,k,:] = ijlk
        return ijlk
    def chain(ijkl, i, j, k, l):
        tmp = transpose23(ijkl, i, j, k, l)
        tmp = transpose12(tmp, i, j, l, k)
        tmp = transpose23(tmp, i, l, j, k)
        tmp = transpose12(tmp, i, l, k, j)
        tmp = transpose23(tmp, i, k, l, j)
        return tmp

# ijkl -> ijlk -> iljk -> ilkj -> iklj -> ikjl
#      -> jikl -> jilk -> jlik -> jlki -> jkli -> jkil
#(ikjl)-> kijl -> kilj -> klij -> klji -> kjli -> kjil
#(iljk)-> lijk -> likj -> lkij -> lkji -> ljki -> ljik
    norb = dm3.shape[0]
    for i in range(norb):
        for k in range(i+1):
            for j in range(k+1):
                for l in range(j+1):
                    tmp = chain(dm4[i,:,j,:,k,:,l,:].copy(), i, j, k, l)
                    tmp = transpose01(tmp, i, k, j, l)
                    tmp = chain(tmp, k, i, j, l)
                    tmp = transpose01(dm4[i,:,j,:,k,:,l,:].copy(), i, j, k, l)
                    tmp = chain(tmp, j, i, k, l)
                    tmp = transpose01(dm4[i,:,l,:,j,:,k,:].copy(), i, l, j, k)
                    tmp = chain(tmp, l, i, j, k)
    return dm4

def reorder_dm12(rdm1, rdm2, inplace=True):
    return reorder_rdm(rdm1, rdm2, inplace)

# <p^+ q r^+ s t^+ u> => <p^+ r^+ t^+ u s q>
# rdm2[p,q,r,s] is <p^+ q r^+ s>
def reorder_dm123(rdm1, rdm2, rdm3, inplace=True):
    rdm1, rdm2 = reorder_rdm(rdm1, rdm2, inplace)
    if not inplace:
        rdm3 = rdm3.copy()
    norb = rdm1.shape[0]
    for q in range(norb):
        rdm3[:,q,q,:,:,:] -= rdm2
        rdm3[:,:,:,q,q,:] -= rdm2
        rdm3[:,q,:,:,q,:] -= rdm2.transpose(0,2,3,1)
        for s in range(norb):
            rdm3[:,q,q,s,s,:] -= rdm1.T
    return rdm1, rdm2, rdm3

# <p^+ q r^+ s t^+ u w^+ v> => <p^+ r^+ t^+ w^+ v u s q>
# rdm2, rdm3 are the (reordered) standard 2-pdm and 3-pdm
def reorder_dm1234(rdm1, rdm2, rdm3, rdm4, inplace=True):
    rdm1, rdm2, rdm3 = reorder_dm123(rdm1, rdm2, rdm3, inplace)
    if not inplace:
        rdm4 = rdm4.copy()
    norb = rdm1.shape[0]
    for q in range(norb):
        rdm4[:,q,:,:,:,:,q,:] -= rdm3.transpose(0,2,3,4,5,1)
        rdm4[:,:,:,q,:,:,q,:] -= rdm3.transpose(0,1,2,4,5,3)
        rdm4[:,:,:,:,:,q,q,:] -= rdm3
        rdm4[:,q,:,:,q,:,:,:] -= rdm3.transpose(0,2,3,1,4,5)
        rdm4[:,:,:,q,q,:,:,:] -= rdm3
        rdm4[:,q,q,:,:,:,:,:] -= rdm3
        for s in range(norb):
            rdm4[:,q,q,s,:,:,s,:] -= rdm2.transpose(0,2,3,1)
            rdm4[:,q,q,:,:,s,s,:] -= rdm2
            rdm4[:,q,:,:,q,s,s,:] -= rdm2.transpose(0,2,3,1)
            rdm4[:,q,:,s,q,:,s,:] -= rdm2.transpose(0,2,1,3)
            rdm4[:,q,:,s,s,:,q,:] -= rdm2.transpose(0,2,3,1)
            rdm4[:,:,:,s,s,q,q,:] -= rdm2
            rdm4[:,q,q,s,s,:,:,:] -= rdm2
            for u in range(norb):
                rdm4[:,q,q,s,s,u,u,:] -= rdm1.T
    return rdm1, rdm2, rdm3, rdm4


#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

'''Functions to quickly simulate the observation of the sky through the year.

We call *scanning strategy* the way the instrument changes its orientation
towards the sky with the passing of time. It is encoded as a set of times and
sky coordinates, and it depends on the way the Earth rotates and the instrument
wheels are operated.

Stripeline offers a set of utilities to generate a scanning strategy, given a
set of instrumental parameters. These are used to determine which configuration
has the most desirable properties in terms of the scientific outcome (sky
coverage, integration time per pixel, etc.).

The function which creates pointing information from a description of the
scanning strategy is :meth:`~stripeline.scanning.generate_pointings`, which
depends on the class :class:`~stripeline.scanning.ScanningStrategy`. This
function is `quick`, in the sense that it does not take into account many
time-dependent effects (like the nutation of the Earth's axis), it just models
the Earth as a sphere rotating around a fixed axis with constant angular
velocity. To make a more realistic simulation, see Sect.
:ref:`accurate-generation-of-pointing-timelines`.

Here is a short example which shows how to use these facilities::

    import stripeline.scanning as sc

    def save_tod(pointings,
                 scanning: sc.ScanningStrategy,
                 dir_vec,
                 index: int):
        # Here you can save/use "pointings", which is a 4xn matrix
        pass

    scanning = sc.ScanningStrategy(wheel3_rpm=1.0,
                                   wheel2_angle0_deg=45.0,
                                   latitude_deg=28.3,
                                   overall_time_s=3600.0,
                                   sampling_frequency_hz=50.0)
    sc.generate_pointings(scanning=scanning,
                          dir_vec=[0., 0., 1.],
                          num_of_chunks=1,
                          tod_callback=save_tod)

This module provides the class :class:`~stripeline.scanning.TodWriter`, which
saves pointing information in FITS files.
'''

import logging as log
import io
import os.path
import sys
from typing import Any
import click
import healpy
import numpy as np
from astropy.io import fits
import yaml

import stripeline.quaternions as q
import stripeline.timetools as timetools

class ScanningStrategy:
    '''Parameters of a sky scanning strategy.

    This class holds together a number of information needed to define a
    strategy to scan the sky. It works like a named tuple, but it allows
    members to be changed after the object has been created.

    The following parameters are accepted:

    - `wheel1_rpm`: rotations per minute of the first wheel
        (focal plane wheel);
    - `wheel3_rpm`: rotations per minute of the third wheel
        (ground wheel, also called «azimuth wheel»);
    - `wheel1_angle0_deg`: angle of the first wheel when the simulation
        starts (degrees);
    - `wheel2_angle0_deg`: angle of the second wheel (elevation wheel)
        when the simulation starts (degrees);
    - `wheel3_angle0_deg`: angle of the third wheel when the simulation
        starts (degrees);
    - `latitude_deg`: latitude (in degrees) of the observing site;
    - `overall_time_s`: overall duration of the observation (in seconds);
    - `sampling_frequency_hz`: sampling frequency of the detector (in Hz).

     The class implements YAML serialization through the methods
     :meth:`~stripeline.scanning.ScanningStrategy.load` and
     :meth:`~stripeline.scanning.ScanningStrategy.save`.'''

    def __init__(self,
                 wheel1_rpm=0.0,
                 wheel3_rpm=0.0,
                 wheel1_angle0_deg=0.0,
                 wheel2_angle0_deg=0.0,
                 wheel3_angle0_deg=0.0,
                 latitude_deg=0.0,
                 overall_time_s=0.0,
                 sampling_frequency_hz=0.0):
        self.wheel1_rpm = wheel1_rpm
        self.wheel3_rpm = wheel3_rpm
        self.wheel1_angle0_deg = wheel1_angle0_deg
        self.wheel2_angle0_deg = wheel2_angle0_deg
        self.wheel3_angle0_deg = wheel3_angle0_deg
        self.latitude_deg = latitude_deg
        self.overall_time_s = overall_time_s
        self.sampling_frequency_hz = sampling_frequency_hz

    def validate(self):
        '''Raise a ValueError if the scanning strategy is invalid.'''

        if not isinstance(self.wheel1_rpm, float) or self.wheel1_rpm < 0.0:
            raise ValueError('invalid value for wheel1_rpm ({0})'
                             .format(self.wheel1_rpm))

        if not isinstance(self.wheel3_rpm, float) or self.wheel3_rpm < 0.0:
            raise ValueError('invalid value for wheel3_rpm ({0})'
                             .format(self.wheel3_rpm))

        if not isinstance(self.wheel1_angle0_deg, float):
            raise ValueError('invalid value for wheel1_angle0_deg ({0})'
                             .format(self.wheel1_angle0_deg))

        if not isinstance(self.wheel2_angle0_deg, float):
            raise ValueError('invalid value for wheel2_angle0_deg ({0})'
                             .format(self.wheel2_angle0_deg))

        if not isinstance(self.wheel3_angle0_deg, float):
            raise ValueError('invalid value for wheel3_angle0_deg ({0})'
                             .format(self.wheel3_angle0_deg))

        if not isinstance(self.latitude_deg, float) or \
                self.latitude_deg < 0.0 or self.latitude_deg > 90.0:
            raise ValueError('invalid value for latitude_deg ({0})'
                             .format(self.latitude_deg))

        if not isinstance(self.sampling_frequency_hz, float) or \
                self.sampling_frequency_hz <= 0.0:
            raise ValueError('invalid value for sampling_frequency_hz ({0})'
                             .format(self.sampling_frequency_hz))

    def save(self, stream):
        '''Write a YAML representation of `self` into the stream.'''

        yaml.dump({'wheel1_rpm': self.wheel1_rpm,
                   'wheel3_rpm': self.wheel3_rpm,
                   'wheel1_angle0_deg': self.wheel1_angle0_deg,
                   'wheel2_angle0_deg': self.wheel2_angle0_deg,
                   'wheel3_angle0_deg': self.wheel3_angle0_deg,
                   'latitude_deg': self.latitude_deg,
                   'overall_time_s': self.overall_time_s,
                   'sampling_frequency_hz': self.sampling_frequency_hz},
                  stream=stream,
                  explicit_start=True,
                  explicit_end=True)

    def load(self, input):
        '''Build a :class:`ScanningStrategy` object from its YAML representation.

        The parameter "input" can either be a file object, a dictionary or a string.'''
        if isinstance(input, dict):
            d = input
        else:
            d = yaml.load(input)

        self.wheel1_rpm = d['wheel1_rpm']
        self.wheel3_rpm = d['wheel3_rpm']
        self.wheel1_angle0_deg = d['wheel1_angle0_deg']
        self.wheel2_angle0_deg = d['wheel2_angle0_deg']
        self.wheel3_angle0_deg = d['wheel3_angle0_deg']
        self.latitude_deg = d['latitude_deg']
        self.overall_time_s = d['overall_time_s']
        self.sampling_frequency_hz = d['sampling_frequency_hz']

def time_to_rot_angle(time_vec: Any, rpm: float) -> Any:
    '''Return a set of angles given a set of times and the RPMs.

    RPM is Rotation Per Minute, of course!'''

    assert rpm >= 0.0

    if rpm == 0.0:
        return np.zeros_like(time_vec)
    else:
        return 2 * np.pi * time_vec * (rpm / 60.0)

def generate_pointings(scanning: ScanningStrategy,
                       dir_vec=[0, 0, 1],
                       num_of_chunks=1,
                       tod_callback=None,
                       time0_s=0.0):
    '''Generate a set of pointing directions.

    Simulate the scanning of the sky with the parameters provided in `scanning`,
    `dir_vec`. The `tod_callback` parameter is a function which is called
    whenever a new chunk of samples has been calculated. (It is fine if it is
    set to ``None``: in this case, pointings will be silently thrown away once
    they have been computed.) The value `time0_s` specifies the time of the
    first sample.

    The callback must accept the following parameters:

    - `pointings`: 4xn matrix containing the time (in seconds), colatitude (in
      radians), longitude (ditto), and polarization angle (ditto), each in its
      own column;
    - `scanning`: copy of the parameter passed to this function;
    - `dir_vec`: copy of the parameter passed to this function;
    - `index`: counter which keeps track of how many times the callback has been
      called, starting from 0.
    '''
    x_vec = np.array([1, 0, 0])
    z_vec = np.array([0, 0, 1])

    chunks = timetools.split_time_range(time_length=scanning.overall_time_s,
                                        num_of_chunks=num_of_chunks,
                                        sampfreq=scanning.sampling_frequency_hz,
                                        time0=time0_s)
    for chunk_idx, cur_chunk in enumerate(chunks):
        start_time, samples_per_chunk = cur_chunk

        time_vec = start_time + \
            np.arange(samples_per_chunk) / scanning.sampling_frequency_hz

        # Determine the angle of each wheel (the second wheel is the simplest)
        wheel1_angle = np.deg2rad(scanning.wheel1_angle0_deg) + time_to_rot_angle(
            time_vec, scanning.wheel1_rpm)
        wheel3_angle = np.deg2rad(scanning.wheel3_angle0_deg) + time_to_rot_angle(
            time_vec, scanning.wheel3_rpm)

        tile_dir = np.reshape(np.tile(dir_vec, time_vec.size), (-1, 3))
        tile_x = np.reshape(np.tile(x_vec, time_vec.size), (-1, 3))
        tile_z = np.reshape(np.tile(z_vec, time_vec.size), (-1, 3))
        # Build the wheel quaternions
        qwheel1 = q.qfromaxisangle(tile_dir, wheel1_angle)
        qwheel2 = np.reshape(
            np.tile(
                q.qfromaxisangle(
                    [x_vec], [np.deg2rad(scanning.wheel2_angle0_deg)]),
                time_vec.size), (-1, 4))
        qwheel3 = q.qfromaxisangle(tile_z, wheel3_angle)

        # This is in the ground's reference frame
        ground_quat = q.qmul(qwheel3, q.qmul(qwheel2, qwheel1))

        # Now we convert from the ground reference frame to the Earth's centre
        location_quat = np.reshape(
            np.tile(
                q.qfromaxisangle(
                    [x_vec], [np.deg2rad(90.0 - scanning.latitude_deg)]),
                time_vec.size), (-1, 4))
        earth_rot_quat = q.qfromaxisangle(tile_z,
                                          2 * np.pi * time_vec / 86400.0)
        quat = q.qmul(earth_rot_quat, q.qmul(location_quat, ground_quat))
        dirs = q.qrotate(tile_dir, quat)
        poldirs = q.qrotate(tile_x, quat)
        theta, phi = healpy.vec2ang(dirs)

        # The north direction for a vector v is just -dv/dtheta, as
        # theta is the colatitude and moves along the meridian
        thetapol, phipol = healpy.vec2ang(dirs)
        northdir = np.column_stack((-np.cos(thetapol) * np.cos(phipol),
                                    -np.cos(thetapol) * np.sin(phipol),
                                    np.sin(thetapol)))

        # The counterclockwise/clockwise measurement of the polarization angle
        # is determined by the ordering of the terms in the call to "np.cross"
        cos_psi = np.clip(np.sum(northdir * poldirs, axis=1), -1.0, 1.0)
        cross = np.cross(northdir, poldirs)
        sin_psi = np.clip(np.sum(cross * cross, axis=1), -1.0, 1.0)
        psi = np.arctan2(sin_psi, cos_psi)
        psi *= np.sign(np.sum(cross * dirs, axis=1))

        if tod_callback is not None:
            tod_callback(pointings=np.column_stack((time_vec, theta, phi, psi)),
                         scanning=scanning,
                         dir_vec=dir_vec,
                         index=chunk_idx)

class TodWriter:
    '''Write a TOD.

    This class has been designed to be used together with
    :meth:`~stripeline.scanning.generate_pointings`.

    You create an instance of the object and then pass it to
    :meth:`~stripeline.scanning.generate_pointings`, like in the following way::

        writer = TodWriter(outdir='/storage',
                           file_name_mask='mytod_{index:04d}.fits.gz')
        generate_pointings(..., tod_callback=writer)

    As you can see, you can use the `index` key in the `file_name_mask`
    parameter to separate the files in chunks. The number of times `writer` is
    called depends on the parameter `num_of_chunks` passed to
    :meth:`~stripeline.scanning.generate_pointings`.
    '''

    def __init__(self,
                 outdir='.',
                 file_name_mask='pointings_{index:04d}.fits'):
        self.outdir = outdir
        self.file_name_mask = file_name_mask

    def __call__(self,
                 pointings,
                 scanning: ScanningStrategy,
                 dir_vec,
                 index: int):
        ''' Save a TOD into a FITS file'''

        file_name = os.path.join(self.outdir,
                                 self.file_name_mask.format(index=index))
        cols = [
            fits.Column(name=name, format=fmt, unit=unit, array=arr)
            for name, fmt, unit, arr in (('TIME', 'D', 's', pointings[:, 0]),
                                         ('THETA', 'D', 'rad',
                                          pointings[:, 1]),
                                         ('PHI', 'D', 'rad', pointings[:, 2]),
                                         ('PSI', 'D', 'rad', pointings[:, 3]))
        ]
        hdu = fits.BinTableHDU.from_columns(cols, name='TOD')
        hdu.header['FSTTIME'] = (
            pointings[0, 0], 'Time of the first sample in the file [s]')
        hdu.header['LSTTIME'] = (
            pointings[-1, 0], 'Time of the last sample in the file [s]')
        hdu.header['DIRX'] = (dir_vec[0], 'X component of the beam axis')
        hdu.header['DIRY'] = (dir_vec[1], 'Y component of the beam axis')
        hdu.header['DIRZ'] = (dir_vec[2], 'Z component of the beam axis')
        hdu.header['SAMPFREQ'] = (scanning.sampling_frequency_hz,
                                  'Sampling frequency [Hz]')
        hdu.header['SITELAT'] = (
            scanning.latitude_deg, 'Latitude of the site [deg]')
        hdu.header['W1RPM'] = (scanning.wheel1_rpm,
                               'Angular speed of wheel 1 [rpm]')
        hdu.header['W3RPM'] = (scanning.wheel3_rpm,
                               'Angular speed of wheel 3 [rpm]')
        hdu.header['W1ANG0'] = (
            scanning.wheel1_angle0_deg, 'Start angle for wheel 1 [deg]')
        hdu.header['W2ANG0'] = (
            scanning.wheel2_angle0_deg, 'Start angle for wheel 2 [deg]')
        hdu.header['W3ANG0'] = (
            scanning.wheel3_angle0_deg, 'Start angle for wheel 3 [deg]')
        hdu.header['TIMELEN'] = (
            scanning.overall_time_s, 'Time span of the *whole* sim [s]')
        hdu.header['TODIDX'] = (index, '0-based index of this file')

        with io.StringIO() as primary_data:
            scanning.save(stream=primary_data)
            raw_bytes = np.array(list(primary_data.getvalue().encode('utf-8')))
            primhdu = fits.PrimaryHDU(data=raw_bytes)

        hdulist = fits.HDUList([primhdu, hdu])
        hdulist.writeto(file_name, clobber=True)
        log.info('file "%s" written successfully', file_name)

@click.command()
@click.argument('output_path')
@click.option('--input-file',
              type=str,
              default=None,
              help='Path to a YAML file containing the parameters of the '
              'scanning strategy')
@click.option('--wheel1-rpm',
              'wheel1_rpm',
              type=float,
              default=None,
              help='Rotations per minute of the first (focal plane) wheel')
@click.option('--wheel3-rpm',
              'wheel3_rpm',
              type=float,
              default=None,
              help='Rotations per minute of the third (ground) wheel')
@click.option('--wheel1-angle0',
              'wheel1_angle0',
              type=float,
              default=None,
              help='Initial angle of the first (focal plane) wheel [deg]')
@click.option('--wheel2-angle0',
              'wheel2_angle0',
              type=float,
              default=None,
              help='Initial angle of the second (elevation) wheel [deg]')
@click.option('--wheel3-angle0',
              'wheel3_angle0',
              type=float,
              default=None,
              help='Initial angle of the third (ground) wheel [deg]')
@click.option('--latitude',
              'latitude',
              type=float,
              default=None,
              help='Latitude of the observing site (North is positive) [deg]')
@click.option('--time',
              'time_length',
              type=float,
              default=None,
              help='Amount of observation time [s]')
@click.option('--samp',
              'sampfreq',
              type=float,
              default=None,
              help='Sampling frequency [Hz]')
@click.option('--num-of-chunks',
              '-n',
              type=int,
              default=1,
              help='Number of chunks for splitting the computation')
@click.option('--direction',
              type=str,
              default='0,0,1',
              help='Pointing direction of the main beam with respect to '
              'the focal plane (3D vector, written as a comma-separated list '
              'of 3 numbers)')
def main(output_path, input_file, wheel1_rpm, wheel3_rpm, wheel1_angle0, wheel2_angle0,
         wheel3_angle0, latitude, time_length, sampfreq, num_of_chunks,
         direction):
    '''This function is called when the script is ran from the command line.'''

    log.basicConfig(
        level=log.INFO, format='%(asctime)s %(levelname)s] %(message)s')

    num_of_samples = int(time_length * sampfreq)
    log.info('%d samples will be processed in %d steps', num_of_samples,
             num_of_chunks)

    try:
        direction = np.array([float(x) for x in direction.split(',')])
        if len(direction) != 3:
            raise ValueError()
    except ValueError:
        log.error('pointing direction must be a comma-separated list of '
                  'three floating-point numbers (es., "0,0,1")')
        sys.exit(1)

    scanning = ScanningStrategy()
    if input_file is not None:
        with open(input_file, 'rt') as f:
            scanning.load(f)

    if wheel1_rpm is not None:
        scanning.wheel1_rpm = wheel1_rpm
    if wheel3_rpm is not None:
        scanning.wheel3_rpm = wheel3_rpm
    if wheel1_angle0 is not None:
        scanning.wheel1_angle0_deg = wheel1_angle0
    if wheel2_angle0 is not None:
        scanning.wheel2_angle0_deg = wheel2_angle0
    if wheel3_angle0 is not None:
        scanning.wheel3_angle0_deg = wheel3_angle0
    if latitude is not None:
        scanning.latitude_deg = latitude
    if time_length is not None:
        scanning.overall_time_s = time_length
    if sampfreq is not None:
        scanning.sampling_frequency_hz = sampfreq

    scanning.validate()

    writer = TodWriter(output_path)
    generate_pointings(scanning=scanning,
                       dir_vec=direction,
                       num_of_chunks=num_of_chunks,
                       tod_callback=writer)

if __name__ == '__main__':
    main()


# coding: utf-8
# pylint: disable=invalid-name, no-member, arguments-differ
""" EIT protocol """
# Copyright (c) Benyuan Liu. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
from __future__ import absolute_import, division, print_function

from dataclasses import dataclass
from typing import Union

import numpy as np

@dataclass
class PyEITProtocol:
    """
    EIT Protocol buid-in protocol object

    """

    ex_mat: np.ndarray
    meas_mat: np.ndarray

    def __post_init__(self) -> None:
        """Checking of the inputs"""
        self.ex_mat = self._check_ex_mat(self.ex_mat)
        self.meas_mat = self._check_meas_mat(self.meas_mat)

    def _check_ex_mat(self, ex_mat: np.ndarray) -> np.ndarray:
        """
        Check/init stimulation

        Parameters
        ----------
        ex_mat : np.ndarray
            stimulation/excitation matrix, of shape (n_exc, 2).
            If single stimulation (ex_line) is passed only a list of length 2
            and np.ndarray of size 2 will be treated.

        Returns
        -------
        np.ndarray
            stimulation matrix

        Raises
        ------
        TypeError
            Only accept, list of length 2, np.ndarray of size 2,
            or np.ndarray of shape (n_exc, 2)
        """
        if isinstance(ex_mat, list) and len(ex_mat) == 2:
            # case ex_line has been passed instead of ex_mat
            ex_mat = np.array([ex_mat]).reshape((1, 2))  # build a 2D array
        elif isinstance(ex_mat, np.ndarray) and ex_mat.size == 2:
            #     case ex_line np.ndarray has been passed instead of ex_mat
            ex_mat = ex_mat.reshape((-1, 2))

        if (
            not isinstance(ex_mat, np.ndarray)
            or ex_mat.ndim != 2
            or ex_mat.shape[1] != 2
        ):
            raise TypeError(
                f"Wrong shape of {ex_mat=} expected an ndarray ; shape (n_exc, 2)"
            )

        return ex_mat

    def _check_meas_mat(self, meas_mat: np.ndarray) -> np.ndarray:
        """
        Check measurement pattern

        Parameters
        ----------
        n_exc : int
            number of excitations/stimulations
        meas_pattern : np.ndarray, optional
           measurements pattern / subtract_row pairs [N, M] to check; shape (n_exc, n_meas_per_exc, 2)

        Returns
        -------
        np.ndarray
            measurements pattern / subtract_row pairs [N, M]; shape (n_exc, n_meas_per_exc, 2)

        Raises
        ------
        TypeError
            raised if meas_pattern is not a np.ndarray of shape (n_exc, : , 2)
        """
        if not isinstance(meas_mat, np.ndarray):
            raise TypeError(
                f"Wrong type of {meas_mat=}, expected an ndarray; shape ({self.n_exc}, n_meas_per_exc, 2)"
            )
        # test shape is something like (n_exc, :, 2)
        if meas_mat.ndim != 3 or meas_mat.shape[::2] != (self.n_exc, 2):
            raise TypeError(
                f"Wrong shape of {meas_mat=}: {meas_mat.shape=}, expected an ndarray; shape ({self.n_exc}, n_meas_per_exc, 2)"
            )

        return meas_mat

    @property
    def n_exc(self) -> int:
        """
        Returns
        -------
        int
            number of excitation
        """
        return self.ex_mat.shape[0]

    @property
    def n_meas(self) -> int:
        """
        Returns
        -------
        int
            number of measurements per excitations
        """
        return self.meas_mat.shape[1]

    @property
    def n_meas_tot(self) -> int:
        """
        Returns
        -------
        int
            total amount of measurements
        """
        return self.n_meas * self.n_exc

    @property
    def n_el(self) -> int:
        """
        Returns
        -------
        int
            number of electrodes used in the excitation and
        """
        return max(max(self.ex_mat.flatten()), max(self.meas_mat.flatten())) + 1

def create(
    n_el: int = 16,
    dist_exc: Union[int, list[int]] = 1,
    step_meas: int = 1,
    parser_meas: Union[str, list[str]] = "std",
) -> PyEITProtocol:
    """
    Return an EIT protocol, comprising an excitation and a measuremnet pattern

    Parameters
    ----------
    n_el : int, optional
        number of total electrodes, by default 16
    dist_exc : Union[int, list[int]], optional
        distance (number of electrodes) of A to B, by default 1
        For 'adjacent'- or 'neighbore'-mode (default) use `1` , and
        for 'apposition'-mode use `n_el/2`. (see `build_exc_pattern`)
        if a list of integer is passed the excitation will bee stacked together.
    step_meas : int, optional
    measurement method (two adjacent electrodes are used for measuring), by default 1 (adjacent).
        (see `build_meas_pattern`)
    parser_meas : Union[str, list[str]], optional
        parsing the format of each frame in measurement/file, by default 'std'.
        (see `build_meas_pattern`)

    Returns
    -------
    PyEITProtocol
        EIT protocol object

    Raises
    ------
    TypeError
        if dist_exc is not list or an int
    """
    if isinstance(dist_exc, int):
        dist_exc = [dist_exc]

    if not isinstance(dist_exc, list):
        raise TypeError(f"{dist_exc=}; {type(dist_exc)=} should be a list[int]")

    _ex_mat = [build_exc_pattern_std(n_el, dist) for dist in dist_exc]
    ex_mat = np.vstack(_ex_mat)

    meas_mat = build_meas_pattern_std(ex_mat, n_el, step_meas, parser_meas)
    return PyEITProtocol(ex_mat, meas_mat)

def build_meas_pattern_std(
    ex_mat: np.ndarray,
    n_el: int = 16,
    step: int = 1,
    parser: Union[str, list[str]] = "std",
) -> np.ndarray:
    """
    Build the measurement pattern (subtract_row-voltage pairs [N, M])
    for all excitations on boundary electrodes.

    we direct operate on measurements or Jacobian on electrodes,
    so, we can use LOCAL index in this module, do not require el_pos.

    Notes
    -----
    ABMN Model.
    A: current driving electrode,
    B: current sink,
    M, N: boundary electrodes, where v_diff = v_n - v_m.

    Parameters
    ----------
    ex_mat : np.ndarray
        Nx2 array, [positive electrode, negative electrode]. ; shape (n_exc, 2)
    n_el : int, optional
        number of total electrodes, by default 16
    step : int, optional
        measurement method (two adjacent electrodes are used for measuring), by default 1 (adjacent)
    parser : Union[str, list[str]], optional
        parsing the format of each frame in measurement/file, by default 'std'
        if parser contains 'fmmu', or 'rotate_meas' then data are trimmed,
        boundary voltage measurements are re-indexed and rotated,
        start from the positive stimulus electrode start index 'A'.
        if parser contains 'std', or 'no_rotate_meas' then data are trimmed,
        the start index (i) of boundary voltage measurements is always 0.
        if parser contains 'meas_current', the measurements on current carrying
        electrodes are allowed. Otherwise the measurements on current carrying
        electrodes are discarded (like 'no_meas_current' option in EIDORS3D).

    Returns
    -------
    np.ndarray
        measurements pattern / subtract_row pairs [N, M]; shape (n_exc, n_meas_per_exc, 2)
    """
    if not isinstance(parser, list):  # transform parser into list
        parser = [parser]
    meas_current = "meas_current" in parser
    fmmu_rotate = any(p in ("fmmu", "rotate_meas") for p in parser)

    diff_op = []
    for ex_line in ex_mat:
        a, b = ex_line[0], ex_line[1]
        i0 = a if fmmu_rotate else 0
        m = (i0 + np.arange(n_el)) % n_el
        n = (m + step) % n_el
        meas_pattern = np.vstack([n, m]).T

        if not meas_current:
            diff_keep = np.logical_and.reduce((m != a, m != b, n != a, n != b))
            meas_pattern = meas_pattern[diff_keep]

        diff_op.append(meas_pattern)

    return np.array(diff_op)

def build_exc_pattern_std(n_el: int = 16, dist: int = 1) -> np.ndarray:
    """
    Generate scan matrix, `ex_mat` ( or excitation pattern), see notes

    Parameters
    ----------
    n_el : int, optional
        number of electrodes, by default 16
    dist : int, optional
        distance (number of electrodes) of A to B, by default 1
        For 'adjacent'- or 'neighbore'-mode (default) use `1` , and
        for 'apposition'-mode use `n_el/2` (see Examples).

    Returns
    -------
    np.ndarray
        stimulation matrix; shape (n_exc, 2)

    Notes
    -----
        - in the scan of EIT (or stimulation matrix), we use 4-electrodes
        mode, where A, B are used as positive and negative stimulation
        electrodes and M, N are used as voltage measurements.
        - `1` (A) for positive current injection, `-1` (B) for negative current
        sink

    Examples
    --------
        n_el=16
        if mode=='neighbore':
            ex_mat = build_exc_pattern(n_el=n_el)
        elif mode=='apposition':
            ex_mat = build_exc_pattern(dist=n_el/2)

    WARNING
    -------
        `ex_mat` is a local index, where it is ranged from 0...15, within the
        range of the number of electrodes. In FEM applications, you should
        convert `ex_mat` to global index using the (global) `el_pos` parameters.
    """
    return np.array([[i, np.mod(i + dist, n_el)] for i in range(n_el)])