"""Test perfect survey."""
import matplotlib.pyplot as plt

from frbpoppy import CosmicPopulation, Survey, SurveyPopulation, hist
from tests.convenience import plot_aa_style, rel_path

pop = CosmicPopulation.simple(1e4, generate=False)
pop.generate()

survey = Survey('perfect')

survey_pop = SurveyPopulation(pop, survey)

plot_aa_style()
plt.rcParams["figure.figsize"] = (5.75373, 5.75373)

f, axes = plt.subplots(2, 2)
s = survey_pop.frbs
axes[0, 0].step(*hist(s.snr, bin_type='log'), where='mid')
axes[0, 0].set_title('SNR')

axes[1, 0].step(*hist(s.T_sys, bin_type='log'), where='mid')
axes[1, 0].set_title(r'T$_{\text{sys}}$')

axes[0, 1].step(*hist(s.s_peak, bin_type='log'), where='mid')
axes[0, 1].set_title(r'S$_{\text{peak}}$')

axes[1, 1].step(*hist(s.w_arr, bin_type='log'), where='mid')
axes[1, 1].set_title(r'w$_{\text{arr}}$')

for x in [0, 1]:
    for y in [0, 1]:
        axes[x, y].set_xscale('log')
        axes[x, y].set_yscale('log')

plt.tight_layout()
plt.savefig(rel_path('plots/perfect.pdf'))
