package main

import (
	"encoding/csv"
	"fmt"
	"math"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"gonum.org/v1/gonum/mat"
)

const DATA_DIR string = "data"

type layer struct {
	weights *mat.Dense
	biases  *mat.Dense
}

func array_flatten(two_d_array [][]string) []float64 {
	var result []float64
	for _, line := range two_d_array {
		for _, value := range line {
			float_value, _ := strconv.ParseFloat(value, 64)
			result = append(result, float_value)
		}
	}
	return result
}

func load_matrix_from_csv(file_path string) (*mat.Dense, error) {
	f, err := os.Open(file_path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	data, err := csv.NewReader(f).ReadAll()
	return mat.NewDense(len(data), len(data[0]), array_flatten(data)), err
}

func load_test_data(dir string) (*mat.Dense, []layer, error) {
	inputs, err := load_matrix_from_csv(filepath.Join(dir, "inputs.csv"))
	if err != nil { return nil, nil, err }
	var layers []layer
	var weights *mat.Dense
	var biases *mat.Dense
	var path string

	for n := 1; n < 100; n++ {
		path = filepath.Join(dir, fmt.Sprintf("weights%02d.csv", n))
		_, err = os.Stat(path)
		if err != nil { break }
		weights, err = load_matrix_from_csv(path)
		if err != nil { return nil, nil, err }

		path = filepath.Join(dir, fmt.Sprintf("biases%02d.csv", n))
		_, err = os.Stat(path)
		if err != nil { break }
		biases, err = load_matrix_from_csv(path)
		if err != nil { return nil, nil, err }

		layers = append(layers, layer{weights, biases})
	}
	return inputs, layers, nil
}

func sigmoid_scalar(x float64) float64 {
	return 1 / (1 + math.Exp(-x))
}

func sigmoid_element(i, j int, value float64) float64 {
	return sigmoid_scalar(value)
}

func sigmoid_matrix(x mat.Matrix) *mat.Dense {
	var y mat.Dense
	y.Apply(sigmoid_element, x)
	return &y
}

func layer_func(inputs *mat.Dense, weights *mat.Dense, biases *mat.Dense, activation func(x mat.Matrix) *mat.Dense) *mat.Dense {
	var output mat.Dense
	output.Mul(weights, inputs)
	output.Add(&output, biases)
	return sigmoid_matrix(&output)
}

func feed_forward(x *mat.Dense, layers []layer) *mat.Dense {
	var y *mat.Dense = x
	for _, l := range layers {
		y = layer_func(y, l.weights, l.biases, sigmoid_matrix)
	}
	return y
}

func time_feed_forward(n int) {
	inputs, layers, _ := load_test_data(DATA_DIR)
	t0 := time.Now()
	for i := 0; i < n; i++ { feed_forward(inputs, layers) }
	elapsed := time.Since(t0)
	fmt.Printf("%.5f\n", elapsed.Seconds())
}

func examples() {
	x := mat.NewDense(2, 1, []float64{0, 2})
	w1 := mat.NewDense(2, 2, []float64{1, 0, 2, 4})
	b1 := mat.NewDense(2, 1, []float64{0, 1})
	w2 := mat.NewDense(2, 2, []float64{1, 0, 2, 4})
	b2 := mat.NewDense(2, 1, []float64{-0.5, 1})
	layers := []layer{{w1, b1}, {w2, b2}}

	pre := "     "
	fmt.Printf("x1 = %v\n", mat.Formatted(x, mat.Prefix(pre)))

	fmt.Printf("w1 = %v\n", mat.Formatted(w1, mat.Prefix(pre)))
	fmt.Printf("b1 = %v\n", mat.Formatted(b1, mat.Prefix(pre)))
	fmt.Println("σ(w1 * x1 + b1) = ")
	x2 := layer_func(x, w1, b1, sigmoid_matrix)
	fmt.Printf("x2 = %v\n\n", mat.Formatted(x2, mat.Prefix(pre)))

	fmt.Printf("w2 = %v\n", mat.Formatted(w2, mat.Prefix(pre)))
	fmt.Printf("b2 = %v\n", mat.Formatted(b2, mat.Prefix(pre)))
	fmt.Println("σ(w2 * x2 + b2) = ")
	x3 := feed_forward(x, layers)
	fmt.Printf("x3 = %v\n", mat.Formatted(x3, mat.Prefix(pre)))
}

func main() {
	n, err := strconv.Atoi(os.Args[1])
    if err != nil {
        fmt.Println(err)
        os.Exit(2)
    }
	time_feed_forward(n)
}