cbpred/cmd/cbpred-export/main.go

package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"github.com/banbox/banbot/biz"
	"github.com/banbox/banbot/config"
	"github.com/banbox/banbot/core"
	"github.com/banbox/banbot/orm"
	"github.com/banbox/banexg/errs"
)

var columns = []string{
	"timestamp", "open", "high", "low", "close", "volume", "quote", "buy_volume", "trade_num",
}

func main() {
	configPath := flag.String("config", "", "banbot config path")
	dataDir := flag.String("data-dir", "", "banbot data directory")
	pairs := flag.String("pairs", "", "comma-separated symbols")
	startMS := flag.Int64("start-ms", 0, "inclusive start timestamp in milliseconds")
	endMS := flag.Int64("end-ms", 0, "exclusive end timestamp in milliseconds")
	timeframe := flag.String("timeframe", "8h", "bar timeframe (8h or 1d)")
	outDir := flag.String("out-dir", "", "CSV output directory")
	skipUnavailable := flag.Bool("skip-unavailable", true, "skip pairs unavailable on the exchange and absent from the local cache")
	flag.Parse()
	if (*timeframe != "8h" && *timeframe != "1d") || *configPath == "" || *dataDir == "" || *pairs == "" || *outDir == "" || *startMS <= 0 || *endMS <= *startMS {
		flag.Usage()
		os.Exit(2)
	}

	args := &config.CmdArgs{
		Configs:   config.ArrString{*configPath},
		DataDir:   *dataDir,
		NoDefault: true,
		RawPairs:  *pairs,
		TimeStart: strconv.FormatInt(*startMS, 10),
		TimeEnd:   strconv.FormatInt(*endMS, 10),
		LogLevel:  "error",
	}
	if err := biz.SetupComs(args); err != nil {
		fail(err)
	}
	defer core.RunExitCalls()
	if err := os.MkdirAll(*outDir, 0o755); err != nil {
		fail(err)
	}

	for _, rawPair := range strings.Split(*pairs, ",") {
		pair := strings.TrimSpace(rawPair)
		if pair == "" {
			continue
		}
		exs, err := orm.GetExSymbolCur(pair)
		if err != nil {
			if *skipUnavailable && isUnavailablePair(err) {
				continue
			}
			fail(err)
		}
		_, rows, err := orm.GetSeries(exs, "1h", *startMS, *endMS, 0, false)
		if err != nil {
			fail(err)
		}
		bars := make([]*orm.SeriesOHLCV, 0, len(rows))
		for _, row := range rows {
			bar, err := row.OHLCV(exs)
			if err != nil {
				fail(err)
			}
			bars = append(bars, bar)
		}
		bars = aggregateHourly(bars, *timeframe)
		name := fmt.Sprintf("%d_%s.csv", exs.ID, strings.NewReplacer("/", "_", ":", "_").Replace(pair))
		if err := writeCSV(filepath.Join(*outDir, name), bars); err != nil {
			fail(err)
		}
	}
}

func isUnavailablePair(err *errs.Error) bool {
	return err != nil && err.Code == errs.CodeNoMarketForPair
}

// aggregateHourly emits only complete UTC buckets with exactly one bar per hour.
func aggregateHourly(rows []*orm.SeriesOHLCV, timeframe string) []*orm.SeriesOHLCV {
	const hourMS int64 = 60 * 60 * 1000
	hours := 8
	if timeframe == "1d" {
		hours = 24
	}
	stepMS := int64(hours) * hourMS
	ordered := append([]*orm.SeriesOHLCV(nil), rows...)
	sort.Slice(ordered, func(i, j int) bool { return ordered[i].Time < ordered[j].Time })
	var output []*orm.SeriesOHLCV
	for first := 0; first < len(ordered); {
		bucket := ordered[first].Time / stepMS * stepMS
		end := first + 1
		for end < len(ordered) && ordered[end].Time/stepMS*stepMS == bucket {
			end++
		}
		complete := end-first == hours
		for i := first; complete && i < end; i++ {
			complete = ordered[i].Time == bucket+int64(i-first)*hourMS
		}
		if complete {
			bar := *ordered[first]
			bar.Time, bar.EndMS, bar.TimeFrame, bar.Closed = bucket, bucket+stepMS, timeframe, true
			for _, row := range ordered[first+1 : end] {
				bar.High = max(bar.High, row.High)
				bar.Low = min(bar.Low, row.Low)
				bar.Close = row.Close
				bar.Volume += row.Volume
				bar.Quote += row.Quote
				bar.BuyVolume += row.BuyVolume
				bar.TradeNum += row.TradeNum
			}
			output = append(output, &bar)
		}
		first = end
	}
	return output
}

func writeCSV(path string, rows []*orm.SeriesOHLCV) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()
	writer := csv.NewWriter(file)
	if err := writer.Write(columns); err != nil {
		return err
	}
	for _, bar := range rows {
		values := []string{
			strconv.FormatInt(bar.Time, 10),
			strconv.FormatFloat(bar.Open, 'f', -1, 64),
			strconv.FormatFloat(bar.High, 'f', -1, 64),
			strconv.FormatFloat(bar.Low, 'f', -1, 64),
			strconv.FormatFloat(bar.Close, 'f', -1, 64),
			strconv.FormatFloat(bar.Volume, 'f', -1, 64),
			strconv.FormatFloat(bar.Quote, 'f', -1, 64),
			strconv.FormatFloat(bar.BuyVolume, 'f', -1, 64),
			strconv.FormatInt(bar.TradeNum, 10),
		}
		if err := writer.Write(values); err != nil {
			return err
		}
	}
	writer.Flush()
	return writer.Error()
}

func fail(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}