package collectionx

import "github.com/samber/lo"

type CountMapN[T comparable] struct {
	N   int
	Map map[T][]int
}

func NewCountMapN[T comparable](n int) *CountMapN[T] {
	return &CountMapN[T]{
		N:   n,
		Map: make(map[T][]int),
	}
}

func (cm *CountMapN[T]) Inc(key T, ith int, value int) {
	m, ok := cm.Map[key]
	if !ok {
		m = make([]int, cm.N)
		cm.Map[key] = m
	}

	m[ith] += value
}

func (cm *CountMapN[T]) Get(key T, ith int) int {
	m, ok := cm.Map[key]
	if !ok {
		return 0
	}

	return m[ith]
}

func (cm *CountMapN[T]) Sum(key T) int {
	m, ok := cm.Map[key]
	if !ok {
		return 0
	}
	return lo.Sum(m)
}