package helper

type CountReward struct {
	Counts  []int
	Rewards []float64
}

func (c *CountReward) ResetTo(size int) {
	if len(c.Counts) > size {
		c.Counts = make([]int, size)
	}
	c.Counts = c.Counts[:size]

	if len(c.Rewards) > size {
		c.Rewards = make([]float64, size)
	}
	c.Rewards = c.Rewards[:size]
}

func (c *CountReward) Count(res *[]int) {
	if res == nil {
		r := make([]int, len(c.Counts))
		res = &r
	}
	if len(c.Counts) < len(*res) {
		*res = append(*res, len(c.Counts)-len(*res))
	}

	(*res) = (*res)[:len(c.Counts)]
	copy(*res, c.Counts)
}

func (c *CountReward) Reward(res *[]float64) {
	if res == nil {
		r := make([]float64, len(c.Rewards))
		res = &r
	}
	if len(c.Rewards) < len(*res) {
		*res = append(*res, make([]float64, len(c.Rewards)-len(*res))...)
	}
	(*res) = (*res)[:len(c.Rewards)]
	copy(*res, c.Rewards)
}