import numpy as np
import matplotlib.pyplot as plt
import glob


f_size = 20
fig_r = (8, 5)
m = 1.3
fig_s = [k * m for k in fig_r]

# rates, vio, vio_err = np.genfromtxt('S_value_with_rates_expriement.dat').T
# rate_th, vio_th, _, _ = np.genfromtxt('simulation_rates_vs_S').T
# fig1 = plt.figure('rates vs violation', figsize=fig_s)
# plt.errorbar(rates / .81, vio, yerr=vio_err,
#              markersize=12,
#              fmt='o',
#              linewidth=3,
#              label=lab1)
# plt.plot(rate_th / .81, vio_th, linewidth=3, label=lab2)
# plt.xlim([0, 15000])
# plt.xlabel('generated pair rate (pairs/s)', fontsize=f_size)
# plt.ylabel('CHSH parameter S', fontsize=f_size)
# plt.gca().tick_params(labelsize=f_size)
# plt.legend(fontsize=f_size)
# plt.tight_layout()
# plt.savefig('vio_vs_rates.pdf')

def accumulate_and_plot_g2(folder):
    file_list = glob.glob(folder + '/*_corr24.g2')
    time, pairs, _, _ = np.genfromtxt(file_list[0]).T
    print(file_list[0])
    for file in file_list[1:]:
        print(file)
        data = np.genfromtxt(file)
        pairs += data[:, 1]
    plt.figure(figsize=fig_s)
    plt.errorbar(time - 500, pairs, yerr=np.sqrt(pairs), fmt='o',
                 markersize=12,
                 linewidth=3)
    plt.gca().tick_params(labelsize=f_size)

    # plt.plot(time,pairs, '-.')
    plt.xlabel('detection time difference (ns)', fontsize=f_size)
    plt.ylabel('coincidences', fontsize=f_size)
    plt.xlim(000, 1000)
    # print()
    # np.savetxt('summedG2.g2', np.column_stack((time, pairs)), header='time pairs')
    plt.tight_layout()
    plt.savefig('cavity_output.pdf')
    plt.show()


def statistics_of_g2(filename):
    data = np.genfromtxt(filename)
    time = data[:, 0]
    pairs = data[:, 1]
    time_start = 720
    time_end = 970
    time_idx_start = int(np.argwhere(time == time_start))
    time_idx_end = int(np.argwhere(time == time_end))
    print(np.sum(pairs[time_idx_start:time_idx_end]))
    # integral_pairs = np.sum(pairs[0])


def main():
    accumulate_and_plot_g2('good_data_2/*/')
    # statistics_of_g2('summedG2.g2')


if __name__ == '__main__':
    main()
