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 = False

draw = True

mode = "full_points"

min_bin_scaling = 100

set_size = 4e5

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)


if load_set:

    with open(r"set.pkl", "rb") as input_file:
        set_dic = pkl.load(input_file)
    
    part_set = (set_dic["x_part"], set_dic["y_part"])
    counter = set_dic["counter"]

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

if draw:
    modl.draw_plots(part_set = part_set, counter = counter, mode = mode, min_bin_scaling = min_bin_scaling)
    

print("Run finished")