import numpy as np
import plotille
import tqdm
from IPython import embed
from scipy.linalg import pinv2
from sklearn.datasets import make_circles, make_classification, make_moons


class ELM:
    def __init__(self, x, y, hidden_dim=None):
        self.input_size = x.shape[-1]
        if hidden_dim is None:
            self.hidden_size = int(x.shape[0] / 7) 
        else:
            self.hidden_size = hidden_dim
        self.input_weights = np.random.normal(size=[self.input_size, self.hidden_size])
        self.biases = np.random.normal(size=[self.hidden_size])
        self.output_weights = np.dot(pinv2(self.hidden_nodes(x)), y)

    def hidden_nodes(self, X):
        G = np.dot(X, self.input_weights) + self.biases
        relu = lambda x_: np.maximum(x_, 0, x_)
        H = relu(G)
        return H

    def __call__(self, x):
        H = self.hidden_nodes(x)
        y = np.dot(H, self.output_weights)
        return y


if __name__ == '__main__':

    X, y = make_moons(n_samples=500, noise=0.3, random_state=0)
    X = (X - X.min(0) + 1e-6) / (X.max(0) - X.min(0) + 1e-6)

    n_hiddens = list(range(2, 32))
    losses = list()
    for n_hidden in tqdm.tqdm(n_hiddens):
        mask = np.random.random(X.shape[0]) < 0.1 
        train_x, train_y = X[mask], y[mask]
        test_x, test_y = X[~mask], y[~mask]
        model = ELM(train_x, train_y, n_hidden)
        loss = (0.5 * (test_y - model(test_x)) ** 2).mean()
        losses.append(loss)

    print(plotille.plot(n_hiddens, losses))

    model = ELM(train_x, train_y)
    loss = (0.5 * (test_y - model(test_x)) ** 2).mean()
    embed(); exit()