quant-source/wqb-demo/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/semaphore"
)

// AlphaConfig Alpha因子配置
type AlphaConfig struct {
	Name       string                 `json:"name"`
	Expression string                 `json:"expression"`
	Settings   map[string]interface{} `json:"settings,omitempty"`
}

// SimulationResult 模拟结果
type SimulationResult struct {
	Name       string      `json:"name"`
	Expression string      `json:"expression"`
	AlphaID    string      `json:"alpha_id,omitempty"`
	Success    bool        `json:"success"`
	Error      string      `json:"error,omitempty"`
	ResultData interface{} `json:"result_data,omitempty"`
}

// WorldQuantBrainAPI WorldQuant Brain API客户端
type WorldQuantBrainAPI struct {
	BaseURL         string
	CredentialsFile string
	HTTPClient      *http.Client
	Username        string
	Password        string
}

// NewWorldQuantBrainAPI 创建新的API客户端
func NewWorldQuantBrainAPI(credentialsFile string) *WorldQuantBrainAPI {
	if credentialsFile == "" {
		credentialsFile = "brain_credentials.txt"
	}

	// 扩展文件路径
	fullPath := expandUser(credentialsFile)

	return &WorldQuantBrainAPI{
		BaseURL:         "https://api.worldquantbrain.com",
		CredentialsFile: fullPath,
		HTTPClient:      &http.Client{Timeout: 30 * time.Second},
	}
}

// expandUser 扩展用户主目录路径
func expandUser(path string) string {
	if strings.HasPrefix(path, "~/") {
		home, err := os.UserHomeDir()
		if err != nil {
			return path
		}
		return filepath.Join(home, path[2:])
	}
	return path
}

// Login 登录认证
func (wq *WorldQuantBrainAPI) Login() error {
	// 读取凭证文件
	data, err := os.ReadFile(wq.CredentialsFile)
	if err != nil {
		return fmt.Errorf("读取凭证文件失败: %v", err)
	}

	content := strings.TrimSpace(string(data))

	var credentials []string

	// 方法1: 尝试标准JSON解析
	if err := json.Unmarshal([]byte(content), &credentials); err == nil && len(credentials) >= 2 {
		// JSON解析成功
	} else {
		// 方法2: 尝试处理单引号JSON
		singleQuoteContent := strings.ReplaceAll(content, "'", "\"")
		if err := json.Unmarshal([]byte(singleQuoteContent), &credentials); err == nil && len(credentials) >= 2 {
			// 单引号JSON解析成功
		} else {
			// 方法3: 尝试其他文本格式
			credentials = wq.parseTextCredentials(content)
		}
	}

	if len(credentials) < 2 {
		return fmt.Errorf("无法解析凭证文件，需要用户名和密码。当前内容: %s", content)
	}

	wq.Username = strings.TrimSpace(credentials[0])
	wq.Password = strings.TrimSpace(credentials[1])

	// 创建认证请求
	req, err := http.NewRequest("POST", wq.BaseURL+"/authentication", nil)
	if err != nil {
		return fmt.Errorf("创建请求失败: %v", err)
	}
	req.SetBasicAuth(wq.Username, wq.Password)

	// 发送认证请求
	resp, err := wq.HTTPClient.Do(req)
	if err != nil {
		return fmt.Errorf("认证请求失败: %v", err)
	}
	defer resp.Body.Close()

	fmt.Printf("认证状态: %d\n", resp.StatusCode)

	// 修复：检查 200 和 201 状态码都表示成功
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("登录失败，状态码: %d, 响应: %s", resp.StatusCode, string(body))
	}

	// 读取成功响应
	body, _ := io.ReadAll(resp.Body)
	fmt.Printf("登录成功响应: %s\n", string(body))

	fmt.Println("登录成功!")
	return nil
}

func (wq *WorldQuantBrainAPI) parseTextCredentials(content string) []string {
	panic("unimplemented")
}

// SimulateAlpha 模拟单个Alpha因子
func (wq *WorldQuantBrainAPI) SimulateAlpha(config AlphaConfig) SimulationResult {
	if wq.HTTPClient == nil {
		return SimulationResult{
			Name:    config.Name,
			Success: false,
			Error:   "请先调用Login方法登录",
		}
	}

	name := config.Name
	if name == "" {
		name = "未命名Alpha"
	}

	fmt.Printf("\n🚀 开始模拟: %s\n", name)
	fmt.Printf("📝 表达式: %s\n", config.Expression)

	// 构建模拟数据
	simulationData := map[string]interface{}{
		"type":     "REGULAR",
		"settings": config.Settings,
		"regular":  config.Expression,
	}

	jsonData, err := json.Marshal(simulationData)
	if err != nil {
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   fmt.Sprintf("序列化模拟数据失败: %v", err),
		}
	}

	// 发送模拟请求
	req, err := http.NewRequest("POST", wq.BaseURL+"/simulations", strings.NewReader(string(jsonData)))
	if err != nil {
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   fmt.Sprintf("创建请求失败: %v", err),
		}
	}
	req.SetBasicAuth(wq.Username, wq.Password)
	req.Header.Set("Content-Type", "application/json")

	resp, err := wq.HTTPClient.Do(req)
	if err != nil {
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   fmt.Sprintf("模拟请求失败: %v", err),
		}
	}
	defer resp.Body.Close()

	fmt.Printf("📤 模拟提交状态: %d\n", resp.StatusCode)

	if resp.StatusCode != http.StatusOK {
		// 尝试读取错误响应体
		body, _ := io.ReadAll(resp.Body)
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   fmt.Sprintf("提交失败: %d, 响应: %s", resp.StatusCode, string(body)),
		}
	}

	// 获取进度URL
	progressURL := resp.Header.Get("Location")
	if progressURL == "" {
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   "未找到进度URL",
		}
	}

	fmt.Printf("⏳ %s 进度URL获取成功\n", name)

	// 轮询进度
	var resultData map[string]interface{}
	for {
		progressReq, err := http.NewRequest("GET", progressURL, nil)
		if err != nil {
			return SimulationResult{
				Name:    name,
				Success: false,
				Error:   fmt.Sprintf("创建进度请求失败: %v", err),
			}
		}
		progressReq.SetBasicAuth(wq.Username, wq.Password)

		progressResp, err := wq.HTTPClient.Do(progressReq)
		if err != nil {
			return SimulationResult{
				Name:    name,
				Success: false,
				Error:   fmt.Sprintf("进度请求失败: %v", err),
			}
		}

		retryAfter := progressResp.Header.Get("Retry-After")
		if retryAfter == "" || retryAfter == "0" {
			// 模拟完成
			if err := json.NewDecoder(progressResp.Body).Decode(&resultData); err != nil {
				progressResp.Body.Close()
				return SimulationResult{
					Name:    name,
					Success: false,
					Error:   fmt.Sprintf("解析结果失败: %v", err),
				}
			}
			progressResp.Body.Close()
			break
		}

		progressResp.Body.Close()

		// 等待指定时间
		waitSeconds, _ := strconv.ParseFloat(retryAfter, 64)
		fmt.Printf("⏰ %s 等待 %.1f 秒...\n", name, waitSeconds)
		time.Sleep(time.Duration(waitSeconds * float64(time.Second)))
	}

	// 提取Alpha ID
	alphaID, ok := resultData["alpha"].(string)
	if !ok {
		return SimulationResult{
			Name:    name,
			Success: false,
			Error:   "未找到Alpha ID",
		}
	}

	fmt.Printf("✅ %s 模拟完成! Alpha ID: %s\n", name, alphaID)

	return SimulationResult{
		Name:       name,
		Expression: config.Expression,
		AlphaID:    alphaID,
		Success:    true,
		ResultData: resultData,
	}
}

// BatchSimulateAlphas 批量测试多个Alpha因子
func (wq *WorldQuantBrainAPI) BatchSimulateAlphas(configFile string, maxWorkers int) []SimulationResult {
	// 读取Alpha配置
	configs, err := wq.loadAlphaConfigs(configFile)
	if err != nil {
		log.Printf("❌ 读取配置文件失败: %v", err)
		return nil
	}

	fmt.Printf("📁 读取到 %d 个Alpha配置\n", len(configs))

	// 登录
	if err := wq.Login(); err != nil {
		log.Printf("❌ 登录失败: %v", err)
		return nil
	}

	var results []SimulationResult
	var mu sync.Mutex
	var wg sync.WaitGroup

	// 使用信号量控制并发数量
	sem := semaphore.NewWeighted(int64(maxWorkers))
	ctx := context.Background()

	for _, config := range configs {
		wg.Add(1)

		go func(c AlphaConfig) {
			defer wg.Done()

			// 获取信号量
			if err := sem.Acquire(ctx, 1); err != nil {
				log.Printf("❌ 获取信号量失败: %v", err)
				return
			}
			defer sem.Release(1)

			// 执行模拟
			result := wq.SimulateAlpha(c)

			// 安全地添加结果
			mu.Lock()
			results = append(results, result)
			mu.Unlock()
		}(config)
	}

	wg.Wait()

	// 打印汇总结果
	wq.printSummary(results)
	return results
}

// loadAlphaConfigs 加载Alpha配置
func (wq *WorldQuantBrainAPI) loadAlphaConfigs(configFile string) ([]AlphaConfig, error) {
	data, err := os.ReadFile(configFile)
	if err != nil {
		return nil, fmt.Errorf("读取配置文件 %s 失败: %v", configFile, err)
	}

	var configs []AlphaConfig
	if err := json.Unmarshal(data, &configs); err != nil {
		return nil, fmt.Errorf("解析配置文件 %s 失败: %v", configFile, err)
	}

	return configs, nil
}

// printSummary 打印测试结果汇总
func (wq *WorldQuantBrainAPI) printSummary(results []SimulationResult) {
	fmt.Println("\n" + strings.Repeat("=", 50))
	fmt.Println("📊 Alpha测试结果汇总")
	fmt.Println(strings.Repeat("=", 50))

	successCount := 0
	for _, result := range results {
		if result.Success {
			successCount++
		}
	}
	failedCount := len(results) - successCount

	fmt.Printf("✅ 成功: %d 个\n", successCount)
	fmt.Printf("❌ 失败: %d 个\n", failedCount)

	fmt.Println("\n成功详情:")
	for _, result := range results {
		if result.Success {
			fmt.Printf("  ✓ %s: %s\n", result.Name, result.AlphaID)
		}
	}

	if failedCount > 0 {
		fmt.Println("\n失败详情:")
		for _, result := range results {
			if !result.Success {
				fmt.Printf("  ✗ %s: %s\n", result.Name, result.Error)
			}
		}
	}
}

// Close 关闭客户端
func (wq *WorldQuantBrainAPI) Close() {
	// 如果有需要关闭的资源，可以在这里处理
}

// 保存结果到文件
func saveResultsToFile(results []SimulationResult, filename string) error {
	data, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(filename, data, 0644)
}

func main() {
	// 创建API实例
	wqAPI := NewWorldQuantBrainAPI("")

	// 确保在程序结束时关闭连接
	defer wqAPI.Close()

	fmt.Println("🎯 开始批量测试Alpha因子...")
	startTime := time.Now()

	// 批量测试所有Alpha
	results := wqAPI.BatchSimulateAlphas("alphas.json", 2) // 并发数量，根据API限制调整

	duration := time.Since(startTime)
	fmt.Printf("\n⏱️  总耗时: %.2f 秒\n", duration.Seconds())

	fmt.Println(results)
}