#!/usr/bin/env python3
"""
# -*- coding: utf-8 -*-
#
# Copyright 2021 Michael Büsch <m@bues.ch>
#
# Licensed under the Apache License version 2.0
# or the MIT license, at your option.
# SPDX-License-Identifier: Apache-2.0 OR MIT
#
"""

import mlplib as mlp
import numpy as np

def make_y(x):
    """Generate the true net output corresponding to the net input layer X.
    """
    def classify(x):
        if x <= 100.0:
            return np.array([1.0, 0.0, 0.0])
        elif x <= 500.0:
            return np.array([0.0, 1.0, 0.0])
        elif x <= 1000.0:
            return np.array([0.0, 0.0, 1.0])
        else:
            return np.array([0.0, 0.0, 0.0])
    nr_samples = x.shape[0]
    y = np.zeros((nr_samples, 3))
    for i in range(nr_samples):
        y[i] = classify(x[i, 0])
    return y

def make_x(nr_samples, nr_inputs):
    """Generate random net input data.
    """
    return mlp.random((nr_samples, nr_inputs)) * 1000.0

def make_net():
    """Construct the MLP network architecture.
    """
    nr_inputs = 1   # Number of inputs.
    layout = (6,    # First hidden layer size.
              12,   # Second hidden layer size.
              6,    # Third hidden layer size.
              3,)   # Output layer size.
    params = mlp.Parameters(
        weights=mlp.init_layers_weights(nr_inputs, layout), # Random weights init.
        biases=mlp.init_layers_biases(layout),              # Neutral bias init.
        actvns=(
            mlp.Sigmoid(),  # First hidden layer activation.
            mlp.Sigmoid(),  # Second hidden layer activation.
            mlp.Sigmoid(),  # Third hidden layer activation.
            mlp.Sigmoid(),  # Output layer activation.
            #mlp.Softmax(),  # Output layer activation.
        ),
    )
    return params

def train_net(params, show_progress):
    print("Synthesizing training data...")
    train_x = make_x(nr_samples=100000,
                     nr_inputs=params.nr_inputs)
    train_y = make_y(train_x)

    print("Synthesizing dev/test data...")
    test_x = make_x(nr_samples=100,
                    nr_inputs=params.nr_inputs)
    test_y = make_y(test_x)

    print("Training the network...")
    nr_iterations = 200
    minibatch_size = 128
    optimizer = mlp.AlphaDecaySimple(
        mlp.Adam(params, alpha=0.001, beta1=0.9, beta2=0.999),
        decay_rate=0.01)
    lossfn = mlp.MSE()
    train_loss = []
    for i in range(nr_iterations):
        show_progress(i, train_loss)
        for tx, ty in mlp.minibatches(train_x, train_y, minibatch_size):
            gradients, yh = mlp.backward_prop(tx, ty, params, lossfn)
            train_loss.append(lossfn.fn(yh, ty))
            optimizer.apply(i, gradients)

    train_yh = mlp.forward_prop(train_x, params)
    test_yh = mlp.forward_prop(test_x, params)
    test_loss = lossfn.fn(test_yh, test_y)

    # Calculate the percentage of matches in the test set.
    print("X", test_x)
    print("YH", mlp.collapse_bool_nodes(test_yh))
    print("Y", mlp.collapse_bool_nodes(test_y))
    train_match = mlp.proportion_equal(mlp.collapse_bool_nodes(train_yh),
                                       mlp.collapse_bool_nodes(train_y))
    test_match = mlp.proportion_equal(mlp.collapse_bool_nodes(test_yh),
                                      mlp.collapse_bool_nodes(test_y))

    print(params)
    print(f"Train set loss:  {train_loss[-1]:.6f}")
    print(f"Train set match: {train_match*100.0:.2f} %")
    print(f"Test set loss:   {test_loss:.6f}")
    print(f"Test set match:  {test_match*100.0:.2f} %")

    return train_loss

def example():
    mlp.seed(142)

    import matplotlib.pyplot as plt
    plt.ion()
    fig, ax = plt.subplots(1)
    fig.tight_layout()
    ax.set_title("Loss during training (live view...)")

    def show_progress(i, loss_values):
        if i % 4 == 0:
            print(f"Training epoch {i}...")
            ax.clear()
            ax.plot(loss_values)
            plt.pause(0.01)

    params = make_net()
    loss_values = train_net(params, show_progress)

    #TODO use plt.pause?
    #TODO scatterplot

    plt.ioff()
    ax.set_title("Loss during training")
    plt.show()

if __name__ == "__main__":
    example()

# vim: ts=4 sw=4 expandtab