import ROOT
#from ROOT import TTree, TFile, Double
import numpy as np
from pdg_const import pdg
import matplotlib
matplotlib.use("Qt5Agg")
import matplotlib.pyplot as plt
import pickle as pkl
import sys
import time
from helperfunctions import display_time
import cmath as c
import raremodel as rm

    
modl = rm.model()

load_set = True

draw = False

mode = "slim_points"

modl.mode = mode

min_bin_scaling = 100

set_size = 1e3

x_min = 211.0

x_max= 4781.0

jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale = pdg["jpsi"]
modl.add_resonance(jpsi_mass, jpsi_width, jpsi_phase, jpsi_scale)

psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale = pdg["psi2s"]
modl.add_resonance(psi2s_mass, psi2s_width, psi2s_phase, psi2s_scale)

modl.add_cusp(3525, 3e-7, 200, 7)

modl.normalize_pdf()

if load_set:

    with open(r"./data/set_{0}_range{1}-{2}.pkl".format(int(set_size), int(x_min), int(x_max)), "rb") as input_file:
        set_dic = pkl.load(input_file)
    
    part_set = (set_dic["x_part"], set_dic["y_part"])
    counter_tot = set_dic["counter_tot"]

else:
    x_part, y_part, counter = modl.generate_points(set_size, mode = mode, min_bin_scaling = min_bin_scaling, verbose = 1)
    part_set = (x_part, y_part)

if draw:
    modl.draw_plots(part_set = part_set, counter = counter_tot, mode = mode, min_bin_scaling = min_bin_scaling)
    
cusp_amp_scan_y = []
cusp_amp_scan_x = np.linspace(0, 2*modl.cusp_amp, 100)

print("Likelihood scan starting")

counter = 0

for i in cusp_amp_scan_x:
    cusp_amp_scan_y.append(modl.log_likelihood(x_part = set_dic["x_part"], cusp_amp = i))
    counter +=1
    print("{0}{1}".format(counter, "%"))

print("Scan finished")
print("Plotting...")

plt.clf()

plt.title("Cusp amp log_likelihood scan")

#plt.yscale("log")

#plt.ylim(0, 2*self._mean)

plt.grid()

plt.plot(cusp_amp_scan_x, cusp_amp_scan_y, label = "log_likelihood")

plt.legend()

plt.savefig("./cusp_amp_scan.png")

print(" pdf_and_parts.png created")

print(modl.log_likelihood(x_part = set_dic["x_part"], cusp_amp = modl.cusp_amp))

scan_set = (cusp_amp_scan_x, cusp_amp_scan_y)
            
pkl.dump( scan_set, open("scan_set.pkl" , "wb" ) )

print("Run finished")