# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_skymaps.ipynb (unless otherwise specified).

__all__ = ['valid', 'plot_week', 'SquareWCS']

# Cell
import numpy as np
import os, sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

import healpy as hp
from healpy.rotator import Rotator
from astropy.coordinates import SkyCoord
from astropy.wcs import WCS
from utilities import healpix as hpx

from .config import *
from .data_man import get_week_map
from .sources import findsource
valid = Config().valid;

# Cell
def plot_week(week=None, mjd=None, nside=32, **kwargs):
    """
    Make an AIT plot of the given week's photon data

    Combine all energies for now

    - week -- the week number from 9
    - mjd -- [None] If set, derive the week from it
    - nside [32] -- HEALPix nside to project data before plotting.
    - kwargs -- args for healpix.ait_plot

    """
    from utilities import healpix as hpx
    from .config import mission_week
    if mjd is not None: week=mission_week(mjd)

    pmap = get_week_map(week, nside=nside)

    kw = dict(log=True, tick_labels=False, vmin=5, vmax=1e3,
             cblabel=f'counts per nside={nside} pixel')

    t = hpx.HPmap(pmap, f'week_{week:d}', #\n{utc}',
                  **kwargs)#f'week_{week:d}\n{utc}', nest=True)
    t.ait_plot(**kw)
    return plt.gcf()

# Cell

class SquareWCS(WCS):
    """
    Create and use a WCS object

    - center : a SkyCoord that will be the center
    - size   : width and height of the display
    - pixsize [0.1] : pixel size
    - frame [None] : The frame is taken from the center SkyCoord, unless specified here --  only accept "galactic" or "fk5"
    - proj ["ZEA"] : projection to use
    """

    def __init__(self, center, size, pixsize=0.1, frame=None, proj='ZEA'):
        """

        """
        assert isinstance(center, SkyCoord), 'Expect SkyCoord'

        frame = frame or center.frame.name
        if frame=='galactic':
            lon, lat = center.galactic.l.deg, center.galactic.b.deg
            lon_name,lat_name = 'GLON','GLAT'
            self.axis_labels='$l$', '$b$'
        elif frame=='fk5':
            lon,lat = center.fk5.ra.deg, center.fk5.dec.deg
            lon_name, lat_name = 'RA--', 'DEC-'
            self.axis_labels = 'RA', 'Dec'
        else:
            raise Exception(f'Expect frame to be "galactic" or "fk5", not {frame}')

        nx=ny=naxis = int(size/pixsize) | 1 # make odd so central pixel has source in middle
        self.center = center
        self.frame=frame
        self.galactic = frame=='galactic'
        super().__init__(
                         dict(
            NAXIS1=nx, CTYPE1=f'{lon_name}-{proj}', CUNIT1='deg', CRPIX1=nx//2+1, CRVAL1=lon, CDELT1=-pixsize,
            NAXIS2=ny, CTYPE2=f'{lat_name}-{proj}', CUNIT2='deg', CRPIX2=ny//2+1, CRVAL2=lat, CDELT2=pixsize, )
              )

    def _make_grid(self):
        # get coordinates of every pixel`
        nx, ny = self.array_shape
        pixlists = list(range(1,nx+1)),list(range(1,ny+1))
        cgrid = self.pixel_to_world(*np.meshgrid(*pixlists) )
        if not self.galactic:
            cgrid = cgrid.galactic
        lon, lat = (cgrid.l.deg, cgrid.b.deg)
        return lon, lat

    def plot(self, hmap, name='' , log=False, cmap='jet', colorbar=True,
             cblabel='', vmin=None, vmax=None, cb_kw={},
             annotator=None):

        import healpy as hp
        from matplotlib import colors

        wcs = self
        grid = self._make_grid();
        nside = hp.get_nside(hmap)

        # lon, lat = grid.l.deg, grid.b.deg
        ipix = hp.ang2pix(nside, *grid, lonlat=True)

        fig = plt.figure(figsize=(6,6))
        fig.add_subplot(111, projection=self)
        ax = fig.axes[0]

        if log:
            norm = colors.LogNorm(vmin=vmin,vmax=vmax)
            vmin=vmax=None
        else:
            norm = None
        ipix = hp.ang2pix(nside, *grid, lonlat=True)
        im = ax.imshow(hmap[ipix], cmap=cmap, origin='lower', norm=norm, vmin=vmin);

        nx, ny = wcs.array_shape
        ax.set(xlabel=self.axis_labels[0], xlim=(-0.5, nx-0.5),
               ylabel=self.axis_labels[1], ylim=(-0.5, ny-0.5),
              title=f'centered on {name}')
        ax.grid();
        if colorbar:
            ticklabels = cb_kw.pop('ticklabels', None)
            cb_kw.update(label=cblabel,)
            cb = plt.colorbar(im, ax=ax, **cb_kw)
            if ticklabels is not None:
                cb.ax.set_yticklabels(ticklabels)
        if annotator is not None:
            annotator(ax, self.frame)
        return fig