powerball_ai.py

# powerball_ai.py
import pandas as pd
import numpy as np
from typing import List, Optional, Dict, Tuple
from functools import lru_cache

WHITE_MAX = 69
PB_MAX = 26

class LotteryAnalyzer:
    def __init__(self):
        self.df: Optional[pd.DataFrame] = None
        self.white_cols: List[str] = []
        self.pb_col: Optional[str] = None
        self.features_white: Optional[pd.DataFrame] = None
        self.features_pb: Optional[pd.DataFrame] = None
        self.scores_combined: Optional[pd.DataFrame] = None
        self.weights = (0.3, 0.3, 0.2, 0.2)  # freq, recent, gap, performance

    def load_csv(self, path: str) -> "LotteryAnalyzer":
        df = pd.read_csv(path)
        cols_lower = [c.lower().strip() for c in df.columns]

        if 'winning numbers' in cols_lower:
            wn_col = df.columns[cols_lower.index('winning numbers')]
            if 'draw date' in cols_lower:
                date_col = df.columns[cols_lower.index('draw date')]
                df['date'] = pd.to_datetime(df[date_col], errors='coerce')
            else:
                df['date'] = pd.NaT

            def parse_row(s):
                parts = str(s).split()
                if len(parts) >= 6:
                    return parts[:5] + [parts[5]]
                return [np.nan]*6

            parsed = df[wn_col].astype(str).apply(parse_row).apply(pd.Series)
            parsed.columns = ['n1','n2','n3','n4','n5','pb']
            df = pd.concat([df, parsed], axis=1)
            self.white_cols = ['n1','n2','n3','n4','n5']
            self.pb_col = 'pb'
        else:
            if set(['n1','n2','n3','n4','n5','pb']).issubset(set(cols_lower)):
                mapping = {c: df.columns[cols_lower.index(c)] for c in ['n1','n2','n3','n4','n5','pb']}
                df = df.rename(columns={mapping[c]: c for c in mapping})
                self.white_cols = ['n1','n2','n3','n4','n5']
                self.pb_col = 'pb'
            else:
                raise ValueError(
                    "无法识别输入文件列。请确保 CSV 包含其中一种格式：\n"
                    " - 列名 'Winning Numbers'（NY Open Data），或\n"
                    " - 单独列名 n1..n5,pb。"
                )

        # 强制数值类型并验证范围
        for c in self.white_cols + [self.pb_col]:
            df[c] = pd.to_numeric(df[c], errors='coerce').astype('Int64')
        df = df[df['date'] >= '2015-10-07']  # 过滤 2015-10-07 前的旧矩阵数据
        for c in self.white_cols:
            if not df[c].isna().all() and not ((df[c] >= 1) & (df[c] <= WHITE_MAX)).all():
                raise ValueError(f"白球列 {c} 包含超出 1-{WHITE_MAX} 的值")
        if not df[self.pb_col].isna().all() and not ((df[self.pb_col] >= 1) & (df[self.pb_col] <= PB_MAX)).all():
            raise ValueError(f"红球列包含超出 1-{PB_MAX} 的值")

        if 'date' not in df.columns:
            possible_dates = [c for c in df.columns if 'date' in c.lower()]
            if possible_dates:
                df['date'] = pd.to_datetime(df[possible_dates[0]], errors='coerce')
            else:
                df['date'] = pd.NaT

        df = df.sort_values('date', na_position='last').reset_index(drop=True)
        self.df = df
        print(f"Loaded {len(self.df)} draws from {path}")
        return self

    def backtest_strategy(self, strategy: str = 'hybrid', n_tickets: int = 15, window: int = 50, constraints: Optional[Dict] = None, random_state: Optional[int] = None, step: int = 10) -> Tuple[pd.DataFrame, pd.Series, pd.Series]:
        if self.df is None:
            raise ValueError("请先 load_csv() 加载数据。")
        rng = np.random.default_rng(random_state)
        records = []
        white_perf = pd.Series(0, index=range(1, WHITE_MAX+1), dtype=int)
        pb_perf = pd.Series(0, index=range(1, PB_MAX+1), dtype=int)
        if len(self.df) < window + 1:
            raise ValueError("数据不足以进行回测。请减少 window 或加载更多数据。")
        for t in range(window, len(self.df), step):  # 使用步长减少迭代
            train = self.df.iloc[:t].copy()
            test_row = self.df.iloc[t]
            la_temp = LotteryAnalyzer()
            la_temp.df = train
            la_temp.white_cols = self.white_cols
            la_temp.pb_col = self.pb_col
            la_temp.compute_features(recent_window=50, backtest_window=window, use_backtest=False)
            picks = la_temp.generate_picks(n_tickets=n_tickets, strategy=strategy, constraints=constraints, random_state=int(rng.integers(0, 1000000000)))
            truth_whites = set([test_row[c] for c in self.white_cols if pd.notna(test_row[c])])
            truth_pb = test_row[self.pb_col]
            for _, row in picks.iterrows():
                whites = set([row['n1'], row['n2'], row['n3'], row['n4'], row['n5']])
                pb = row['pb']
                pb_match = int(pb == truth_pb)
                white_matches = len(whites & truth_whites)
                records.append({"t": t, "white_matches": white_matches, "pb_match": pb_match})
                for w in whites:
                    if w in truth_whites:
                        white_perf[w] += 1
                if pb_match:
                    pb_perf[pb] += 1
        out = pd.DataFrame(records)
        summary = out.groupby(["white_matches", "pb_match"]).size().rename("count").reset_index()
        return summary, white_perf, pb_perf

    @lru_cache(maxsize=128)  # 添加缓存
    def _compute_features_cached(self, recent_window: int, data_hash: int) -> Tuple[pd.DataFrame, pd.DataFrame]:
        total = len(self.df)

        # 频次计数（白球 1..69）
        white_counts = pd.Series(0, index=range(1, WHITE_MAX+1), dtype=int)
        for c in self.white_cols:
            vc = self.df[c].value_counts().reindex(range(1, WHITE_MAX+1)).fillna(0).astype(int)
            white_counts = white_counts.add(vc, fill_value=0).astype(int)
        pb_counts = self.df[self.pb_col].value_counts().reindex(range(1, PB_MAX+1)).fillna(0).astype(int)

        # 近期窗口
        window = min(recent_window, total)
        recent = self.df.tail(window)
        rec_w = pd.Series(0, index=range(1, WHITE_MAX+1), dtype=int)
        for c in self.white_cols:
            vc = recent[c].value_counts().reindex(range(1, WHITE_MAX+1)).fillna(0).astype(int)
            rec_w = rec_w.add(vc, fill_value=0).astype(int)
        rec_pb = recent[self.pb_col].value_counts().reindex(range(1, PB_MAX+1)).fillna(0).astype(int)

        # 优化后的 gap 计算（向量化）
        gap_w = pd.Series(np.inf, index=range(1, WHITE_MAX+1))
        gap_pb = pd.Series(np.inf, index=range(1, PB_MAX+1))
        last_seen_w = pd.Series(np.nan, index=range(1, WHITE_MAX+1))
        last_seen_pb = pd.Series(np.nan, index=range(1, PB_MAX+1))
        for num in range(1, WHITE_MAX+1):
            mask = self.df[self.white_cols].eq(num).any(axis=1)
            if mask.any():
                last_seen_w[num] = self.df.index[mask][::-1][0]
        for num in range(1, PB_MAX+1):
            mask = self.df[self.pb_col].eq(num)
            if mask.any():
                last_seen_pb[num] = self.df.index[mask][::-1][0]
        gap_w = (len(self.df) - last_seen_w).fillna(len(self.df)).astype(int)
        gap_pb = (len(self.df) - last_seen_pb).fillna(len(self.df)).astype(int)

        fw = pd.DataFrame({'freq': white_counts, 'recent': rec_w, 'gap': gap_w, 'performance': pd.Series(0, index=range(1, WHITE_MAX+1), dtype=int)})
        fpb = pd.DataFrame({'freq': pb_counts, 'recent': rec_pb, 'gap': gap_pb, 'performance': pd.Series(0, index=range(1, PB_MAX+1), dtype=int)})
        return fw, fpb

    def compute_features(self, recent_window: int = 50, backtest_window: int = 50, use_backtest: bool = True) -> "LotteryAnalyzer":
        if self.df is None:
            raise ValueError("请先 load_csv() 加载数据。")
        
        # 计算数据哈希用于缓存
        data_hash = hash(self.df.to_string())
        fw, fpb = self._compute_features_cached(recent_window, data_hash)

        # 回测性能
        if use_backtest:
            _, white_perf, pb_perf = self.backtest_strategy(strategy='hybrid', n_tickets=5, window=backtest_window, step=50)
            fw['performance'] = white_perf
            fpb['performance'] = pb_perf
        else:
            fw['performance'] = pd.Series(0, index=range(1, WHITE_MAX+1), dtype=int)
            fpb['performance'] = pd.Series(0, index=range(1, PB_MAX+1), dtype=int)

        # 向量化归一化
        def normalize_df(df):
            df = df.astype(float)
            result = df.copy()
            for col in df.columns:
                s = df[col]
                if s.max() == s.min():
                    result[col] = 0.0
                else:
                    result[col] = (s - s.min()) / (s.max() - s.min())
            return result

        fw_norm = normalize_df(fw)
        fpb_norm = normalize_df(fpb)

        w_freq, w_recent, w_gap, w_perf = self.weights
        white_score = fw_norm['freq']*w_freq + fw_norm['recent']*w_recent + fw_norm['gap']*w_gap + fw_norm['performance']*w_perf
        pb_score = fpb_norm['freq']*w_freq + fpb_norm['recent']*w_recent + fpb_norm['gap']*w_gap + fpb_norm['performance']*w_perf

        self.features_white = fw.assign(freq_norm=fw_norm['freq'], recent_norm=fw_norm['recent'],
                                       gap_norm=fw_norm['gap'], performance_norm=fw_norm['performance'], score=white_score)
        self.features_pb = fpb.assign(freq_norm=fpb_norm['freq'], recent_norm=fpb_norm['recent'],
                                      gap_norm=fpb_norm['gap'], performance_norm=fpb_norm['performance'], score=pb_score)

        whites_df = self.features_white.reset_index().rename(columns={'index':'number'}).assign(type='white')
        pbs_df = self.features_pb.reset_index().rename(columns={'index':'number'}).assign(type='pb')
        self.scores_combined = pd.concat([whites_df[['number','score','type']], pbs_df[['number','score','type']]], ignore_index=True)
        print("Computed features for white balls and Powerball")
        return self

    def rank_numbers(self, topn: Optional[int] = None) -> pd.DataFrame:
        if self.scores_combined is None:
            raise ValueError("请先调用 compute_features() 来计算特征。")
        df = self.scores_combined.sort_values('score', ascending=False).reset_index(drop=True)
        if topn:
            return df.head(topn)
        return df

    def generate_picks(self, n_tickets: int = 5, strategy: str = 'hybrid',
                       constraints: Optional[Dict] = None, random_state: Optional[int] = None) -> pd.DataFrame:
        if self.scores_combined is None:
            self.compute_features()
        rng = np.random.default_rng(random_state)

        w_scores = self.features_white['score'].reindex(range(1, WHITE_MAX+1)).fillna(0).astype(float).values
        p_scores = self.features_pb['score'].reindex(range(1, PB_MAX+1)).fillna(0).astype(float).values

        if strategy == 'hot':
            w_probs = np.where(w_scores<=0, 1e-6, w_scores)
            p_probs = np.where(p_scores<=0, 1e-6, p_scores)
        elif strategy == 'overdue':
            w_probs = self.features_white['gap'].reindex(range(1, WHITE_MAX+1)).fillna(0).values
            p_probs = self.features_pb['gap'].reindex(range(1, PB_MAX+1)).fillna(0).values
            w_probs = np.where(w_probs<=0, 1e-6, w_probs)
            p_probs = np.where(p_probs<=0, 1e-6, p_probs)
        elif strategy == 'random':
            w_probs = np.ones(WHITE_MAX)
            p_probs = np.ones(PB_MAX)
        else:  # hybrid
            w_probs = np.where(w_scores<=0, 1e-6, w_scores)
            p_probs = np.where(p_scores<=0, 1e-6, p_scores)

        w_probs = w_probs / w_probs.sum()
        p_probs = p_probs / p_probs.sum()

        picks = []
        tries = 0
        max_tries = 5000
        while len(picks) < n_tickets and tries < max_tries:
            tries += 1
            whites = rng.choice(np.arange(1, WHITE_MAX+1), size=5, replace=False, p=w_probs)
            whites = np.sort(whites).tolist()
            pb = int(rng.choice(np.arange(1, PB_MAX+1), p=p_probs))
            if self._satisfy_constraints(whites, pb, constraints):
                picks.append(whites + [pb])

        df_picks = pd.DataFrame(picks, columns=['n1','n2','n3','n4','n5','pb'])
        return df_picks

    def _satisfy_constraints(self, whites: List[int], pb: int, constraints: Optional[Dict]) -> bool:
        if not constraints:
            return True
        total = sum(whites)
        odd = sum(1 for w in whites if w % 2 == 1)
        spread = max(whites) - min(whites)
        if 'sum_range' in constraints:
            lo, hi = constraints['sum_range']
            if not (lo <= total <= hi):
                return False
        if 'odd_even' in constraints:
            lo, hi = constraints['odd_even']
            if not (lo <= odd <= hi):
                return False
        if 'min_spread' in constraints:
            if spread < constraints['min_spread']:
                return False
        return True

    def save_picks_csv(self, df_picks: pd.DataFrame, path: str = 'recommended_picks.csv'):
        df_picks.to_csv(path, index=False)
        print(f"Saved picks to {path}")