import sys
from pathlib import Path
from timeit import timeit

import numpy as np


THIS_DIR = Path(__file__).parent
TEST_DATA_DIR = Path(THIS_DIR, '..', 'data')


def csv_to_array(file_path: Path, dtype=np.float32) -> np.ndarray:
    return np.loadtxt(file_path, dtype=dtype, delimiter=',')


def load_test_data() -> tuple[np.ndarray, list[tuple[np.ndarray, np.ndarray]]]:
    inputs = csv_to_array(Path(TEST_DATA_DIR, 'inputs.csv'))
    weights_iter = (csv_to_array(p) for p in sorted(TEST_DATA_DIR.glob('weights*.csv')))
    biases_iter = (csv_to_array(p) for p in sorted(TEST_DATA_DIR.glob('biases*.csv')))
    return inputs, list(zip(weights_iter, biases_iter))


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def layer_func(input_vector: np.ndarray, weight_matrix: np.ndarray, bias_vector: np.ndarray) -> np.ndarray:
    return sigmoid(np.matmul(weight_matrix, input_vector) + bias_vector)


def feed_forward(x: np.ndarray, *layers: tuple[np.ndarray, np.ndarray]) -> np.ndarray:
    for w, b in layers:
        x = layer_func(x, w, b)
    return x


def main(n: int) -> None:
    t = timeit(
        stmt='feed_forward(inp, *layers)',
        setup='from __main__ import feed_forward, load_test_data;' + 
              'inp, layers = load_test_data()', 
        number=n
    )
    print(round(t, 5))


if __name__ == '__main__':
    main(int(sys.argv[1]))