import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

from nnfwtbn import Variable, correlation_matrix
from nnfwtbn import toydata

df = toydata.get()

v_higgs_m = Variable(r"$m^H$", "higgs_m", "GeV")
v_jet_1_pt = Variable(r"$p_{\mathrm{T}}^{j_1}$", "jet_1_pt", "GeV")
v_jet_2_pt = Variable(r"$p_{\mathrm{T}}^{j_2}$", "jet_2_pt", "GeV")
v_m_jj = Variable(r"$m^{jj}$", "m_jj", "GeV")
v_jet_1_eta = Variable(r"$\eta^{j_1}$" ,"jet_1_eta") 
v_jet_2_eta = Variable(r"$\eta^{j_2}$" ,"jet_2_eta")
v_tau_pt = Variable(r"$p_{\mathrm{T}}^{\tau}$", "tau_pt", "GeV")
v_lep_pt = Variable(r"$p_{\mathrm{T}}^{\ell}$", "lep_pt", "GeV")

fig, axes = plt.subplots(figsize=(5, 4.5))
correlation_matrix(df, [v_jet_1_pt, v_jet_2_pt, v_m_jj, v_higgs_m,
                        v_tau_pt, v_lep_pt, v_jet_1_eta, v_jet_2_eta],
                  figure=fig, axes=axes)

	$p_{\mathrm{T}}^{j_1}$	$p_{\mathrm{T}}^{j_2}$	$m^{jj}$	$m^H$	$p_{\mathrm{T}}^{\tau}$	$p_{\mathrm{T}}^{\ell}$	$\eta^{j_1}$	$\eta^{j_2}$
$p_{\mathrm{T}}^{j_1}$	1.000000	0.319978	0.211387	-0.007761	-0.065800	-0.071961	0.011964	0.013058
$p_{\mathrm{T}}^{j_2}$	0.319978	1.000000	0.368783	0.159946	-0.010026	-0.000068	0.016535	0.004242
$m^{jj}$	0.211387	0.368783	1.000000	0.456126	0.015672	0.000365	0.002525	0.004646
$m^H$	-0.007761	0.159946	0.456126	1.000000	0.237944	-0.011807	0.002297	0.015111
$p_{\mathrm{T}}^{\tau}$	-0.065800	-0.010026	0.015672	0.237944	1.000000	-0.120521	0.001576	0.006336
$p_{\mathrm{T}}^{\ell}$	-0.071961	-0.000068	0.000365	-0.011807	-0.120521	1.000000	-0.009070	0.006785
$\eta^{j_1}$	0.011964	0.016535	0.002525	0.002297	0.001576	-0.009070	1.000000	0.569375
$\eta^{j_2}$	0.013058	0.004242	0.004646	0.015111	0.006336	0.006785	0.569375	1.000000