from nilmtk.electric import align_two_meters
import numpy as np

import math

def tp_tn_fp_fn(states_pred, states_ground):
    tp = np.sum(np.logical_and(states_pred == 1, states_ground == 1))
    fp = np.sum(np.logical_and(states_pred == 1, states_ground == 0))
    fn = np.sum(np.logical_and(states_pred == 0, states_ground == 1))
    tn = np.sum(np.logical_and(states_pred == 0, states_ground == 0))
    return tp, tn, fp, fn

def recall_precision_accuracy_f1(pred, ground):
    aligned_meters = align_two_meters(pred, ground)

    threshold = ground.on_power_threshold()
    chunk_results = []
    sum_samples = 0.0
    for chunk in aligned_meters:
        sum_samples += len(chunk)
        pr = np.array([0 if (p)<threshold else 1 for p in chunk.iloc[:,0]])
        gr = np.array([0 if p<threshold else 1 for p in chunk.iloc[:,1]])

        tp, tn, fp, fn = tp_tn_fp_fn(pr,gr)
        p = sum(pr)
        n = len(pr) - p

        chunk_results.append([tp,tn,fp,fn,p,n])

    if sum_samples == 0:
        return None
    else:
        [tp,tn,fp,fn,p,n] = np.sum(chunk_results, axis=0)

        res_recall = recall(tp,fn)
        res_precision = precision(tp,fp)
        res_f1 = f1(res_precision,res_recall)
        res_accuracy = accuracy(tp,tn,p,n)

        # If value is NaN return None instead so JSON is valid
        res_recall = None if math.isnan(res_recall) else res_recall
        res_precision = None if math.isnan(res_precision) else res_precision
        res_f1 = None if math.isnan(res_f1) else res_f1
        res_accuracy = None if math.isnan(res_accuracy) else res_accuracy

        return (res_recall,res_precision,res_accuracy,res_f1)

def relative_error_total_energy(pred, ground):
    aligned_meters = align_two_meters(pred, ground)
    chunk_results = []
    sum_samples = 0.0
    for chunk in aligned_meters:
        chunk.fillna(0, inplace=True)
        sum_samples += len(chunk)
        E_pred = sum(chunk.iloc[:,0])
        E_ground = sum(chunk.iloc[:,1])

        chunk_results.append([
                            E_pred,
                            E_ground
                            ])
    if sum_samples == 0:
        return None
    else:
        [E_pred, E_ground] = np.sum(chunk_results,axis=0)
        return abs(E_pred - E_ground) / float(max(E_pred,E_ground))

def mean_absolute_error(pred, ground):
    aligned_meters = align_two_meters(pred, ground)
    total_sum = 0.0
    sum_samples = 0.0
    for chunk in aligned_meters:
        chunk.fillna(0, inplace=True)
        sum_samples += len(chunk)
        total_sum += sum(abs((chunk.iloc[:,0]) - chunk.iloc[:,1]))
    if sum_samples == 0:
        return None
    else:
        return total_sum / sum_samples


def recall(tp,fn):
    return tp/float(tp+fn)

def precision(tp,fp):
    return tp/float(tp+fp)

def f1(prec,rec):
    return 2 * (prec*rec) / float(prec+rec)

def accuracy(tp, tn, p, n):
    return (tp + tn) / float(p + n)

# Normalized Aboluste Distance
def nad(pred, ground):
    aligned_meters = align_two_meters(pred, ground)

    nominator = 0.0
    denominator = 0.0
    sum_samples = 0.0
    for chunk in aligned_meters:
        chunk.fillna(0, inplace=True)
        sum_samples += len(chunk)

        nominator += sum(abs((chunk.iloc[:,0]) - chunk.iloc[:,1]))
        denominator += sum(abs(chunk.iloc[:,1]))
    if sum_samples == 0:
        return None
    else:
        return np.sqrt(nominator / denominator)

# Mean Squared Error  np.mean(np.square(y_predict - y))
def mean_square_error(pred, ground):
    aligned_meters = align_two_meters(pred, ground)
    total_sum = 0.0
    sum_samples = 0.0
    for chunk in aligned_meters:
        chunk.fillna(0, inplace=True)
        sum_samples += len(chunk)
        total_sum += sum(np.square((chunk.iloc[:,0]) - chunk.iloc[:,1]))
    if sum_samples == 0:
        return None
    else:
        return total_sum / sum_samples

# TODO: add in code and edit JSON format
# Lungu's Disaggregation Accuracy for Appliance
def disaggregation_accuracy(pred, ground):

    aligned_meters = align_two_meters(pred, ground)

    nominator = 0.0
    denominator = 0.0
    sum_samples = 0.0
    for chunk in aligned_meters:
        chunk.fillna(0, inplace=True)
        sum_samples += len(chunk)

        nominator += np.linalg.norm(chunk.iloc[:,0] - chunk.iloc[:,1], ord=1)
        denominator += np.linalg.norm(chunk.iloc[:,0], ord=1)
    if sum_samples == 0:
        return None
    else:
        return  1 - (float(nominator) / (2*denominator))
    #     denominator += 2 * np.linalg.norm(chunk.iloc[:,0], ord=1)
    # if sum_samples == 0:
    #     return None
    # else:
    #     return  1 - (float(nominator) / denominator)
