#!/usr/bin/env python
# coding: utf-8

# # Check SkyPointSource
# 
# A quick queck if `SkyPointSource` gives correct results or if it's buggy.
# 
# Ticket: https://github.com/gammapy/gammapy/issues/2367

# In[1]:


import numpy as np
import astropy.units as u


# In[2]:


import warnings
warnings.filterwarnings('ignore')


# In[3]:


import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')


# In[4]:


from gammapy.maps import Map, MapAxis
from gammapy.modeling.models import SkyPointSource, ConstantModel, SkyModel
from gammapy.cube import MapEvaluator


# In[19]:


# Test case setup
energy_axis = MapAxis.from_edges([1, 10], unit="TeV", name="energy", interp="log")
exposure = Map.create(binsz=1, proj="AIT", unit="cm2 s", axes=[energy_axis])
exposure.data = np.ones_like(exposure.data)

spatial_model = SkyPointSource(0 * u.deg, 0 * u.deg, frame="icrs")
spectral_model = ConstantModel("1 cm-2 s-1 TeV-1")
model = SkyModel(spatial_model=spatial_model, spectral_model=spectral_model)
evaluator = MapEvaluator(model=model, exposure=exposure)


# In[20]:


# Compute flux error image for different sky positions
# If algorithm is correct, should get a constant image with value 1
def make_ratio_map(skydir, binsz, proj, npix=None):
    ratio_map = Map.create(skydir=skydir, binsz=binsz, proj=proj, npix=npix)
    coord = ratio_map.geom.get_coord()

    for idx_lat in range(coord.shape[0]):
        for idx_lon in range(coord.shape[1]):
            spatial_model.lon_0.value = coord.lon.value[idx_lat, idx_lon]
            spatial_model.lat_0.value = coord.lat.value[idx_lat, idx_lon]
            flux = evaluator.compute_flux()
            ratio = float(np.nansum(flux) / spectral_model.integral(1 * u.TeV, 10 * u.TeV))
            ratio_map.data[idx_lat, idx_lon] = ratio
    
    return ratio_map


# In[21]:


# Check for errors related to WCS distortions in all-sky maps
ratio_map = make_ratio_map(skydir=(0, 0), binsz=3, proj="CAR", npix=None)
print(ratio_map.data.mean(), ratio_map.data.min(), ratio_map.data.max())
ratio_map.plot(add_cbar=True);


# In[22]:


# Check for errors related to pixel contains bugs in the implementation
ratio_map = make_ratio_map(skydir=(90, 60), binsz=0.11, proj="CAR", npix=(200, 200))
print(ratio_map.data.mean(), ratio_map.data.min(), ratio_map.data.max())
ratio_map.plot(add_cbar=True);


# In[18]:


# Check for errors related to pixel contains bugs at the pixel edges and corner
ratio_map = make_ratio_map(skydir=(95, 5), binsz=0.01, proj="CAR", npix=(100, 100))
print(ratio_map.data.mean(), ratio_map.data.min(), ratio_map.data.max())
ratio_map.plot(add_cbar=True);


# In[ ]: