# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------
"""
Groups together code used for creating a NuPIC model and dealing with IO.
(This is a component of the One Hot Gym Anomaly Tutorial.)
"""
import importlib
import sys
import csv
import datetime
import numpy as np
import matplotlib.pyplot as plt
import datetime
import pickle
import time
from numpy import genfromtxt



from nupic.data.inference_shifter import InferenceShifter
from nupic.frameworks.opf.model_factory import ModelFactory
from nupic.algorithms import anomaly_likelihood

import nupic_anomaly_output


DESCRIPTION = (
    "Starts a NuPIC model from the model params returned by the swarm\n"
    "and pushes each line of input from the gym into the model. Results\n"
    "are written to an output file (default) or plotted dynamically if\n"
    "the --plot option is specified.\n"
)
DATA_DIR = "."
MODEL_PARAMS_DIR = "./model_params"
# '7/2/10 0:00'
DATE_FORMAT = "%m/%d/%y %H:%M"


#Function that loads data into a numpy matrix from csv file.
#These files are generated by running ./garch/garch_long.py
def getData(datatype="D1", size="100"):
    D = genfromtxt('./' + datatype + '_' + size + '.csv', delimiter=',')
    D_no_anomalies = genfromtxt('./' + datatype + '_unpolluted_' + size + '.csv', delimiter=',')
    D_truth = genfromtxt('./' + datatype + '_truth_' + size + '.csv', delimiter=',')
    return D, D_no_anomalies, D_truth

#Saves a csv-fil of the numpy matrix "data" 
def saveCsvOfDataWithTimestamp(data):
        
    dt = datetime.datetime(2010, 12, 1)
    end = datetime.datetime(2010, 12, 30, 23, 59, 59)
    step = datetime.timedelta(seconds=5)

    timestamps = []

    while dt < end:
        timestamps.append(dt.strftime('%Y/%m/%d %H:%M:%S'))
        dt += step

    timestamps = timestamps[:len(data)]
    data = np.append(data, np.array(timestamps))

    joineddata = [timestamps, data]
    joineddata = zip(*joineddata)

    with open("Datetime.csv","w") as csvFile:
        Fileout = csv.writer(csvFile, delimiter=',')
        for jd in joineddata:
            Fileout.writerow(jd)
    np.savetxt("3000_LS_timeserie.csv", data.astype(float), delimiter=",")


#Creates a model with the parameters specified in the model parameters.
def createModel(modelParams):
    """
    Given a model params dictionary, create a CLA Model. Automatically enables
    inference for kw_energy_consumption.
    :param modelParams: Model params dict
    :return: OPF Model object
    """
    model = ModelFactory.create(modelParams)
    model.enableInference({"predictedField": "kw_energy_consumption"})
    return model


#Fetches the model parameters from file.
def getModelParamsFromType(dataType):
    """
    Given a gym name, assumes a matching model params python module exists within
    the model_params directory and attempts to import it.
    :param gymName: Gym name, used to guess the model params module name.
    :return: OPF Model params dictionary
    """
    importName = "model_params." + str(dataType) + "_model_params"
    print "Importing model params from %s" % importName
    try:
        importedModelParams = importlib.import_module(importName).MODEL_PARAMS
    except ImportError:
        raise Exception("No model params exist for '%s'. Run swarm first!"
                                        % dataName)
    return importedModelParams


#HTM-model capable of performing anomaly detection
def HTM(data, model):


    shifter = InferenceShifter()
    
    counter = 0
    anomalyLikelihood = anomaly_likelihood.AnomalyLikelihood(claLearningPeriod=None, learningPeriod=100, estimationSamples=10, historicWindowSize=800, reestimationPeriod=1)
    
    actuals = np.zeros(len(data))
    predictions = np.zeros(len(data))
    anomalyScores = np.zeros(len(data))
    anomalyLikelihoods = np.zeros(len(data))

    #For every timestep in time serie.
    for t in range(len(data)):
        counter += 1
        if (counter % 100 == 0):
            pass
            #print "Read %i lines..." % counter
        consumption = float(data[t])
        result = model.run({
            "kw_energy_consumption": consumption
        })

        result = shifter.shift(result)

        prediction = result.inferences["multiStepBestPredictions"][1]
        anomalyScore = result.inferences["anomalyScore"]
        actuals[t] = data[t]
        predictions[t] = prediction
        anomalyScores[t] = anomalyScore 
        
        likelihood = anomalyLikelihood.anomalyProbability(
            consumption, anomalyScore
        )
        anomalyLikelihoods[t] = likelihood


    return predictions, anomalyScores, anomalyLikelihoods


#Runs the anomaly detection and calculates the score.
def runAnomalyDetection(D, D_truth, D_unpolluted, dataType, dataSize):
    FP = 0.
    FN = 0.
    TP = 0.
    TN = 0.

    FN_ALO = 0.
    TP_ALO = 0.
    FN_LS = 0.
    TP_LS = 0.
    FN_TC = 0.
    TP_TC = 0.
    FN_LTO = 0.
    TP_LTO = 0.
    FN_SALO = 0.
    TP_SALO = 0.

    epsilon = 0.00001

    reaction_counter = 0.
    tot_reaction = 0.

    print "Creating model from %s..." % dataType
    model = createModel(getModelParamsFromType(dataType))

    print("Running HTM-AD")
    guessed_anomalies = np.zeros((len(D), len(D[0])))
    all_preds = np.zeros((len(D), len(D[0])))
    all_scores = np.zeros((len(D), len(D[0])))
    all_likelihoods = np.zeros((len(D), len(D[0])))
    for i in range(len(D)):
        if dataTypes == "D3":
            model = createModel(getModelParamsFromType(dataType))
        predictions, anomalyScores, anomalyLikelihoods = HTM(D[i], model)

        all_preds[i] = predictions
        all_scores[i] = anomalyScores
        all_likelihoods[i] = anomalyLikelihoods

        for t in range(int(0.1*len(D[i])), len(D[i])):
            if anomalyLikelihoods[t] > (1 - epsilon):
                guessed_anomalies[i][t] = 1 #This is an anomaly guess

        #Evaluate
        in_anomaly_window = False
        anomaly_window_type = 0
        has_flagged_anomaly_window = False
        for t in range(1, len(D[i])):

            #Exiting an anomaly window. Check if it has been flagged.
            if D_truth[i][t] == 0 and in_anomaly_window:
                in_anomaly_window = False
                if not has_flagged_anomaly_window: 
                    FN += 1 #Failed to flag the entire anomaly window.
                    if int(anomaly_window_type) == 1:
                        FN_ALO += 1 #Failed to flag the entire anomaly window.
                    elif int(anomaly_window_type) == 2:
                        FN_LS += 1 #Failed to flag the entire anomaly window.
                    elif int(anomaly_window_type) == 3:
                        FN_TC += 1 #Failed to flag the entire anomaly window.
                    elif int(anomaly_window_type) == 4:
                        FN_LTO += 1 #Failed to flag the entire anomaly window.
                    elif int(anomaly_window_type) == 5:
                        FN_SALO += 1 #Failed to flag the entire anomaly window.
                    else:
                        print("ERROR ANOMALY WINDOW  HAD NO TYPE!")
                        print(int(anomaly_window_type))
                else:
                    has_flagged_anomaly_window = False

                reaction_counter = 0.
                anomaly_window_type = 0
            #Entering an anomaly window.
            if D_truth[i][t] != 0 and not in_anomaly_window:
                in_anomaly_window = True
                anomaly_window_type = D_truth[i][t]

            if guessed_anomalies[i][t] == 0 and not in_anomaly_window:
                TN += 1 #Correct to not guess for an anomaly.
            if guessed_anomalies[i][t] != 0 and in_anomaly_window and not has_flagged_anomaly_window:
                TP += 1 #Correct guess, within anomaly window.
                tot_reaction += reaction_counter #Add reaction counter to av
                reaction_counter = 0.
                has_flagged_anomaly_window = True
                if int(anomaly_window_type) == 1:
                    TP_ALO += 1 
                elif int(anomaly_window_type) == 2:
                    TP_LS += 1 
                elif int(anomaly_window_type) == 3:
                    TP_TC += 1 
                elif int(anomaly_window_type) == 4:
                    TP_LTO += 1 
                elif int(anomaly_window_type) == 5:
                    TP_SALO += 1 
                else:
                    print("ERROR ANOMALY WINDOW  HAD NO TYPE!")
                    print(int(anomaly_window_type))
            if guessed_anomalies[i][t] != 0 and not in_anomaly_window:
                FP += 1 #Erroneous guess, outside anomaly window.
            if guessed_anomalies[i][t] == 0 and in_anomaly_window and not has_flagged_anomaly_window:
                reaction_counter += 1 #Failed to react to an anomaly which was present in the current timestep.

    print("|" + "TP: " + str(TP) + "|TN: " + str(TN) +  "|" + "FP: " + str(FP) + "|" + "FN: " + str(FN) + "|")
    precision = TP / (TP + FP)
    recall = TP / (TP + FN)

    if precision == 0 and recall == 0:
        print("Somethings is wrong. No TP's")
        precision = 0.0000001
        recall = 0.00000001
    
    if(TP == 0):
        avg_reaction = 1232153463
    else:
        avg_reaction = tot_reaction/(TP)

    F1 = 2* (precision * recall) / (precision + recall)
    print("Anomalies:   " + str(TP + FN))
    print("Guesses:     " + str(TP + FP))
    print("precision:   " + str(precision))
    print("recall:      " + str(recall))
    print("F1:          " + str(F1))
    print("Reaction:    " + str(avg_reaction))
    print("TP_ALO:      " + str(TP_ALO))
    print("FN_ALO:      " + str(FN_ALO))
    print("TP_LS:       " + str(TP_LS))
    print("FN_LS:       " + str(FN_LS))
    print("TP_TC:       " + str(TP_TC))
    print("FN_TC:       " + str(FN_TC))
    print("TP_LTO:      " + str(TP_LTO))
    print("FN_LTO:      " + str(FN_LTO))
    print("TP_SALO:     " + str(TP_SALO))
    print("FN_SALO:     " + str(FN_SALO))

    with open(str(dataType) + "_" + str(dataSize) + '_result_htm.txt', 'a') as the_result_file:
        the_result_file.write("\n\n\n")
        the_result_file.write(datetime.datetime.now().strftime("%Y-%m-%d %H:%M"))
        the_result_file.write("\n")
        the_result_file.write("|" + "TP: " + str(TP) + "|TN: " + str(TN) +  "|" + "FP: " + str(FP) + "|" + "FN: " + str(FN) + "|")
        the_result_file.write("\nAnomalies:   " + str(TP + FN))
        the_result_file.write("\nGuesses:     " + str(TP + FP))
        the_result_file.write("\nprecision:   " + str(precision))
        the_result_file.write("\nrecall:      " + str(recall))
        the_result_file.write("\nF1:          " + str(F1))
        the_result_file.write("\nReaction:    " + str(avg_reaction))
        the_result_file.write("\nTP_ALO:      " + str(TP_ALO))
        the_result_file.write("\nFN_ALO:      " + str(FN_ALO))
        the_result_file.write("\nTP_LS:       " + str(TP_LS))
        the_result_file.write("\nFN_LS:       " + str(FN_LS))
        the_result_file.write("\nTP_TC:       " + str(TP_TC))
        the_result_file.write("\nFN_TC:       " + str(FN_TC))
        the_result_file.write("\nTP_LTO:      " + str(TP_LTO))
        the_result_file.write("\nFN_LTO:      " + str(FN_LTO))
        the_result_file.write("\nTP_SALO:     " + str(TP_SALO))
        the_result_file.write("\nFN_SALO:     " + str(FN_SALO))
    return guessed_anomalies, all_preds, all_scores, all_likelihoods

def plotAnomalyDetection(D, D_truth, D_unpolluted, D_guesses, D_preds, D_anomaly_scores, D_anomaly_likelihoods):

    show_this_many_timesteps = min(len(D[0]), 3000)
    differencer = len(D[0]) - show_this_many_timesteps

    for sample_nr in range(len(D)):
        plt.figure(1)

        all_gi = np.argwhere(D_guesses[sample_nr][-show_this_many_timesteps:] != 0) #all guesses
        all_ai = np.argwhere(D_truth[sample_nr][-show_this_many_timesteps:] != 0)   #all actuall anomalies
        plt.subplot(311)
        plt.title("Actuals and predictions")

        #plt.plot(D_truth[sample_nr][-show_this_many_timesteps:], 'b')
        line_actual, = plt.plot(D[sample_nr][-show_this_many_timesteps:], 'r', label="True Values")
        line_preds, = plt.plot(D_preds[sample_nr][-show_this_many_timesteps:], 'g', label="HTM predictions")
        for gi in all_gi:
            dot_wrong, = plt.plot(gi, D[sample_nr][differencer + gi], 'yo', label="Incorrect Guess")
            #plt.axvspan(ai-0.1, ai+0.1, facecolor='#990000', alpha=0.8)
        for ai in all_ai:
            dot_anom, = plt.plot(ai, D[sample_nr][differencer + ai], 'ro', label="Anomaly")
            if ai in all_gi:
                dot_correct, = plt.plot(ai, D[sample_nr][differencer +ai], 'ko', label="Correct Guess")

        plt.legend(handles=[line_preds, line_actual, dot_wrong, dot_anom, dot_correct])

        plt.subplot(312)
        plt.title("Residuals")
        line_resi, = plt.plot(abs(D[sample_nr][-show_this_many_timesteps:]-D_preds[sample_nr][-show_this_many_timesteps:]), label="Residuals")
        plt.legend(handles=[line_resi])

        plt.subplot(313)
        plt.title("anomalyScores and anomalyLikelihoods")
        line_pred_err, = plt.plot(D_anomaly_scores[sample_nr][-show_this_many_timesteps:], "r", label="Prediction Error")
        line_anom_likel, = plt.plot((D_anomaly_likelihoods)[sample_nr][-show_this_many_timesteps:], "purple", label="Anomaly Likelihood")
        for gi in all_gi:
            plt.plot(gi, D_anomaly_likelihoods[sample_nr][differencer + gi], 'yo', label="Incorrect Guess")
            #plt.axvspan(ai-0.1, ai+0.1, facecolor='#990000', alpha=0.8)
        for ai in all_ai:
            plt.plot(ai, D_anomaly_likelihoods[sample_nr][differencer + ai], 'ro', label="Anomaly")
            if ai in all_gi:
                plt.plot(ai, D_anomaly_likelihoods[sample_nr][differencer + ai], 'ko', label="Correct Guess")

        plt.legend(handles=[line_pred_err, line_anom_likel])

        plt.show()





def runModel(dataName, plot=False):
    """
    DEPRECATED DESCRIPTION!
    Assumes the gynName corresponds to both a like-named model_params file in the
    model_params directory, and that the data exists in a like-named CSV file in
    the current directory.
    :param gymName: Important for finding model params and input CSV file
    :param plot: Plot in matplotlib? Don't use this unless matplotlib is
    installed.
    """
    
    inputData = dataName + ".csv"
    runIoThroughNupic(inputData, model, dataName, plot)



if __name__ == "__main__":
    print DESCRIPTION
    plot = False
    args = sys.argv[1:]
    if "--plot" in args:
        plot = True

    N = 1 # number of series to evaluate.

    dataTypes = ["D1", "D2", "D4"]
    dataSizes = ["5000"]
    for dt in dataTypes:
        for ds in dataSizes:
            #Get data from file.
            print("Starting Dataset:" + dt + "_" + ds)
            D, D_no_anomalies, D_truth = getData(dt, ds)
            start_htm = time.time()
            htm_guesses, htm_preds, htm_scores, htm_likelihoods = runAnomalyDetection(D[0:N], D_truth[0:N], D_no_anomalies[0:N], dt, ds)
            end_htm = time.time()
            np.savetxt(str(dt) + "_" + str(ds) + "_guesses_htm.csv", htm_guesses, delimiter=",")
            print('Elapsed time for HTM: ' + str(end_htm-start_htm))

            if plot:
                plotAnomalyDetection(D[0:N], D_truth[0:N], D_no_anomalies[0:N], htm_guesses[0:N], htm_preds[0:N], htm_scores[0:N], htm_likelihoods[0:N])