#from ROOT import TTree, TFile, Double
import numpy as np
from pdg_const import pdg
import Tkinter
import matplotlib
matplotlib.use("agg")
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 ROOT


modl = rm.model()

load_set = False

draw = True

mode = "fast_binned"

modl.mode = mode

set_size = 1e5
nonres_set_size = 44000

# x_min = 3150.0
# x_max= 3650.0

x_min = modl.x_min
x_max = modl.x_max

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(3550, 6e-7, 200, 7)

modl.normalize_pdf()

print(modl.total_scale_amp)

if load_set:

    if mode == "true_data":

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

    else:

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

        counter = set_dic["counter_tot"]

else:
    part_set = modl.generate_points(set_size, x_min = x_min, x_max = x_max, mode = mode, verbose = 1)

if draw:
    modl.draw_plots(part_set = part_set, x_min = x_min, x_max = x_max, mode = mode)

    plt.clf()

    bin_mean = part_set["bin_mean"]
    bin_true_height = part_set["bin_true_height"]
    nbins = part_set["nbins"]

    plt.hist(bin_mean, bins=nbins, range=(x_min, x_max), weights = bin_true_height, label = "toy data binned")

    plt.savefig("./plots/fast_binned/test.png")

# 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")