// 修复后的banbot完整策略代码
package llmjob

import (
	"fmt"
	"github.com/banbox/banbot/btime"
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/core"
	"github.com/banbox/banbot/orm/ormo"
	"github.com/banbox/banbot/strat"
	"github.com/banbox/banexg/log"
	ta "github.com/banbox/banta"
	"github.com/sasha-s/go-deadlock"
	"go.uber.org/zap"
	"math"
)

/*
根据AtrBase基于不同品种权重，限制范围后更好
从4h根据行情自动调整加仓阈值和止盈阈值后：
之前：夏普0.93，显示回撤35.8%，最大回撤95.4%，最大资金占用3.2W，总收益119.6%
之后：夏普1.38，显示回撤5.66%，最大回撤36.2%，最大资金占用1.64W，总收益79.4%

尝试使用公式0.383*pow(1+target_qty, x)加仓，符合测试的数据，但效果更差
*/

// State1525S 存储每个品种的策略状态
type State1525S struct {
	avgPrice    float64 // 平均持仓成本
	positionQty float64 // 当前持仓数量
	tpLimit     float64
	score       float64 // 下跌行情力度分数，值域[0,6]
}

func J1525(pol *config.RunPolicyConfig) *strat.TradeStrat {
	// 定义可调超参数
	smaPeriod := pol.DefInt("sma_period", 18, core.PNorm(5, 50))
	targetPercRaw := pol.Def("target_perc", -0.14, core.PNorm(-0.3, -0.05))
	takeProfitRaw := pol.Def("take_profit", 0.09, core.PNorm(0.05, 0.3))
	targetQty := pol.Def("target_qty", 0.65, core.PNorm(0.2, 1.0))
	// 加仓次数，超过20后变差；低于19后有段时间不稳定
	maxAdd := pol.DefInt("max_add", 19, core.PNorm(1, 30))
	maxOpenPair := pol.DefInt("max_open_pair", 25, core.PNorm(1, 30)) // 最大同时交易品种

	openPairs := make(map[*strat.StratJob]int)
	pairLock := deadlock.Mutex{}
	lenAtr := 20
	baseAtrLen := int(math.Round(float64(lenAtr) * 4.3))

	targetPerc := targetPercRaw
	takeProfit := takeProfitRaw

	printMS := int64(0)
	maxPairs := 0
	// 加仓倍率，基于targetQty： 1, 0.6, 1 1.7 3 4.7 7.7 12.8 21.1 35 57.5 95 157 256 385 464 767
	return &strat.TradeStrat{
		WarmupNum:       max(smaPeriod, baseAtrLen),
		HedgeOff:        true, // 单向持仓模式
		RunTimeFrames:   []string{"5m"},
		BatchInfo:       true,
		OrderOnRotation: "open",
		OnStartUp: func(s *strat.StratJob) {
			// 初始化策略状态
			s.More = &State1525S{}
		},
		OnPairInfos: func(s *strat.StratJob) []*strat.PairSub {
			return []*strat.PairSub{
				{"_cur_", "4h", 20},
			}
		},
		OnBar: func(s *strat.StratJob) {
			e := s.Env
			state, _ := s.More.(*State1525S)
			closePrice := e.Close.Get(0)

			// 不利行情下调整参数
			targetPercJob, takeProfitJob := targetPercRaw, takeProfitRaw
			if state.score >= 5 {
				targetPercJob = targetPercRaw * 2
				takeProfitJob = takeProfitRaw * 0.6
			} else if state.score >= 4 {
				targetPercJob = targetPercRaw * 1.7
				takeProfitJob = takeProfitRaw * 0.7
			} else if state.score >= 3 {
				targetPercJob = targetPercRaw * 1.3
				takeProfitJob = takeProfitRaw * 0.8
			}
			targetPercCur := (targetPercJob + targetPerc) / 2
			takeProfitCur := (takeProfitJob + takeProfit) / 2

			// 计算SMA
			ma := ta.SMA(e.Close, smaPeriod)
			atr := ta.ATR(e.High, e.Low, e.Close, lenAtr)
			atrBase := ta.Lowest(ta.Highest(atr, lenAtr), baseAtrLen).Get(0)
			openRate := 0.006 * closePrice / atrBase
			if openRate < 1 {
				openRate = math.Sqrt(openRate)
			}
			openRate = max(0.7, min(1.5, openRate))

			// 计算当前持仓盈亏比例
			var pnl float64
			if state.positionQty > 0 && state.avgPrice > 0 {
				pnl = (closePrice - state.avgPrice) / state.avgPrice
			}

			// 计算止盈价格
			if state.positionQty > 0 && state.avgPrice > 0 {
				curTakeProfit := state.avgPrice * (1 + takeProfitCur)
				if curTakeProfit != state.tpLimit {
					state.tpLimit = curTakeProfit
					// 设置限价单止盈
					s.CloseOrders(&strat.ExitReq{
						Tag:   "Take Profit",
						Limit: state.tpLimit,
					})
				}
			}

			// 首次入场条件：价格上穿SMA且无持仓
			if e.Close.Cross(ma) == 1 && state.positionQty == 0 {
				pairLock.Lock()
				_, allow := openPairs[s]
				if !allow && len(openPairs) < maxOpenPair {
					openPairs[s] = 1
					maxPairs = max(maxPairs, len(openPairs))
					allow = true
				}
				pairLock.Unlock()
				if allow {
					s.OpenOrder(&strat.EnterReq{
						Tag:      "Long",
						CostRate: openRate,
					})
				}
			}

			// 加仓条件：亏损达到目标比例且有持仓
			if state.positionQty > 0 && pnl <= targetPercCur && len(s.LongOrders) < maxAdd {
				// 计算加仓数量，确保是整数
				odNum := len(s.LongOrders)
				addQty := math.Floor(state.positionQty * targetQty)
				pairLock.Lock()
				openPairs[s] = odNum
				pairLock.Unlock()
				s.OpenOrder(&strat.EnterReq{
					Tag:    fmt.Sprintf("%v", odNum),
					Amount: addQty,
				})
				if odNum > 5 {
					printMS = e.TimeStop
				}
			}
		},
		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
				}
			}
			ma25 := ta.SMA(e.Close, 25).Get(0)
			ma99 := ta.SMA(e.Close, 99).Get(0)
			if ma99 > ma25 {
				score += 1
			}
			state, _ := s.More.(*State1525S)
			state.score = score
		},
		OnBatchInfos: func(tf string, jobs map[string]*strat.JobEnv) {
			jobNum := len(jobs)
			if jobNum < 3 {
				return
			}
			downNum, downHeavyNum := float64(0), float64(0)
			timeStop := int64(0)
			var last *strat.JobEnv
			for _, j := range jobs {
				state, _ := j.Job.More.(*State1525S)
				if state.score >= 5 {
					downHeavyNum += 1
				}
				if state.score >= 3 {
					downNum += 1
				}
				timeStop = j.Env.TimeStop
				last = j
			}
			downRate := downNum / float64(len(jobs))
			downHeavyRate := downHeavyNum / float64(len(jobs))
			if printMS > 0 {
				printMS = 0
				dateStr := btime.ToDateStr(timeStop, core.DefaultDateFmt)
				fmt.Printf("%v downRate: %.3f downHeavyRate: %.3f\n", dateStr, downRate*100, downHeavyRate*100)
			}
			if downHeavyRate >= 0.8 {
				// 极端下跌
				targetPerc = targetPercRaw * 2
				takeProfit = takeProfitRaw * 0.5
			} else if downRate >= 0.85 && downHeavyRate >= 0.6 {
				// 严重下跌
				targetPerc = targetPercRaw * 1.8
				takeProfit = takeProfitRaw * 0.5
			} else if downRate >= 0.7 && downRate <= 0.85 {
				// 市场中度下跌
				targetPerc = targetPercRaw * 1.6
				takeProfit = takeProfitRaw * 0.65
			} else if downRate >= 0.5 && downRate <= 0.7 {
				// 轻度下跌
				targetPerc = targetPercRaw * 1.3
				takeProfit = takeProfitRaw * 0.8
			} else {
				// 轻度或未下跌
				targetPerc = targetPercRaw
				takeProfit = takeProfitRaw
			}
			pairLock.Lock()
			totNum := 0
			for _, num := range openPairs {
				totNum += num
			}
			pairLock.Unlock()
			if totNum > 0 {
				last.Job.Strat.WriteOutput(fmt.Sprintf("pairNum: %v", totNum), true)
			}
		},
		OnOrderChange: func(s *strat.StratJob, od *ormo.InOutOrder, chgType int) {
			state, _ := s.More.(*State1525S)

			// 当订单入场时更新持仓状态
			if chgType == strat.OdChgEnterFill && od.EnterTag != "" {
				// 计算新的平均成本
				if state.positionQty == 0 {
					// 首次开仓
					state.avgPrice = od.InitPrice
					state.positionQty = od.HoldAmount()
				} else {
					// 加仓，更新平均成本
					oldQty := state.positionQty
					oldAvgPrice := state.avgPrice
					newQty := od.HoldAmount()
					newAvgPrice := od.InitPrice

					// 计算新的平均成本
					totalQty := oldQty + newQty
					totalCost := oldQty*oldAvgPrice + newQty*newAvgPrice
					state.avgPrice = totalCost / totalQty
					state.positionQty = totalQty
				}
			}

			// 当订单完全平仓时重置持仓状态
			if chgType == strat.OdChgExitFill && od.ExitTag != "" && od.Status == ormo.InOutStatusFullExit {
				state.avgPrice = 0
				state.positionQty = 0
				state.tpLimit = 0
				pairLock.Lock()
				delete(openPairs, s)
				pairLock.Unlock()
			}
		},
		OnShutDown: func(s *strat.StratJob) {
			if maxPairs == 0 {
				return
			}
			log.Info("max open pairs", zap.Int("num", maxPairs))
			maxPairs = 0
		},
	}
}

// State1525Sim 存储每个品种的策略状态
type State1525Sim struct {
	totAmount float64
	totCost   float64
	tpLimit   float64 // 止盈出场价格
	odNum     int
	upFlag    float64 // 1: ma20 > ma99 && ma20 > ma20[2]
	score     float64 // 下跌行情力度分数，值域[0,6]
}

func J1525Sim(pol *config.RunPolicyConfig) *strat.TradeStrat {
	// 最多30个品种，单笔30，初始3W
	// 定义可调超参数
	smaPeriod := pol.DefInt("sma_period", 18, core.PNorm(5, 50))
	targetPercRaw := pol.Def("target_perc", -0.14, core.PNorm(-0.3, -0.05))
	takeProfitRaw := pol.Def("take_profit", 0.09, core.PNorm(0.05, 0.3))
	targetQty := pol.Def("target_qty", 0.6, core.PNorm(0.2, 1.0))
	// 加仓次数，超过20后变差；低于19后有段时间不稳定
	maxAdd := pol.DefInt("max_add", 19, core.PNorm(1, 30))
	startLvlRaw := pol.DefInt("start_lvl", 3, core.PNorm(1, 15))
	maxOpenPair := pol.DefInt("max_open_pair", 20, core.PNorm(1, 30)) // 最大同时交易品种

	openPairs := make(map[*strat.StratJob]int)
	pairLock := deadlock.Mutex{}
	lenAtr := 20
	baseAtrLen := int(math.Round(float64(lenAtr) * 4.3))

	targetPerc := targetPercRaw
	takeProfit := takeProfitRaw
	startLvl := startLvlRaw

	// 加仓倍率，基于targetQty： 1, 0.6, 1 1.7 3 4.7 7.7 12.8 21.1 35 57.5 95 157 256 385 464 767
	return &strat.TradeStrat{
		WarmupNum:       max(smaPeriod, baseAtrLen),
		HedgeOff:        true, // 单向持仓模式
		RunTimeFrames:   []string{"5m"},
		BatchInfo:       true,
		StakeRate:       0.25,
		OrderOnRotation: "open", // 品种切换时允许开单
		OnStartUp: func(s *strat.StratJob) {
			s.More = &State1525Sim{}
		},
		OnPairInfos: func(s *strat.StratJob) []*strat.PairSub {
			return []*strat.PairSub{
				{"_cur_", "4h", 20},
			}
		},
		OnBar: func(s *strat.StratJob) {
			e := s.Env
			state, _ := s.More.(*State1525Sim)
			closePrice := e.Close.Get(0)

			// 不利行情下调整参数
			targetPercJob, takeProfitJob := targetPercRaw, takeProfitRaw
			if state.score >= 5 {
				targetPercJob = targetPercRaw * 2
				takeProfitJob = takeProfitRaw * 0.6
			} else if state.score >= 4 {
				targetPercJob = targetPercRaw * 1.7
				takeProfitJob = takeProfitRaw * 0.7
			} else if state.score >= 3 {
				targetPercJob = targetPercRaw * 1.3
				takeProfitJob = takeProfitRaw * 0.8
			}
			targetPercCur := (targetPercJob + targetPerc) / 2
			takeProfitCur := (takeProfitJob + takeProfit) / 2

			// 计算SMA
			ma := ta.SMA(e.Close, smaPeriod)
			atr := ta.ATR(e.High, e.Low, e.Close, lenAtr)
			atrBase := ta.Lowest(ta.Highest(atr, lenAtr), baseAtrLen).Get(0)
			openRate := 0.006 * closePrice / atrBase
			if openRate < 1 {
				openRate = math.Sqrt(openRate)
			}
			openRate = max(0.7, min(1.5, openRate))

			// 检查止盈触发
			if state.tpLimit > 0 && state.tpLimit <= e.High.Get(0) {
				state.totAmount = 0
				state.totCost = 0
				state.tpLimit = 0
				state.odNum = 0
				pairLock.Lock()
				delete(openPairs, s)
				pairLock.Unlock()
			}

			// 计算当前持仓盈亏比例、止盈
			var pnl, avgPrice float64
			if state.totAmount > 0 {
				avgPrice = state.totCost / state.totAmount
				pnl = closePrice/avgPrice - 1
				curTP := avgPrice * (1 + takeProfitCur)
				if curTP != state.tpLimit {
					state.tpLimit = curTP
					// 设置限价单止盈
					s.CloseOrders(&strat.ExitReq{
						Tag:   "Take Profit",
						Limit: curTP,
					})
				}
			}

			// 首次入场条件：价格上穿SMA且无持仓
			if e.Close.Cross(ma) == 1 && state.totAmount == 0 && openRate > 0 {
				state.odNum = 1
				state.totAmount = openRate
				state.totCost = closePrice * openRate
			}

			// 加仓条件：亏损达到目标比例且有持仓
			if avgPrice > 0 && pnl <= targetPercCur && state.odNum < maxAdd {
				amt := state.totAmount * targetQty
				state.totAmount += amt
				state.totCost += amt * closePrice
				curStart := startLvl
				if state.upFlag < 0 {
					// 下降趋势上升3个阶梯，和全局阈值取平均
					curStart = (startLvl + startLvlRaw + 3) / 2
				}
				if state.odNum >= curStart {
					pairLock.Lock()
					_, allow := openPairs[s]
					if !allow && len(openPairs) < maxOpenPair {
						allow = true
					}
					if allow {
						openPairs[s] = len(s.LongOrders)
					}
					pairLock.Unlock()
					if allow {
						s.OpenOrder(&strat.EnterReq{
							Tag:      fmt.Sprintf("%v", state.odNum),
							CostRate: amt,
						})
					}
				}
				state.odNum += 1
			}
		},
		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
				}
			}
			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.(*State1525Sim)
			state.score = score
			state.upFlag = 0
			if ma20V > ma99 && ma20V > ma20.Get(2) {
				state.upFlag = 1
			} else if ma7 < min(ma20V, ma99) {
				state.upFlag = -1
			}
		},
		OnBatchInfos: func(tf string, jobs map[string]*strat.JobEnv) {
			jobNum := len(jobs)
			if jobNum < 3 {
				return
			}
			downNum, downHeavyNum, sumUp := float64(0), float64(0), float64(0)
			var last *strat.JobEnv
			for _, j := range jobs {
				state, _ := j.Job.More.(*State1525Sim)
				if state.score >= 5 {
					downHeavyNum += 1
				}
				if state.score >= 3 {
					downNum += 1
				}
				sumUp += max(0, state.upFlag)
				last = j
			}
			downRate := downNum / float64(len(jobs))
			downHeavyRate := downHeavyNum / float64(len(jobs))
			sumUp /= float64(len(jobs))
			if downHeavyRate >= 0.8 {
				// 极端下跌
				targetPerc = targetPercRaw * 2
				takeProfit = takeProfitRaw * 0.5
			} else if downRate >= 0.85 && downHeavyRate >= 0.6 {
				// 严重下跌
				targetPerc = targetPercRaw * 1.8
				takeProfit = takeProfitRaw * 0.5
			} else if downRate >= 0.7 && downRate <= 0.85 {
				// 市场中度下跌
				targetPerc = targetPercRaw * 1.6
				takeProfit = takeProfitRaw * 0.65
			} else if downRate >= 0.5 && downRate <= 0.7 {
				// 轻度下跌
				targetPerc = targetPercRaw * 1.3
				takeProfit = takeProfitRaw * 0.8
			} else {
				// 轻度或未下跌
				targetPerc = targetPercRaw
				takeProfit = takeProfitRaw
			}
			// 调整最小开仓级别
			if sumUp > 0.7 {
				startLvl = startLvlRaw
			} else if sumUp > 0.5 && downRate <= 0.5 {
				startLvl = startLvlRaw + 2
			} else if sumUp > 0.3 && downRate <= 0.7 {
				startLvl = startLvlRaw + 4
			} else {
				startLvl = startLvlRaw + 6
			}
			pairLock.Lock()
			totNum := 0
			for _, num := range openPairs {
				totNum += num
			}
			pairLock.Unlock()
			if totNum > 0 {
				last.Job.Strat.WriteOutput(fmt.Sprintf("pairNum: %v", totNum), true)
			}
		},
	}
}