crpto_predictor/api/models.py

"""
数据库模型和操作类 - MySQL
"""
import pymysql
import os
from typing import Optional, List, Dict
from datetime import datetime
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 表"""
    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;
        ''')
        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:
        """
        获取历史交易记录
        
        Args:
            page: 页码
            limit: 每页条数
            month: 月份筛选 (YYYY-MM)
            day: 日期筛选 (YYYY-MM-DD)
        """
        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:
    """预测记录模型"""
    
    @staticmethod
    def get_prediction(date: str, hour: int, model_name: str) -> Optional[Dict]:
        """获取指定时间和模型的预测记录"""
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT * FROM model_predictions 
                WHERE date = %s AND hour = %s AND model_name = %s
            """, (date, hour, model_name))
            row = cursor.fetchone()
            return _dict_from_row(row)
    
    @staticmethod
    def save_prediction(data: Dict) -> int:
        """
        保存预测记录
        
        Args:
            data: 预测数据字典
        
        Returns:
            插入的记录ID
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            sql = """
                INSERT INTO model_predictions 
                (timestamp, timestamp_et, date, hour, model_name, 
                 prediction, probability, actual_result, is_correct,
                 open_price, close_price)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON DUPLICATE KEY UPDATE
                    prediction = VALUES(prediction),
                    probability = VALUES(probability),
                    actual_result = VALUES(actual_result),
                    is_correct = VALUES(is_correct),
                    open_price = VALUES(open_price),
                    close_price = VALUES(close_price)
            """
            cursor.execute(sql, (
                data.get('timestamp'),
                data.get('timestamp_et'),
                data.get('date'),
                data.get('hour'),
                data.get('model_name'),
                data.get('prediction'),
                data.get('probability'),
                data.get('actual_result'),
                data.get('is_correct'),
                data.get('open_price'),
                data.get('close_price')
            ))
            conn.commit()
            return cursor.lastrowid or 0
    
    @staticmethod
    def update_actual_result(date: str, hour: int, model_name: str,
                             actual_result: int, close_price: float) -> bool:
        """更新实际结果"""
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            # 先获取预测值
            cursor.execute("""
                SELECT prediction FROM model_predictions 
                WHERE date = %s AND hour = %s AND model_name = %s
            """, (date, hour, model_name))
            row = cursor.fetchone()
            
            if not row:
                return False
            
            prediction = row['prediction']
            is_correct = 1 if prediction == actual_result else 0
            
            cursor.execute("""
                UPDATE model_predictions 
                SET actual_result = %s, is_correct = %s, close_price = %s
                WHERE date = %s AND hour = %s AND model_name = %s
            """, (actual_result, is_correct, close_price, date, hour, model_name))
            conn.commit()
            return True
    
    @staticmethod
    def get_monthly_stats(month: str) -> Dict:
        """
        获取月度统计
        
        Args:
            month: 月份 (YYYY-MM)
        
        Returns:
            各模型的统计数据
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            sql = """
                SELECT 
                    model_name,
                    COUNT(*) as total,
                    SUM(CASE WHEN is_correct = 1 THEN 1 ELSE 0 END) as correct,
                    SUM(CASE WHEN is_correct = 0 THEN 1 ELSE 0 END) as incorrect
                FROM model_predictions 
                WHERE DATE_FORMAT(CONVERT_TZ(timestamp, '+00:00', '+08:00'), '%%Y-%%m') = %s
                  AND actual_result IS NOT NULL
                GROUP BY model_name
            """
            cursor.execute(sql, (month,))
            
            rows = cursor.fetchall()
            result = {}
            
            for row in rows:
                row_dict = _dict_from_row(row)
                model = row_dict['model_name']
                total = row_dict['total']
                correct = row_dict['correct'] or 0
                incorrect = row_dict['incorrect'] or 0
                accuracy = (correct / total * 100) if total > 0 else 0
                
                result[model] = {
                    "total": total,
                    "correct": correct,
                    "incorrect": incorrect,
                    "accuracy": round(accuracy, 2)
                }
            
            return result
    
    @staticmethod
    def get_daily_stats(month: str, page: int = 1, limit: int = 10) -> Dict:
        """
        获取每日统计
        
        Args:
            month: 月份 (YYYY-MM)
            page: 页码
            limit: 每页条数
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            
            sql = """
                SELECT 
                    DATE(CONVERT_TZ(timestamp, '+00:00', '+08:00')) as date_cst,
                    model_name,
                    COUNT(*) as total,
                    SUM(CASE WHEN is_correct = 1 THEN 1 ELSE 0 END) as correct,
                    SUM(CASE WHEN is_correct = 0 THEN 1 ELSE 0 END) as incorrect
                FROM model_predictions 
                WHERE DATE_FORMAT(CONVERT_TZ(timestamp, '+00:00', '+08:00'), '%%Y-%%m') = %s
                  AND actual_result IS NOT NULL
                GROUP BY date_cst, model_name
                ORDER BY date_cst DESC
            """
            cursor.execute(sql, (month,))
            
            rows = cursor.fetchall()
            
            # 重组数据: 按日期合并不同模型的统计
            daily_data = {}
            for row in rows:
                row_dict = _dict_from_row(row)
                date = str(row_dict['date_cst'])  # CST 日期
                if date not in daily_data:
                    daily_data[date] = {"date": date}
                
                model = row_dict['model_name']
                total = row_dict['total']
                correct = row_dict['correct'] or 0
                incorrect = row_dict['incorrect'] or 0
                accuracy = (correct / total * 100) if total > 0 else 0
                
                daily_data[date][model] = {
                    "total": total,
                    "correct": correct,
                    "incorrect": incorrect,
                    "accuracy": round(accuracy, 2)
                }
            
            # 排序并分页
            sorted_dates = sorted(daily_data.keys(), reverse=True)
            total_count = len(sorted_dates)
            offset = (page - 1) * limit
            paginated_dates = sorted_dates[offset:offset + limit]
            
            data = [daily_data[d] for d in paginated_dates]
            
            return {
                "data": 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(date: str) -> List[Dict]:
        """
        获取指定日期的每小时详情
        
        Args:
            date: 日期 (YYYY-MM-DD, CST 北京时间)
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            # 按 CST 日期筛选，并获取 CST 小时
            cursor.execute("""
                SELECT 
                    *,
                    HOUR(CONVERT_TZ(timestamp, '+00:00', '+08:00')) as hour_cst,
                    CONVERT_TZ(timestamp, '+00:00', '+08:00') as timestamp_cst_dt
                FROM model_predictions 
                WHERE DATE(CONVERT_TZ(timestamp, '+00:00', '+08:00')) = %s
                ORDER BY hour_cst, model_name
            """, (date,))
            
            rows = cursor.fetchall()
            
            # 按 CST 小时重组数据
            hourly_data = {}
            for row in rows:
                row_dict = _dict_from_row(row)
                hour_cst = row_dict['hour_cst']  # 使用 CST 小时
                if hour_cst not in hourly_data:
                    # 将 timestamp 转换为 CST 时间字符串
                    timestamp_cst = convert_to_cst(row_dict.get('timestamp', ''))
                    hourly_data[hour_cst] = {
                        "hour": hour_cst,  # 返回 CST 小时
                        "timestamp_et": row_dict['timestamp_et'],
                        "timestamp_cst": timestamp_cst,
                        "open_price": row_dict['open_price'],
                        "close_price": row_dict['close_price'],
                        "actual_result": row_dict['actual_result'],
                        "models": {}
                    }
                
                model = row_dict['model_name']
                hourly_data[hour_cst]["models"][model] = {
                    "prediction": row_dict['prediction'],
                    "probability": row_dict['probability'],
                    "is_correct": row_dict['is_correct']
                }
            
            # 按 CST 小时排序返回
            return [hourly_data[h] for h in sorted(hourly_data.keys())]
    
    @staticmethod
    def get_missing_hours(date: str, model_name: str) -> List[int]:
        """
        获取指定日期缺少预测的小时列表
        
        Args:
            date: 日期 (YYYY-MM-DD)
            model_name: 模型名称
        
        Returns:
            缺少预测的小时列表
        """
        with get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT hour FROM model_predictions 
                WHERE date = %s AND model_name = %s
            """, (date, model_name))
            
            rows = cursor.fetchall()
            existing_hours = {row['hour'] for row in rows}
            all_hours = set(range(24))
            
            return sorted(list(all_hours - existing_hours))