package eas

import (
	"fmt"

	"gitlab.alibaba-inc.com/pai_biz_arch/pairec/algorithm/eas/tf"
	"gitlab.alibaba-inc.com/pai_biz_arch/pairec/algorithm/response"
)

type tfResponse struct {
	score          float64
	scoreArr       map[string]float64
	multiValModule bool
}

func (r *tfResponse) GetScore() float64 {
	return r.score
}

func (r *tfResponse) GetScoreMap() map[string]float64 {
	return r.scoreArr
}

func (r *tfResponse) GetModuleType() bool {
	return r.multiValModule
}

func tfMutValResponseFunc(data interface{}) (ret []response.AlgoResponse, err error) {
	resp, ok := data.(*tf.PredictResponse)
	if !ok {
		err = fmt.Errorf("invalid data type, %v", data)
		return
	}
	var response []map[string]float64
	for name, arrayProto := range resp.GetOutputs() {
		for i, val := range arrayProto.FloatVal {
			if i >= len(response) {
				response = append(response, map[string]float64{name: float64(val)})
			} else {
				response[i][name] = float64(val)
			}
		}
	}
	for _, v := range response {
		ret = append(ret, &tfResponse{scoreArr: v, multiValModule: true})
	}
	return
}

func tfResponseFunc(data interface{}) (ret []response.AlgoResponse, err error) {
	resp, ok := data.(*tf.PredictResponse)
	if !ok {
		err = fmt.Errorf("invalid data type, %v", data)
		return
	}
	for _, arrayProto := range resp.GetOutputs() {
		for _, val := range arrayProto.FloatVal {
			ret = append(ret, &tfResponse{score: float64(val)})
		}
		break
	}
	return
}
