import pandas as pd
import numpy as np

###SCRIPT DESCRIPTION###
# This script provides statistical analysis for LSTM labeled data.
###SCRIPT INPUT###
# The .csv file given to this script should be equivalent to the labeled output generated by the corresponding
# data-generation.py script. This means a "Format" column must be present with a unique identifier for each
# unique format. The column "true_label" should hold the integer value of the true class of the sample, the
# column "label" should hold the integer value of the predicted class of the sample.
###SCRIPT OUTPUT###
# This script provides the following values in a .csv file
#   - Classwise precision, accuracy and recall for each format
#   - Formatwise accuracy
###SCRIPT BEGIN####

input_file = input("Path to test results file: ")
output_file = input("Path to output file: ")

# Load data
data = pd.read_csv(input_file)

# Initial column names and row list
rows = []
cols = ['Format ID', 'Format Example', 'Sample Count']
#each class has a column for precision, recall and accuracy
for cl in data['true_label'].unique():
    cols.append(str(cl) + "_precision")
    cols.append(str(cl) + "_recall")
    cols.append(str(cl) + "_accuracy")
#add format accuracy at the end
cols.append("Format Accuracy")

#for each unique format ID
for format in data['Format'].unique():
    #create a subset containing only entries from this format
    subset = data[data['Format'] == format]
    #find the number of rows this format has
    n = subset.shape[0]
    #get one example of the format
    example = subset['Date'].iloc[0]

    row = [format, example, n]

    # for each class that truly exists
    for cl in data['true_label'].unique():
        #create subset with all samples in the format that have this class
        class_subset = subset[subset['true_label'] == cl]
        #create subset with all samples in the format that are predicted as this class
        predicted_subset = subset[subset['label'] == cl]
        #create subset with all samples in the format that are not predicted as this class
        negative_subset = subset[subset['label'] != cl]
        #get indices of rows where this class was correctly classified
        correct = np.where(class_subset['true_label'] == class_subset['label'])

        #get amount of real, predicted and correctly predicted values of this class
        real = class_subset.shape[0]
        predicted = predicted_subset.shape[0]
        correctly_predicted = len(correct[0])
        true_negatives = negative_subset[negative_subset['true_label'] != cl].shape[0]

        #precision = True Positives / Predicted
        precision = (correctly_predicted / predicted) if predicted > 0 else "N/A" if real == 0 else 0
        #recall = True Positives / Real Samples
        recall = (correctly_predicted / real) if real > 0 else "N/A"
        #accuracy = True Positives + True Negatives / All Format Samples
        accuracy = (correctly_predicted + true_negatives) / n
        
        #Add formatwise precision, recall and accuracy to the row
        row += [precision, recall, accuracy]

    #Add format accuracy to the row (all matching entries / all entries)
    acc = subset[subset['label'] == subset['true_label']]
    row.append(acc.shape[0] / n)
    rows.append(row)

#output dataframe
df=pd.DataFrame(rows, columns=cols)
df.to_csv(output_file, index=False)