import numpy as np
from test_generic import *

def train1(x, model_manager, sequence_len):
    models = model_manager.extract_features_td(x, sequence_len)
    models, arr_min, arr_max = rescale(models)

    number_of_features = models.shape[1]
    network = TRBM(number_of_features,number_of_features*2)
    network.train(models, epochs=10, learning_rate=0.1)
    return network, arr_min, arr_max

def train2(x, model_manager, sequence_len):
    models = model_manager.extract_features_td(x, sequence_len)
    models, arr_min, arr_max = rescale(models)

    number_of_features = models.shape[1]
    network = TRBM(number_of_features,number_of_features*3)
    network.train(models, epochs=10, learning_rate=0.1)
    return network, arr_min, arr_max

def train_with_models(models):
    if models.shape[0] == 0:
        return None

    number_of_features = models.shape[1]
    network = DBN(number_of_features, [number_of_features*2])
    network.train(models, epochs=2, learning_rate=0.1)
    return network

def retrain(x, arr_min, arr_max, network, bad_memory, model_manager, sequence_len, similarity_threshold=1e-3, use_bad_memory=False):
    models = model_manager.extract_features_td(x, sequence_len)
    models,_,_ = rescale(models, arr_min, arr_max)
    if use_bad_memory and bad_memory != None:
        number_of_models = models.shape[0]
        for i in xrange(number_of_models):
            sample = bad_memory.sample_network()
            dist = distance(models[i,:], sample)
            if dist > similarity_threshold:
                model = models[i][np.newaxis]
                network.train(model, epochs=10, learning_rate=0.1)
    else:
        network.train(models, epochs=2, learning_rate=0.1)
    return network

def test(y, arr_min, arr_max, network, model_manager, sequence_len, threshold=-1):
    number_of_sequences = len(y) - sequence_len
    test_models = model_manager.extract_features_td(y, sequence_len)
    test_models,_,_ = rescale(test_models, arr_min, arr_max)

    initial_data = np.zeros((sequence_len-2,1))
    for i in xrange(sequence_len-2):
        initial_data[i] = np.random.rand()

    network.initialise(initial_data)

    distances = list()
    suspicious_data = list()
    for i in xrange(0,number_of_sequences):
        sample = network.sample_network(test_models[i])
        dist = distance(sample, test_models[i])

        if threshold > 0. and dist > threshold:
            suspicious_data.append(test_models[i])
        distances.append(dist)

    suspicious_data = np.array(suspicious_data)
    return distances, suspicious_data

sequence_len = 10
model_manager = ModelFitLibrary()
number_of_test_sequences = 90

#print 'Not using bad memory'
generic_feature_bad_memory_test(train1, retrain, train_with_models, test, model_manager, sequence_len, number_of_test_sequences, False)

# print 'Using bad memory'
# generic_feature_bad_memory_test(train1, retrain, train_with_models, test, model_manager, sequence_len, number_of_test_sequences, True)

#print 'Testing with hidden neurons = 2 * visible neurons'
#generic_feature_test(train1, retrain, train_with_models, test, model_manager, sequence_len, number_of_test_sequences)

#print 'Testing with hidden neurons = 3 * visible neurons'
#generic_feature_bad_memory_test(train2, retrain, train_with_models, test, model_manager, sequence_len, number_of_test_sequences)

# train_data_front = np.genfromtxt('../../logs/laser_front.log')
# test_data_front = np.genfromtxt('../../logs/laser_front_test.log')
# anomalous_data_front = np.genfromtxt('../../logs/laser_front_faulty.log')
# retrain_data_front = np.genfromtxt('../../logs/laser_front_retrain.log')

# train_data = train_data_front.mean(axis=1)
# test_data = test_data_front.mean(axis=1)
# anomalous_data = anomalous_data_front.mean(axis=1)
# retrain_data = retrain_data_front.mean(axis=1)

# # train_data = train_data_front.max(axis=1)
# # test_data = test_data_front.max(axis=1)
# # anomalous_data = anomalous_data_front.max(axis=1)
# # retrain_data = retrain_data_front.max(axis=1)

# number_of_test_sequences =  max(test_data_front.shape[0], anomalous_data.shape[0]) - sequence_len
# generic_sensor_feature_test(train_data, retrain_data, test_data, anomalous_data, train1, retrain, train_with_models, test, model_manager, sequence_len, number_of_test_sequences, False)