// Package ml machine learning operations for flow
package ml

import (
	"flow/registry"
	"math"
	"math/rand"

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

// Matrix wrapper
type Matrix = mat.Matrix

// New registry
func New() *registry.R {

	r := registry.New()

	registry.Describer(
		r.Add(matNew).Inputs("rows", "columns", "data"),
		r.Add(
			normFloat,
			matMul,
			matTranspose,
			matSigmoid,
			matSigmoidD,
			toFloatArr,
		),
	).Description("gonum functions").
		Tags("gonum").
		Extra("style", registry.M{"color": "#953"})

	return r
}

func normFloat(n int) []float64 {
	data := make([]float64, n)
	for i := range data {
		data[i] = rand.NormFloat64()
	}
	return data
}

func matNew(r, c int, data []float64) Matrix {
	return mat.NewDense(r, c, data)
}

func matMul(a Matrix, b Matrix) Matrix {
	var c mat.Dense
	c.Mul(a, b)
	return &c
}

func matTranspose(a Matrix) Matrix {
	return a.T()
}

func matSigmoid(a Matrix) Matrix {
	// Should be a vector perhaps
	r, c := a.Dims()
	ret := mat.NewDense(r, c, nil)
	for c--; c >= 0; c-- {
		for r--; r >= 0; r-- {
			ret.Set(r, c, activator(a.At(r, c)))
		}
	}
	return ret
}

func matSigmoidD(a Matrix) Matrix {
	r, c := a.Dims()
	ret := mat.NewDense(r, c, nil)
	for c--; c >= 0; c-- {
		for r--; r >= 0; r-- {
			ret.Set(r, c, derivative(a.At(r, c)))
		}
	}
	return ret
}
func toFloatArr(a Matrix) []float64 {
	r, c := a.Dims()
	ret := make([]float64, r*c)
	for i := 0; i < c; i++ {
		for j := 0; j < r; j++ {
			ret[i*r+j] = a.At(j, i)
		}
	}
	return ret

}

func activator(v float64) float64 {
	return 1 / (1 + math.Exp(-v))
}
func derivative(v float64) float64 {
	return v * (1 - v)
}