from file_reader import XMLReader
import glob
import matplotlib.pyplot as plt
import my_ROOT
import numpy as np
import os
import sys

TH1D = my_ROOT.TH1D

batch_directory = sys.argv[1]

xml_reader = XMLReader(';')


def exhaust_directory(directory, histogram, validity_condition, boundary_condition, quantization):
    number_of_profiles = 0
    for filepath in glob.glob(os.path.join(directory, xml_reader.pattern)):
        profiles = xml_reader.extract_profiles_from_file(filepath)
        profiles = filter(validity_condition, profiles)  # filter to obtain good quality profiles
        for profile in profiles:
            profile.bring_to_zero()
        # select region of interest
        profiles = [profile.create_copy_between_boundaries(*boundary_condition(profile)) for profile in profiles]
        for profile in profiles:
            histogram.Fill(quantization(profile))
        number_of_profiles += len(profiles)

    return number_of_profiles


def plot_quantization_for_all(particle_type, validity_condition, boundary_condition, quantization, histogram_parameters,
                              normalize_to=1):
    histogram = TH1D('histogram_%s' % particle_type, particle_type, *histogram_parameters)
    print '%s: processed %d profiles' % (particle_type,
                                         exhaust_directory(os.path.join(batch_directory, particle_type), histogram,
                                                           validity_condition, boundary_condition, quantization))
    xs, ys = histogram.get_bin_centers_x(), histogram.get_bin_contents()
    ys = ys / np.sum(ys) * normalize_to
    plt.plot(xs, ys, label=particle_type)

    underflow, overflow = histogram.get_underflow_and_overflow_bin_contents()
    print 'underflow: %f, overflow: %f' % (underflow, overflow)


def plot_quantization_different_extraction_voltages(directory, validity_condition, boundary_condition, quantization,
                                                    histogram_parameters, normalize_to=1):
    histogram = TH1D(directory, os.path.split(directory)[1], *histogram_parameters)
    print '%s: processed %d profiles' % (os.path.split(directory)[1], exhaust_directory(directory, histogram,
                                                                                        validity_condition,
                                                                                        boundary_condition,
                                                                                        quantization))
    xs, ys = histogram.get_bin_centers_x(), histogram.get_bin_contents()
    ys = ys / np.sum(ys) * normalize_to
    plt.plot(xs, ys, label=os.path.split(directory)[1])

    underflow, overflow = histogram.get_underflow_and_overflow_bin_contents()
    print 'underflow: %f, overflow: %f' % (underflow, overflow)


validity_condition = lambda profile: profile.has_sufficient_height(absolute_difference=2.5)
boundary_condition = lambda profile: profile.boundaries_for_height(0.25)
quantization = lambda profile: profile.excess_kurtosis
histogram_parameters = (100, -2.0, 0.0)

# plot_quantization_for_all('electrons', validity_condition=validity_condition, boundary_condition=boundary_condition,
#                           quantization=quantization, histogram_parameters=histogram_parameters)
# plot_quantization_for_all('ions', validity_condition=validity_condition, boundary_condition=boundary_condition,
#                           quantization=quantization, histogram_parameters=histogram_parameters)

for EV in sys.argv[2:]:
    plot_quantization_different_extraction_voltages(os.path.join(sys.argv[1], EV),
                                                    validity_condition=validity_condition,
                                                    boundary_condition=boundary_condition,
                                                    quantization=quantization,
                                                    histogram_parameters=histogram_parameters)

plt.legend()
plt.show()