import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy
import os
import itertools
import math

# TO BE CHANGED --> trial for new REAL data; select code blocks for respective runs
trial = 'trial_12_05_20' #real

# trial_list = ['train_12_28_20', 'train_12_29_20'] #286 load size
# label_list = ['B50, smooth', 'B16, smooth'] #286 load size
# load_size = 286
# results_real = 'trial_12_05_20'
trial_list = ['train_12_23_20', 'train_12_26_20'] #321 load size
label_list = ['B1, sharp', 'B50, sharp'] #321 load size
load_size = 321
results_real = 'train_12_23_20'

# ------------------------------------------------------------------------------------------------- #

# numPackets directories, both original and cropped, and output folder

# FILE PATHS:
real_results_dir = '/home/tom_phelan_ext/Documents/microstructure_analysis/grains2packets/numPackets_results_real/' + results_real + '/'
orig_imgs_dir = '/home/tom_phelan_ext/Documents/microstructure_analysis/grains2packets/numPackets/' + trial +'/'
cropped_imgs_dir = '/home/tom_phelan_ext/Documents/microstructure_analysis/grains2packets/numPackets_cropped/' + trial + '_cropped/'
fake_imgs_dir = '/home/tom_phelan_ext/Documents/microstructure_analysis/grains2packets/numPackets_FAKE/'
output_folder = '/home/tom_phelan_ext/Documents/microstructure_analysis/grains2packets/numPackets_hists/'

# ------------------------------------------------------------------------------------------------- #

def read_csv_files(imgs_dir, img_type):
    # append each image's numPackets data to array
    data = []
    print("Reading csv files of " + img_type + " images...")
    csv_files = os.listdir(imgs_dir)
    count = 0
    for csv in csv_files:
        # open csv file for reading numPackets data, line by line
        with open(imgs_dir + csv, 'r') as readfile:
            last_line = len(readfile.readlines())
            # file must be-reopened
            readfile = open(imgs_dir + csv, 'r')
            for num in itertools.islice(readfile, 2, last_line): # omit headers
                data.append(num[:-1]) # omit \n new line characters
            # parse string list to int list
            data = list(map(int, data))
            count += 1
    print("numPackets data for " + str(count) + " " + img_type + " images collected.")
    print()
    return numpy.array(data)

def create_histogram(array, label, bins):
    y2, binEdges2 = numpy.histogram(array, bins=bins, density=True)
    bincenters2 = 0.5*(binEdges2[1:]+binEdges2[:-1])
    plt.plot(bincenters2, y2, '-', label=label)

def finish_histogram():
    plt.xlabel('ln(Packets per Grain)')
    plt.ylabel('Frequency')
    plt.title('Distribution of numPackets in Grain (load_size=' + str(load_size) + ')')
    plt.legend(loc='upper left')
    plt.savefig(output_folder + 'numPackets_' + str(load_size) + '_hist.png')
    plt.clf()
    print("numPackets log histogram saved.")
# ------------------------------------------------------------------------------------------------- #

# numPacket array of results real data
numPackets_results = read_csv_files(imgs_dir=real_results_dir, img_type='real (results,' + results_real + ')') 
numPackets_results = numpy.log(numPackets_results)
# generating histogram of results real data
y, binEdges = numpy.histogram(numPackets_results, density=True)
print(binEdges)
bincenters = 0.5*(binEdges[1:]+binEdges[:-1])
plt.plot(bincenters, y, label='results real', linestyle='dashed', linewidth=4, color='g')

if (load_size == 321):
    # set bins to same as load size 286 run
    binEdges = [0, 0.69057533, 1.38115066, 2.07172598, 2.76230131, 3.45287664, 4.14345197, 4.83402729, 5.52460262, 6.21517795, 6.90575328]
    # read in original and cropped image data
    numPackets = read_csv_files(imgs_dir=orig_imgs_dir, img_type="original")
    numPackets_cropped = read_csv_files(imgs_dir=cropped_imgs_dir, img_type="cropped")   
    # take log of each array
    numPackets = numpy.log(numPackets)
    numPackets_cropped = numpy.log(numPackets_cropped)
    # cropped image data added to histogram + results real data
    create_histogram(array=numPackets, label='real', bins=binEdges)
    create_histogram(array=numPackets_cropped, label='cropped', bins=binEdges)

# fake trials added to histogram
i=0
for trial in trial_list:
    numPackets_FAKE = read_csv_files(imgs_dir=fake_imgs_dir + trial_list[i] + '/', img_type=trial_list[i] + ' (fake)')
    numPackets_FAKE = numpy.log(numPackets_FAKE)
    # create histogram of current fake trial
    create_histogram(array=numPackets_FAKE, label=label_list[i], bins=binEdges)
    i += 1

finish_histogram()
print()
print("Program completed.")
print()
