# Loading data
import pandas as pd
import numpy as np

# Output of weights
import re
from tensorflow import transpose

# Custom activation function
from tensorflow import constant as k
from tensorflow.math import add, divide, maximum, minimum, multiply, square

# Computing size
from tensorflow.keras.backend import get_value

# Building models
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Concatenate
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Convolution2D
from tensorflow.keras.layers import Flatten


# We need something that's close to logistic, but cheaper to compute.
# (12.0+x+50.0*x/(x*x+10.0))/24.0, clamped between 0 and 1
# as it would otherwise exceed this range at +- 4.6 or so
def truncated_pade_logistic(x):
    val = add(k(12.0),
              add(x, multiply(k(50.0),
                              divide(x, add(square(x), k(10.0))))))
    return minimum(k(1.0), maximum(k(0.0), divide(val, k(24.0))))

def make_model(size, magic=[12, 9, 9]):
    # Our model for the stacks, a small CNN
    stack_shape = (size, size, 8)
    stack_input = Input(shape=stack_shape)
    stack_model = Convolution2D(magic[0], kernel_size=(3, 3), strides=(1, 1),
                                padding='valid', activation="relu",
                                use_bias=True)(stack_input)
    stack_model = Flatten()(stack_model)
    stack_model = Model(inputs=stack_input, outputs=stack_model)
    # The overall model
    flats_input = Input(shape=(2,))
    combn_input = Concatenate()([stack_model.output, flats_input])
    model = Dense(magic[1], activation="relu", use_bias=True)(combn_input)
    model = Dense(magic[2], activation="relu", use_bias=True)(model)
    model = Dense(1, activation=truncated_pade_logistic, use_bias=True)(model)
    model = Model(inputs=[stack_model.input, flats_input], outputs=model)
    model.compile(optimizer='adam',
                  # loss='binary_crossentropy',
                  # loss='categorical_crossentropy',
                  loss='mean_squared_error',
                  # loss='mean_absolute_error',
                  metrics=['accuracy'])
    model.summary()
    return model


def load_data(size):
    shape = (-1, size, size, 8)

    tr_fn = "training-"+str(size)+".csv"
    training_csv = pd.read_csv(tr_fn)  #.head(100000)
    training_data = training_csv
    training_stack_input = np.array(training_data.iloc[:, 2:-1]).reshape(shape, order='F')
    training_flats_input = np.array(training_data.iloc[:, 0:2])
    training_input = [training_stack_input, training_flats_input]
    training_outcome = training_data.iloc[:, -1:]

    val_fn = "validation-"+str(size)+".csv"
    val_data = pd.read_csv(val_fn).tail(20000)
    val_stack_input = np.array(val_data.iloc[:, 2:-1]).reshape(shape, order='F')
    val_flats_input = np.array(val_data.iloc[:, 0:2])
    val_input = [val_stack_input, val_flats_input]
    val_outcome = val_data.iloc[:, -1:]

    return [(training_input, training_outcome), (val_input, val_outcome)]

def train(size, model, data, iterations=1, epochs=10):
    (training_input, training_outcome), (val_input, val_outcome) = data
    results = []
    for i in range(0, iterations):
        print("Iteration {0}/{1}".format(i+1, iterations))
        model.fit(training_input, training_outcome, epochs=epochs,
                  validation_data=(val_input, val_outcome),
                  verbose=True)
        v_loss, v_accuracy = model.evaluate(val_input, val_outcome,
                                            verbose=False)
        t_loss, t_accuracy = model.evaluate(training_input, training_outcome,
                                            verbose=False)
        results += [((v_loss, t_loss), (v_accuracy, t_accuracy))]
    return (results, model)


def model_size(model):
    return sum([np.prod(get_value(w).shape) for w in model.trainable_weights])


def magic_search(data):
    results = []
    for width in range(12, 16):
        for dense1 in range(8, 24):
            for dense2 in range(6, dense1+1):
                m = make_model(5, [width, dense1, dense2])
                if model_size(m) <= 2000:
                    results += [([width, dense1, dense2],
                                 train(5, m, data, iterations=1, epochs=20))]
    print("\nSummary")
    for r in results:
        print("Parameters {0}: {1}".format(r[0], r[1][0][-1:][0]))
    return results


def write_weights(model):
    def fix(string):
        string = string.replace("[", "{").replace("]", "}")
        string = re.sub(r'}\n', '},\n', string)
        string = re.sub(r'([0-9]+)\n', r'\1,\n', string)
        string = re.sub(r'([0-9]+) ', r'\1, ', string)
        return string
    f = open("weights.txt", "w")
    conv2d_weights = np.array(transpose(model.trainable_variables[0], perm=[3, 1, 0, 2]))
    conv2d_biases = np.array(model.trainable_variables[1])
    dense1_weights = np.array(transpose(model.trainable_variables[2], perm=[1, 0]))
    dense1_biases = np.array(model.trainable_variables[3])
    dense2_weights = np.array(transpose(model.trainable_variables[4], perm=[1, 0]))
    dense2_biases = np.array(model.trainable_variables[5])
    output_weights = np.array(transpose(model.trainable_variables[6], perm=[1, 0]))
    output_biases = np.array(model.trainable_variables[7])
    for v in [conv2d_weights, conv2d_biases,
              dense1_weights, dense1_biases,
              dense2_weights, dense2_biases,
              output_weights, output_biases]:
        f.write(fix(str(v))+"\n\n")
    f.close()


data = load_data(5)
model = make_model(5, [12, 9, 9])
results, model = train(5, model, data, iterations=10, epochs=20)
print("\nScores")
for i in range(len(results)):
    print("Iteration {0}: {1}".format(i+1, results[i]))

write_weights(model)

#results = magic_search(data)