import os
import shutil
import numpy as np
from keras import Input
from keras.layers import Dense
from keras.models import Model
from keras.optimizers import RMSprop
from keras import callbacks
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

TENSORBOARD_PATH = './logs'
reset=True
print('Resetting:', reset)
if reset:
    while os.path.isdir(TENSORBOARD_PATH):
        shutil.rmtree(TENSORBOARD_PATH)

def build_model(shape, name=None):
    x = Input(shape)
    y = Dense(128, activation='relu')(x)
    model = Model(inputs=x, outputs=y, name=name)
    optimizer = RMSprop(lr=0.01, rho=0.9, epsilon=1e-08, decay=0.0)
    model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer)
    return model

x, y = load_digits(return_X_y=True)
x, x_test, y, y_test = train_test_split(x, y)

samples, features = x.shape
models = [build_model([features], name=name) for name in ('alpha', 'beta')]

for model in models:
    # Training "fresh" models.
    model.fit(x, y,
              epochs=2,
              batch_size=None,
              validation_data=(x_test, y_test),
              callbacks=[callbacks.TensorBoard('./logs/' + model.name)])
    model.save_weights('./%s.hdf5'%(model.name))

for model in models:
    # Training "trained" models.
    model.load_weights('./%s.hdf5'%(model.name))
    model.fit(x, y,
              epochs=4,
              batch_size=None,
              validation_data=(x_test, y_test),
              callbacks=[callbacks.TensorBoard('./logs/' + model.name)],
              initial_epoch=2)
    model.save_weights('./%s.hdf5'%(model.name))


for model in models:
    # Training "fresh, but trained" models.
    model = Model(model.inputs, model.outputs, name=model.name)
    optimizer = RMSprop(lr=0.001, rho=0.9, epsilon=1e-08, decay=0.0)
    model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer)
    model.load_weights('./%s.hdf5'%(model.name))
    model.fit(x, y,
              epochs=6,
              batch_size=None,
              validation_data=(x_test, y_test),
              callbacks=[callbacks.TensorBoard('./logs/' + model.name)],
              initial_epoch=6)