package main

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/james-bowman/sparse"
	"gonum.org/v1/gonum/mat"
)

var gNumTestIterations = 30

type SparseMatrixTiming struct {
	Label string `json:"label"`
	Rows  int    `json:"rows"`
	Cols  int    `json:"cols"`
	NNZ   int    `json:"nnz"`

	MatMulRuns    int   `json:"matmul_runs"`
	MatMulTotalNs int64 `json:"matmul_total_ns"`
	MatMulAvgNs   int64 `json:"matmul_avg_ns"`

	SpMVRuns    int   `json:"spmv_runs"`
	SpMVTotalNs int64 `json:"spmv_total_ns"`
	SpMVAvgNs   int64 `json:"spmv_avg_ns"`
}

type SparseBenchmarkCase struct {
	Timing SparseMatrixTiming
	Matrix *sparse.CSR
}

// The "SuiteSparse" collection of matrices come in a format called
// "market matrix" and so we parse and load them to a sparse.COO format here.
func matrixLoadMarket(path string) *sparse.CSR {
	f, err := os.Open(path)
	if err != nil {
		panic(err)
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)

	// skip header
	scanner.Scan()

	// skip comments
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "%") {
			break
		}
	}

	fields := strings.Fields(scanner.Text())

	rows, _ := strconv.Atoi(fields[0])
	cols, _ := strconv.Atoi(fields[1])
	nnz, _ := strconv.Atoi(fields[2])

	r := make([]int, nnz)
	c := make([]int, nnz)
	v := make([]float64, nnz)

	i := 0

	for scanner.Scan() {
		f := strings.Fields(scanner.Text())

		ri, _ := strconv.Atoi(f[0])
		ci, _ := strconv.Atoi(f[1])
		val, _ := strconv.ParseFloat(f[2], 64)

		r[i] = ri - 1
		c[i] = ci - 1
		v[i] = val
		i++
	}

	coo := sparse.NewCOO(rows, cols, r, c, v)

	return coo.ToCSR()
}

// small helper just to strip path and file ending
func matrixGetLabel(path string) string {
	base := filepath.Base(path)          // f. ex "1138_bus.mtx"
	ext := filepath.Ext(base)            // ".mtx"
	out := strings.TrimSuffix(base, ext) // just "1138_bus"
	return out
}

func getSparseBenchmarkCase(path string) SparseBenchmarkCase {
	var out SparseBenchmarkCase

	loadMarketStart := time.Now()
	label := matrixGetLabel(path)
	m := matrixLoadMarket(path)
	loadMarketEnd := time.Since(loadMarketStart)
	loadedMarketMillisecs := float64(loadMarketEnd.Milliseconds())
	fmt.Printf("Loaded market matrix %s in %.4f ms \n", label, loadedMarketMillisecs)

	out.Matrix = m
	out.Timing.Label = label

	rows, cols := out.Matrix.Dims()
	nnz := m.NNZ()

	out.Timing.Rows = rows
	out.Timing.Cols = cols
	out.Timing.NNZ = nnz

	return out
}

func timeSparseMatmuls(bcase *SparseBenchmarkCase) {
	mCSR := bcase.Matrix
	var warm sparse.CSR
	for i := 0; i < 3; i += 1 {
		warm.Mul(mCSR, mCSR)
	}

	// The CSR x CSR matrix multiplication is supposed to have a specific optimised
	// path for matmuls
	var out sparse.CSR
	numberOfMults := gNumTestIterations

	fmt.Printf("NNZ before matmuls: %d\n", mCSR.NNZ())

	timeBegin := time.Now()
	for i := 0; i < numberOfMults; i += 1 {
		out.Mul(mCSR, mCSR)
	}
	timeElapsed := time.Since(timeBegin)

	fmt.Printf("NNZ after matmuls: %d\n", out.NNZ())

	if warm.NNZ() != out.NNZ() {
		panic("Sparsity pattern changed unexpectedly after matmul!")
	}

	bcase.Timing.MatMulRuns = numberOfMults
	bcase.Timing.MatMulTotalNs = timeElapsed.Nanoseconds()

	timeAvgNS := timeElapsed.Nanoseconds() / int64(numberOfMults)
	bcase.Timing.MatMulAvgNs = timeAvgNS
}

func timeSparseMatVec(bcase *SparseBenchmarkCase) {
	A := bcase.Matrix
	rows, cols := A.Dims()

	x := mat.NewVecDense(cols, nil)
	y := mat.NewVecDense(rows, nil)

	for i := 0; i < cols; i++ {
		x.SetVec(i, 1.0)
	}

	// warmup
	A.MulVecTo(y.RawVector().Data, false, x.RawVector().Data)
	norm := y.Norm(2)

	if norm < 1e-12 {
		panic("Norm of resulting SpMV is zero. Something went wrong.\n")
	}

	numberOfRuns := gNumTestIterations

	timeBegin := time.Now()

	for i := 0; i < numberOfRuns; i++ {
		// fmt.Printf("spmv iteration %d\n", i)
		A.MulVecTo(y.RawVector().Data, false, x.RawVector().Data)
	}

	timeElapsed := time.Since(timeBegin)

	total := timeElapsed.Nanoseconds()
	avg := total / int64(numberOfRuns)
	bcase.Timing.SpMVTotalNs = total
	bcase.Timing.SpMVAvgNs = avg
	bcase.Timing.SpMVRuns = numberOfRuns
}

func timeNanoToMS(timeNS int64) float64 {
	return float64(timeNS) / float64(1e6)
}

func doTimings(path string) {
	bcase := getSparseBenchmarkCase(path)
	rows := bcase.Timing.Rows
	cols := bcase.Timing.Cols

	if rows != cols {
		panic("Sparse benchmark assumes square matrix")
	}

	doMatMul := true
	if doMatMul {
		fmt.Printf("Timing sparse matrix %s with size: %d x %d \n", bcase.Timing.Label, rows, cols)
		fmt.Printf("Matmul:\n")
		timeSparseMatmuls(&bcase)
		avgMatMulTimeMS := timeNanoToMS(bcase.Timing.MatMulAvgNs)
		fmt.Printf("Avg matmul time for %s: %.4f ms \n", bcase.Timing.Label, avgMatMulTimeMS)
	}

	doSpMV := true
	if doSpMV {
		fmt.Printf("SpMV:\n")
		timeSparseMatVec(&bcase)
		avgSpMVTimeMS := timeNanoToMS(bcase.Timing.SpMVAvgNs)
		fmt.Printf("Avg SpMV time for %s: %.4f ms\n", bcase.Timing.Label, avgSpMVTimeMS)
	}
}

func main() {
	mainBegin := time.Now()

	thermal2Path := "suitesparse_test_matrices/thermal2.mtx"
	doTimings(thermal2Path)

	mainElapsed := time.Since(mainBegin)

	fmt.Printf("Program finished in %.4f seconds \n", float64(mainElapsed.Milliseconds())/1e3)
}