factor_engine/gp/mining.py

"""Main GP factor mining loop using qlib data."""

from __future__ import annotations

import pickle
import random
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm

from data_pipeline.init_qlib import init_qlib
from data_pipeline.load_data import load_instruments
from factor_engine.gp.config import load_gp_config
from factor_engine.gp.evolution import crossover, mutate, tournament_selection, tree_too_large
from factor_engine.gp.fitness import FitnessConfig, calculate_fitness, factor_report, sampled_spearman_corr_torch
from factor_engine.gp.operators import generate_random_tree
from factor_engine.gp.qlib_engine import QlibTensorDataEngine


def _append_csv(path: Path, rows: list[dict]):
    df = pd.DataFrame(rows)
    df.to_csv(path, mode="a" if path.exists() else "w", header=not path.exists(), index=False, encoding="utf-8-sig")


def _save_population(path: Path, population, generation_offset: int, cfg: dict):
    with open(path, "wb") as f:
        pickle.dump({"population": population, "generation_offset": generation_offset, "config": cfg}, f)


def _load_population(path: Path):
    if not path.exists():
        return None, 0
    with open(path, "rb") as f:
        payload = pickle.load(f)
    return payload["population"], int(payload.get("generation_offset", 0))


def _export_ml_features(engine, target, top_trees, output_path: Path):
    if not top_trees:
        print("No orthogonal factors selected for ML export.")
        return

    dates_col = np.repeat(engine.times, len(engine.symbols))
    symbols_col = np.tile(engine.symbols, len(engine.times))
    ml_df = pd.DataFrame({"date": dates_col, "symbol": symbols_col})
    ml_df["target_return"] = target.detach().cpu().numpy().flatten(order="F")

    formulas = {}
    for i, tree in enumerate(top_trees, start=1):
        factor = tree.evaluate(engine)
        ml_df[f"factor_{i}"] = factor.detach().cpu().numpy().flatten(order="F")
        formulas[f"factor_{i}"] = str(tree)

    ml_df = ml_df.replace([np.inf, -np.inf], np.nan).dropna().reset_index(drop=True)
    ml_df.to_csv(output_path, index=False, encoding="utf-8-sig")
    ml_df.to_parquet(output_path.with_suffix(".parquet"), index=False)
    pd.Series(formulas, name="formula").to_csv(output_path.with_name("factor_formulas.csv"), header=True)
    print(f"ML features saved: {output_path} ({len(top_trees)} factors)")


def _select_low_corr_trees(engine, candidate_pairs, train_mask, top_k, threshold):
    selected_trees, selected_factors = [], []
    for fit, tree in candidate_pairs:
        if len(selected_trees) >= top_k:
            break
        factor = tree.evaluate(engine)
        duplicate = False
        for old in selected_factors:
            c = sampled_spearman_corr_torch(factor, old, mask=train_mask, max_points=30_000)
            if not np.isnan(c) and abs(c) >= threshold:
                duplicate = True
                break
        if not duplicate:
            selected_trees.append(tree.clone())
            selected_factors.append(factor.detach().clone())
    return selected_trees


def run_gp_mining(
    gp_config_path: str | None = None,
    base_config_path: str | None = None,
) -> Path:
    """Run GP factor mining end-to-end on qlib data."""
    cfg_bundle = load_gp_config(gp_config_path, base_config_path)
    gp_cfg = cfg_bundle["gp"]
    base_cfg = cfg_bundle["base"]
    out_dir: Path = cfg_bundle["output_dir"]
    out_dir.mkdir(parents=True, exist_ok=True)

    seed = int(gp_cfg.get("seed", 42))
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    device = gp_cfg.get("device", "auto")
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    init_qlib(base_config_path)
    market = base_cfg["qlib"].get("market", "csi300")
    instruments = load_instruments(market)

    engine = QlibTensorDataEngine(
        instruments=instruments,
        start_time=base_cfg["data"]["start_time"],
        end_time=base_cfg["data"]["end_time"],
        freq=base_cfg["data"].get("freq", "day"),
        forward_steps=gp_cfg.get("forward_steps", base_cfg["label"].get("forward_days", 5)),
        splits=base_cfg.get("splits"),
        device=device,
        qlib_fields=base_cfg["data"].get("fields"),
    )

    fitness_cfg = FitnessConfig(
        min_stocks=gp_cfg.get("min_stocks", 50),
        min_tree_nodes=gp_cfg.get("min_tree_nodes", 3),
        depth_penalty=gp_cfg.get("depth_penalty", 0.006),
        node_penalty=gp_cfg.get("node_penalty", 0.0012),
        duplicate_penalty=gp_cfg.get("duplicate_penalty", 0.05),
        elite_corr_fatal=gp_cfg.get("elite_corr_fatal", 0.95),
        elite_corr_hard=gp_cfg.get("elite_corr_hard", 0.85),
        elite_corr_mid=gp_cfg.get("elite_corr_mid", 0.75),
        elite_corr_soft=gp_cfg.get("elite_corr_soft", 0.65),
        export_corr_threshold=gp_cfg.get("export_corr_threshold", 0.80),
    )

    pop_size = gp_cfg.get("population_size", 200)
    max_init_depth = gp_cfg.get("max_init_depth", 4)
    max_tree_depth = gp_cfg.get("max_tree_depth", 8)
    max_tree_nodes = gp_cfg.get("max_tree_nodes", 60)
    generations = gp_cfg.get("generations_per_run", 10)
    elite_size = gp_cfg.get("elite_size", 20)
    top_k = gp_cfg.get("top_k_export", 30)

    population_path = out_dir / "population.pkl"
    factor_zoo_path = out_dir / "factor_zoo.csv"
    ml_feature_path = out_dir / "ML_Features_qlib.csv"

    target = engine.get_data("target_return")
    train_mask = engine.get_data("train_mask")
    valid_mask = engine.get_data("valid_mask")
    test_mask = engine.get_data("test_mask")

    population, generation_offset = _load_population(population_path)
    if population is None:
        population = [generate_random_tree(1, max_init_depth) for _ in range(pop_size)]
        generation_offset = 0
    else:
        population = population[:pop_size] + [
            generate_random_tree(1, max_init_depth) for _ in range(max(0, pop_size - len(population)))
        ]

    formula_seen = set()
    if factor_zoo_path.exists():
        zoo = pd.read_csv(factor_zoo_path)
        if "formula" in zoo.columns:
            formula_seen = set(zoo["formula"].astype(str))

    elite_cache = []
    max_elite_cache = gp_cfg.get("max_elite_cache", 60)

    for local_gen in range(1, generations + 1):
        global_gen = generation_offset + local_gen
        print(f"\n========== Generation {global_gen} ==========")

        fitnesses, records = [], []
        for tree in tqdm(population, desc="Evaluating"):
            if tree_too_large(tree, max_tree_depth, max_tree_nodes):
                fitnesses.append(-999.0)
                continue
            factor = tree.evaluate(engine)
            fit, is_rpt, oos_rpt = calculate_fitness(
                tree, factor, target, train_mask, test_mask, fitness_cfg, formula_seen, elite_cache
            )
            fitnesses.append(fit)
            records.append({
                "generation": global_gen,
                "fitness": fit,
                "is_icir": is_rpt["icir"],
                "is_ic_mean": is_rpt["ic_mean"],
                "oos_icir": oos_rpt["icir"],
                "formula": str(tree),
                "depth": tree.get_depth(),
                "nodes": tree.get_size(),
            })

        fitnesses = np.array(fitnesses, dtype=float)
        sorted_idx = np.argsort(fitnesses)[::-1]
        best_tree = population[int(sorted_idx[0])].clone()
        best_factor = best_tree.evaluate(engine)

        best_is = factor_report(best_factor, target, train_mask, fitness_cfg.min_stocks)
        best_valid = factor_report(best_factor, target, valid_mask, fitness_cfg.min_stocks)
        best_test = factor_report(best_factor, target, test_mask, fitness_cfg.min_stocks)

        if not np.isnan(best_is["icir"]):
            if not any(
                not np.isnan(c := sampled_spearman_corr_torch(best_factor, old, train_mask))
                and abs(c) >= fitness_cfg.elite_corr_hard
                for old in elite_cache
            ):
                elite_cache.append(best_factor.detach().clone())
                if len(elite_cache) > max_elite_cache:
                    elite_cache.pop(0)

        print(f"Best fitness={fitnesses[sorted_idx[0]]:.4f} | IS ICIR={best_is['icir']:.4f} | VALID ICIR={best_valid['icir']:.4f} | TEST ICIR={best_test['icir']:.4f}")
        print(f"Formula: {best_tree}")

        _append_csv(factor_zoo_path, sorted(records, key=lambda x: x["fitness"], reverse=True)[:top_k])

        new_population = [population[int(i)].clone() for i in sorted_idx[:elite_size]]
        new_population += [generate_random_tree(1, max_init_depth) for _ in range(int(pop_size * gp_cfg.get("random_immigrant_rate", 0.15)))]

        while len(new_population) < pop_size:
            r = random.random()
            if r < gp_cfg.get("crossover_rate", 0.5):
                child = crossover(tournament_selection(population, fitnesses), tournament_selection(population, fitnesses))
            elif r < gp_cfg.get("crossover_rate", 0.5) + gp_cfg.get("mutation_rate", 0.35):
                child = mutate(tournament_selection(population, fitnesses), max_init_depth)
            else:
                child = tournament_selection(population, fitnesses)
            if tree_too_large(child, max_tree_depth, max_tree_nodes):
                child = generate_random_tree(1, max_init_depth)
            new_population.append(child)

        population = new_population
        _save_population(population_path, population, global_gen, gp_cfg)
        if device == "cuda":
            torch.cuda.empty_cache()

    if factor_zoo_path.exists():
        zoo = pd.read_csv(factor_zoo_path).drop_duplicates("formula").sort_values("fitness", ascending=False)
        zoo.to_csv(out_dir / "top_factors.csv", index=False, encoding="utf-8-sig")

        final_pairs = []
        for tree in tqdm(population, desc="Final export"):
            if tree_too_large(tree, max_tree_depth, max_tree_nodes):
                continue
            factor = tree.evaluate(engine)
            fit, _, _ = calculate_fitness(tree, factor, target, train_mask, test_mask, fitness_cfg)
            final_pairs.append((fit, tree.clone()))

        top_trees = _select_low_corr_trees(
            engine,
            sorted(final_pairs, key=lambda x: x[0], reverse=True),
            train_mask,
            top_k,
            fitness_cfg.export_corr_threshold,
        )
        if top_trees:
            _export_ml_features(engine, target, top_trees, ml_feature_path)

    print(f"\nGP mining complete. Outputs: {out_dir}")
    return out_dir