package utils

// Author: Weisen Pan
// Date: 2023-10-24


import (
	"fmt"
	"math"
	"sort"
	"sync"

	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"

	"github.com/hkust-adsl/kubernetes-scheduler-simulator/knets_pkg/api/v1alpha1"
	simontype "github.com/hkust-adsl/kubernetes-scheduler-simulator/knets_pkg/type"
	gpushareutils "github.com/hkust-adsl/kubernetes-scheduler-simulator/knets_pkg/type/open-gpu-share/utils"
)

const (
	Q1LackBoth  = "q1_lack_both"
	Q2LackGpu   = "q2_lack_gpu"
	Q3Satisfied = "q3_satisfied"
	Q4LackCpu   = "q4_lack_cpu"
	XLSatisfied = "xl_satisfied"
	XRLackCPU   = "xr_lack_cpu"
	NoAccess    = "no_access"
)

var FragRatioDataMap = map[string]int{
	Q1LackBoth:  0,
	Q2LackGpu:   1,
	Q3Satisfied: 2,
	Q4LackCpu:   3,
	XLSatisfied: 4,
	XRLackCPU:   5,
	NoAccess:    6,
}

var GpuNumTypeList = []string{"PureCpu", "ShareGpu", "OneGpu", "TwoGpu", "FourGpu", "EightGpu", "Others"}

type FragRatio struct {
	Data []float64
}

type FragAmount struct {
	NodeName string
	Data     []float64
}

func NewFragAmount(nodeName string, data []float64) FragAmount {
	fragAmount := FragAmount{NodeName: nodeName, Data: make([]float64, len(data))}
	copy(fragAmount.Data, data)
	return fragAmount
}

func (fr FragRatio) AddRatio(fragType string, freq float64) error {
	if freq < 0 || freq > 1 {
		return fmt.Errorf("bad freq")
	}
	if index, ok := FragRatioDataMap[fragType]; !ok {
		return fmt.Errorf("bad fragType")
	} else {
		fr.Data[index] += freq
		return nil
	}
}

func (fa FragAmount) AddGamma(faOther FragAmount, gamma float64) error {
	if len(fa.Data) == 0 {
		fa.Data = make([]float64, len(FragRatioDataMap))
		for i := 0; i < len(FragRatioDataMap); i++ {
			fa.Data[i] = 0
		}
	}
	if len(fa.Data) != len(faOther.Data) {
		return fmt.Errorf("this (%d) does not match the other (%d)", len(fa.Data), len(faOther.Data))
	}
	for i := 0; i < len(fa.Data); i++ {
		fa.Data[i] += gamma * faOther.Data[i]
	}
	return nil
}

func (fa FragAmount) AddFragAmount(faOther FragAmount) error {
	return fa.AddGamma(faOther, 1.0)
}

func (fa FragAmount) AddByFragType(fragType string, amount float64) error {
	if amount < 0 {
		return fmt.Errorf("bad freq")
	}
	if index, ok := FragRatioDataMap[fragType]; !ok {
		return fmt.Errorf("bad fragType")
	} else {
		fa.Data[index] += amount
		return nil
	}
}

func (fr FragRatio) Repr() (outStr string) {
	outStr += "["
	for i, v := range fr.Data {
		if i > 0 {
			outStr += ", "
		}
		outStr += fmt.Sprintf("%4.1f%%", 100*v)
	}
	outStr += "]"
	return outStr
}

func (fa FragAmount) Repr() (outStr string) {
	outStr += fa.NodeName
	outStr += ": ["
	for i, v := range fa.Data {
		if i > 0 {
			outStr += ", "
		}
		outStr += fmt.Sprintf("%6.1f", v)
	}
	outStr += "]"
	return outStr
}

func NodeGpuFragRatio(nodeRes simontype.NodeResource, typicalPods simontype.TargetPodList) FragRatio {
	data := make([]float64, len(FragRatioDataMap))
	fragRatio := FragRatio{data}
	for _, pod := range typicalPods {
		freq := pod.Percentage
		if freq < 0 || freq > 1 {
			log.Errorf("pod %v has bad freq: %f\n", pod.TargetPodResource, freq)
			continue
		}
		fragType := GetNodePodFrag(nodeRes, pod.TargetPodResource)
		log.Debugf("nodeRes: %s; pod: %s => fragType: %s (freq: %.2f)\n", nodeRes.Repr(), pod.TargetPodResource.Repr(), fragType,