import pandas as pd
import numpy as np

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

# 30 10(11) 5
# 27 17 4

def make_model(size, magic=6):
    # Our model for the stacks, a small CNN
    stack_shape = (size, size, 8)
    stack_input = Input(shape=stack_shape)
    stack_model = Convolution2D(magic, kernel_size=(3, 3), strides=(1, 1),
                                padding='valid', activation="relu",
                                use_bias=True)(stack_input)
    stack_model = Flatten()(stack_model)
    # stack_model = Dense(magic, activation="relu", use_bias=True)(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+2, activation="relu", use_bias=True)(combn_input)
    model = Dense(magic, activation="relu", use_bias=True)(model)
    model = Dense(2, activation="sigmoid", 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',
                  metrics=['accuracy'])

    model.summary()
    return model

def train(size, model, reps):
    shape = (-1, size, size, 8)

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

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

    results = []
    for i in range(0, reps):
        print("Iteration {0}/{1}".format(i+1, reps))
        model.fit(training_input, training_outcome, epochs=5,
                  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))]

    print("\n\n")
    for i in range(0, reps):
        print("Iteration {0}/{1}: {2}".format(i+1, reps, results[i]))

    return model

# model = train(5, make_model(5), 50)
# print(model.trainable_variables)