#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import opengate as gate
from test049_pet_digit_blurring_helpers import *
import opengate.contrib.spect_ge_nm670 as gate_spect
import matplotlib.pyplot as plt

paths = gate.get_default_test_paths(__file__, "")

"""
Check the options in DigitizerAdderActor
- TimeDifference: to store the max time difference within all hits in one single
- NumberOfHits: store the number of hits in one single
"""

# create the simulation
sim = gate.Simulation()

# units
m = gate.g4_units("m")
nm = gate.g4_units("nm")
mm = gate.g4_units("mm")
cm = gate.g4_units("cm")
Bq = gate.g4_units("Bq")
sec = gate.g4_units("s")
min = gate.g4_units("min")

# verbose
ui = sim.user_info
ui.visu = False
ui.g4_verbose = False
ui.g4_verbose_level = 1
ui.number_of_threads = 1
ui.random_seed = 123456

# world size
world = sim.world
world.size = [2.5 * m, 2.5 * m, 2.5 * m]
world.material = "G4_AIR"

# spect head without collimator
spect, crystal = gate_spect.add_ge_nm67_spect_head(
    sim, "spect", collimator_type=False, debug=False
)
spect.translation = [0, 0, -15 * cm]

# spect digitizer
hc = sim.add_actor("DigitizerHitsCollectionActor", f"Hits")
hc.mother = crystal.name
hc.output = ""  # No output
hc.attributes = [
    "PostPosition",
    "TotalEnergyDeposit",
    "PreStepUniqueVolumeID",
    "GlobalTime",
]
sc = sim.add_actor("DigitizerAdderActor", f"Singles")
sc.mother = hc.mother
sc.input_digi_collection = hc.name
sc.time_difference = True
sc.number_of_hits = True
sc.output = paths.output / "test051_singles.root"

# physics
p = sim.get_physics_user_info()
p.physics_list_name = "G4EmStandardPhysics_option3"
p.enable_decay = True
sim.set_cut("world", "all", 100 * mm)
sim.set_cut("spect", "all", 0.1 * mm)

# source of Ac225 ion
s1 = sim.add_source("GenericSource", "source")
s1.particle = "ion 89 225"  # Ac225
s1.position.type = "sphere"
s1.position.radius = 1 * nm
s1.position.translation = [0, 0, 0]
s1.direction.type = "iso"
s1.activity = 1e4 * Bq / ui.number_of_threads

# add stat actor
s = sim.add_actor("SimulationStatisticsActor", "stats")
s.track_types_flag = True
s.output = paths.output / "test051_stats.txt"

# go
output = sim.start(start_new_process=True)

# check stats
print()
gate.warning("Check stats")
stats = output.get_actor("stats")
stats_ref = gate.read_stat_file(paths.output_ref / "test051_stats.txt")
is_ok = gate.assert_stats(stats, stats_ref, tolerance=0.04)

# check root
print()
gate.warning("Check root time difference")
root1, n1 = gate.open_root_as_np(paths.output_ref / "test051_singles.root", "Singles")
root2, n2 = gate.open_root_as_np(sc.output, "Singles")

# time difference
td1 = root1["TimeDifference"]
td2 = root2["TimeDifference"]
td1 = td1[td1 > 1 * sec] / min
td2 = td2[td2 > 1 * sec] / min
is_ok = (
    gate.check_diff_abs(
        len(td1),
        len(td2),
        20,
        f"Number of time diff larger than 1 sec (~{len(td1)/n1*100:.2f}%):",
    )
    and is_ok
)
is_ok = (
    gate.check_diff_abs(np.mean(td1), np.mean(td2), 30, f"Time diff mean in minutes:")
    and is_ok
)

print()
gate.warning("Check root nb of hits")
nh1 = root1["NumberOfHits"]
nh2 = root2["NumberOfHits"]
is_ok = (
    gate.check_diff(np.mean(nh1), np.mean(nh2), 6, f"Number of hits in mean:") and is_ok
)

# plot
f, ax = plt.subplots(1, 2, figsize=(25, 10))
gate.plot_hist(ax[0], td1, f"Time diff ref (minutes)")
gate.plot_hist(ax[0], td2, f"Time diff (minutes)")
gate.plot_hist(ax[1], nh1, f"Nb of hits ref")
gate.plot_hist(ax[1], nh2, f"Nb of hits")

fn = paths.output / "test051_singles.png"
plt.savefig(fn)
print(f"Plot in {fn}")

gate.test_ok(is_ok)
