# github.com/mrbid
import sys
import os
import random
import struct
import numpy as np
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from time import time_ns
from sys import exit
from os import mkdir
from os.path import isdir
from os.path import isfile

# disable warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # cpu training only

# print everything / no truncations
np.set_printoptions(threshold=sys.maxsize)

# hyperparameters
project = "Heads_Simple"
model_name = 'keras_model'
optimiser = 'adam'
inputsize = 2
outputsize = 6
activator = 'tanh'
layers = 8
layer_units = 256
batches = 256
epoches = 3
use_bias = True

# load options
argc = len(sys.argv)
if argc >= 2:
    layers = int(sys.argv[1])
    print("layers:", layers)
if argc >= 3:
    layer_units = int(sys.argv[2])
    print("layer_units:", layer_units)
if argc >= 4:
    batches = int(sys.argv[3])
    print("batches:", batches)
if argc >= 5:
    activator = sys.argv[4]
    print("activator:", activator)
if argc >= 6:
    optimiser = sys.argv[5]
    print("optimiser:", optimiser)
if argc >= 7 and sys.argv[6] == '1':
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    print("CPU_ONLY: 1")
if argc >= 8:
    epoches = int(sys.argv[7])
    print("epoches:", epoches)

# make sure save dir exists
model_name = 'models/' + activator + '_' + optimiser + '_' + str(layers) + '_' + str(layer_units) + '_' + str(batches) + '_' + str(epoches)
outdir_name = model_name + '_pd'


##########################################
#   LOAD DATASET
##########################################
print("\n--Loading Dataset")
st = time_ns()

tss = 632847695
# tss = int(os.stat("train_y.dat").st_size / 24)
print("Dataset Size:", "{:,}".format(tss))

if isfile("train_x.npy"):
    # train_x = np.load("train_x.npy")[:1000000]
    # train_y = np.load("train_y.npy")[:1000000]
    train_x = np.load("train_x.npy")
    train_y = np.load("train_y.npy")
else:
    load_x = []
    with open("train_x.dat", 'rb') as f:
        load_x = np.fromfile(f, dtype=np.float32)
    train_x = np.reshape(load_x, [tss, inputsize])

    load_y = []
    with open("train_y.dat", 'rb') as f:
        load_y = np.fromfile(f, dtype=np.float32)
    train_y = np.reshape(load_y, [tss, outputsize])
    
    np.save("train_x.npy", train_x)
    np.save("train_y.npy", train_y)

timetaken = (time_ns()-st)/1e+9
print("Time Taken:", "{:.2f}".format(timetaken), "seconds")

# train_x = np.reshape(train_x, [tss, inputsize])
# train_y = np.reshape(train_y, [tss, outputsize])

# print(train_x.shape)
# print(train_x)
# print(train_y.shape)
# print(train_y)
# exit()


##########################################
#   TRAIN
##########################################
print("\n--Training Model")

# construct neural network
model = Sequential()
model.add(Dense(layer_units, activation=activator, use_bias=use_bias, input_dim=inputsize))
if layers > 0:
    for x in range(layers):
        model.add(Dense(layer_units, use_bias=use_bias, activation=activator))
model.add(Dense(outputsize, use_bias=use_bias))

# output summary
model.summary()

if optimiser == 'adam':
    optim = keras.optimizers.Adam(learning_rate=0.001)
elif optimiser == 'sgd':
    lr_schedule = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=0.3, decay_steps=epoches*tss, decay_rate=0.1)
    #lr_schedule = keras.optimizers.schedules.ExponentialDecay(initial_learning_rate=0.1, decay_steps=epoches*tss, decay_rate=0.01)
    optim = keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.0, nesterov=False)
elif optimiser == 'momentum':
    optim = keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=False)
elif optimiser == 'nesterov':
    optim = keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=True)
elif optimiser == 'nadam':
    optim = keras.optimizers.Nadam(learning_rate=0.001)
elif optimiser == 'adagrad':
    optim = keras.optimizers.Adagrad(learning_rate=0.001)
elif optimiser == 'rmsprop':
    optim = keras.optimizers.RMSprop(learning_rate=0.001)
elif optimiser == 'adadelta':
    optim = keras.optimizers.Adadelta(learning_rate=0.001)
elif optimiser == 'adamax':
    optim = keras.optimizers.Adamax(learning_rate=0.001)
elif optimiser == 'ftrl':
    optim = keras.optimizers.Ftrl(learning_rate=0.001)

model.compile(optimizer=optim, loss='mean_squared_error', metrics=['accuracy'])

# train network
history = model.fit(train_x, train_y, epochs=epoches, batch_size=batches, shuffle=True)
model_name = model_name + "_" + "[{:.2f}]".format(history.history['accuracy'][-1])
timetaken = (time_ns()-st)/1e+9
print("\nTime Taken:", "{:.2f}".format(timetaken), "seconds")


##########################################
#   EXPORT
##########################################

# export model
print("\n--Exporting Model")
if not isdir('models'): mkdir('models')
st = time_ns()

# save keras model
model.save(model_name + ".keras")

# save prediction csv
predict_points = 200000
predict_x = np.empty([predict_points, 2], float)
rseed = random.uniform(0,1)
for i in range(predict_points): predict_x[i] = [rseed, i/predict_points]
p = model.predict(predict_x)
if not isdir(outdir_name): mkdir(outdir_name)
np.asarray(p).tofile(outdir_name + '/head_sample.csv', sep=',')

# export PLY file of prediction
# https://gist.github.com/Shreeyak/9a4948891541cb32b501d058db227fff
fid = open(outdir_name + "/head_sample.ply", 'wb')
fid.write(bytes('ply\n', 'utf-8'))
fid.write(bytes('format binary_little_endian 1.0\n', 'utf-8'))
fid.write(bytes('element vertex ' + str(predict_points) + '\n', 'utf-8'))
fid.write(bytes('property float x\n', 'utf-8'))
fid.write(bytes('property float y\n', 'utf-8'))
fid.write(bytes('property float z\n', 'utf-8'))
fid.write(bytes('property uchar red\n', 'utf-8'))
fid.write(bytes('property uchar green\n', 'utf-8'))
fid.write(bytes('property uchar blue\n', 'utf-8'))
fid.write(bytes('end_header\n', 'utf-8'))
for i in range(predict_points):
    fid.write(bytearray(struct.pack("fffccc",p[i][0],p[i][1],p[i][2],
        int(round(p[i][3]*255)).to_bytes(1, 'little'),
        int(round(p[i][4]*255)).to_bytes(1, 'little'),
        bytes([np.uint8(p[i][5]*255)]),
        )))
fid.close()

# print timing
timetaken = (time_ns()-st)/1e+9
print("\nTime Taken:", "{:.2f}".format(timetaken), "seconds\n")

# save weights for C array
print("")
print("Exporting weights...")
st = time_ns()
li = 0
f = open("models/" + project + "_layers.h", "w")
f.write("#ifndef " + project + "_layers\n#define " + project + "_layers\n\n// accuracy: " + "{:.8f}".format(history.history['accuracy'][-1]) + "\n// loss: " + "{:.8f}".format(history.history['loss'][-1]) + "\n\n")
if f:
    for layer in model.layers:
        total_layer_weights = layer.get_weights()[0].transpose().flatten().shape[0]
        total_layer_units = layer.units
        layer_weights_per_unit = total_layer_weights / total_layer_units
        print("+ Layer:", li)
        print("Total layer weights:", total_layer_weights)
        print("Total layer units:", total_layer_units)
        print("Weights per unit:", int(layer_weights_per_unit))

        f.write("const float " + project + "_layer" + str(li) + "[] = {")
        isfirst = 0
        wc = 0
        bc = 0
        if layer.get_weights() != []:
            for weight in layer.get_weights()[0].transpose().flatten():
                wc += 1
                if isfirst == 0:
                    f.write(str(weight))
                    isfirst = 1
                else:
                    f.write("," + str(weight))
                if wc == layer_weights_per_unit:
                    f.write(", /* bias */ " + str(layer.get_weights()[1].transpose().flatten()[bc]))
                    wc = 0
                    bc += 1
        f.write("};\n\n")
        li += 1
f.write("#endif\n")
f.close()

# print timing
timetaken = (time_ns()-st)/1e+9
print("\nTime Taken:", "{:.2f}".format(timetaken), "seconds\n")
