crpto_predictor/api/models.py

"""
数据库模型和操作类 - MySQL
"""
import pymysql
import os
from typing import Optional, List, Dict
from datetime import datetime, timedelta
from contextlib import contextmanager
from zoneinfo import ZoneInfo
from dotenv import load_dotenv

load_dotenv()

# 时区定义
TZ_UTC = ZoneInfo('UTC')
TZ_ET = ZoneInfo('America/New_York')  # 美东时间 (Polymarket)
TZ_CST = ZoneInfo('Asia/Shanghai')     # 中国标准时间 (前端显示)


def convert_to_cst(timestamp_val, source_tz: str = 'UTC') -> str:
    """
    将时间戳转换为 CST 时间字符串
    
    Args:
        timestamp_val: datetime 对象或 ISO 格式时间戳字符串
        source_tz: 输入时区 - 'UTC'(默认，币安数据), 'ET'(Polymarket), 'CST'(已是北京时间)
    
    Returns:
        CST 时间字符串 (YYYY-MM-DD HH:MM:SS)
    """
    if not timestamp_val:
        return ""
    try:
        # 如果是 datetime 对象 (MySQL 返回)
        if isinstance(timestamp_val, datetime):
            dt = timestamp_val
        else:
            # 字符串格式
            dt = datetime.fromisoformat(str(timestamp_val).replace('Z', '+00:00'))
        
        # 根据来源时区设置时区信息
        if dt.tzinfo is None:
            if source_tz == 'ET':
                dt = dt.replace(tzinfo=TZ_ET)
            elif source_tz == 'CST':
                # 已经是 CST，直接格式化返回
                return dt.strftime('%Y-%m-%d %H:%M:%S')
            else:  # 默认 UTC
                dt = dt.replace(tzinfo=TZ_UTC)
        
        # 如果来源已经是 CST，直接格式化
        if source_tz == 'CST':
            return dt.strftime('%Y-%m-%d %H:%M:%S')
        
        # 转换为 CST
        dt_cst = dt.astimezone(TZ_CST)
        return dt_cst.strftime('%Y-%m-%d %H:%M:%S')
    except Exception:
        return str(timestamp_val)


@contextmanager
def get_db_connection():
    """获取 MySQL 数据库连接"""
    conn = pymysql.connect(
        host=os.getenv('MYSQL_HOST'),
        port=int(os.getenv('MYSQL_PORT', 3306)),
        user=os.getenv('MYSQL_USER'),
        password=os.getenv('MYSQL_PASSWORD'),
        database=os.getenv('MYSQL_DB'),
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )
    try:
        yield conn
    finally:
        conn.close()


def _dict_from_row(row):
    """将数据库行转换为字典"""
    if row is None:
        return None
    if isinstance(row, dict):
        return row
    return dict(row)


def init_model_predictions_table():
    """初始化 model_predictions 表"""
    # 兼容旧表，但主要使用 hourly_model_predictions
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS model_predictions (
                id INT AUTO_INCREMENT PRIMARY KEY,
                timestamp DATETIME NOT NULL,
                timestamp_et VARCHAR(50),
                date VARCHAR(10),
                hour INT,
                model_name VARCHAR(50) NOT NULL,
                prediction INT,
                probability DOUBLE,
                actual_result INT,
                is_correct INT,
                open_price DOUBLE,
                close_price DOUBLE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE KEY uk_date_hour_model (date, hour, model_name),
                INDEX idx_date (date),
                INDEX idx_model (model_name)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        ''')
        
        # 确保 hourly_model_predictions 表存在
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS hourly_model_predictions (
                id INT AUTO_INCREMENT PRIMARY KEY,
                timestamp DATETIME NOT NULL,
                timestamp_et VARCHAR(50),
                date VARCHAR(10),
                hour INT,
                
                open_price DOUBLE,
                close_price DOUBLE,
                actual_result INT COMMENT '0=Down, 1=Up',
                
                dlinear_prediction INT,
                dlinear_prob DOUBLE,
                dlinear_correct INT,
                
                nbeats_prediction INT,
                nbeats_prob DOUBLE,
                nbeats_correct INT,
                
                ttm_prediction INT,
                ttm_prob DOUBLE,
                ttm_correct INT,
                
                ensemble_prediction INT,
                ensemble_votes VARCHAR(50),
                ensemble_correct INT,
                
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                
                UNIQUE KEY uk_date_hour (date, hour),
                INDEX idx_timestamp (timestamp)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
        ''')
        conn.commit()


# 初始化表
init_model_predictions_table()


class TradeModel:
    """交易记录模型"""
    
    @staticmethod
    def get_current_trade() -> Optional[Dict]:
        """获取最新的交易记录"""
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT * FROM trades 
                ORDER BY timestamp DESC 
                LIMIT 1
            """)
            row = cursor.fetchone()
            result = _dict_from_row(row)
            if result and result.get('timestamp'):
                # 交易记录时间已经是 CST，直接格式化
                result['timestamp_cst'] = convert_to_cst(result['timestamp'], source_tz='CST')
            return result
    
    @staticmethod
    def get_history(page: int = 1, limit: int = 10, 
                    month: Optional[str] = None,
                    day: Optional[str] = None) -> Dict:
        """
        获取历史交易记录
        """
        offset = (page - 1) * limit
        
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # 构建查询条件
            where_clauses = []
            params = []
            
            if month:
                where_clauses.append("DATE_FORMAT(timestamp, '%Y-%m') = %s")
                params.append(month)
            
            if day:
                where_clauses.append("DATE(timestamp) = %s")
                params.append(day)
            
            where_sql = " AND ".join(where_clauses) if where_clauses else "1=1"
            
            # 查询数据
            sql = f"SELECT * FROM trades WHERE {where_sql} ORDER BY timestamp DESC LIMIT %s OFFSET %s"
            cursor.execute(sql, params + [limit, offset])
            rows = cursor.fetchall()
            
            # 查询总数
            count_sql = f"SELECT COUNT(*) as cnt FROM trades WHERE {where_sql}"
            cursor.execute(count_sql, params)
            result = cursor.fetchone()
            total = result['cnt'] if result else 0
            
            # 转换每条记录的时间戳
            data_list = []
            for row in rows:
                item = _dict_from_row(row)
                if item and item.get('timestamp'):
                    # 交易记录时间已经是 CST，直接格式化
                    item['timestamp_cst'] = convert_to_cst(item['timestamp'], source_tz='CST')
                data_list.append(item)
            
            return {
                "data": data_list,
                "page": page,
                "limit": limit,
                "total": total,
                "total_pages": (total + limit - 1) // limit if limit > 0 else 0
            }


class PredictionModel:
    """预测记录模型 (V2 - 宽表结构)"""
    
    @staticmethod
    def save_hourly_prediction(data: Dict) -> int:
        """
        保存每小时综合预测记录 (宽表)
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            sql = """
                INSERT INTO hourly_model_predictions 
                (timestamp, timestamp_et, date, hour, open_price, close_price, actual_result,
                 dlinear_prediction, dlinear_prob, dlinear_correct,
                 nbeats_prediction, nbeats_prob, nbeats_correct,
                 ttm_prediction, ttm_prob, ttm_correct,
                 ensemble_prediction, ensemble_votes, ensemble_correct)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    open_price = VALUES(open_price),
                    close_price = VALUES(close_price),
                    actual_result = VALUES(actual_result),
                    dlinear_prediction = VALUES(dlinear_prediction),
                    dlinear_prob = VALUES(dlinear_prob),
                    dlinear_correct = VALUES(dlinear_correct),
                    nbeats_prediction = VALUES(nbeats_prediction),
                    nbeats_prob = VALUES(nbeats_prob),
                    nbeats_correct = VALUES(nbeats_correct),
                    ttm_prediction = VALUES(ttm_prediction),
                    ttm_prob = VALUES(ttm_prob),
                    ttm_correct = VALUES(ttm_correct),
                    ensemble_prediction = VALUES(ensemble_prediction),
                    ensemble_votes = VALUES(ensemble_votes),
                    ensemble_correct = VALUES(ensemble_correct)
            """
            cursor.execute(sql, (
                data.get('timestamp'),
                data.get('timestamp_et'),
                data.get('date'),
                data.get('hour'),
                data.get('open_price'),
                data.get('close_price'),
                data.get('actual_result'),
                data.get('dlinear_prediction'),
                data.get('dlinear_prob'),
                data.get('dlinear_correct'),
                data.get('nbeats_prediction'),
                data.get('nbeats_prob'),
                data.get('nbeats_correct'),
                data.get('ttm_prediction'),
                data.get('ttm_prob'),
                data.get('ttm_correct'),
                data.get('ensemble_prediction'),
                str(data.get('ensemble_votes')) if data.get('ensemble_votes') else None,
                data.get('ensemble_correct')
            ))
            conn.commit()
            return cursor.lastrowid or 0

    @staticmethod
    def update_actual_result_v2(date: str, hour: int, actual_result: int, close_price: float) -> bool:
        """
        更新实际结果，并重新计算各模型正确性 (V2)
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # 获取该小时记录
            cursor.execute("""
                SELECT * FROM hourly_model_predictions 
                WHERE date = %s AND hour = %s
            """, (date, hour))
            row = cursor.fetchone()
            
            if not row:
                return False
            
            # 计算正确性
            updates = {
                'actual_result': actual_result,
                'close_price': close_price
            }
            
            # DLinear
            if row['dlinear_prediction'] is not None:
                updates['dlinear_correct'] = 1 if row['dlinear_prediction'] == actual_result else 0
                
            # N-BEATS
            if row['nbeats_prediction'] is not None:
                updates['nbeats_correct'] = 1 if row['nbeats_prediction'] == actual_result else 0
                
            # TTM
            if row['ttm_prediction'] is not None:
                updates['ttm_correct'] = 1 if row['ttm_prediction'] == actual_result else 0
                
            # Ensemble
            if row['ensemble_prediction'] is not None:
                updates['ensemble_correct'] = 1 if row['ensemble_prediction'] == actual_result else 0
            
            # 构建更新SQL
            set_clause = ", ".join([f"{k} = %s" for k in updates.keys()])
            sql = f"UPDATE hourly_model_predictions SET {set_clause} WHERE date = %s AND hour = %s"
            
            cursor.execute(sql, list(updates.values()) + [date, hour])
            conn.commit()
            return True

    @staticmethod
    def get_monthly_stats_v2(month: str) -> Dict:
        """
        获取月度统计 (V2 - 从宽表)
        修改：不使用 SQL 进行时区转换和格式化，而是查询 UTC 时间范围，然后在 Python 中处理
        """
        # 计算 UTC 时间范围
        try:
            year, mon = map(int, month.split('-'))
            start_date_cst = datetime(year, mon, 1, tzinfo=TZ_CST)
            # 下个月第一天
            if mon == 12:
                end_date_cst = datetime(year + 1, 1, 1, tzinfo=TZ_CST)
            else:
                end_date_cst = datetime(year, mon + 1, 1, tzinfo=TZ_CST)
            
            # 转换为 UTC
            start_dt_utc = start_date_cst.astimezone(TZ_UTC)
            end_dt_utc = end_date_cst.astimezone(TZ_UTC)
            
            start_str = start_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
            end_str = end_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
            
        except Exception as e:
            print(f"日期解析错误: {e}")
            return {}

        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # 统计各个模型的准确率
            models = ['dlinear', 'nbeats', 'ttm', 'ensemble']
            result = {}
            
            for model in models:
                col_pred = f"{model}_prediction"
                col_correct = f"{model}_correct"
                
                # 查询 UTC 时间范围内的所有数据
                sql = f"""
                    SELECT 
                        COUNT({col_pred}) as total,
                        SUM(CASE WHEN {col_correct} = 1 THEN 1 ELSE 0 END) as correct,
                        SUM(CASE WHEN {col_correct} = 0 THEN 1 ELSE 0 END) as incorrect
                    FROM hourly_model_predictions 
                    WHERE timestamp >= %s AND timestamp < %s
                      AND actual_result IS NOT NULL
                """
                cursor.execute(sql, (start_str, end_str))
                row = cursor.fetchone()
                
                if row:
                    total = int(row['total'])
                    correct = int(row['correct'] or 0)
                    incorrect = int(row['incorrect'] or 0)
                    accuracy = (correct / total * 100) if total > 0 else 0
                    
                    # 格式化键名以匹配前端
                    model_key = {
                        'dlinear': 'DLinear', 
                        'nbeats': 'N-BEATS', 
                        'ttm': 'TTM-R2', 
                        'ensemble': 'Ensemble'
                    }.get(model, model)
                    
                    result[model_key] = {
                        "total": total,
                        "correct": correct,
                        "incorrect": incorrect,
                        "accuracy": round(accuracy, 2)
                    }
            
            return result

    @staticmethod
    def get_daily_stats_v2(month: str, page: int = 1, limit: int = 10) -> Dict:
        """
        获取每日统计 (V2 - 从宽表)
        修改：不使用 SQL 进行时区转换，查询 UTC 范围，然后在 Python 中转换
        """
        # 计算 UTC 时间范围
        try:
            year, mon = map(int, month.split('-'))
            start_date_cst = datetime(year, mon, 1, tzinfo=TZ_CST)
            if mon == 12:
                end_date_cst = datetime(year + 1, 1, 1, tzinfo=TZ_CST)
            else:
                end_date_cst = datetime(year, mon + 1, 1, tzinfo=TZ_CST)
            
            start_dt_utc = start_date_cst.astimezone(TZ_UTC)
            end_dt_utc = end_date_cst.astimezone(TZ_UTC)
            
            start_str = start_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
            end_str = end_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
        except Exception:
            return {"data": [], "page": page, "limit": limit, "total": 0, "total_pages": 0}

        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # 获取范围内的所有原始数据
            # 我们需要 timestamp, actual_result, 和各模型的 prediction/correct
            # 为了按 CST 日期分组，我们需要取出每一行，然后在 Python 中转换时间
            
            sql = """
                SELECT 
                    timestamp,
                    actual_result,
                    dlinear_prediction, dlinear_correct,
                    nbeats_prediction, nbeats_correct,
                    ttm_prediction, ttm_correct,
                    ensemble_prediction, ensemble_correct
                FROM hourly_model_predictions
                WHERE timestamp >= %s AND timestamp < %s
                  AND actual_result IS NOT NULL
                ORDER BY timestamp DESC
            """
            cursor.execute(sql, (start_str, end_str))
            rows = cursor.fetchall()
            
            # 在 Python 中处理分组
            daily_stats = {} # { "YYYY-MM-DD": { "date": "...", "dlinear": {...}, ... } }
            
            for row in rows:
                # 转换 timestamp 到 CST
                ts_utc = row['timestamp']
                # 确保是带时区的 datetime
                if isinstance(ts_utc, str):
                    ts_utc = datetime.fromisoformat(str(ts_utc).replace('Z', '+00:00'))
                if ts_utc.tzinfo is None:
                    ts_utc = ts_utc.replace(tzinfo=TZ_UTC)
                
                ts_cst = ts_utc.astimezone(TZ_CST)
                date_str = ts_cst.strftime('%Y-%m-%d')
                
                if date_str not in daily_stats:
                    daily_stats[date_str] = {
                        "date": date_str,
                        "raw_data": {m: {"total": 0, "correct": 0} for m in ['dlinear', 'nbeats', 'ttm', 'ensemble']}
                    }
                
                # 累加统计
                models = ['dlinear', 'nbeats', 'ttm', 'ensemble']
                for model in models:
                    pred = row[f"{model}_prediction"]
                    correct_flag = row[f"{model}_correct"]
                    
                    if pred is not None:
                        daily_stats[date_str]['raw_data'][model]['total'] += 1
                        if correct_flag == 1:
                            daily_stats[date_str]['raw_data'][model]['correct'] += 1

            # 整理为最终格式
            sorted_dates = sorted(daily_stats.keys(), reverse=True)
            total_count = len(sorted_dates)
            
            # 分页
            offset = (page - 1) * limit
            paginated_dates = sorted_dates[offset:offset + limit]
            
            result_data = []
            for date_key in paginated_dates:
                day_item = daily_stats[date_key]
                formatted_item = {"date": day_item["date"]}
                
                for model in ['dlinear', 'nbeats', 'ttm', 'ensemble']:
                    raw = day_item['raw_data'][model]
                    total = int(raw['total'])
                    correct = int(raw['correct'])
                    accuracy = (correct / total * 100) if total > 0 else 0
                    
                    model_key = {
                        'dlinear': 'DLinear', 
                        'nbeats': 'N-BEATS', 
                        'ttm': 'TTM-R2', 
                        'ensemble': 'Ensemble'
                    }.get(model, model)
                    
                    formatted_item[model_key] = {
                        "total": total,
                        "correct": correct,
                        "accuracy": round(accuracy, 2)
                    }
                result_data.append(formatted_item)
            
            return {
                "data": result_data,
                "page": page,
                "limit": limit,
                "total": total_count,
                "total_pages": (total_count + limit - 1) // limit if limit > 0 else 0
            }

    @staticmethod
    def get_hourly_details_v2(date: str) -> List[Dict]:
        """
        获取指定日期的每小时详情 (V2)
        修改：根据传入的 CST 日期计算 UTC 范围，查询该范围内的数据
        """
        try:
            # date 是 'YYYY-MM-DD' (CST)
            date_cst_start = datetime.strptime(date, '%Y-%m-%d').replace(tzinfo=TZ_CST)
            date_cst_end = date_cst_start + timedelta(days=1)
            
            start_dt_utc = date_cst_start.astimezone(TZ_UTC)
            end_dt_utc = date_cst_end.astimezone(TZ_UTC)
            
            start_str = start_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
            end_str = end_dt_utc.strftime('%Y-%m-%d %H:%M:%S')
        except Exception as e:
            print(f"日期处理错误: {e}")
            return []

        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT *
                FROM hourly_model_predictions 
                WHERE timestamp >= %s AND timestamp < %s
                ORDER BY timestamp ASC
            """, (start_str, end_str))
            
            rows = cursor.fetchall()
            result = []
            
            for row in rows:
                # 计算 CST 小时用于显示排序 (虽然 SQL 已经按 timestamp 排序了，但为了确保前端显示 hour 字段正确)
                ts_utc = row['timestamp']
                if isinstance(ts_utc, str):
                    ts_utc = datetime.fromisoformat(str(ts_utc).replace('Z', '+00:00'))
                if ts_utc.tzinfo is None:
                    ts_utc = ts_utc.replace(tzinfo=TZ_UTC)
                ts_cst = ts_utc.astimezone(TZ_CST)
                hour_cst = ts_cst.hour
                
                hour_data = {
                    "hour": hour_cst,
                    "timestamp_et": row['timestamp_et'],
                    "timestamp_cst": convert_to_cst(row.get('timestamp', '')),
                    "open_price": row['open_price'],
                    "close_price": row['close_price'],
                    "actual_result": row['actual_result'],
                    "models": {
                        "DLinear": {
                            "prediction": row['dlinear_prediction'],
                            "probability": row['dlinear_prob'],
                            "is_correct": row['dlinear_correct']
                        },
                        "N-BEATS": {
                            "prediction": row['nbeats_prediction'],
                            "probability": row['nbeats_prob'],
                            "is_correct": row['nbeats_correct']
                        },
                        "TTM-R2": {
                            "prediction": row['ttm_prediction'],
                            "probability": row['ttm_prob'],
                            "is_correct": row['ttm_correct']
                        },
                        "Ensemble": {
                            "prediction": row['ensemble_prediction'],
                            "is_correct": row['ensemble_correct'],
                            "votes": row['ensemble_votes']
                        }
                    }
                }
                result.append(hour_data)
                
            return result

    @staticmethod
    def save_prediction(data: Dict) -> int:
        """兼容旧接口 - 如果需要的话，暂时保留，或者可以将旧表写入也放在这里"""
        # 旧表写入逻辑... 略，或保持不变
        # 为了完整性，我们可以保留对旧表的操作，或者逐步弃用。
        # 这里为了安全，保留原方法的签名，实际调用旧表逻辑。
        # 但我们主要逻辑将切换到 save_hourly_prediction
        # 见 prediction_service.py 的修改
        
        # ... (keep old logic or redirect) ...
        # 这里我们假设 prediction_service.py 会改为调用 save_hourly_prediction
        # 如果还有地方调用 save_prediction (比如测试代码), 我们可以保留它
        return 0 
    
    @staticmethod
    def update_actual_result(date: str, hour: int, model_name: str,
                             actual_result: int, close_price: float) -> bool:
        """兼容旧接口"""
        return False
    
    @staticmethod
    def get_missing_hours(date: str, model_name: str = 'Ensemble') -> List[int]:
        """
        获取指定日期缺少预测的小时列表 (V2)
        model_name 参数现在主要用于兼容，实际检查整行是否存在
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT hour FROM hourly_model_predictions 
                WHERE date = %s
            """, (date,))
            
            rows = cursor.fetchall()
            existing_hours = {row['hour'] for row in rows}
            all_hours = set(range(24))
            
            return sorted(list(all_hours - existing_hours))