#!python
# -*- coding: utf-8 -*-

import sys
import garf
import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import uproot
import ntpath
import click

# -----------------------------------------------------------------------------
CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"])


@click.command(context_settings=CONTEXT_SETTINGS)
@click.argument("data")
@click.option("--rr", required=True, help="Russian Roulette used during simulation")
def garf_plot_training_dataset(data, rr):
    """
    \b
    Display info about the training dataset.
    <DATA> : dataset in root format
    """
    filename = data
    data, theta, phi, E, w = garf.load_training_dataset(filename)

    ## histo binning
    b = 100

    f, ax = plt.subplots(2, 2, figsize=(10, 10))

    n, bins, patches = ax[0, 0].hist(theta, b, density=True, facecolor="g", alpha=0.35)
    n, bins, patches = ax[0, 1].hist(phi, b, density=True, facecolor="g", alpha=0.35)
    n, bins, patches = ax[1, 0].hist(
        E * 1000, b, density=True, facecolor="b", alpha=0.35
    )
    n, bins, patches = ax[1, 1].hist(
        w[w > 0], b, density=False, facecolor="r", alpha=0.35
    )

    ax[0, 0].set_xlabel("Theta angle (deg)")
    ax[0, 1].set_xlabel("Phi angle (deg)")
    ax[1, 0].set_xlabel("Energy (keV)")
    ax[1, 1].set_xlabel("Energy window #")

    rr = float(rr)
    print("Russian Roulette is", rr)
    garf.print_training_dataset_info(data, rr)

    plt.tight_layout()
    plt.show()


# -----------------------------------------------------------------------------
if __name__ == "__main__":
    garf_plot_training_dataset()
