#! env python

import argparse
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.integrate import simps

import packet_trace
import features
import model
import trainer
import defaults
import transform

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-b', '--benign', type = int, default = 78, help = 'number of benign packets')
    parser.add_argument('-a', '--anomalous', type = int, default = 22, help = 'number of anomalous packets')
    parser.add_argument('-g', '--group', type = int, default = 0, help = 'index of feature group for which to calculate roc curves')
    parser.add_argument('-f', '--full-run', action = 'store_true', help = 'perform a full run of the test and display roc curves')
    args = parser.parse_args()

    m = model.Model()
    t = trainer.Trainer()
    p, c = t.get_features_params()
    pt = packet_trace.PacketTrace(pcap_dir = defaults.test_pcap_dir)

    max_score = m.group_max_score(args.group)
    threshold = 0
    #step = 2.146e-11
    step = max_score / 100000.
    stop = False
    x = list()
    y = list()
    packet_scores = list()

    for idx, packet in enumerate(pt):
        score = m.pdf_group(args.group, packet)
        packet_scores.append(score)

    while not stop:
        true_positives = 0
        true_negatives = 0
        false_positives = 0
        false_negatives = 0
        good_res = []
        bad_res = []

        for idx, score in enumerate(packet_scores):
            if score is not None:

                if args.full_run:
                    if idx < args.benign:
                        if score < threshold:
                            false_positives += 1
                        else:
                            true_negatives += 1
                    else:
                        if score < threshold:
                            true_positives += 1
                        else:
                            false_negatives += 1
                else:
                    if idx < args.benign:
                        good_res.append(score)
                    else:
                        bad_res.append((idx - args.benign, score))

        if args.full_run:
            fpr = float(false_positives) / (false_positives + true_negatives)
            tpr = float(true_positives) / (true_positives + false_negatives)
            x.append(fpr)
            y.append(tpr)
        else:
            idx_anom_under = list()
            for idx_anom, s in bad_res:
                if s < min(good_res):
                    idx_anom_under.append(idx_anom)
            print 'Maximum score for this feature group is', max_score
            print 'Minimum score for bening examples', min(good_res)
            print 'Minimum score for anomalous examples', min(bad_res)
            print 'Number of anomalies under min bening score', len(idx_anom_under)
            print 'Found anomalies indices', ' '.join([ str(sv) for sv in idx_anom_under])
            print bad_res

        if args.full_run:
            threshold += step
            if threshold > max_score:
                stop = True
            #if fpr > 0.05:
            #    stop = True
        else:
            stop = True

    if args.full_run:
        #print 'AUC', simps(y, x)
        plt.plot(x, y, linewidth = 2)
        plt.show()

if __name__ == '__main__':
    main()
