{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2e64c206",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import logging\n",
    "from pathlib import Path\n",
    "import numba as nb\n",
    "\n",
    "import awkward as ak\n",
    "import click\n",
    "import h5py\n",
    "import numpy as np\n",
    "import vector\n",
    "\n",
    "#from src.data.cms.convert_to_h5 import MIN_JETS, N_JETS, N_FJETS\n",
    "\n",
    "vector.register_awkward()\n",
    "\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "146f4ece",
   "metadata": {},
   "outputs": [],
   "source": [
    "HIGGS_MASS = 125.0\n",
    "MIN_JETS = 6\n",
    "N_JETS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0e95337e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a function that loads jets from hhh_test.h5\n",
    "def load_jets(in_file):\n",
    "    # load jets from the h5\n",
    "    pt = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"pt\"])\n",
    "    eta = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"eta\"])\n",
    "    phi = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"phi\"])\n",
    "    btag = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"btag\"])\n",
    "    mass = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"mass\"])\n",
    "    mask = ak.Array(in_file[\"INPUTS\"][\"Jets\"][\"MASK\"])\n",
    "\n",
    "    jets = ak.zip(\n",
    "        {\n",
    "            \"pt\": pt,\n",
    "            \"eta\": eta,\n",
    "            \"phi\": phi,\n",
    "            \"btag\": btag,\n",
    "            \"mass\": mass,\n",
    "            \"mask\": mask\n",
    "        },\n",
    "        with_name=\"Momentum4D\",\n",
    "    )\n",
    "    \n",
    "    return jets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9e14d4d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "JET_ASSIGNMENTS = {}\n",
    "for nj in range(MIN_JETS, N_JETS + 1):\n",
    "    a = list(itertools.combinations(range(nj), 2))\n",
    "    b = np.array([(i, j, k) for i, j, k in itertools.combinations(a, 3) if len(set(i + j + k)) == MIN_JETS])\n",
    "    JET_ASSIGNMENTS[nj] = b\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bee5d870",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_file = \"//Users/billyli/UCSD/hhh/reports/bv2/hhh_test.h5\"\n",
    "in_file = h5py.File(in_file)\n",
    "jets = load_jets(in_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "01febbff",
   "metadata": {},
   "outputs": [],
   "source": [
    "nj = 6\n",
    "mjj = (jets[:, JET_ASSIGNMENTS[nj][:, :, 0]] + jets[:, JET_ASSIGNMENTS[nj][:, :, 1]]).mass\n",
    "chi2 = ak.sum(np.square(mjj - HIGGS_MASS), axis=-1)\n",
    "chi2_argmin = ak.argmin(chi2, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "664a1ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi2_argmin = chi2_argmin.to_numpy(allow_missing=False).reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fa25c028",
   "metadata": {},
   "outputs": [],
   "source": [
    "chi2 = ak.to_numpy(chi2, allow_missing=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c9d7f1be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  872.79700636],\n",
       "       [ 9803.28624733],\n",
       "       [35882.93413304],\n",
       "       ...,\n",
       "       [ 2491.51387976],\n",
       "       [49792.71390584],\n",
       "       [13780.4842438 ]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chi2_sum = np.take_along_axis(chi2, chi2_argmin, axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "65107a86",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([35882.93413304])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chi2[2,chi2_argmin[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b7a85e9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}