#!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()
