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]
	for idx := range c.Counts {
		c.Counts[idx] = 0
	}

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

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) CountSum() (i int) {
	for _, v := range c.Counts {
		i = i + v
	}
	return i
}
func (c *CountReward) CountReward() (i int) {
	for _, v := range c.Counts {
		i = i + v
	}
	return i
}

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)
}