# cnn_train.py, train a small CNN to recognise winning Tak positions
#
# Copyright (C) 2021, tslil clingman
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


# Loading data
import pandas as pd
import numpy as np

# Output of weights
from tensorflow import transpose

# Building models
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense


def load_data(size):
    output_len = 2
    # Load training data
    tr_fn = "data/training-"+str(size)+".csv"
    training_csv = pd.read_csv(tr_fn)
    training_data = training_csv
    training_input = np.array(training_data.iloc[:, 0:-output_len])
    training_outcome = np.array(training_data.iloc[:, -output_len:])
    # Load validation data
    val_fn = "data/validation-"+str(size)+".csv"
    val_data = pd.read_csv(val_fn)
    val_input = np.array(val_data.iloc[:, 0:-output_len])
    val_outcome = np.array(val_data.iloc[:, -output_len:])
    return ((training_input, training_outcome), (val_input, val_outcome))


def train(size, model, data, iterations=1, epochs=10, batch=None):
    (tra_input, tra_outcome), (val_input, val_outcome) = data
    results = []
    for i in range(0, iterations):
        print("Iteration {0}/{1}".format(i+1, iterations))
        model.fit(tra_input, tra_outcome, epochs=epochs,
                  validation_data=(val_input, val_outcome),
                  verbose=True, batch_size=batch)
        val_res = model.evaluate(val_input, val_outcome, verbose=False)
        tra_res = model.evaluate(tra_input, tra_outcome, verbose=False)
        results.append((tra_res, val_res))
        print(val_res)
        write_weights(model, i+1, (tra_res, val_res))
    print("\nScores")
    for i, data in enumerate(results):
        print(f"Iteration {i}: {data}")
    return results


def make_model(size, magic):
    inputs = Input(shape=(size * size + 3,))
    model = inputs
    model = Dense(magic, activation="relu", use_bias=True)(model)
    model = Dense(2, activation="relu", use_bias=True)(model)
    model = Model(inputs=inputs, outputs=model)
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    model.summary()
    return model


def write_weights(model, iteration, performance):
    def fix(val):
        string = str(np.array(val).tolist())
        string = string.replace("[", "{").replace("]", "}")
        return string
    dense1_weights = transpose(model.trainable_variables[0], perm=[1, 0])
    dense1_biases = model.trainable_variables[1]
    output_weights = transpose(model.trainable_variables[2], perm=[1, 0])
    output_bias = model.trainable_variables[3]
    # Prepare output
    to_output = [("dense1_weights[DENSE_NUM][INP_NUM]",dense1_weights)
                 ("dense1_biases[DENSE_NUM]", dense1_biases),
                 ("output_weights[2][DENSE_NUM]", output_weights),
                 ("output_bias[2]", output_bias)]
    # Write to file
    f = open("weights-"+str(iteration)+".txt", "w")
    f.write("/*\n")
    model.summary(print_fn=lambda l: f.write(" * "+l+"\n"))
    f.write(" * "+str(performance)+"\n*/\n\n")
    f.write("#include \"weights.h\"\n\n")
    for (name, val) in to_output:
        f.write("const float "+name+" =\n"+fix(val)+";\n\n")
    f.close()


data = load_data(5)
model = make_model(5, 64)

print("Before training", model.evaluate(data[1][0], data[1][1], verbose=False, batch_size=16))
results = train(5, model, data, iterations=1, epochs=10, batch=None)