package datascience

// Pivot transforma datos del formato largo al formato ancho (pivot table).
// Agrupa por index, expande los valores unicos de columns como nuevas columnas
// y agrega values con la funcion indicada.
// Funciones de agregacion soportadas: sum, count, mean, min, max, first, last.
// Valores numericos faltantes se rellenan con 0.
func Pivot(rows []map[string]any, index, columns, values, agg string) []map[string]any {
	// Mantener orden de aparicion de index y column values
	indexOrder := []any{}
	seenIndex := map[any]bool{}
	colOrder := []any{}
	seenCols := map[any]bool{}

	for _, row := range rows {
		idx := row[index]
		col := row[columns]
		if !seenIndex[idx] {
			seenIndex[idx] = true
			indexOrder = append(indexOrder, idx)
		}
		if !seenCols[col] {
			seenCols[col] = true
			colOrder = append(colOrder, col)
		}
	}

	// Acumular: groups[indexVal][colVal] = lista de valores
	type key struct{ idx, col any }
	groups := map[key][]any{}
	for _, row := range rows {
		idx := row[index]
		col := row[columns]
		val := row[values]
		if val != nil {
			k := key{idx, col}
			groups[k] = append(groups[k], val)
		}
	}

	aggregate := func(vals []any, fn string) any {
		if len(vals) == 0 {
			return 0
		}
		switch fn {
		case "count":
			return len(vals)
		case "first":
			return vals[0]
		case "last":
			return vals[len(vals)-1]
		}
		// Funciones numericas: sum, mean, min, max
		toFloat := func(v any) float64 {
			switch n := v.(type) {
			case float64:
				return n
			case float32:
				return float64(n)
			case int:
				return float64(n)
			case int64:
				return float64(n)
			case int32:
				return float64(n)
			}
			return 0
		}
		sum := 0.0
		mn := toFloat(vals[0])
		mx := toFloat(vals[0])
		for _, v := range vals {
			f := toFloat(v)
			sum += f
			if f < mn {
				mn = f
			}
			if f > mx {
				mx = f
			}
		}
		switch fn {
		case "sum":
			return sum
		case "mean":
			return sum / float64(len(vals))
		case "min":
			return mn
		case "max":
			return mx
		}
		return sum
	}

	result := make([]map[string]any, 0, len(indexOrder))
	for _, idx := range indexOrder {
		record := map[string]any{index: idx}
		for _, col := range colOrder {
			k := key{idx, col}
			vals := groups[k]
			if len(vals) > 0 {
				record[col.(string)] = aggregate(vals, agg)
			} else {
				record[col.(string)] = 0
			}
		}
		result = append(result, record)
	}
	return result
}