import numpy as np
from pdg_const import pdg
import matplotlib
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
import scipy.integrate as integrate

modl = rm.model()

load_set = False
draw = True
fit = False

mode = "slim_points"

modl.mode = mode

set_size = 4e5
nonres_set_size = 44000
# nonres_set_size = 1000
nr_of_toys = 1

# x_min = 3150.0
# x_max= 3650.0

x_max = modl.x_max
x_min = modl.x_min

modl.add_nonres()

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

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

modl.add_cusp(3550, 3e-7, 200, 7, 20000)

# print(integrate.quad(modl.total_pdf, modl.x_min**2, modl.x_max**2, limit = 3000))

# modl.param_val[0] = 1e7

# modl.normalize_pdf()

# print(modl.param_val[0])

# print(integrate.quad(modl.total_pdf, modl.x_min**2, modl.x_max**2, limit = 250))

modl.param_list()

modl.mode = "no_data"

modl.draw_plots(part_set = 1, x_min = x_min, x_max = x_max, mode = "no_data")

modl.mode = mode
# print(modl.total_scale_amp)

if load_set:

    if mode == "true_data":

        with open(r"./data/true_data/true_data_toy_{0}_range({1}-{2}).pkl".format(0, int(x_min), int(x_max)), "rb") as input_file:
            part_set = pkl.load(input_file)

    elif mode == "slim_points":

        with open(r"./data/slim_points/slim_points_toy_0_range({0}-{1}).pkl".format(int(x_min), int(x_max)), "rb") as input_file:
            part_set = pkl.load(input_file)

else:
    start = time.time()
    part_set = modl.generate_points(set_size, x_min = x_min, x_max = x_max, mode = mode, verbose = 1, nr_of_toys = nr_of_toys, nonres_set_size = nonres_set_size)
    print(time.time() - start)
if draw:
    modl.draw_plots(part_set = part_set, x_min = x_min, x_max = x_max, mode = mode)

if fit:
    modl.fit_pdf_to_data(part_set)

print()

print("Run finished")