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

# ## Imports 

# In[14]:


import os
import urllib.request

import numpy as np
import matplotlib.pyplot as plt
import astropy.io as aio
import astropy.units as u

import lenstools as lt

get_ipython().run_line_magic('matplotlib', 'inline')


# # One map sample 

# ## Header 

# In[2]:


with aio.fits.open("hlsp_frontier_model_abell2744_cats_v4.1_kappa.fits") as hdu:
    head = hdu[0].header

head


# ## Read one kappa map 

# In[3]:


def loadConvergenceMap(fname):
    with aio.fits.open(fname) as hdu:
        angle = 0.3*hdu[0].header["NAXIS1"]*u.arcsec
        kappa = hdu[0].data
    return angle,kappa


# In[4]:


conv = lt.image.convergence.ConvergenceMap.load("hlsp_frontier_model_abell2744_cats_v4.1_kappa.fits",format=loadConvergenceMap)


# In[5]:


conv.info


# ## Visualize log(kappa) 

# In[6]:


fig,ax = plt.subplots(figsize=(12,8))
lt.image.convergence.ConvergenceMap(np.log(conv.data),angle=conv.side_angle).visualize(colorbar=True,cbar_label=r"$\log(\kappa)$",fig=fig,ax=ax)


# ## 2D angular power spectrum 

# In[45]:


ell_bins = np.linspace(5000,5000*100,100)
ell,Pell = conv.powerSpectrum(ell_bins)


# In[48]:


fig,ax = plt.subplots(figsize=(12,8))

ax.plot(ell,ell*(ell+1)*Pell/(2*np.pi))

ax.set_xlabel(r"$\ell$",fontsize=16)
ax.set_ylabel(r"$\ell(\ell+1)P_\ell/2\pi$",fontsize=16)


# # Ensemble of kappa maps

# ## Download kappa maps

# In[15]:


def downloadMaps(ns,dest):
    url_root = "https://archive.stsci.edu/pub/hlsp/frontier/abell2744/models/cats/v4.1/range/"
    map_name = "hlsp_frontier_model_abell2744_cats-map{0:03d}_v4.1_kappa.fits"
    for n in ns:
        url = url_root + map_name.format(n)
        print("[+] Downloading: "+url)
        urllib.request.urlretrieve(url,os.path.join(dest,map_name.format(n)))


# In[34]:


downloadMaps(range(200),"hlsp_maps/range")


# ## Kappa power spectrum measurements 

# In[25]:


def measureKappaPowerSpectrum(ell_bins,ns,map_root="hlsp_maps/range/hlsp_frontier_model_abell2744_cats-map{0:03d}_v4.1_kappa.fits"):
    ps = np.zeros((len(ns),len(ell_bins)-1))
    for i,n in enumerate(ns):
        map_fname = map_root.format(n)
        print("[+] Measuring kappa power spectrum in map: "+map_fname)
        conv = lt.image.convergence.ConvergenceMap.load(map_fname,format=loadConvergenceMap)
        ell,pn = conv.powerSpectrum(ell_bins)
        ps[i,:] = pn
    
    return ell,ps


# In[50]:


ell,kappa_ps = measureKappaPowerSpectrum(ell_bins,range(200))


# In[51]:


fig,ax = plt.subplots(figsize=(12,8))

for ps in kappa_ps:
    ax.plot(ell,ell*(ell+1)*ps/(2*np.pi),color="blue",alpha=0.1)
ax.plot(ell,ell*(ell+1)*kappa_ps.mean(0)/(2*np.pi),color="red",linewidth=2)

ax.set_xlabel(r"$\ell$",fontsize=16)
ax.set_ylabel(r"$\ell(\ell+1)P_\ell/2\pi$",fontsize=16)


# In[ ]: