{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from herschelhelp_internal import git_version\n",
    "print(\"This notebook was run with herschelhelp_internal version: \\n{}\".format(git_version()))\n",
    "import datetime\n",
    "print(\"This notebook was executed on: \\n{}\".format(datetime.datetime.now()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline\n",
    "#%config InlineBackend.figure_format = 'svg'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.rc('figure', figsize=(10, 6))\n",
    "\n",
    "import os\n",
    "import time\n",
    "import gc\n",
    "import glob\n",
    "from matplotlib_venn import venn2\n",
    "from astropy import units as u\n",
    "from astropy.coordinates import SkyCoord\n",
    "from astropy.table import Column, QTable, Table, join, vstack, hstack, setdiff, unique\n",
    "import numpy as np\n",
    "import pymoc \n",
    "from pymoc import MOC\n",
    "from datetime import datetime\n",
    "\n",
    "\n",
    "from herschelhelp_internal.masterlist import merge_catalogues, nb_merge_dist_plot, specz_merge\n",
    "from herschelhelp_internal.utils import coords_to_hpidx, ebv, gen_help_id, inMoc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DMU_DIR = '../'\n",
    "FIELD = 'Herschel-Stripe-82'\n",
    "SUFFIX = '20210513' #datetime.today().strftime('%Y%m%d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_sample(catalogue, allbands, mocs_list, mcuts, ndet=2):\n",
    "    \"\"\"\n",
    "    param.\n",
    "    ---------------\n",
    "    catalogue: astropy.table\n",
    "        HELP catalogue \n",
    "    allbands: [[..., str]]\n",
    "        Names of bands used in the selections: [[opt_bands], [nir_bands], [mir_bands]]\n",
    "    mocs_list: : [..., str] \n",
    "        list containing the mocs fits files.\n",
    "    mcuts: dict\n",
    "        magnitude cuts limits for each band.\n",
    "    ndet: int\n",
    "        number of detections required per group of bands (opt/nir/mid). Default 2; it requires 2 detection on each to pass the selection.\n",
    "    \n",
    "    output\n",
    "    ---------------\n",
    "    nb_band: astropy.Table \n",
    "        Table with flags for each source indicating whether the source has at least \"ndet\" detections above the magnitude cut.\n",
    "        Flags: optband=1, nirband=2, mirband=4, ndet_total = (optband + nirband + mirband)\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "    from astropy.table import Column, Table, join\n",
    "    import time\n",
    "    from pymoc import MOC\n",
    "    \n",
    "    \n",
    "    nb_band = {}\n",
    "    nb_band['help_id'] = catalogue['help_id']\n",
    "    nb = ['optband', 'nirband', 'mirband']\n",
    "    nx = ['optvalue', 'nirvalue', 'mirvalue']\n",
    "    \n",
    "    r_filter = np.zeros(len(catalogue))\n",
    "\n",
    "    # Name of all moc files \n",
    "    mocs = [x[12:] for x in mocs_list]\n",
    "    path_mocs = mocs_list[0][:12]\n",
    "                  \n",
    "\n",
    "    for idx, bands in enumerate(allbands):\n",
    "        nb_band[nb[idx]] = np.zeros(len(catalogue))\n",
    "        nb_filter = np.zeros(len(catalogue))\n",
    "\n",
    "        for i, filt in enumerate(bands):\n",
    "            # mocs in that filter\n",
    "            mocs_filt = [moc for moc in mocs if filt in moc]\n",
    "            ind = np.zeros(len(catalogue))\n",
    "            for moc_ in mocs_filt:\n",
    "                # Get the name of the filter and survey\n",
    "                m = moc_[:-18]\n",
    "                \n",
    "                # Read in the moc for the specific survey     \n",
    "                moc = MOC()\n",
    "                moc.read(path_mocs+moc_)\n",
    "                ind_moc = inMoc(catalogue['ra'], catalogue['dec'],moc)\n",
    "                \n",
    "                # Source 'detection'\n",
    "                ind_filter = 1 * (catalogue['m_'+filt] <= mcuts[m]) # Sources with magnitude below the magnitud cut defined for that survey\n",
    "                ind_filter = ind_filter * ind_moc # Sources within the moc survey\n",
    "                ind += ind_filter\n",
    "\n",
    "            nb_filter += ind \n",
    "            # Count sources detected on the 'r' filter\n",
    "            if '_r' in filt:\n",
    "                r_filter += ind                   \n",
    "                \n",
    "        # sources detected in at least two filters of those bands (opt, nir, mir)\n",
    "        if idx==2:\n",
    "            n = 2\n",
    "        else:\n",
    "            n = 1\n",
    "        # add flag: 1 * opt ; 2 * nir ; 4 * mir    \n",
    "        nb_band[nb[idx]] += (idx+n) * (nb_filter >= ndet)  \n",
    "    \n",
    "    nb_band = Table(nb_band)\n",
    "    \n",
    "    # at least 1det in r band\n",
    "    has_r_det = 1 * (r_filter >= 1)\n",
    "    nb_band['optband'] = nb_band['optband'] * has_r_det\n",
    "            \n",
    "    total = np.zeros(len(catalogue))\n",
    "    for i in range(len(allbands)):\n",
    "        total += nb_band[nb[i]]\n",
    "\n",
    "    nb_band['ndet_total'] = total\n",
    "    \n",
    "    return nb_band\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def select_sample_subfiles(file_list, allbands, mocs_list, mcut, ndet=2):\n",
    "    \n",
    "    nb = ['optband', 'nirband', 'mirband']\n",
    "    names_det = ['help_id'] + nb[:len(allbands)] + ['ndet_total']\n",
    "    nb_band = Table(names=names_det, dtype=('str','float','float','float'))  \n",
    "    \n",
    "    # read sub_files to calculate snr per each band\n",
    "    ################################################\n",
    "    print('Reading sub_catalogues')\n",
    "    start = time.time()  \n",
    "    for idx, sub_file in enumerate(file_list):\n",
    "        print('subfile: ', idx)\n",
    "        start_sub = time.time() \n",
    "        \n",
    "        sub_catalogue = Table.read(sub_file)\n",
    "        sub_nb_band = select_sample(sub_catalogue, allbands, mocs_list, mcut, ndet=ndet)\n",
    "        nb_band = vstack([nb_band, sub_nb_band])\n",
    "        \n",
    "        end_sub = time.time()\n",
    "        time_ = end_sub - start_sub\n",
    "        print('select_sample subcatalogue time: ', time_)\n",
    "        \n",
    "        del sub_catalogue\n",
    "        del sub_nb_band\n",
    "        gc.collect()\n",
    "    \n",
    "    end = time.time()\n",
    "    time_all = end - start\n",
    "    print('select_sample total time: ', time_all)    \n",
    "    \n",
    "    return nb_band     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compare_cats_venndiagram(cat1, cat2, ind_catalogue, ind_names, cat1name, cat2name, save=True):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    for i, ind_cat in enumerate(ind_catalogue):\n",
    "#             print('joining catalogues')\n",
    "        plt.figure(figsize=(10,10))\n",
    "        set1 = set(cat1.tolist())\n",
    "        set2 = set(cat2[ind_cat].tolist())\n",
    "\n",
    "        out = venn2([set1, set2], (cat1name, cat2name))\n",
    "\n",
    "        for text in out.set_labels:\n",
    "            text.set_fontsize(20)\n",
    "        plt.title(ind_names[i])\n",
    "        \n",
    "        if save==True:\n",
    "            plt.savefig('./data/figs/'+ind_names[i]+'.png', format='png')\n",
    "            plt.savefig('./data/figs/'+ind_names[i]+'.pdf', format='pdf')\n",
    "        plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. SELECTION FUNCTION\n",
    "\n",
    "### 1.1 Read in files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### MOCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moc_list = glob.glob('./mocs/*.fits')\n",
    "len(moc_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Masterlist / Subfiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_subfiles = True\n",
    "\n",
    "if use_subfiles == False:\n",
    "    masterfile = DMU_DIR + 'dmu1/dmu1_ml_Herschel-Stripe-82/data/master_catalogue_herschel-stripe-82_{}.fits'.format(SUFFIX) # too big\n",
    "    masterlist = Table.read(masterfile, memmap=True)\n",
    "\n",
    "else:\n",
    "    file_list = glob.glob(DMU_DIR + 'dmu1/dmu1_ml_Herschel-Stripe-82/data/tiles/sub_catalogue_herschel-stripe-82_{}_*.fits'.format(SUFFIX))\n",
    "    print('There are {} subfiles'.format(len(file_list)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Bands definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HELP bands\n",
    "############\n",
    "\n",
    "u = [                        'sdss_u']\n",
    "g = ['suprime_g', 'decam_g', 'sdss_g']\n",
    "r = ['suprime_r', 'decam_r', 'sdss_r']\n",
    "i = ['suprime_i', 'decam_i', 'sdss_i']\n",
    "z = ['suprime_z', 'decam_z', 'sdss_z']\n",
    "y = ['suprime_y', 'decam_y']\n",
    "\n",
    "y_nir = ['vista_y',  'ukidss_y']\n",
    "h = ['vista_h',  'ukidss_h']\n",
    "j = ['vista_j',  'ukidss_j']\n",
    "k = ['vista_ks', 'ukidss_k']\n",
    "\n",
    "helpbands = [u + g + r + i + z + y] + [y_nir + h + j + k] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Magnitude cuts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SDSS\n",
    "magcuts_sdss = {'sdss_u': 22.8,\n",
    " 'sdss_g': 23.7,\n",
    " 'sdss_r': 23.5,\n",
    " 'sdss_i': 22.9,\n",
    " 'sdss_z': 21.5}\n",
    "\n",
    "# HSC\n",
    "magcuts_hsc = {\n",
    "    'suprime_g': 25.4,\n",
    "    'suprime_r': 25.4, \n",
    "    'suprime_i': 25.1,\n",
    "    'suprime_z': 24.2,\n",
    "    'suprime_y': 23.7}\n",
    "\n",
    "# DECam: DECaLS + DES\n",
    "magcuts_des = {'decam_g_des': 23.3,\n",
    " 'decam_r_des': 23.1,\n",
    " 'decam_i_des': 22.7,\n",
    " 'decam_z_des': 21.8,\n",
    " 'decam_y_des': 20.6}\n",
    "\n",
    "magcuts_decals =  {'decam_g_decals': 23.9, 'decam_r_decals': 23.5, 'decam_z_decals': 22.5}\n",
    "\n",
    "# VISTA: VHS + VICS82\n",
    "magcuts_vhs = {'vista_y_vhs': 20.1,\n",
    " 'vista_h_vhs': 19.6,\n",
    " 'vista_j_vhs': 19.9,\n",
    " 'vista_ks_vhs': 19.6}\n",
    "\n",
    "magcuts_vics = {'vista_j_vics': 21.0,\n",
    " 'vista_ks_vics': 21.0}\n",
    "\n",
    "# UKIDSS - LAS\n",
    "magcuts_ukidss = {'ukidss_h':19.0, 'ukidss_j':19.0, 'ukidss_k':19.0, 'ukidss_y':20.2}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mcuts = magcuts_sdss.copy()\n",
    "\n",
    "mcuts.update(magcuts_hsc)\n",
    "mcuts.update(magcuts_des)\n",
    "mcuts.update(magcuts_decals)\n",
    "mcuts.update(magcuts_vhs)\n",
    "mcuts.update(magcuts_vics)\n",
    "mcuts.update(magcuts_ukidss)\n",
    "mcuts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Flagged catalogue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if use_subfiles == False:\n",
    "    ind_cat = select_sample(masterlist, helpbands, moc_list, mcuts)\n",
    "\n",
    "else:\n",
    "    ind_cat = select_sample_subfiles(file_list, helpbands, moc_list, mcuts)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind_cat.write('./data/ind_cat_sf_{}.fits'.format(SUFFIX),overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## II Compare with Ldust catalogue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ldustfile = DMU_DIR + 'dmu28/dmu28_Herschel-Stripe-82/data/Ldust/old/herschel_stripe_82_Ldust_prediction_all.fits'\n",
    "#ldust_cat = Table.read(ldustfile, memmap=True)\n",
    "#ldust_cat['id'].name = 'help_id'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# indexs catalogues based on flags\n",
    "#ind_catalogue_ldustcriteria = (ind_catalogue['ndet_total']==3) | (ind_catalogue['ndet_total']==7)\n",
    "# ind_catalogue_iraccriteria = ind_catalogue['ndet_total']>=5\n",
    "\n",
    "\n",
    "# Compare sources after selection criteria with ldust catalogue \n",
    "# Plotting Venn diagrams \n",
    "#########################################\n",
    "#ind_catalogues = np.array([\n",
    "#    ind_catalogue_ldustcriteria,\n",
    "#     ind_catalogue_iraccriteria\n",
    "#    ])\n",
    "\n",
    "#ind_names =  np.array([\n",
    "#    '_ldustcriteria_{}'.format(SUFFIX),\n",
    "#     cat_name+'_iraccriteria'\n",
    "#])\n",
    "\n",
    "# Venn Diagrams\n",
    "#compare_cats_venndiagram(ldust_cat['help_id'], ind_catalogue['help_id'], ind_catalogues, ind_names, 'ldust', 'sf')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## III Selected catalogue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Define the final sample\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hs82_cat_selected = masterlist[ind_catalogue_ldustcriteria]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# len(hs82_cat_selected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hs82_cat_selected.write('./data/hs82_selected_sources_sf_{}.fits'.format(SUFFIX),overwrite=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (herschelhelp_internal)",
   "language": "python",
   "name": "helpint"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
