import sys,os
#import uproot4 as uproot
#import awkward as ak
import json
import numpy as np
import matplotlib.pyplot as plt
#from fcc_python_tools.locations import loc
#from fcc_python_tools import kinematics_flat
from particle import literals as lp
#from fcc_python_tools import plotting
import tensorflow as tf
import zfit
import random
import matplotlib.ticker as tick
from scipy.stats import norm
from scipy.stats import linregress
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.axes_grid1 import ImageGrid
from scipy.optimize import curve_fit
from seaborn import regplot

from matplotlib import rc
rc('font',**{'family':'serif','serif':['Roman']})
rc('text', usetex=True)

parameters = {'axes.labelsize': 20,
          'axes.titlesize': 20, 'xtick.labelsize':15,'ytick.labelsize':15, 'figure.titlesize':23, 'legend.fontsize':15 }
plt.rcParams.update(parameters)
import ROOT
rand=0

npsigmass=np.loadtxt("data1.txt")

low = 4.6
high = 6

bins=int((6-4.6)/0.03)
bin_w = (high - low)/bins
sig_hist=ROOT.TH1D("sig mass", "sig mass", bins, 4.6, 6.0)
for i in npsigmass:
    sig_hist.Fill(i, 175.547295/len(npsigmass))
sig_hist.SetFillColor(906)
stack_hist=ROOT.THStack("stack histo", f"Invariant B0 mass with sel solutions and natural number of event")
stack_hist.Add(sig_hist)
legend=ROOT.TLegend(0.5,0.25,0.89,0.6)
legend.AddEntry(sig_hist,"sig","f")
legend.SetFillStyle(0)
ttext=ROOT.TPaveText(0.3,0.65,0.89,0.89,"NDC")
ttext.SetTextSize(0.04)
ttext.AddText(f"PV (XX"+"\\mu m)"+f" \\& SV \\& TV XX\mu m, XX\\mu m) SMEARED")
ttext.AddText("Probability\tto\tidentify\ta\t\\pi^0"+f"=XX")
ttext.SetFillColor(0)
ttext.SetBorderSize(0)
c1=ROOT.TCanvas()
stack_hist.Draw("HIST")
xaxis=stack_hist.GetXaxis()
xaxis.SetTitle("m(K^{*} [3\\pi]_\\tau [3\\pi]_\\tau) [GeV/c^{2}]")
yaxis=stack_hist.GetYaxis()
yaxis.SetTitle(f"Candidates / ({bin_w:.2f}"+ "GeV/c^{2})")
legend.Draw()
ttext.Draw()
c1.Draw()

mass = ROOT.RooRealVar("mass", "mass", 5., 6.)
bins=int((6-5)/0.03)

signal_mass = ROOT.RooDataSet.from_numpy(data={"mass": npsigmass}, variables=[mass])

low = 4.6
high = 6

rand+=1

#signal fit
#crystal ball parameters
mu=ROOT.RooRealVar(f"mu_{rand}", f"mu_{rand}", lp.B_0.mass/1000., low, high)
sigma = ROOT.RooRealVar(f"sigma_{rand}", f"sigma_{rand}", 0.5, 0., 2.)
alphaR = ROOT.RooRealVar(f"alphaR_{rand}", f"alphaR_{rand}", 0.2, 0., 5.)
nR = ROOT.RooRealVar(f"nR_{rand}", f"nR_{rand}", 10., 0., 200.)
alphaL = ROOT.RooRealVar(f"alphaL_{rand}", f"alphaL_{rand}", 0.2, 0., 5.)
nL = ROOT.RooRealVar(f"nL_{rand}", f"nL_{rand}", 10., 0., 200.)
# Core gaussian parameters
sigmaG1 = ROOT.RooRealVar(f"sigmaG1_{rand}", f"sigmaG1_{rand}", 0.015, 0., 1)
# Fractions
frac = ROOT.RooRealVar(f"frac_{rand}", f"frac_{rand}", 0.5, 0., 1.)

#signal model
CB = ROOT.RooCrystalBall(f"CB_{rand}",f"CB_{rand}",x=mass,x0=mu,sigmaLR=sigma,alphaR=alphaR, nR=nR,alphaL=alphaL, nL=nL )
Gau = ROOT.RooGaussian(f"gau_{rand}",f"gau_{rand}",_x=mass,_mean=mu,_sigma=sigmaG1)
sig_fit= ROOT.RooAddPdf( f"sig_fit_{rand}",f"sig_fit_{rand}",CB, Gau,frac)

sig_fit.fitTo(signal_mass)#,ROOT.RooFit.Extended(),ROOT.RooFit.NumCPU(7))
c22=ROOT.TCanvas()
Mframe = mass.frame(bins)
signal_mass.plotOn(Mframe)
sig_fit.plotOn(Mframe)
Mframe.Draw()
c22.Draw('mass')