package kk_adjacency_matrix

import (
	"fmt"
	"slices"
	"sync"
)

type vertexEqualFunc[V any] func(a, b V) (equal bool)

type edgeEqualFunc[E any] func(a, b E) (equal bool)

type AdjacencyMatrix[E, V any] struct {
	matrix [][]E
	// zero weight means no edge or edge is unreachable
	zeroWeight *E
	// compare two weight
	edgeEqualFunc edgeEqualFunc[E]
	//the index of vertex indicates the index of the matrix's row and col
	vertexList []V
	// find index in vertexList by data
	findIndexFunc func(vertexList []V, data V) (index int, found bool)
	// compare two vertex
	// use to filter duplicate vertex and find the index of vertex in vertexList
	vertexEqualFunc vertexEqualFunc[V]
	mutex           sync.RWMutex
}

func NewAdjacencyMatrix[E, V any](
	zeroWeight *E,
	vertexEqualFunc vertexEqualFunc[V],
	edgeEqualFunc edgeEqualFunc[E]) (am *AdjacencyMatrix[E, V]) {
	if zeroWeight == nil {
		panic("zeroWeight is nil")
	}
	if vertexEqualFunc == nil {
		panic("vertexEqualFunc is nil")
	}
	if edgeEqualFunc == nil {
		panic("edgeEqualFunc is nil")
	}
	am = &AdjacencyMatrix[E, V]{
		matrix:          make([][]E, 0),
		zeroWeight:      zeroWeight,
		edgeEqualFunc:   edgeEqualFunc,
		vertexList:      make([]V, 0),
		vertexEqualFunc: vertexEqualFunc,
		mutex:           sync.RWMutex{},
	}
	am.findIndexFunc = func(vertexList []V, data V) (index int, found bool) {
		index = slices.IndexFunc(vertexList, func(v V) bool {
			return am.vertexEqualFunc(v, data)
		})
		return index, index != -1
	}
	return am
}
func (am *AdjacencyMatrix[E, V]) AddNodes(dataList ...V) {
	am.mutex.Lock()
	defer am.mutex.Unlock()
	//remove deduplicate data in dataList
	dataList = slices.CompactFunc(dataList, func(a, b V) bool {
		return am.vertexEqualFunc(a, b)
	})
	//add if not exist
	for _, v := range dataList {
		if _, found := am.findIndexFunc(am.vertexList, v); found {
			continue
		}
		am.vertexList = append(am.vertexList, v)
	}

	moreRowCount := len(am.vertexList) - len(am.matrix)
	// add row
	am.matrix = append(am.matrix, make([][]E, moreRowCount)...)
	// add col for each row
	// The new column length should be equal to the number of rows after adding the new row
	newColCount := len(am.matrix)
	for i := 0; i < len(am.matrix); i++ {
		// add col
		// extend the capacity of existing rows
		if cap(am.matrix[i]) >= newColCount {
			am.matrix[i] = am.matrix[i][:newColCount]
		} else {
			am.matrix[i] = append(am.matrix[i], make([]E, newColCount-len(am.matrix[i]))...)
		}
		am.matrix[i][newColCount-1] = *am.zeroWeight
	}
}
func (am *AdjacencyMatrix[E, V]) RemoveNode(data V) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	index, found := am.findIndexFunc(am.vertexList, data)
	if found {
		// remove data
		am.vertexList = slices.DeleteFunc(am.vertexList, func(v V) bool {
			return am.vertexEqualFunc(v, data)
		})
		// remove col
		for i := 0; i < len(am.matrix); i++ {
			am.matrix[i] = slices.Delete(am.matrix[i], index, index+1)
		}
		// remove row
		am.matrix = slices.Delete(am.matrix, index, index+1)
	}
}
func (am *AdjacencyMatrix[E, V]) AddEdge(fromNode, toNode V, weight E) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	fromIndex, fromFound := am.findIndexFunc(am.vertexList, fromNode)
	toIndex, toFound := am.findIndexFunc(am.vertexList, toNode)
	if fromFound && toFound {
		am.matrix[fromIndex][toIndex] = weight
	}
}
func (am *AdjacencyMatrix[E, V]) RemoveEdge(fromNode, toNode V) {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	fromIndex, fromFound := am.findIndexFunc(am.vertexList, fromNode)
	toIndex, toFound := am.findIndexFunc(am.vertexList, toNode)
	if fromFound && toFound {
		am.matrix[fromIndex][toIndex] = *am.zeroWeight
	}
}
func (am *AdjacencyMatrix[E, V]) CheckEdge(fromNode, toNode V) (reachable bool, weight E) {
	am.mutex.RLock()
	defer am.mutex.RUnlock()
	fromNodeIndex, fromFound := am.findIndexFunc(am.vertexList, fromNode)
	toNodeIndex, toFound := am.findIndexFunc(am.vertexList, toNode)
	if fromFound && toFound {
		weight = am.matrix[fromNodeIndex][toNodeIndex]
		if am.edgeEqualFunc(weight, *am.zeroWeight) {
			reachable = false
		} else {
			reachable = true
		}
	}
	return reachable, weight
}
func (am *AdjacencyMatrix[E, V]) GetMatrix() (matrix [][]E) {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	return am.matrix
}
func (am *AdjacencyMatrix[E, V]) GetVertexList() (vertexList []V) {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	return am.vertexList
}
func (am *AdjacencyMatrix[E, V]) MatrixString() string {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	s := ""
	for i := 0; i < len(am.matrix); i++ {
		for j := 0; j < len(am.matrix[i]); j++ {
			s += fmt.Sprintf("%v	 ", am.matrix[i][j])
		}
		s += "\n"
	}
	return s
}