stock-analysis/internal/backtest/engine.go

package backtest

import (
	"context"
	"fmt"
	"sort"

	"stock-analysis/internal/model"
	"stock-analysis/internal/storage"
)

// Params 回测运行参数。
type Params struct {
	InitialCash float64 // 初始资金（元）
	Commission  float64 // 单边佣金比例（含佣金，默认 0.00025）
	Slippage    float64 // 滑点比例（默认 0）
	StartDate   string  // 回测起始日（含），空则使用数据首日
	EndDate     string  // 回测结束日（含），空则使用数据末日
}

// DefaultParams 返回推荐默认参数。
func DefaultParams() Params {
	return Params{
		InitialCash: 100000,
		Commission:  0.00025,
	}
}

// DataFeed 从存储加载全量K线，构建按日期索引的行情数据。
type DataFeed struct {
	quotes map[string][]Quote // date → 当日全部转债行情
	dates  []string           // 升序交易日列表
	hist   *History
}

// NewDataFeed 从存储加载全部转债K线并构建回测数据源。
func NewDataFeed(ctx context.Context, store storage.KlineStore) (*DataFeed, error) {
	all, err := store.LoadAllKlines(ctx)
	if err != nil {
		return nil, fmt.Errorf("load klines: %w", err)
	}

	quotes := make(map[string][]Quote)
	for code, ks := range all {
		for _, k := range ks {
			quotes[k.Date] = append(quotes[k.Date], Quote{
				Code:         code,
				Price:        k.Close,
				Open:         k.Open,
				High:         k.High,
				Low:          k.Low,
				Volume:       k.Volume,
				ConvertValue: k.ConvertValue,
				PremiumRate:  k.PremiumRate,
			})
		}
	}

	dates := make([]string, 0, len(quotes))
	for d := range quotes {
		dates = append(dates, d)
	}
	sort.Strings(dates)

	return &DataFeed{quotes: quotes, dates: dates, hist: &History{klines: all}}, nil
}

// Dates 返回升序交易日列表。
func (d *DataFeed) Dates() []string {
	return d.dates
}

// QuotesOn 返回指定交易日的全部转债行情。
func (d *DataFeed) QuotesOn(date string) []Quote {
	return d.quotes[date]
}

// History 返回全量历史数据访问器。
func (d *DataFeed) History() *History {
	return d.hist
}

// Filter 只保留指定代码的数据，用于按用户选择的转债子集回测。
func (d *DataFeed) Filter(codes []string) {
	keep := make(map[string]bool, len(codes))
	for _, c := range codes {
		keep[c] = true
	}

	// 过滤 quotes。
	for date, qs := range d.quotes {
		var filtered []Quote
		for _, q := range qs {
			if keep[q.Code] {
				filtered = append(filtered, q)
			}
		}
		d.quotes[date] = filtered
	}

	// 过滤 history.klines。
	filteredKlines := make(map[string][]model.Kline)
	for code, ks := range d.hist.klines {
		if keep[code] {
			filteredKlines[code] = ks
		}
	}
	d.hist.klines = filteredKlines
}

// Engine 事件驱动回测引擎。
type Engine struct {
	params   Params
	strategy Strategy
	account  *Account
	filter   *FundamentalFilter
	fundMap  map[string]*model.StockFundamental
}

// NewEngine 创建回测引擎。
func NewEngine(params Params, strategy Strategy) *Engine {
	return &Engine{params: params, strategy: strategy}
}

// SetFundamentalFilter 设置基本面过滤器（可选）。
// stock2bond 为正股代码到转债代码的映射；fundamentals 为正股代码到基本面数据的映射。
func (e *Engine) SetFundamentalFilter(filter *FundamentalFilter, fundamentals map[string]*model.StockFundamental) {
	e.filter = filter
	e.fundMap = fundamentals
}

// Run 执行回测，返回结果（含净值曲线与绩效指标）。
func (e *Engine) Run(ctx context.Context, feed *DataFeed) (*Result, error) {
	e.account = NewAccount(e.params)

	for _, date := range feed.dates {
		if e.params.StartDate != "" && date < e.params.StartDate {
			continue
		}
		if e.params.EndDate != "" && date > e.params.EndDate {
			break
		}

		universe := feed.quotes[date]
			if e.filter != nil {
				universe = e.filter.Apply(ctx, universe, e.fundMap)
			}
			targets, err := e.strategy.OnBar(ctx, date, universe, feed.hist)
		if err != nil {
			return nil, fmt.Errorf("strategy %s on %s: %w", e.strategy.Name(), date, err)
		}
		if targets != nil {
			e.account.Rebalance(date, targets, feed.hist.CloseOn)
		}
		e.account.MarkToMarket(date, feed.hist.CloseOn)
	}

	curve := e.account.Curve()
	start, end := "", ""
	if len(curve) > 0 {
		start, end = curve[0].Date, curve[len(curve)-1].Date
	}
	return &Result{
		Strategy: e.strategy.Name(),
		Start:    start,
		End:      end,
		Metrics:  ComputeMetrics(curve),
		Curve:    curve,
	}, nil
}

// Result 回测输出。
type Result struct {
	Strategy string
	Start    string
	End      string
	Metrics  Metrics
	Curve    []EquityPoint
}

// EquityPoint 净值曲线点。
type EquityPoint struct {
	Date   string
	Value  float64 // 组合总市值
	Return float64 // 累计收益率
}