package his

import (
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/core"
	"github.com/banbox/banbot/strat"
	ta "github.com/banbox/banta"
	"github.com/sasha-s/go-deadlock"
	"math"
)

/*
需要多品种综合市场信息，品种数不可少于30
*/

// ZoneTradeState 用于存储每个交易区间的状态信息
type ZoneTradeState struct {
	LastTradeTime2 int64 // 买入区间2的最后交易时间
	LastTradeTime3 int64 // 买入区间3的最后交易时间
	LastTradeTime4 int64 // 买入区间4的最后交易时间
	LastSellTime   int64 // 最后卖出时间
	atrBase        float64
	score          float64
	upFlag         float64
}

func J3782(pol *config.RunPolicyConfig) *strat.TradeStrat {
	// 定义可调超参数
	maPeriod := pol.DefInt("maPeriod", 100, core.PNorm(20, 200))
	baseRetracement := pol.Def("baseRetracement", 0.08, core.PNorm(0.05, 0.2))
	retracementRate := pol.Def("retracementRate", 1.8, core.PNorm(1.5, 2.5))
	profitRate := pol.Def("profitRate", 0.03, core.PNorm(0.015, 0.06))
	tradeDelay := pol.Def("tradeDelay", 40, core.PNorm(10, 100))
	slRate := pol.Def("slRate", 63, core.PNorm(5, 35))

	// 计算各回撤比例
	retracement2 := baseRetracement
	retracement3 := baseRetracement * retracementRate
	retracement4 := baseRetracement * retracementRate * 1.33

	lenAtr := 20
	baseAtrLen := int(math.Round(float64(lenAtr) * 4.3))
	rates := make([]float64, 0, 10)

	// 最大交易数量
	maxTrades2 := 2
	maxTrades3 := 3
	maxTrades4 := 4

	openPairs := 0.0
	var pairLock deadlock.Mutex
	addThres, firstRate := 1.0, 1.0
	avgAtrBase := 2.5
	maxPair := 3.0
	return &strat.TradeStrat{
		WarmupNum:     maPeriod,
		BatchInfo:     true,
		BatchInOut:    true,
		RunTimeFrames: []string{"15m"},
		OnStartUp: func(s *strat.StratJob) {
			// 初始化状态信息
			s.More = &ZoneTradeState{}
		},
		OnPairInfos: func(s *strat.StratJob) []*strat.PairSub {
			return []*strat.PairSub{
				{"_cur_", "4h", baseAtrLen * 2},
			}
		},
		OnBar: func(s *strat.StratJob) {
			e := s.Env
			state, _ := s.More.(*ZoneTradeState)

			// 计算MA100
			ma100 := ta.SMA(e.Close, maPeriod)
			ma100Val := ma100.Get(0)

			// 计算交易价格水平
			curRate := max(addThres, 1+state.score/10)
			itemRate := max(0.7, min(math.Sqrt(state.atrBase/avgAtrBase), 1.4))
			buyPrice2 := ma100Val * (1 - retracement2*curRate*itemRate*firstRate)
			buyPrice3 := ma100Val * (1 - retracement3*curRate*itemRate)
			buyPrice4 := ma100Val * (1 - retracement4*curRate*itemRate)
			sellPrice := ma100Val * (1 + profitRate)

			// 获取当前价格
			lowPrice := e.Low.Get(0)
			highPrice := e.High.Get(0)

			// 检查交易间隔条件
			currentBarTime := e.TimeStop

			lastEnt := max(state.LastTradeTime2, state.LastTradeTime3, state.LastTradeTime4)
			canTrade := e.BarCount(lastEnt) >= 10
			canTrade2 := canTrade && e.BarCount(state.LastTradeTime2) >= tradeDelay
			canTrade3 := canTrade && e.BarCount(state.LastTradeTime3) >= tradeDelay
			canTrade4 := canTrade && e.BarCount(state.LastTradeTime4) >= tradeDelay
			canSell := e.BarCount(state.LastSellTime) >= tradeDelay

			// 获取当前总开仓数量，与原始策略的strategy.opentrades保持一致
			openTrades := len(s.LongOrders) + len(s.ShortOrders)

			if lowPrice <= buyPrice4 && openTrades < maxTrades4 && canTrade4 {
				pairLock.Lock()
				allow := false
				if openPairs < maxPair {
					openPairs += 3
					allow = true
				}
				pairLock.Unlock()
				if allow {
					s.OpenOrder(&strat.EnterReq{
						Tag:      "BuyZone4",
						CostRate: 3,
						Log:      true,
					})
					state.LastTradeTime4 = currentBarTime
				}
			}
			if lowPrice <= buyPrice3 && openTrades < maxTrades3 && canTrade3 {
				pairLock.Lock()
				allow := false
				if openPairs < maxPair {
					openPairs += 1.5
					allow = true
				}
				pairLock.Unlock()
				if allow {
					s.OpenOrder(&strat.EnterReq{
						Tag:      "BuyZone3",
						CostRate: 1.5,
						Log:      true,
					})
					state.LastTradeTime3 = currentBarTime
				}
			}
			if lowPrice <= buyPrice2 && openTrades < maxTrades2 && canTrade2 {
				pairLock.Lock()
				allow := false
				if openPairs < maxPair {
					openPairs += 1
					allow = true
				}
				pairLock.Unlock()
				if allow {
					s.OpenOrder(&strat.EnterReq{
						Tag:      "BuyZone2",
						CostRate: 1,
						Log:      true,
					})
					state.LastTradeTime2 = currentBarTime
				}
			}

			// 检查卖出条件
			if s.OrderNum > 0 {
				curPrice := e.Close.Get(0)
				slRange := slRate * itemRate * curPrice / 100
				avgPrice := s.PositionAvgPrice(core.OdDirtLong, "")
				if curPrice+slRange <= avgPrice {
					s.CloseOrders(&strat.ExitReq{
						Tag: "StopLoss",
						Log: true,
					})
					state.LastSellTime = currentBarTime
				} else {
					exitPrice := avgPrice*(1+profitRate)*0.4 + sellPrice*0.6
					if openTrades > 0 && highPrice >= exitPrice && canSell {
						s.CloseOrders(&strat.ExitReq{
							Tag: "SellAll",
							Log: true,
						})
						state.LastSellTime = currentBarTime
					}
				}
			}
		},
		OnInfoBar: func(s *strat.StratJob, e *ta.BarEnv, pair, tf string) {
			score := float64(0)
			if ta.RSI(e.Close, 14).Get(0) < 30 {
				score += 1
			}
			roc := ta.ROC(e.Close, 20).Get(0)
			if roc < -20 {
				score += 2
			} else if roc < -10 {
				score += 1
			}
			plusDi, minusDi := ta.PluMinDI(e.High, e.Low, e.Close, 14)
			adx := ta.ADX(e.High, e.Low, e.Close, 14).Get(0)
			if minusDi.Get(0) > plusDi.Get(0) {
				if adx > 40 {
					score += 2
				} else if adx > 25 {
					score += 1
				}
			}
			atr := ta.ATR(e.High, e.Low, e.Close, lenAtr)
			atrBase := ta.Lowest(ta.Highest(atr, lenAtr), baseAtrLen).Get(0)
			rates = append(rates, atrBase/(e.Close.Get(0)*0.01))
			//openRate := 0.006 * closePrice / atrBase
			//if openRate < 1 {
			//	openRate = math.Sqrt(openRate)
			//}
			//openRate = max(0.7, min(1.5, openRate))
			ma7 := ta.SMA(e.Close, 7).Get(0)
			ma20 := ta.SMA(e.Close, 20)
			ma20V := ma20.Get(0)
			ma99 := ta.SMA(e.Close, 99).Get(0)
			if ma99 > ma20V {
				score += 1
			}
			state, _ := s.More.(*ZoneTradeState)
			state.score = score
			state.upFlag = 0
			state.atrBase = atrBase / (e.Close.Get(0) * 0.01)
			if ma20V > ma99 && ma20V > ma20.Get(2) {
				state.upFlag = 1
			} else if ma7 < min(ma20V, ma99) {
				state.upFlag = -1
			}
		},
		OnBatchJobs: func(jobs []*strat.StratJob) {
			pairLock.Lock()
			openPairs = 0
			pairLock.Unlock()
		},
		OnBatchInfos: func(tf string, jobs map[string]*strat.JobEnv) {
			if len(jobs) < 3 {
				return
			}
			downNum, downHeavyNum := float64(0), float64(0)
			atrRates := make([]float64, 0, len(jobs))
			for _, j := range jobs {
				state, _ := j.Job.More.(*ZoneTradeState)
				if state.score >= 5 {
					downHeavyNum += 1
				}
				if state.score >= 3 {
					downNum += 1
				}
				atrRates = append(atrRates, state.atrBase)
			}
			jobNum := float64(len(jobs))
			downRate := downNum / jobNum
			downHeavyRate := downHeavyNum / jobNum
			avgAtrBase = atrRates[len(jobs)/2]
			if downHeavyRate >= 0.8 {
				// 极端下跌
				addThres, firstRate = 1.6, 1.2
				maxPair = jobNum * 0.5
			} else if downRate >= 0.85 && downHeavyRate >= 0.6 {
				// 严重下跌
				addThres, firstRate = 1.45, 1.15
				maxPair = jobNum * 0.55
			} else if downRate >= 0.7 && downRate <= 0.85 {
				// 市场中度下跌
				addThres, firstRate = 1.3, 1.1
				maxPair = jobNum * 0.6
			} else if downRate >= 0.5 && downRate <= 0.7 {
				// 轻度下跌
				addThres, firstRate = 1.15, 1.05
				maxPair = jobNum * 0.7
			} else {
				// 轻度或未下跌
				maxPair = jobNum
				addThres, firstRate = 1, 1
			}
		},
	}
}