"""数据处理模块 - 使用 Polars 进行高效数据处理"""

import polars as pl
from pydantic import BaseModel, Field
from typing import List, Optional
import json

class TweetFeatures(BaseModel):
    """推文特征模型"""
    tweet_id: str
    airline: str
    text: str
    sentiment: str
    confidence: float
    negativereason: Optional[str] = None
    retweet_count: int
    user_name: str
    
    class Config:
        """配置"""
        from_attributes = True

class SentimentPrediction(BaseModel):
    """情感预测结果模型"""
    tweet_id: str
    predicted_sentiment: str
    confidence_score: float
    risk_level: str
    risk_factors: List[str]

class PolarsDataProcessor:
    """使用 Polars 进行数据处理的类"""
    
    def __init__(self, data_path: str):
        """初始化"""
        self.data_path = data_path
        self.df = None
    
    def load_data(self) -> pl.DataFrame:
        """加载数据"""
        self.df = pl.read_csv(self.data_path)
        return self.df
    
    def preprocess_data(self) -> pl.DataFrame:
        """预处理数据"""
        if self.df is None:
            self.load_data()
        
        # 选择相关列
        processed_df = self.df.select([
            "tweet_id",
            "airline",
            "text",
            "airline_sentiment",
            "airline_sentiment_confidence",
            "negativereason",
            "retweet_count",
            "name"
        ])
        
        # 重命名列
        processed_df = processed_df.rename({
            "airline_sentiment": "sentiment",
            "airline_sentiment_confidence": "confidence",
            "name": "user_name"
        })
        
        # 处理缺失值
        processed_df = processed_df.fill_null({
            "negativereason": "Unknown"
        })
        
        return processed_df
    
    def extract_features(self, df: pl.DataFrame) -> pl.DataFrame:
        """提取特征"""
        # 文本长度特征
        df = df.with_columns([
            pl.col("text").str.lengths().alias("text_length"),
            pl.col("text").str.count_matches(r"\b(?:delay|cancel|lost|baggage|late)\b", case_sensitive=False).alias("negative_keywords_count")
        ])
        
        return df
    
    def convert_to_pydantic(self, df: pl.DataFrame) -> List[TweetFeatures]:
        """转换为 Pydantic 模型列表"""
        records = df.to_dicts()
        return [TweetFeatures(**record) for record in records]
    
    def batch_process(self, batch_size: int = 1000) -> List[SentimentPrediction]:
        """批量处理数据"""
        processed_df = self.preprocess_data()
        features_df = self.extract_features(processed_df)
        
        # 模拟预测结果
        predictions = []
        for row in features_df.iter_rows(named=True):
            prediction = SentimentPrediction(
                tweet_id=str(row["tweet_id"]),
                predicted_sentiment=row["sentiment"],
                confidence_score=row["confidence"],
                risk_level="high" if row["sentiment"] == "negative" else "low",
                risk_factors=["negative sentiment"] if row["sentiment"] == "negative" else []
            )
            predictions.append(prediction)
        
        return predictions

# 示例用法
if __name__ == "__main__":
    processor = PolarsDataProcessor("../data/Tweets.csv")
    df = processor.load_data()
    print(f"Loaded {len(df)} tweets")
    
    processed_df = processor.preprocess_data()
    print(f"Processed data shape: {processed_df.shape}")
    
    features_df = processor.extract_features(processed_df)
    print(f"Features extracted: {features_df.columns}")
    
    samples = processor.convert_to_pydantic(features_df.head(5))
    print(f"Converted to Pydantic models: {len(samples)}")
    
    predictions = processor.batch_process(batch_size=10)
    print(f"Generated {len(predictions)} predictions")
    print(json.dumps(predictions[0].model_dump(), ensure_ascii=False, indent=2))