visualization_tgstat.py

import os
from PIL import Image, ImageDraw, ImageFont
from datetime import datetime
from config import OUTPUT_DIR, AVATAR_DIR
from utils import create_rounded_avatar
from wordcloud import WordCloud
import jieba
from collections import Counter
import numpy as np
import pandas as pd


class TelegramTGStatStyleVisualizer:
    def __init__(self, analyzer):
        self.analyzer = analyzer
        
        self.width = 1500
        self.height = 1300
        self.bg_color = '#FFFFFF'
        self.canvas = Image.new('RGB', (self.width, self.height), self.bg_color)
        self.draw = ImageDraw.Draw(self.canvas)

        self.font_path = os.path.join(os.path.dirname(__file__), 'ZhuZiAWan-2.ttc')
        
        self.font_title = ImageFont.truetype(self.font_path, 52)
        self.font_subtitle = ImageFont.truetype(self.font_path, 32)
        self.font_label = ImageFont.truetype(self.font_path, 28)
        self.font_note = ImageFont.truetype(self.font_path, 20)
        self.font_small = ImageFont.truetype(self.font_path, 16)

        self.colors = {
            'primary': '#2B6CB0',
            'secondary': '#10B981',
            'accent': '#F59E0B',
            'danger': '#EF4444',
            'purple': '#8B5CF6',
            'text_dark': '#1F2937',
            'text_medium': '#4B5563',
            'text_light': '#6B7280',
            'border': '#D1D5DB',
            'background': '#F9FAFB',
            'card_bg': '#FFFFFF'
        }

    def draw_user_info(self, user_data):
        
        avatar_size = 120
        avatar_x = 50
        avatar_y = 40

        
        avatar_rel_path = user_data.get('avatar_path')
        avatar_path = None
        if avatar_rel_path:
            
            potential_path = os.path.join(AVATAR_DIR, avatar_rel_path)
            if os.path.exists(potential_path) and os.path.isfile(potential_path):
                avatar_path = potential_path
            else:
                
                base_name = os.path.splitext(avatar_rel_path)[0]
                for ext in ['.jpg', '.png', '.jpeg']:
                    alt_path = os.path.join(AVATAR_DIR, base_name + ext)
                    if os.path.exists(alt_path) and os.path.isfile(alt_path):
                        avatar_path = alt_path
                        break

        avatar_bytes = None
        if avatar_path:
            try:
                with open(avatar_path, 'rb') as f:
                    avatar_bytes = f.read()
            except Exception as e:
                print(f"读取头像失败: {e}")

        avatar_img = create_rounded_avatar(
            avatar_bytes,
            size=avatar_size
        )
        self.canvas.paste(avatar_img, (avatar_x, avatar_y), avatar_img)

        
        name = user_data.get('first_name') or f"User {user_data.get('user_id', '未知')}"
        username = user_data.get('username')
        if username:
            name += f" (@{username})"
        
        self.draw.text((200, 50), name, font=self.font_title, fill=self.colors['primary'])
        self.draw.text((200, 110), f"ID: {user_data.get('user_id', '-')}", font=self.font_note, fill=self.colors['text_light'])
        self.draw.text((200, 135), f"DC: {user_data.get('dc_id', '-')}", font=self.font_note, fill=self.colors['text_light'])

    def draw_stat_block(self, x, y, title, value, subtitle=None, color='#2B6CB0', width=300, height=160):
        
        
        self.draw.rounded_rectangle(
            (x, y, x + width, y + height), 
            radius=12, 
            fill=self.colors['card_bg'], 
            outline=self.colors['border'], 
            width=2
        )
        
        
        self.draw.text((x + 20, y + 20), str(value), font=self.font_title, fill=color)
        
        self.draw.text((x + 20, y + 85), title, font=self.font_label, fill=self.colors['text_dark'])
        
        if subtitle:
            self.draw.text((x + 20, y + 120), subtitle, font=self.font_note, fill=self.colors['text_light'])

    def draw_all_stats(self):
        
        messages = self.analyzer.messages_df
        total_msgs = len(messages)
        unique_users = messages['sender_id'].nunique()
        total_chats = messages['chat_id'].nunique()
        private_chats = messages[messages['chat_type'] == 'private']['chat_id'].nunique()
        group_chats = messages[messages['chat_type'].isin(['group', 'supergroup'])]['chat_id'].nunique()
        
        
        if not messages.empty:
            date_range = (messages['timestamp'].max() - messages['timestamp'].min()) / (24 * 3600)
            daily_avg = int(total_msgs / max(date_range, 1))
        else:
            daily_avg = 0

        
        self.draw_stat_block(50, 200, '消息总数', total_msgs, f'日均 {daily_avg} 条')
        self.draw_stat_block(370, 200, '独立用户数', unique_users, color=self.colors['secondary'])
        self.draw_stat_block(690, 200, '活跃群组', group_chats, color=self.colors['accent'])
        self.draw_stat_block(1010, 200, '私聊对象数', private_chats, color=self.colors['purple'])

    def draw_top_groups_section(self):
        
        section_y = 400
        section_height = 400
        
        
        self.draw.rounded_rectangle(
            (50, section_y, self.width - 50, section_y + section_height),
            radius=12,
            fill=self.colors['card_bg'],
            outline=self.colors['border'],
            width=2
        )
        
        
        self.draw.text((70, section_y + 20), "最活跃群聊 TOP5", font=self.font_subtitle, fill=self.colors['primary'])
        
        
        top_groups = self.get_top_user_groups()
        if not top_groups:
            self.draw.text((70, section_y + 70), "暂无群聊数据", font=self.font_note, fill=self.colors['text_light'])
            return
        
        
        item_height = 50
        start_y = section_y + 70
        max_count = max([g['count'] for g in top_groups]) if top_groups else 1
        
        for i, group in enumerate(top_groups[:5]):
            y_pos = start_y + i * item_height
            
            
            avatar_size = 32
            avatar_x = 70
            avatar_y = y_pos + 4
            
            group_avatar = self.get_group_avatar(group['chat_id'])
            if group_avatar:
                try:
                    self.canvas.paste(group_avatar, (avatar_x, avatar_y), group_avatar)
                except:
                    pass
            
            
            group_name = group['name'][:30] + ('...' if len(group['name']) > 30 else '')
            self.draw.text((avatar_x + avatar_size + 15, y_pos + 2), 
                          f"#{i+1} {group_name}", 
                          font=self.font_label, 
                          fill=self.colors['text_dark'])
            
            
            count_text = f"{group['count']} 条消息"
            self.draw.text((avatar_x + avatar_size + 15, y_pos + 22), 
                          count_text, 
                          font=self.font_small, 
                          fill=self.colors['text_light'])
            
            
            progress_width = 200
            progress_height = 6
            progress_x = self.width - 300
            progress_y = y_pos + 15
            
            
            self.draw.rounded_rectangle(
                (progress_x, progress_y, progress_x + progress_width, progress_y + progress_height),
                radius=3,
                fill=self.colors['background']
            )
            
            
            fill_width = int((group['count'] / max_count) * progress_width)
            if fill_width > 0:
                color = [self.colors['primary'], self.colors['secondary'], self.colors['accent'], 
                        self.colors['purple'], self.colors['danger']][i % 5]
                self.draw.rounded_rectangle(
                    (progress_x, progress_y, progress_x + fill_width, progress_y + progress_height),
                    radius=3,
                    fill=color
                )

    def get_top_user_groups(self):
        
        if self.analyzer.messages_df is None or self.analyzer.messages_df.empty:
            return []
        
        
        messages = self.analyzer.messages_df
        user_id = messages['sender_id'].iloc[0] if not messages.empty else None
        
        if not user_id:
            return []
        
        
        group_messages = messages[
            (messages['chat_type'].isin(['group', 'supergroup'])) &
            (messages['sender_id'] == user_id)
        ]
        
        if group_messages.empty:
            return []
        
        
        group_stats = group_messages.groupby(['chat_id', 'chat_title']).size().reset_index(name='count')
        group_stats = group_stats.sort_values('count', ascending=False)
        
        
        result = []
        for _, row in group_stats.head(5).iterrows():
            result.append({
                'chat_id': row['chat_id'],
                'name': row['chat_title'] or f"群聊 {row['chat_id']}",
                'count': row['count']
            })
        
        return result

    def get_group_avatar(self, chat_id):
        
        avatar_filename = f"chat_{chat_id}.jpg"
        avatar_path = os.path.join(AVATAR_DIR, avatar_filename)
        
        if os.path.exists(avatar_path):
            try:
                with open(avatar_path, 'rb') as f:
                    avatar_bytes = f.read()
                return create_rounded_avatar(avatar_bytes, size=32)
            except:
                pass
        
        
        return create_rounded_avatar(None, size=32)

    def draw_activity_analysis(self):
        
        section_y = 720
        section_height = 180
        
        
        self.draw.rounded_rectangle(
            (50, section_y, self.width - 50, section_y + section_height),
            radius=12,
            fill=self.colors['card_bg'],
            outline=self.colors['border'],
            width=2
        )
        
        
        self.draw.text((70, section_y + 20), "聊天活动分析", font=self.font_subtitle, fill=self.colors['primary'])
        
        if self.analyzer.messages_df is None or self.analyzer.messages_df.empty:
            return
        
        messages = self.analyzer.messages_df
        
        
        total_chars = messages['text'].fillna('').str.len().sum()
        avg_msg_length = int(total_chars / len(messages)) if len(messages) > 0 else 0
        
        
        reply_count = messages['reply_to_msg_id'].notna().sum()
        reply_rate = int((reply_count / len(messages)) * 100) if len(messages) > 0 else 0
        
        
        media_count = messages['has_media'].sum() if 'has_media' in messages.columns else 0
        media_rate = int((media_count / len(messages)) * 100) if len(messages) > 0 else 0
        
        
        if 'timestamp' in messages.columns:
            messages['hour'] = pd.to_datetime(messages['timestamp'], unit='s').dt.hour
            most_active_hour = messages['hour'].mode().iloc[0] if not messages['hour'].empty else 12
            time_period = "上午" if 6 <= most_active_hour < 12 else "下午" if 12 <= most_active_hour < 18 else "晚上"
        else:
            time_period = "未知"
        
        
        stats_y = section_y + 65
        col_width = (self.width - 140) // 4
        
        stats_data = [
            ("平均消息长度", f"{avg_msg_length} 字符", self.colors['primary']),
            ("回复消息比例", f"{reply_rate}%", self.colors['secondary']),
            ("媒体消息比例", f"{media_rate}%", self.colors['accent']),
            ("最活跃时段", time_period, self.colors['purple'])
        ]
        
        for i, (label, value, color) in enumerate(stats_data):
            x_pos = 70 + i * col_width
            
            
            self.draw.text((x_pos, stats_y), value, font=self.font_label, fill=color)
            
            self.draw.text((x_pos, stats_y + 35), label, font=self.font_note, fill=self.colors['text_medium'])

    def draw_top_reply_users(self):
        
        if self.analyzer.top_reply_users is None or self.analyzer.top_reply_users.empty:
            return

        section_y = 920
        section_height = 320
        
        
        self.draw.rounded_rectangle(
            (50, section_y, 650, section_y + section_height),
            radius=12,
            fill=self.colors['card_bg'],
            outline=self.colors['border'],
            width=2
        )
        
        self.draw.text((70, section_y + 20), "互动最多的用户", font=self.font_subtitle, fill=self.colors['primary'])
        
        for i, (_, user) in enumerate(self.analyzer.top_reply_users.head(3).iterrows()):
            y_pos = section_y + 60 + i * 25
            text = f"#{i+1}  {user.get('user_name', '未知用户')}  ({user.get('reply_count', 0)} 次回复)"
            self.draw.text((70, y_pos), text, font=self.font_note, fill=self.colors['text_dark'])

    def draw_wordcloud(self):
        
        section_y = 920
        section_height = 320
        
        
        self.draw.rounded_rectangle(
            (670, section_y, self.width - 50, section_y + section_height),
            radius=12,
            fill=self.colors['card_bg'],
            outline=self.colors['border'],
            width=2
        )
        
        self.draw.text((690, section_y + 20), "消息词云", font=self.font_subtitle, fill=self.colors['primary'])
        
        try:
            messages = self.analyzer.messages_df
            if messages is None or messages.empty:
                return
                
            
            text = ' '.join(messages['text'].dropna().astype(str).tolist())
            if not text.strip():
                return
                
            words = jieba.lcut(text)
            
            stop_words = {'的', '了', '是', '我', '你', '他', '她', '它', '在', '有', '和', '就', '不', '这', '那', '也', '都', '要', '会', '可以', '没有', '什么', '怎么', '为什么', '因为', '所以'}
            word_freq = Counter(w for w in words if len(w) > 1 and w not in stop_words and not w.startswith('@'))
            
            if not word_freq:
                return
                
            
            wordcloud = WordCloud(
                font_path=self.font_path, 
                background_color='white', 
                width=760, 
                height=260,
                max_words=100,
                colormap='Set3'
            ).generate_from_frequencies(word_freq)
            
            wordcloud_img = wordcloud.to_image()
            
            self.canvas.paste(wordcloud_img, (690, section_y + 50))
        except Exception as e:
            print(f"生成词云失败: {e}")

    def draw_footer(self):
        
        footer_y = self.height - 50
        now = datetime.now().strftime('%Y-%m-%d %H:%M')
        
        
        self.draw.text((self.width - 300, footer_y), f"生成时间: {now}", 
                      font=self.font_note, fill=self.colors['text_light'])
        
        
        self.draw.text((50, footer_y), "Powered by Telestat", 
                      font=self.font_note, fill=self.colors['text_light'])

    def save(self, output_path):
        
        self.canvas.save(output_path, quality=95)

    def generate(self, output_path):
        
        
        if self.analyzer.users_df is None or self.analyzer.users_df.empty:
            raise ValueError("用户数据为空，无法生成可视化")

        user_data = self.analyzer.users_df.iloc[0].to_dict()
        
        
        self.draw_user_info(user_data)
        self.draw_all_stats()
        self.draw_top_groups_section()
        self.draw_activity_analysis()
        self.draw_top_reply_users()
        self.draw_wordcloud()
        self.draw_footer()
        
        self.save(output_path)