from read_data import get_data_from_csv, split_data
from neural_network_exercice import MyNeuralNetwork
import numpy as np

if __name__ == "__main__":
    steps = 100000
    all_data = get_data_from_csv()
    test_data, train_data = split_data(all_data)
    print(f"we have {len(train_data)} training examples")
    print(f"we have {len(test_data)} test examples")
    model = MyNeuralNetwork(
        input_dimension=13,
        hidden_dimension=32,
        relu_to_use=1,
        learning_rate=0.05,
        total_num_steps=steps,
    )

    mse = []
    for example in test_data:
        X = example[:-1]
        Y = example[-1:]
        P = model(X)
        mse.append((P - Y) ** 2)
    mse_train = []
    for example in train_data:
        X = example[:-1]
        Y = example[-1:]
        P = model(X)
        mse_train.append((P - Y) ** 2)
    print(
        f"\npre-training MSE: {np.mean(mse):.3f} (test) {np.mean(mse_train):.3f} (train)"
    )
    batch_size = 2
    for step in range(steps):
        index = np.random.choice(train_data.shape[0], batch_size, replace=False)
        batch = np.copy(train_data[index, :])
        X = batch[:, :-1].astype(np.float32)
        Y = batch[:, -1:].astype(np.float32)
        model(np.transpose(X), np.transpose(Y))
    print()
    mse = []
    for example in test_data:
        X = example[:-1]
        Y = example[-1:]
        P = model(X)
        mse.append((P - Y) ** 2)
    mse_train = []
    for example in train_data:
        X = example[:-1]
        Y = example[-1:]
        P = model(X)
        mse_train.append((P - Y) ** 2)
    print(
        f"post-training MSE: {np.mean(mse):.3f} (test) {np.mean(mse_train):.3f} (train)"
    )