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

# Imports
import opengate as gate
import opengate_core as g4
import opengate.contrib.spect_ge_nm670 as gate_spect
import uproot
import sys
import matplotlib.pyplot as plt
import numpy as np
import os
import gatetools.phsp as phsp
from scipy import stats


"""
File: GE_Healthcare_NM_670_Discovery_Simulations.py

This code uses the [opengate package](https://github.com/OpenGATE/opengate) for simulation applications of the GE Healthcare Discovery NM 670 SPECT/CT detector.
It makes use of the pre-defined spect_ge_nm670.py file which is part of the previously mentioned package in the folder contrib.

Usage:
- The following settings can be used to adjust the simulation parameters
    - g_energy: Sets the energy of the source in kilo electron volt
    - g_activity: Sets the activity of the source in Becquerel
    - g_runtime: Sets the runtime of the simulation in secounds
    - g_channels: The number of channels that sould be monitored for the simulation. Increasing this value leads to a better spectral resolution.
    - g_waterbox: If this setting is True a 15cmx15cmx15cm waterbox is located around the source as a scatter medium.
    - g_visu: This setting creates a visualization of the setup if set to True.
    - g_distance: This value indicates the distance between the source and the detector crystal.

Date: September 2023
"""

# Set path to data, output and mac
paths = gate.get_default_test_paths(__file__, "gate_test000_hits_collection")


# Settings
key_of_interest = "KineticEnergy" # Use key "KineticEnergy" to use for the spectrum.
scalings = 1000  # Scaled to keV. For different energies has to be scaled differently.


# Colors:
blue_transparent = [0, 0, 1, 0.5] 


def format_magnitude(value, unit="Bq"):
    """
    Formats a given numerical value with a specified unit using SI unit prefixes.

    Parameters:
    - value (float or int): The numerical value to be formatted.
    - unit (str, optional): The unit to be appended to the formatted value (default is "Bq").

    Returns:
    - str: A formatted string representing the value with the appropriate SI unit prefix, e.g., "1.2 kBq".

    Example:
    >>> format_magnitude(1200, "Bq")
    '1.2 kBq'
    >>> format_magnitude(4500000, "Bq")
    '4.5 MBq'
    >>> format_magnitude(0, "Bq")
    '0 Bq'
    >>> format_magnitude(2500, "eV")
    '2.5 keV'

    Notes:
    - SI unit prefixes (k, M, G, T, P, E, Z, Y) are used to scale the value for readability.
    - The function automatically determines the appropriate prefix based on the magnitude of the input value.
    - If the input value is zero, the function returns '0' followed by the specified unit.

    Reference:
    - SI unit prefixes: https://en.wikipedia.org/wiki/Metric_prefix
    """

    prefixes = ['', 'k', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y']

    if value == 0:
        return f"0 {unit}"

    magnitude = int(value).bit_length() // 10

    if magnitude < 0:
        magnitude = 0
    elif magnitude > len(prefixes) - 1:
        magnitude = len(prefixes) - 1

    scaled_value = value / (1e3 ** magnitude)

    # Adjust formatting based on the last digit
    if scaled_value % 1 == 0:
        formatted_value = f"{scaled_value:.0f}"
    else:
        formatted_value = f"{scaled_value:.1f}"

    return f"{formatted_value} {prefixes[magnitude]}{unit}"




def create_simulation(g_energy, g_activity, g_runtime = 100, g_channels = 512, g_waterbox = False, g_visu = False, g_distance = 15, g_collimator_type = "lehr"):
    
    # Create the simulation object.
    sim = gate.Simulation()

    # Set the main options of the simulaiton.
    ui = sim.user_info
    ui.g4_verbose = False
    ui.visu = g_visu
    ui.start_new_process=True
    ui.visu_type = "vrml"
    ui.number_of_threads = 1
    ui.check_volumes_overlap = False

    # Setup of the Geant4 units.
    m = gate.g4_units("m")
    cm = gate.g4_units("cm")
    eV = gate.g4_units("eV")
    mm = gate.g4_units("mm")
    Bq = gate.g4_units("Bq")
    sec = gate.g4_units("second")

    # Define the world volume of size 1mx1mx1m.
    world = sim.world
    world.size = [1 * m, 1 * m, 1 * m]

    # Set the path to the Material database.
    sim.add_material_database(paths.data / "GateMaterials.db")


    # Based on the setting g_waterbox add a waterbox as a scatter medium around the source.
    if g_waterbox:
        waterbox = sim.add_volume("Box", "waterbox")
        waterbox.size = [15 * cm, 15 * cm, 15 * cm]
        waterbox.material = "G4_WATER"
        waterbox.color = blue_transparent
    


    # Add the GE Healthcare NM 670 Discovery spect head into the simulation.
    spect, crystal = gate_spect.add_ge_nm67_spect_head(sim, "spect", collimator_type=g_collimator_type, debug=False)
    psd = 6.11 * cm
    spect.translation = [0, 0, -(g_distance * cm + psd)]



    # Create the physics list.
    p = sim.get_physics_user_info()
    p.physics_list_name = "G4EmStandardPhysics_option4"
    p.enable_decay = True
    cuts = p.production_cuts
    cuts.world.gamma = 0.01 * mm
    cuts.world.electron = 0.01 * mm
    cuts.world.positron = 1 * mm
    cuts.world.proton = 1 * mm


    # Add the Tc-99m point source.
    source_Tc_99m = sim.add_source("GenericSource", "Tc-99m")
    source_Tc_99m.particle = "gamma"
    source_Tc_99m.energy.type = "gauss"
    source_Tc_99m.energy.mono = g_energy * eV  # Tc-99m gamma energy
    source_Tc_99m.activity = g_activity * Bq  # Activity of the point source
    source_Tc_99m.position.type = "point"
    source_Tc_99m.direction.type = "iso"

    # Track the hits on the crystal by adding a hits collection actor.
    hc = sim.add_actor("DigitizerHitsCollectionActor", "Hits")
    hc.mother = [crystal.name] 


    # Set the output path such that it contains the simulation settings.
    c_waterbox = ""
    if g_waterbox:
        c_waterbox = "_waterbox"
    hc.output = str(paths.output / f"test000_NM_670_{format_magnitude(g_energy, 'eV')}_{format_magnitude(g_activity)}_{g_runtime}s") + str(c_waterbox) + str(".root") 
    

    # Prepare the output path for the file.
    output_path = str(paths.output / f"test000_NM_670_{format_magnitude(g_energy, 'eV')}_{format_magnitude(g_activity, 'Bq')}_{g_runtime}s") + str(c_waterbox)

    
    # Set the attributes of interest
    hc.attributes = ["Direction", "EventDirection", "EventID", "EventKineticEnergy", "EventPosition", "GlobalTime", "HitUniqueVolumeID", "KineticEnergy", "LocalTime", "ParticleName", "Position", "PostDirection", "PostKineticEnergy", "PostPosition", "PostStepUniqueVolumeID", "PostStepVolumeCopyNo", "PreDirection", "PreDirectionLocal", "PreKineticEnergy", "PrePosition", "PreStepUniqueVolumeID", "PreStepVolumeCopyNo", "ProcessDefinedStep", "RunID", "ThreadID", "TimeFromBeginOfEvent", "TotalEnergyDeposit", "TrackCreatorProcess", "TrackID", "TrackProperTime", "TrackVertexKineticEnergy", "TrackVertexMomentumDirection", "TrackVertexPosition", "TrackVolumeCopyNo", "TrackVolumeInstanceID", "TrackVolumeName", "Weight"]

    
    # Create the run time intervals. In this case we only need one interval from 0 to g_runtime in secounds.
    sim.run_timing_intervals = [[0, g_runtime * sec]]
    

    # Start a new process to not interfere with other simulations that might have run previously.
    output = sim.start(start_new_process = True)


    # Get the Hits actor to further process the hits on the detector.
    hc_file = output.get_actor("Hits").user_info.output  


    # Prepare the data for graphical representation.
    hits = uproot.open(hc_file)["Hits"]
    number_of_hits = hits.num_entries
    hits = hits.arrays(library="numpy")
    hits_data = hits[key_of_interest] * scalings

    
    # --------- Plot ------------
    f, ax = plt.subplots(1, 1, figsize=(16, 9))
    plt.hist(hits_data, bins=g_channels, label="Simulation")
    ax.set_xlabel("KineticEnergy [keV]")
    ax.set_ylabel("Counts")
    plt.suptitle(f"Simulation of a point source with {format_magnitude(g_energy, 'eV')} measured with the GE Discovery NM 670 detector.")
    ax.plot([], [], ' ', label=f"Energy: {format_magnitude(g_energy, 'eV')}\nActivity: {format_magnitude(g_activity, 'Bq')}\nDistance detector <-> source: {g_distance} cm\nRuntime: {g_runtime} s\nChannels: {g_channels}\nHits: {number_of_hits}")
    ax.legend()
    plt.savefig(output_path + str(".svg"))


    # Write Values in txt:
    with open(output_path + str(".txt"), "w") as file:
        file.write(f"# Number of hits: {number_of_hits}\n")
        for v in np.sort(hits_data):
            file.write(f"{v}\n")




if __name__ == "__main__":

    # Settings:
    g_energy = 140.5e3 #[eV]
    g_activity = 3.7e5 #[Bq]
    g_runtime = 60 #[s]
    g_channels = 512 
    g_waterbox = False
    g_visu = False
    g_distance = 15 #[cm]


    for energy, activity, runtime, waterbox in [[g_energy, g_activity, g_runtime, g_waterbox]]:
        print()
        print("***************")
        print(f"Start simulation with energy: {format_magnitude(energy, 'eV')}, activity: {format_magnitude(activity, 'Bq')}")
        create_simulation(energy, activity, runtime, g_channels, waterbox)