import matplotlib.pyplot as plt
from BH_mass_function import gene_BHBH, dl, solve_z
from cal_likelihood import fac_s_eff_v
import pickle
import glob
import random
import scipy.optimize as op
import time
a, mbh_max, mbh_min = 2.35, 80., 5.
event_rate_list = []
zs_all_list = []
#%%
sce = 2
test = gene_BHBH(h0=70, rho0=0, m_type='BHBH', scenario=sce)
event_rate, zs_detected, masses, rhos_detected = test.mc_year_rate(
a=a, mbh_max=mbh_max, mbh_min=mbh_min, itera=5)
zs_all, chirp_mass_all, m1_all, m2_all = zs_detected, masses[:,
0], masses[:,
1], masses[:,
2]
index = np.arange(len(m1_all))
idx = random.sample(index, 5000)
m1 = m1_all[idx]
m2 = m2_all[idx]
rho = rhos_detected[idx]
zs = zs_all[idx]
#save_filename = 'simulated_{0}_low_end.txt'.format(test.m_type)
#save_file = open(save_filename,'w')
#save_file.write("#zs, m1, m2, SNR, total_rate = {0:.3f} \n".format(event_rate))
import astropy.io.fits as pyfits
import matplotlib.pyplot as plt
import glob
from BH_mass_function import gene_BHBH, dl, solve_z
import random
import time
solve_z = np.vectorize(solve_z)
#a0, a1, mbh_max, mbh_min = 2.35, 0.1, 80., 5.
a0, a1, mbh_max, mbh_min = 2.35, 0.7, 80., 5. #mode1
#filename = 'test3_sim_a_{0}_max_{1}_min_{2}'.format(round(a,2), round(mbh_max,1), round(mbh_min,1))
#if_file = glob.glob(filename)
seed = 1
np.random.seed(seed)
test = gene_BHBH(h0=70)
#event_rate0, zs_detected0, masses0, rhos_detected0 = test.mc_year_rate(a=a, mbh_max=mbh_max, mbh_min=mbh_min, ev_type = 'const')
#event_rate1, zs_detected1, masses1, rhos_detected1 = test.mc_year_rate(a=[2.35,0.1], mbh_max=mbh_max, mbh_min=mbh_min, ev_type = 'mode0')
event_rate2, zs_detected2, masses2, rhos_detected2 = test.mc_year_rate(
a=[a0, a1], mbh_max=mbh_max, mbh_min=mbh_min, ev_type='mode1')
dl_zs = dl(zs_detected2)
zs_all, chirp_mass_all, m1_all, m2_all, lumi_dis_all = zs_detected2, masses2[:,
0], masses2[:,
1], masses2[:,
2], dl_zs
#
##%% Test if the data slope follows the redshift trend.
#plt.hist(m1_all[zs_all==2.141], m1_all[zs_all==5.531],normed=True, log=True)
#plt.show()