{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hep/davide/miniconda3/envs/root_env/lib/python2.7/site-packages/root_numpy/_tree.py:5: RuntimeWarning: numpy.dtype size changed, may indicate binary incompatibility\n",
      "  from . import _librootnumpy\n"
     ]
    }
   ],
   "source": [
    "import root_numpy as rn\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import ROOT as r\n",
    "import pickle\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_name='DplusMuNu_Full'\n",
    "file_path='/disk/lhcb_data/davide/HCAL_project_full_event/'+file_name+'.root'\n",
    "tree_name='ntuple/DecayTree'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = r.TFile(file_path)\n",
    "t = f.Get(tree_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_tracker_info(particle):\n",
    "    \n",
    "    t.SetBranchStatus(\"*\",0)\n",
    "    \n",
    "    t.SetBranchStatus(particle+\"_L0Calo_HCAL_xProjection\",1)\n",
    "    t.SetBranchStatus(particle+\"_L0Calo_HCAL_yProjection\",1)\n",
    "    t.SetBranchStatus(particle+\"_L0Calo_HCAL_realET\",1)\n",
    "    t.SetBranchStatus(particle+\"_L0Calo_HCAL_region\", 1)\n",
    "\n",
    "    x_projections=[]\n",
    "    y_projections=[]\n",
    "    real_ET=[]\n",
    "    region=[]\n",
    "\n",
    "    for i, event in enumerate(t):\n",
    "        \n",
    "        #if i > 2:\n",
    "        #    break\n",
    "\n",
    "        region.append(getattr(event,particle+\"_L0Calo_HCAL_region\"))\n",
    "        x_projections.append(getattr(event,particle+\"_L0Calo_HCAL_xProjection\"))\n",
    "        y_projections.append(getattr(event,particle+\"_L0Calo_HCAL_yProjection\"))\n",
    "        real_ET.append(getattr(event,particle+\"_L0Calo_HCAL_realET\"))\n",
    "        \n",
    "    true_events ={'true_x':np.array(x_projections),'true_y':np.array(y_projections),'true_ET':np.array(real_ET),'region':np.array(region)}\n",
    "    \n",
    "    return true_events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_events_piplus = get_tracker_info(\"piplus\")\n",
    "true_events_piplus0 = get_tracker_info(\"piplus0\")\n",
    "true_events_Kminus = get_tracker_info(\"Kminus\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(465693,)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_events_piplus0['region'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(24401,)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_events_piplus0['region'][np.where(true_events_piplus0['region']<0)].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1131.9902543165215, 507.3645231436275, 1)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "true_events_Kminus['true_x'][0],true_events_Kminus['true_y'][0],true_events_Kminus['region'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame.from_dict(true_events_piplus)\n",
    "df.to_csv('/disk/lhcb_data/davide/HCAL_project_full_event/csv/piplus/MCtrue_piplus.csv', index=False)\n",
    "\n",
    "df = pd.DataFrame.from_dict(true_events_piplus0)\n",
    "df.to_csv('/disk/lhcb_data/davide/HCAL_project_full_event/csv/piplus0/MCtrue_piplus0.csv', index=False)\n",
    "\n",
    "df = pd.DataFrame.from_dict(true_events_Kminus)\n",
    "df.to_csv('/disk/lhcb_data/davide/HCAL_project_full_event/csv/Kminus/MCtrue_Kminus.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = r.TBrowser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}