import numpy as np
import util
import sys
from random import random

sys.path.append('../linearclass')

### NOTE : You need to complete logreg implementation first!

from logreg import LogisticRegression

# Character to replace with sub-problem letter in plot_path/save_path
WILDCARD = 'X'
# Ratio of class 1 to class 0
kappa = 0.1

def confusion_matrix(y_predicted, y_real, verbose=False):
    """ Calculates number of True positive, true negative, false positives and false negatives

    Args:
        y_predicted: predictions generated by the model
        y_real: real labels
        verbose: if true, print TP, FP, TN, FN values

    Returns: TP, FP, TN, FN
    """
    # First find the indices at which the true value is = 1
    TP_vec = y_predicted[y_real == 1]
    # Now evaluate whether the model estimates the probability over 0.5, obtaining indices
    TP_tup = np.where(TP_vec >= 0.5)
    # Results are stored in a ([...]) tuple, so we extract the list with [0] and count # items
    TP = len(TP_tup[0])

    FP_vec = y_predicted[y_real == 0]
    FP_tup = np.where(FP_vec >= 0.5)
    FP = len(FP_tup[0])

    TN_vec = y_predicted[y_real == 0]
    TN_tup = np.where(TN_vec < 0.5)
    TN = len(TN_tup[0])

    FN_vec = y_predicted[y_real == 1]
    FN_tup = np.where(FN_vec < 0.5)
    FN = len(FN_tup[0])

    if verbose==True:
        print(f"TP: {TP}, FP: {FP}, TN: {TN}, FN: {FN}")

    return TP, FP, TN, FN

def confusion_accuracy(TP, FP, TN, FN):
    """ Reports accuracy metrics for the predictions

    Args:
        TP, FP, TN, FN of model
    """
    A = (TP + TN) / (TP + FP + TN + FN)
    A0 = (TN) / (TN + FP)
    A1 = (TP) / (TP + FN)
    Balanced_A = (1 / 2) * (A0 + A1)

    print(f"Accuracy: {A}")
    print(f"Accuracy on 0: {A0}")
    print(f"Accuracy on 1: {A1}")
    print(f"Balanced Accuracy: {Balanced_A}")


def main(train_path, validation_path, save_path):
    """Problem 2: Logistic regression for imbalanced labels.

    Run under the following conditions:
        1. naive logistic regression
        2. upsampling minority class

    Args:
        train_path: Path to CSV file containing training set.
        validation_path: Path to CSV file containing validation set.
        save_path: Path to save predictions.
    """
    output_path_naive = save_path.replace(WILDCARD, 'naive')
    output_path_upsampling = save_path.replace(WILDCARD, 'upsampling')

    # *** START CODE HERE ***
    # Part (b): Vanilla logistic regression1
    # Make sure to save predicted probabilities to output_path_naive using np.savetxt()

    # Load training and validation data
    x_train, y_train = util.load_dataset(train_path, add_intercept=True)
    x_val, y_val = util.load_dataset(validation_path, add_intercept=True)

    # Train logistic regression on training sample, then generate predictions on validation
    model = LogisticRegression()
    model.fit(x_train, y_train)
    y_val_hat = model.predict(x_val)

    # To obtain TP, FP, TN, FN, use our confusion_matrix helper function
    TP, FP, TN, FN = confusion_matrix(y_val_hat, y_val, verbose=True)

    # Calculating accuracy
    confusion_accuracy(TP, FP, TN, FN)

    # Plotting
    output_path_naive_img = output_path_naive[:-3] + 'png'
    util.plot(x_val, y_val, model.theta, output_path_naive_img)

    # Saving predicted probabilities to output_path_naive
    np.savetxt(output_path_naive, y_val_hat)

    # Part (d): Upsampling minority class
    # Make sure to save predicted probabilities to output_path_upsampling using np.savetxt()
    # Repeat minority examples 1 / kappa times
    print('-'*30)

    # Calculating p, which is portion of minority (positive) examples
    p = y_train.sum() / len(y_train)

    # Calculating k
    k = p / (1-p)

    # Adding (1/k)-1 repetitions of positive examples to training data
    x_pos = x_train[y_train==1]
    x_rep = np.repeat(x_pos, (1/k)-1, axis=0)
    x_mod = np.vstack((x_train, x_rep))

    # Adding (1/k)-1 repetitions of positive examples to y vector
    y_pos = y_train[y_train==1]
    y_rep = np.repeat(y_pos, (1/k)-1)
    y_mod = np.append(y_train, y_rep)

    # Building and fitting model, then making predictions on original validation data
    upsample_model = LogisticRegression()
    upsample_model.fit(x_mod, y_mod)
    y_mod_hat = upsample_model.predict(x_val)

    # Calculating confusion matrix and reporting accuracy metrics
    TP_mod, FP_mod, TN_mod, FN_mod = confusion_matrix(y_mod_hat, y_val, verbose=True)
    confusion_accuracy(TP_mod, FP_mod, TN_mod, FN_mod)

    # Plotting the re-weighted model
    output_path_upsampling_img = output_path_upsampling[:-3] + 'png'
    util.plot(x_val, y_val, upsample_model.theta, output_path_upsampling_img)

    # Saving predictions using re-weighted logistic model
    np.savetxt(output_path_upsampling, y_mod_hat)

    # *** END CODE HERE

if __name__ == '__main__':
    main(train_path='train.csv',
        validation_path='validation.csv',
        save_path='imbalanced_X_pred.txt')
