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

# # Introduction
# 
# QA plots for the TPC detector

# In[1]:


# imports to write dynamic markdown contents
import os
from IPython.display import display, Markdown, Latex
from IPython.display import HTML


# In[2]:


# turn off/on code for the result HTML page

display(Markdown('*For the result HTML page:* '))
    
HTML('''<script>
code_show=false; 
function code_toggle() {
 if (code_show){
 $('div.input').hide();
 } else {
 $('div.input').show();
 }
 code_show = !code_show
} 
$( document ).ready(code_toggle);
</script>
<form action="javascript:code_toggle()"><input type="submit" value="Click here to toggle on/off the raw code."></form>''')


# 
# # Initialization

# In[3]:


import uproot


import matplotlib.pyplot as plt
plt.style.use('seaborn-white')

import numpy as np


# # 2D Field map export

# In[4]:


MapFile2D = '/cvmfs/eic.opensciencegrid.org/ecce/gcc-8.3/release/release_prop/prop.7/share/calibrations//Field/Map/sPHENIX.2d.root'

# 2D field map scaling on its working point
# https://github.com/ECCE-EIC/macros/blob/prop.7.0/detectors/EICDetector/Fun4All_G4_EICDetector.C#L460 
magfield_2D_rescale = -1.4 / 1.5;


# In[5]:


file = uproot.open(MapFile2D)


# In[6]:


file.keys()


# In[7]:


file['fieldmap'].keys()


# In[8]:


df_fieldmap = file['fieldmap'].arrays(library="pd")

df_fieldmap['br'] = df_fieldmap['br'].apply(lambda x: x * magfield_2D_rescale)
df_fieldmap['bz'] = df_fieldmap['bz'].apply(lambda x: x * magfield_2D_rescale)


# In[9]:


df_fieldmap


# In[10]:


X=df_fieldmap['z'].values
Y=df_fieldmap['r'].values
Z=df_fieldmap.apply(lambda x: np.sqrt(x['br']*x['br'] + x['bz']*x['bz']), axis=1).values

limits =  np.arange(0,2.5,.1);


# ## Check field map plot

# In[11]:


fig, ax = plt.subplots(figsize=(16,6))
ax.set_aspect('equal')
ax.tricontourf(X,Y,Z,limits, cmap='Spectral_r') 
contours = ax.tricontour(X,Y,Z, limits)
ax.clabel(contours, inline=True, fontsize=8)

ax.set_xlabel('z [cm]')
ax.set_ylabel('R [cm]')
ax.set_title("Total B strengh [T] : " + MapFile2D)

fig.show()


# In[12]:


df_fieldmap.to_csv('sPHENIX.2d.csv');


# # 3D field map

# In[ ]:





# In[13]:


MapFile3D = '/cvmfs/eic.opensciencegrid.org/ecce/gcc-8.3/release/release_prop/prop.7/share/calibrations//Field/Map/sphenix3dbigmapxyz.root'


# In[14]:


file = uproot.open(MapFile3D)


# In[15]:


file.keys()


# In[16]:


file['fieldmap'].keys()


# In[17]:


df_fieldmap = file['fieldmap'].arrays(library="pd")


# In[18]:


df_fieldmap


# In[19]:


df_fieldmap_yz = df_fieldmap[df_fieldmap['x'] == 0]


# In[20]:


df_fieldmap_yz


# In[21]:


X=df_fieldmap_yz['z'].values
Y=df_fieldmap_yz['y'].values
Z=df_fieldmap_yz.apply(lambda x: np.sqrt(x['bx']*x['bx'] + x['by']*x['by'] + x['bz']*x['bz']), axis=1).values

limits =  np.arange(0,2.5,.1);


# ## Check field map plot

# In[22]:


fig, ax = plt.subplots(figsize=(16,16))
ax.set_aspect('equal')
ax.tricontourf(X,Y,Z,limits, cmap='Spectral_r') 
contours = ax.tricontour(X,Y,Z, limits)
ax.clabel(contours, inline=True, fontsize=8)

ax.set_xlabel('z [cm]')
ax.set_ylabel('y [cm]')
ax.set_title("Total B strengh [T] @ x=0 : " + MapFile3D)

fig.show()


# In[23]:


df_fieldmap.to_csv('sphenix3dbigmapxyz.csv');


# In[ ]: