package arb

import (
	"fmt"
	"github.com/banbox/banbot/biz"
	"github.com/banbox/banbot/btime"
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/core"
	"github.com/banbox/banbot/data"
	"github.com/banbox/banbot/exg"
	"github.com/banbox/banbot/goods"
	"github.com/banbox/banbot/orm"
	utils2 "github.com/banbox/banbot/utils"
	"github.com/banbox/banexg"
	"github.com/banbox/banexg/errs"
	"github.com/banbox/banexg/log"
	"github.com/banbox/banexg/utils"
	"github.com/sasha-s/go-deadlock"
	"go.uber.org/zap"
	"math/rand"
	"path/filepath"
	"sort"
	"strings"
	"testing"
	"time"
)

func TestPairPick(t *testing.T) {
	all := "ETH,BTC,SOL,DOGE,XRP,WLFI,MYX,ENA,BNB,1000PEPE,LINK,SUI,AVAX,ADA,PUMP,WLD,SOMI,HYPE,BIO,FARTCOIN,NMR,M,PENGU,PYTH,AAVE,LTC,HIFI,TA,WIF,ARB,RNDR,TRUMP,1000BONK,BCH,DOT,UNI,IP,BAKE,MATIC,FIL,EOS,TRX,NEAR,ONDO,XPL,1000SHIB,DOLO,PROMPT,CRV,OP,MEME,ETHFI,HBAR,XLM,ETC,TIA,LDO,SEI,TREE,EIGEN,TAO,XNY,SPK,APT,DAM,W,NAORIS,INJ,SAPIEN,JELLYJELLY,MOODENG,VIRTUAL,RED,TST,QTUM,MERL,TON,F,TUT,PNUT,LPT,SOLV,PENDLE,CFX,POL,KAITO,BERA,ZORA,GALA,FHE,PAXG,ONT,JUP,OG,FET,PUMPBTC,SKL,GAL,ATOM,NEIRO,BOME,MAV,PROVE,ATH,SPX"
	pickA, pickB := 30, 30
	coins := strings.Split(all, ",")
	rand.Shuffle(len(coins), func(i, j int) { coins[i], coins[j] = coins[j], coins[i] })
	groupA := coins[:pickA]
	groupB := coins[pickA : pickA+pickB]
	fmt.Println(strings.Join(groupA, ","))
	fmt.Println(strings.Join(groupB, ","))
}

// 从高频逐笔数据中提取每隔5秒的价格
func TestExtractPrices(t *testing.T) {
	core.SetRunMode(core.RunModeBackTest)
	err := biz.SetupComsExg(&config.CmdArgs{})
	if err != nil {
		panic(err)
	}
	loader, err := data.NewWsDataLoader()
	if err != nil {
		panic(err)
	}
	dataDir := config.GetDataDir()
	outPath := filepath.Join(dataDir, "dump", "group1.gob")
	group1 := "LPT,ATH,MOODENG,RNDR,APT,TRUMP,DOT,HYPE,BAKE,SPX,BOME,TUT,PUMP,BERA,OG,EOS,MERL,TAO,NMR,FET,TIA,FHE,ZORA,LINK,POL,TRX,1000PEPE,ETHFI,PROVE,TREE"
	//group2 := "RED,SUI,VIRTUAL,NEIRO,GALA,TON,PAXG,XNY,CRV,NAORIS,XRP,NEAR,AVAX,ETC,XLM,MATIC,AAVE,SEI,1000SHIB,DAM,SPK,MEME,OP,FIL,KAITO,FARTCOIN,EIGEN,BCH,IP,XPL"
	pairs := strings.Split(group1, ",")
	pairs, err = config.ParsePairs(pairs...)
	if err != nil {
		panic(err)
	}
	result := make(map[string][]float64)
	if utils2.Exists(outPath) {
		err = utils2.DecodeGobFile(outPath, &result)
		if err != nil {
			panic(err)
		}
		cache := make(map[string]int)
		num := 0
		for k, v := range result {
			cache[k] = len(v)
			if len(v) > 0 {
				num += 1
			}
		}
		log.Info("load cache", zap.Int("num", num), zap.Any("keys", cache))
	}
	lock := deadlock.Mutex{}
	year, mon, day := time.Now().Add(-time.Hour * 24).Date()
	endTime := time.Date(year, mon, day, 0, 0, 0, 0, time.UTC)
	year, mon, day = endTime.AddDate(0, 0, -15).Date()
	startTime := time.Date(year, mon, day, 0, 0, 0, 0, time.UTC)
	gapMS := int64(5000) // 采样间隔ms
	startMS := startTime.UnixMilli()
	endMS := endTime.UnixMilli()
	size := int((endMS - startMS) / gapMS)
	pBar := utils2.NewPrgBar(len(pairs)*size, "run")
	defer pBar.Close()
	err = utils2.ParallelRun(pairs, 4, func(i int, pair string) *errs.Error {
		lock.Lock()
		arr, ok := result[pair]
		lock.Unlock()
		if ok && len(arr) > 0 {
			pBar.Add(size)
			return nil
		}
		arr = make([]float64, size+10)
		lastIdx := 0
		saveNum := 0
		curTime := startTime
		success := true
		for curTime.Before(endTime) {
			trades, err := loader.LoadTrades(&data.WsSymbol{
				ExgId:  "binance",
				Market: banexg.MarketLinear,
				WsType: core.WsSubTrade,
				Symbol: pair,
				Date:   curTime.Format(core.DateFmt),
				Hour:   curTime.Hour(),
			})
			if err != nil {
				log.Error("down trades fail", zap.Error(err))
				success = false
				break
			}
			curTime = curTime.Add(time.Hour)
			for _, td := range trades {
				curIdx := int((td.Timestamp - startMS) / gapMS)
				if curIdx > lastIdx {
					prevVal := arr[lastIdx]
					for lastIdx+1 < curIdx {
						lastIdx += 1
						arr[lastIdx] = prevVal
					}
					lastIdx = curIdx
					arr[curIdx] = td.Price
				}
			}
			if lastIdx > saveNum {
				pBar.Add(lastIdx - saveNum)
				saveNum = lastIdx
			}
		}
		if success {
			lock.Lock()
			result[pair] = arr
			lock.Unlock()
		}
		return nil
	})
	err = utils2.EncodeGob(outPath, result)
	if err != nil {
		panic(err)
	}
}

// 测试随机两组品种的平均价格比率（逐秒级别）
func TestDayBatchRates(t *testing.T) {
	// 读取两组30个交易品种的每隔5s的价格序列，计算平均价比率，研究套利空间
	tf := "5s"
	dataDir := config.GetDataDir()
	outPath := filepath.Join(dataDir, "dump", "group1.gob")
	group1 := make(map[string][]float64)
	err := utils2.DecodeGobFile(outPath, &group1)
	if err != nil {
		panic(err)
	}
	outPath = filepath.Join(dataDir, "dump", "group2.gob")
	group2 := make(map[string][]float64)
	err = utils2.DecodeGobFile(outPath, &group2)
	if err != nil {
		panic(err)
	}
	tfSecs := utils.TFToSecs(tf)
	daySize := utils.TFToSecs("1d") / tfSecs
	dataLen := 0
	for _, arr := range group1 {
		dataLen = len(arr)
		break
	}
	tfMSecs := int64(tfSecs * 1000)
	year, mon, day := time.Now().Add(-time.Hour * 24).Date()
	endMS := time.Date(year, mon, day, 0, 0, 0, 0, time.UTC).UnixMilli()
	startMS := endMS - int64(dataLen)*tfMSecs
	dumpDir := filepath.Join(config.GetDataDir(), "dump")
	for offset := 0; offset+daySize < dataLen; offset += daySize {
		offsetMS := startMS + int64(offset)*tfMSecs
		batchA := normInRange(group1, offset, offset+daySize, offsetMS, tfMSecs)
		batchB := normInRange(group2, offset, offset+daySize, offsetMS, tfMSecs)
		rateArr := make([]*TimeFlt, len(batchA))
		for i, va := range batchA {
			rateArr[i] = &TimeFlt{
				Stamp: va.Stamp,
				Val:   va.Val / batchB[i].Val,
			}
		}
		dateStr := btime.ToDateStr(offsetMS, "20060102")
		dumpPath := filepath.Join(dumpDir, dateStr+".html")
		log.Info("dump rate chart", zap.String("title", dateStr))
		err = dumpRateChart(rateArr, tf, dumpPath, dateStr)
		if err != nil {
			panic(err)
		}
	}
}

func normInRange(dataMap map[string][]float64, start, end int, offsetMS, stepMS int64) []*TimeFlt {
	batchA := make(map[string][]float64)
	minLen := -1
	// 对每个品种的价格序列进行归一化
	for key, arr := range dataMap {
		vals := arr[start:end]
		res := make([]float64, len(vals))
		first := float64(0)
		for i, v := range vals {
			if v > 0 {
				if first > 0 {
					res[i] = v / first
				} else {
					res[i] = 1
					first = v
				}
			}
		}
		batchA[key] = res
		if minLen < 0 {
			minLen = len(res)
		} else {
			minLen = min(minLen, len(res))
		}
	}
	// 直接对所有品种价格序列取平均值
	res := make([]*TimeFlt, minLen)
	for i := 0; i < minLen; i++ {
		sumVal := float64(0)
		validNum := 0
		for _, arr := range batchA {
			val := arr[i]
			if val > 0 {
				sumVal += val
				validNum += 1
			}
		}
		item := &TimeFlt{
			Stamp: offsetMS + int64(i)*stepMS,
		}
		if validNum > 0 {
			item.Val = sumVal / float64(validNum)
		}
		res[i] = item
	}
	return res
}

/*
每7日取成交量最高top60
随机分为两组计算平均价格比率

获取7日成交量倒序前30%的品种列表
每天计算：

	计算最近14天1h的相关性矩阵
	按相关性从高到低取top60
	奇数位和偶数位分为两组

对比了按相关性从高到低取top60和随机分配top60，差别不大，简单起见，使用随机方法
*/
func TestDailyGroupArb(t *testing.T) {
	core.SetRunMode(core.RunModeBackTest)
	err := biz.SetupComsExg(&config.CmdArgs{})
	if err != nil {
		panic(err)
	}
	groupSize := 30
	dayTf := "1d"
	dayMSecs := int64(utils.TFToSecs("1d") * 1000)
	year, mon, day := time.Now().Add(-time.Hour * 24).Date()
	endMS := time.Date(year, mon, day, 0, 0, 0, 0, time.UTC).UnixMilli()
	startMS := endMS - dayMSecs*28
	minListMS := endMS - dayMSecs*180
	curMS := startMS
	exsMap := orm.GetExSymbols(core.ExgName, core.Market)
	pairs := make([]string, 0, len(exsMap))
	for _, exs := range exsMap {
		if exs.ListMs > minListMS {
			// 忽略上市不久的品种
			continue
		}
		pairs = append(pairs, exs.Symbol)
	}
	err = orm.BulkDownOHLCV(exg.Default, exsMap, "1h", startMS, endMS, 0, nil)
	if err != nil {
		panic(err)
	}
	var pre30Pairs []string
	var pre30Ms int64
	dumpDir := filepath.Join(config.GetDataDir(), "dump")
	for curMS < endMS {
		if len(pre30Pairs) == 0 || (curMS-pre30Ms)/dayMSecs >= 7 {
			// 获取14日成交量倒序前30%的品种列表
			pre30Pairs = nil
			symbolVols, err := goods.GetSymbolVols(pairs, dayTf, 14, curMS, false)
			if err != nil {
				panic(err)
			}
			sort.Slice(symbolVols, func(i, j int) bool {
				return symbolVols[i].Vol > symbolVols[j].Vol
			})
			cutNum := int(float64(len(symbolVols)) * 0.23)
			for i := 0; i < cutNum; i++ {
				pre30Pairs = append(pre30Pairs, symbolVols[i].Symbol)
			}
			pre30Ms = curMS
			dateStr := btime.ToDateStr(curMS, core.DateFmt)
			log.Info("cut pre 30% pairs by vols", zap.Int("num", len(pre30Pairs)), zap.String("date", dateStr))
		}
		// 按成交量取前面的，然后随机挑选。
		coins := pre30Pairs[:groupSize*2]
		rand.Shuffle(len(coins), func(i, j int) { coins[i], coins[j] = coins[j], coins[i] })
		var pairsA, pairsB []string
		downMap := make(map[int32]*orm.ExSymbol)
		for i := 0; i < groupSize*2; i++ {
			pair := coins[i]
			if i%2 == 0 {
				pairsA = append(pairsA, pair)
			} else {
				pairsB = append(pairsB, pair)
			}
			exs, _ := orm.GetExSymbolCur(pair)
			downMap[exs.ID] = exs
		}
		// 对A和B两组计算1m加权价格
		curEndMS := curMS + dayMSecs
		err = orm.BulkDownOHLCV(exg.Default, downMap, "1m", curMS, curEndMS, 0, nil)
		if err != nil {
			panic(err)
		}
		pricesA, err := GetGroupPrices(pairsA, "1m", curMS, curEndMS)
		if err != nil {
			panic(err)
		}
		pricesB, err := GetGroupPrices(pairsB, "1m", curMS, curEndMS)
		if err != nil {
			panic(err)
		}
		rateArr := make([]*TimeFlt, len(pricesA))
		for i, va := range pricesA {
			rateArr[i] = &TimeFlt{
				Stamp: va.Stamp,
				Val:   va.Val / pricesB[i].Val,
			}
		}
		dateStr := btime.ToDateStr(curMS, "20060102")
		dumpPath := filepath.Join(dumpDir, dateStr+"_r2.html")
		log.Info("dump rate chart", zap.String("title", dateStr))
		err = dumpRateChart(rateArr, "1m", dumpPath, dateStr)
		if err != nil {
			panic(err)
		}
		curMS += dayMSecs
	}
}

func getPairCorrs(curMS int64, pairs []string) []string {
	// 计算最近14天1h的相关性矩阵
	dayMSecs := int64(utils.TFToSecs("1d") * 1000)
	hourMSecs := int64(utils.TFToSecs("1h") * 1000)
	pre14MS := curMS - dayMSecs*14
	names := make([]string, 0, len(pairs))
	dataArr := make([][]float64, 0, len(pairs))
	klineNum := int((curMS - pre14MS) / hourMSecs)
	err := orm.FastBulkOHLCV(exg.Default, pairs, "1h", pre14MS, curMS, 0, func(s string, s2 string, klines []*banexg.Kline, infos []*orm.AdjInfo) {
		if len(klines) == klineNum {
			prices := make([]float64, 0, len(klines))
			for _, b := range klines {
				prices = append(prices, b.Close)
			}
			names = append(names, s)
			dataArr = append(dataArr, prices[:klineNum])
		}
	})
	if err != nil {
		panic(err)
	}
	_, avgs, err_ := utils2.CalcCorrMat(klineNum, dataArr, true)
	if err_ != nil {
		panic(err_)
	}
	// 基于平均相似度降序
	arr := make([]*goods.IdVal, 0, len(avgs))
	lefts := make(map[int]bool)
	for i, avg := range avgs {
		arr = append(arr, &goods.IdVal{Id: i, Val: avg})
		lefts[i] = true
	}
	sort.Slice(arr, func(i, j int) bool {
		return arr[i].Val < arr[j].Val
	})
	var resPairs []string
	for _, v := range arr {
		resPairs = append(resPairs, names[v.Id])
	}
	return resPairs
}