#!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('-n',
              default=1e4,
              help='Number of values to plot.')

def garf_plot_training_dataset(data, n):
    '''
    \b
    Display info about the training dataset.
    <DATA> : dataset in root format
    '''
    filename = data
    data, x, y, theta, phi, E = garf.load_test_dataset(filename)

    print('Number of values: ', len(data))

    n = int(n)
    x = x[:n]
    y = y[:n]
    theta = theta[:n]
    phi = phi[:n]
    E = E[:n]

    print('Plot with ', len(x), 'values.')

    ## 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)
    ax[1,1].scatter(x, y, color='r', alpha=0.35, s=1)

    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('X')
    ax[1,1].set_ylabel('Y')

    plt.tight_layout()
    plt.show()


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