#!/usr/bin/env python3

import argparse
import os
import glob
import json
import numpy as np
import matplotlib.pyplot as plt
import collections

def main():
    args = __parseArguments()
    
    __stats(args["comparisonDir"], args["outputDir"])


def __parseArguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--directory", help="the direcotry with all comparison files", type=str)
    parser.add_argument("-o", "--output", help="Directory to store the stats", type=str)
    args = parser.parse_args()
    
    arguments = {}
    print(args)
    
    arguments["comparisonDir"] = args.directory
    if arguments["comparisonDir"] == None:
         arguments["comparisonDir"] = str(input("Comparison directory: "))
    
    arguments["comparisonDir"] = os.path.abspath(arguments["comparisonDir"])
    
    arguments["outputDir"] = args.output
    if arguments["outputDir"] == None:
        arguments["outputDir"] = str(input("Output directory: "))
        
    arguments["outputDir"] = os.path.abspath(arguments["outputDir"])
    
    return arguments

def __stats(comparisonDir, outputDir):
    stats = __collectStats(comparisonDir)
    
    __writeStats(stats, outputDir)
    
def __collectStats(comparisonDir):
    files = glob.glob(os.path.join(comparisonDir, "*.cmp"))
    
    stats = []
    
    for path in files:
        
        comparison = __readComparison(path)
    
        stats.append(__processSingleInstance(comparison))
            
    return stats
    
    
def __processSingleInstance(comparison):
    instanceStats = {}
    
    degrees = comparison["degrees_of_variables"]
    degreeArr = np.array(list(degrees.values()))
    
    instanceStats["degree_of_variables_mean"] = degreeArr.mean()
    instanceStats["degree_of_variables_median"] = np.median(degreeArr)
    instanceStats["degree_of_variables_std_dev"] = np.std(degreeArr)
    instanceStats["degree_of_variables_max"] = degreeArr.max()
    instanceStats["degree_of_variables_min"] = degreeArr.min()
    instanceStats["variables_per_degree"] = __getVarsPerDegree(degreeArr)
    
    if comparison["minisat_satisfiable"]:
        if __instanceIsFalseNegative(comparison):
            instanceStats["result"] = "false_negative"
        else:
            instanceStats["result"] = "satisfiable"
    else:
        instanceStats["result"] = "unsatisfiable"
    
    return instanceStats

def __instanceIsFalseNegative(comparison):
    return (comparison["minisat_satisfiable"] == True and
            comparison["qubo_satisfiable"] == False)

def __getVarsPerDegree(degreeArr):
    degCount = collections.Counter(degreeArr)
    
    varsPerDegree = {}
    
    for degree in degCount:
        varsPerDegree[degree] = degCount[degree]
        
    return varsPerDegree
    
    
def __readComparison(path):
    cmpFile = open(path, "r")
    comparison = json.load(cmpFile)
    cmpFile.close()
    
    return comparison
    
    
def __writeStats(stats, outputDir):

    fig1 = plt.figure()
    
    data = __seperateMatchesAndFalseNegatives(stats)
    
    ax0 = fig1.add_subplot(141,)
    ax0.boxplot([data["mean"]["satisfiable"],
                 data["mean"]["false_negative"],
                 data["mean"]["unsatisfiable"]])
    ax0.set_title("mean")
    
    ax1 = fig1.add_subplot(142, sharey=ax0)
    ax1.boxplot([data["median"]["satisfiable"],
                 data["median"]["false_negative"],
                 data["median"]["unsatisfiable"]])
    ax1.set_title("median")
    
    ax2 = fig1.add_subplot(143, sharey=ax0)
    ax2.boxplot([data["max"]["satisfiable"],
                 data["max"]["false_negative"],
                 data["max"]["unsatisfiable"]])
    ax2.set_title("max degree")
    
    ax3 = fig1.add_subplot(144, sharey=ax0)
    ax3.boxplot([data["min"]["satisfiable"],
                 data["min"]["false_negative"],
                 data["min"]["unsatisfiable"]])
    ax3.set_title("min degree")
    
    
    fig2 = plt.figure()
    ax4 = fig2.add_subplot(111)
    ax4.boxplot([data["std_dev"]["satisfiable"],
                 data["std_dev"]["false_negative"],
                 data["std_dev"]["unsatisfiable"]])
    ax4.set_title("standard deviation")
    
    
    _BINS_ = 23
    
    fig3 = plt.figure()
    ax5 = fig3.add_subplot(311)
    varsPerDegreeSat = __accumulateVarsPerDegree(data["vars_per_degree"]["satisfiable"])
    ax5.hist(varsPerDegreeSat, density=True, bins=_BINS_)
    
    ax6 = fig3.add_subplot(312, sharex=ax5)
    varsPerDegreeFP = __accumulateVarsPerDegree(data["vars_per_degree"]["false_negative"])
    ax6.hist(varsPerDegreeFP, density=True, bins=_BINS_)
    
    ax7 = fig3.add_subplot(313, sharex=ax6)
    varsPerDegreeUnsat = __accumulateVarsPerDegree(data["vars_per_degree"]["unsatisfiable"])
    ax7.hist(varsPerDegreeUnsat, density=True, bins=_BINS_)
    
    
    plt.setp([ax0, ax1, ax2, ax3, ax4], xticks=[1, 2, 3], xticklabels=["satisfiable",
                                                                       "false negative",
                                                                       "unsatisfiable"])
    plt.setp(ax0.get_xticklabels(), rotation=45)
    plt.setp(ax1.get_xticklabels(), rotation=45)
    plt.setp(ax2.get_xticklabels(), rotation=45)
    plt.setp(ax3.get_xticklabels(), rotation=45)
    plt.setp(ax4.get_xticklabels(), rotation=45)
    fig1.set_size_inches(12, 8)
    fig1.suptitle("Degrees of variables", fontsize=16)
    
    fig2.set_size_inches(4, 8)
    
    fig3.set_size_inches(5, 12)
    
    fig1.savefig(os.path.join(outputDir, "degrees1.png"))
    fig2.savefig(os.path.join(outputDir, "degrees2.png"))
    fig3.savefig(os.path.join(outputDir, "degrees3.png"))
    
    plt.show()
    

def __accumulateVarsPerDegree(listOfVarsPerDegreeDicts):
    accumulated = []
    
    for instance in listOfVarsPerDegreeDicts:
        for degree in instance:
            accumulated += [degree] * instance[degree]
    
    return accumulated
    
def __compressVarsPerDegree(listOfVarsPerDegreeDicts):
    compressed = {}
    
    countOfVars = 0
    
    for instance in listOfVarsPerDegreeDicts:
        for degree in instance:
            if degree in compressed:
                compressed[degree] += float(instance[degree])
            else:
                compressed[degree] = float(instance[degree])
            
            countOfVars += instance[degree]
    
    check = 0
    for degree in compressed:
        compressed[degree] /= countOfVars
        
        check += compressed[degree]
        
    
    print("check: ", check)
    
    return compressed
    

def __seperateMatchesAndFalseNegatives(stats):
    data = {}
    data["mean"] = {"false_negative": [],
                    "satisfiable": [],
                    "unsatisfiable": []}
                    
    data["median"] = {"false_negative": [],
                      "satisfiable": [],
                      "unsatisfiable": []}
                      
    data["std_dev"] = {"false_negative": [],
                       "satisfiable": [],
                       "unsatisfiable": []}
                       
    data["max"] = {"false_negative": [],
                   "satisfiable": [],
                   "unsatisfiable": []}
                   
    data["min"] = {"false_negative": [],
                   "satisfiable": [],
                   "unsatisfiable": []}
                   
    data["vars_per_degree"] = {"false_negative": [],
                               "satisfiable": [],
                               "unsatisfiable": []}
    
    for instance in stats:
        target = instance["result"]
        
        data["mean"][target].append(instance["degree_of_variables_mean"])
        data["median"][target].append(instance["degree_of_variables_median"])
        data["std_dev"][target].append(instance["degree_of_variables_std_dev"])
        data["max"][target].append(instance["degree_of_variables_max"])
        data["min"][target].append(instance["degree_of_variables_min"])
        data["vars_per_degree"][target].append(instance["variables_per_degree"])
    
    return data

if __name__ == "__main__":
    main()