import os

import numpy as np
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping
from sklearn.utils import shuffle

import constants as c
from utils import prepare_data, export_data
from models import target_model, shadow_model, attack_model
from validate import validate_attack_model, naive_solution, plot_roc_curve


def train_target_model(train_size, epochs, random_state=None, path=None, data_path=None):
    """
    Train target/victim model.
    Return: model           | trained model
            predictions     | model predictions on complete dataset
            classifications | classification of predictions as train(1) or test(0)
    """
    x_train, x_test, y_train, y_test = prepare_data(train_size=train_size, random_state=random_state)

    model = target_model()
    model.fit(x_train, y_train, epochs=epochs, validation_split=0.1, verbose=True)
    model.evaluate(x_test, y_test, verbose=True)

    train_predictions = model.predict(x_train)
    test_predictions  = model.predict(x_test)

    x                 = np.concatenate([x_train, x_test], axis=0)
    y                 = np.concatenate([y_train, y_test], axis=0)
    predictions       = np.concatenate([train_predictions, test_predictions], axis=0)
    classifications   = np.concatenate([np.ones(train_predictions.shape[0]), np.zeros(test_predictions.shape[0])], axis=0)

    [x, y, predictions, classifications] = shuffle(x, y, predictions, classifications)

    if path: model.save(path)

    if data_path: export_data(x, y, predictions, classifications, data_path)

    return model, predictions, classifications


def train_shadow_model(train_size, epochs, path=None):
    """
    Train shadow model.
    Assume architecture of target/victim model is known.
    Return: predictions     | model predictions on complete dataset
            classifications | classification of predictions as train(1) or test(0)
    """
    _, predictions, classifications = train_target_model(train_size, epochs, path=path)

    return predictions, classifications


def train_attack_model(shadow_predictions, shadow_classifications, epochs, path=None):
    """
    Train attack model.
    Uses predictions / classifications of shadow models.
    Return: model | trained attack model
    """
    model = attack_model()
    callback = EarlyStopping(monitor='val_acc', min_delta=0, patience=5)
    model.fit(shadow_predictions, shadow_classifications, epochs=epochs, callbacks=[callback], validation_split=0.1, verbose=True)

    if path: model.save(path)

    return model


if __name__ == "__main__":
    EPOCHS = 15
    NUM_MODELS = 1
    NUM_SHADOW_MODELS = 5

    for target_model_index in range(NUM_MODELS):
        print('Training model: ', target_model_index)
        target_model_path = os.path.join(c.MODELS_DIRECTORY, c.TARGET_MODEL_NAME.format(target_model_index=target_model_index))
        target_model_data_path = os.path.join(c.MODELS_DIRECTORY, c.TARGET_DATA_NAME.format(target_model_index=target_model_index))

        _, target_predictions, target_classifications = train_target_model(0.5, EPOCHS, random_state=None, path=target_model_path,
                                                                                                           data_path=target_model_data_path)

        shadow_predictions = []
        shadow_classifications = []

        for shadow_model_index in range(NUM_SHADOW_MODELS):
            print('Training shadow model: ', shadow_model_index)
            shadow_model_path = os.path.join(c.MODELS_DIRECTORY, c.SHADOW_MODEL_NAME.format(target_model_index=target_model_index,
                                                                                        shadow_model_index=shadow_model_index))

            predictions, classifications = train_shadow_model(0.2, EPOCHS, path=shadow_model_path)

            shadow_predictions.append(predictions)
            shadow_classifications.append(classifications)

        shadow_predictions = np.concatenate(shadow_predictions, axis=0)
        shadow_classifications = np.concatenate(shadow_classifications, axis=0)

        print('Training attack model: ', target_model_index)
        attack_model_path = os.path.join(c.MODELS_DIRECTORY, c.ATTACK_MODEL_NAME.format(target_model_index=target_model_index))
        attack_model = train_attack_model(shadow_predictions, shadow_classifications, epochs=10, path=attack_model_path)

        attack_model.evaluate(target_predictions, target_classifications)

        # validation
        attack_predictions = attack_model.predict(target_predictions)
        validate_attack_model(attack_predictions, target_classifications, cutoff=0.5)

        naive_solution(target_predictions, target_classifications, cutoff=0.99)

        # plots
        plot_roc_curve(target_classifications, attack_predictions)