HKDataManagment/main_gui.py

"""
香港股票数据管理系统 - GUI界面

基于PyQt5的图形用户界面，通过按钮控制各个功能模块的执行
"""
import sys
import os
from PyQt5.QtWidgets import (QApplication, QMainWindow, QPushButton, 
                             QTextEdit, QVBoxLayout, QWidget, QProgressBar,
                             QHBoxLayout, QGroupBox, QLabel, QMessageBox, QComboBox)
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QFont

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from UpdateFutuData.KLineUpdater import KLineUpdater
from base.LogHelper import LogHelper
from base import Config
from DataAnalysis import DataExporter, MarketDataCalculator
from DataAnalysis.checktable import StockTableChecker
from UpdateFutuData.ConditionalSelection import FutuStockFilter
from datetime import datetime, timedelta
from tqdm import tqdm
from pathlib import Path
import time

class LogHandler:
    """日志处理器，用于在GUI中显示日志"""
    def __init__(self, text_widget):
        self.text_widget = text_widget
        
    def write(self, message):
        if message.strip():
            self.text_widget.append(f"{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} - {message}")
            
    def flush(self):
        pass

class WorkerThread(QThread):
    """工作线程类，用于执行耗时操作"""
    log_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int)
    finished_signal = pyqtSignal(bool, str)
    
    def __init__(self, task_func, *args, **kwargs):
        super().__init__()
        self.task_func = task_func
        self.args = args
        self.kwargs = kwargs
        
    def run(self):
        try:
            self.log_signal.emit("开始执行任务...")
            result = self.task_func(*self.args, **self.kwargs)
            self.finished_signal.emit(True, "任务执行成功")
        except Exception as e:
            self.log_signal.emit(f"执行失败: {str(e)}")
            self.finished_signal.emit(False, f"任务执行失败: {str(e)}")

class MainWindow(QMainWindow):
    """主窗口类"""
    
    def __init__(self):
        super().__init__()
        self.worker_threads = []
        self.initUI()
        
    def initUI(self):
        """初始化用户界面"""
        self.setWindowTitle('白泽数科数据管理平台')
        self.setGeometry(100, 100, 1000, 800)
        
        # 设置字体
        font = QFont("Microsoft YaHei", 10)
        self.setFont(font)
        
        # 应用白色主题样式
        self.apply_light_theme()
        
        # 创建中央部件和布局
        central_widget = QWidget()
        main_layout = QVBoxLayout()
        
        # 创建标题
        title_label = QLabel("港股数据管理系统")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setStyleSheet("""
            font-size: 18px; 
            font-weight: bold; 
            margin: 10px;
            color: #333333;
            background-color: #e0e0e0;
            padding: 10px;
            border-radius: 5px;
        """)
        main_layout.addWidget(title_label)
        
        # 创建配置选择组
        self.create_config_selection_group(main_layout)
        
        # 创建功能按钮组
        self.create_button_group(main_layout)
        
        # 创建数据导入按钮组
        self.create_import_button_group(main_layout)
        
        # 创建进度条
        self.progress_bar = QProgressBar()
        # self.progress_bar.setVisible(False)
        main_layout.addWidget(self.progress_bar)
        self.progress_bar.setStyleSheet("""
            QProgressBar {
                border: 2px solid grey;
                border-radius: 5px;
                text-align: center;      /* 文本水平居中 */
                color: black;            /* 设置文本颜色，确保与背景对比明显 */
                background-color: #FFFFFF; /* 进度条背景色 */
            }
            QProgressBar::chunk {
                background-color: #05B8CC; /* 进度块颜色 */
                width: 20px;               /* 进度块宽度 */
                margin: 0.5px;             /* 进度块间隔 */
            }
        """)
        
        # 创建日志显示区域
        self.create_log_area(main_layout)
        
        # 设置状态栏
        self.statusBar().showMessage('就绪')
        
        central_widget.setLayout(main_layout)
        self.setCentralWidget(central_widget)
        
    def create_config_selection_group(self, layout):
        """创建配置选择组"""
        config_group = QGroupBox("读取股票列表")
        config_layout = QHBoxLayout()
        
        # 配置文件选择标签
        config_label = QLabel("选择股票列表配置文件:")
        config_layout.addWidget(config_label)
        
        # 配置文件选择下拉框
        self.config_combo = QComboBox()
        self.config_combo.setToolTip("选择要使用的股票列表配置文件")
        
        # 获取config目录下的所有txt文件
        config_dir = Path("config")
        if config_dir.exists():
            txt_files = [f for f in config_dir.glob("*.txt") if f.is_file()]
            for file in txt_files:
                self.config_combo.addItem(file.name, str(file))
        
        # 如果没有找到文件，添加默认选项
        if self.config_combo.count() == 0:
            self.config_combo.addItem("未找到配置文件", "")
        
        config_layout.addWidget(self.config_combo)
        
        # 当前选择显示
        self.selected_config_label = QLabel("当前选择: 无")
        config_layout.addWidget(self.selected_config_label)
        
        # 连接选择变化信号
        self.config_combo.currentIndexChanged.connect(self.on_config_changed)
        
        config_group.setLayout(config_layout)
        layout.addWidget(config_group)
        
    def on_config_changed(self, index):
        """配置文件选择变化事件"""
        if index >= 0:
            file_path = self.config_combo.itemData(index)
            if file_path:
                self.selected_config_label.setText(f"当前选择: {self.config_combo.itemText(index)}")
                self.log_message(f"已选择配置文件: {self.config_combo.itemText(index)}")
            else:
                self.selected_config_label.setText("当前选择: 无")
        
    def apply_light_theme(self):
        """应用白色主题样式"""
        light_stylesheet = """
            QMainWindow {
                background-color: #f5f5f5;
                color: #333333;
            }
            QWidget {
                background-color: #f5f5f5;
                color: #333333;
            }
            QPushButton {
                background-color: #007acc;
                color: #ffffff;
                border: none;
                padding: 8px 16px;
                border-radius: 4px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #005a9e;
            }
            QPushButton:pressed {
                background-color: #004275;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
            QGroupBox {
                font-weight: bold;
                border: 1px solid #cccccc;
                border-radius: 5px;
                margin-top: 1ex;
                padding-top: 10px;
                background-color: #ffffff;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                subcontrol-position: top left;
                padding: 0 8px;
                color: #333333;
            }
            QTextEdit {
                background-color: #ffffff;
                color: #333333;
                border: 1px solid #cccccc;
                border-radius: 3px;
                selection-background-color: #007acc;
                selection-color: #ffffff;
            }
            QLabel {
                color: #333333;
            }
            QComboBox {
                background-color: #ffffff;
                color: #333333;
                border: 1px solid #cccccc;
                border-radius: 3px;
                padding: 5px;
            }
            QComboBox:drop-down {
                subcontrol-origin: padding;
                subcontrol-position: top right;
                width: 20px;
                border-left-width: 1px;
                border-left-color: #cccccc;
                border-left-style: solid;
            }
            QComboBox QAbstractItemView {
                background-color: #ffffff;
                color: #333333;
                selection-background-color: #007acc;
                selection-color: #ffffff;
            }
            QProgressBar {
                border: 1px solid #cccccc;
                border-radius: 3px;
                text-align: center;
                background-color: #ffffff;
                color: #333333;
            }
            QProgressBar::chunk {
                background-color: #007acc;
                width: 10px;
            }
            QStatusBar {
                background-color: #ffffff;
                color: #333333;
                border-top: 1px solid #cccccc;
            }
        """
        self.setStyleSheet(light_stylesheet)
        
    def create_button_group(self, layout):
        """创建功能按钮组"""
        button_group = QGroupBox("功能操作")
        button_layout = QHBoxLayout()
        
        # 更新流通股数量更新流通股数量
        self.btn_float_share = QPushButton('更新流通股数量')
        self.btn_float_share.clicked.connect(self.on_share_clicked)
        self.btn_float_share.setToolTip('更新流通股数量')
        button_layout.addWidget(self.btn_float_share)

        # 检查数据按钮,检测列表中的股票是不是全部下载好了k线数据
        self.btn_check = QPushButton('检查数据完整性')
        self.btn_check.clicked.connect(self.on_check_clicked)
        self.btn_check.setToolTip('检查数据表的完整性')
        button_layout.addWidget(self.btn_check)

        # 更新数据按钮，根据登录OpenD账号的不同，加载不同的股票列表，用于更新K线数据
        self.btn_update = QPushButton('更新K线数据')
        self.btn_update.clicked.connect(self.on_update_clicked)
        self.btn_update.setToolTip('从Futu API获取并更新股票K线数据')
        button_layout.addWidget(self.btn_update)

        # 计算月度平均按钮
        self.btn_calculate = QPushButton('计算月度平均')
        self.btn_calculate.clicked.connect(self.on_calculate_clicked)
        self.btn_calculate.setToolTip('计算股票的月度平均数据')
        button_layout.addWidget(self.btn_calculate)
        
        # 导出数据按钮
        self.btn_export = QPushButton('导出数据')
        self.btn_export.clicked.connect(self.on_export_clicked)
        self.btn_export.setToolTip('导出月度平均数据到CSV文件')
        button_layout.addWidget(self.btn_export)
        
        button_group.setLayout(button_layout)
        layout.addWidget(button_group)

    def create_import_button_group(self, layout):
        """创建数据导入按钮组"""
        import_group = QGroupBox("数据导入")
        import_layout = QHBoxLayout()
        
        # 导入数据按钮
        self.btn_import = QPushButton('导入股票数据')
        self.btn_import.clicked.connect(self.on_import_clicked)
        self.btn_import.setToolTip('从CSV文件导入股票数据到数据库')
        import_layout.addWidget(self.btn_import)
        
        import_group.setLayout(import_layout)
        layout.addWidget(import_group)
        
    def create_log_area(self, layout):
        """创建日志显示区域"""
        log_group = QGroupBox("操作日志")
        log_layout = QVBoxLayout()
        
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        # self.log_text.setMaximumHeight(400)
        log_layout.addWidget(self.log_text)
        
        # 清空日志按钮
        btn_clear = QPushButton('清空日志')
        btn_clear.clicked.connect(self.on_clear_log)
        log_layout.addWidget(btn_clear)
        
        log_group.setLayout(log_layout)
        layout.addWidget(log_group)
        
    def log_message(self, message):
        """添加日志消息"""
        self.log_text.append(f"{datetime.now().strftime('%H:%M:%S')} - {message}")
        
    def on_clear_log(self):
        """清空日志"""
        self.log_text.clear()
        self.log_message("日志已清空")
        
    def set_buttons_enabled(self, enabled):
        """设置按钮启用状态"""
        self.btn_update.setEnabled(enabled)
        self.btn_export.setEnabled(enabled)
        self.btn_calculate.setEnabled(enabled)
        self.btn_check.setEnabled(enabled)
        self.btn_import.setEnabled(enabled)
        self.btn_float_share.setEnabled(enabled)
        
    def on_update_clicked(self):
        """更新数据按钮点击事件"""
        self.log_message("开始更新K线数据...")
        self.set_buttons_enabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setRange(0, 0)  # 不确定进度

        self.statusBar().showMessage('正在更新数据')
        
        # 创建工作线程执行更新任务
        worker = WorkerThread(self.update_data)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)
        
    def on_export_clicked(self):
        """导出数据按钮点击事件"""
        self.log_message("开始导出数据...")
        self.set_buttons_enabled(False)
        
        worker = WorkerThread(self.export_data)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)
        
    def on_calculate_clicked(self):
        """计算月度平均按钮点击事件"""
        self.log_message("开始计算月度平均数据...")
        self.set_buttons_enabled(False)
        
        worker = WorkerThread(self.calculate_monthly_avg)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)
    
    def on_share_clicked(self):
        """更新流通股数量"""
        self.log_message("开始更新流通股数量...")
        self.set_buttons_enabled(False)
        
        worker = WorkerThread(self.update_float_share)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)

    def on_check_clicked(self):
        """检查数据按钮点击事件"""
        self.log_message("开始检查数据完整性...")
        self.set_buttons_enabled(False)
        
        worker = WorkerThread(self.check_data)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)

    def on_import_clicked(self):
        """导入数据按钮点击事件"""
        self.log_message("开始导入股票数据...")
        self.set_buttons_enabled(False)
        
        worker = WorkerThread(self.import_stock_data)
        worker.log_signal.connect(self.log_message)
        worker.finished_signal.connect(self.on_task_finished)
        worker.start()
        self.worker_threads.append(worker)
        
    def on_task_finished(self, success, message):
        """任务完成回调"""
        self.set_buttons_enabled(True)
        # self.progress_bar.setVisible(False)
        
        if success:
            self.log_message(message)
            self.statusBar().showMessage('任务完成')
        else:
            self.log_message(message)
            self.statusBar().showMessage('任务失败')
            QMessageBox.warning(self, "警告", message)
            
    def update_data(self):
        """更新数据任务"""
        try:
            # 获取选择的配置文件路径
            current_index = self.config_combo.currentIndex()
            if current_index < 0:
                self.log_message("错误: 未选择配置文件")
                return False
                
            config_file_path = self.config_combo.itemData(current_index)
            if not config_file_path:
                self.log_message("错误: 配置文件路径无效")
                return False
                
            updater = KLineUpdater()
            self.log_message(f"使用配置文件: {Path(config_file_path).name}")
            
            # 从选择的配置文件读取股票代码
            stock_codes = updater.read_single_account_stock_codes(config_file_path)
            if not stock_codes:
                self.log_message("错误: 无法从配置文件读取股票代码或文件为空")
                return False
                
            self.log_message(f"从配置文件获取到 {len(stock_codes)} 个股票代码")
            
            # # 显示前几个代码作为示例
            # if stock_codes:
            #     sample_codes = stock_codes[:5]
            #     self.log_message(f"示例代码: {', '.join(sample_codes)}")
            
            # 更新数据
            updater.update_kline_data(stock_codes=stock_codes)
            return True
        except Exception as e:
            self.log_message(f"更新数据失败: {str(e)}")
            raise
            
    def export_data(self):
        """导出数据任务"""
        try:
            self.log_message("开始导出月度平均数据...")
            # 使用现有的导出函数
            exporter = DataExporter(Config.ConfigInfo.db_hk_kline_1d)
            
            # 根据导出时间命名
            target_table_name = 'hk_monthly_avg_' + datetime.now().strftime("%Y%m%d")
            target_file_name = 'hk_monthly_avg_' + datetime.now().strftime("%Y%m%d") + ".csv"
            
            success = exporter.export_data(
                monthly_table=target_table_name,
                csv_file=target_file_name
            )
            
            if success:
                self.log_message(f"数据导出成功: {target_file_name}")
                return True
            else:
                self.log_message("数据导出失败")
                return False
        except Exception as e:
            self.log_message(f"导出数据失败: {str(e)}")
            raise
            
    def calculate_monthly_avg(self):
        """计算月度平均任务"""
        try:
            self.log_message("开始计算月度平均数据...")
            calculator = MarketDataCalculator(Config.ConfigInfo.db_hk_kline_1d)
            
            # 移除人民币交易的股票：股票名称最后一个字符为R，误删除的从配置文件读回来
            reserved_codes = calculator.read_stock_codes_list(Path.cwd() / "config" / "Reservedcode.txt")
            remove_codes = calculator.read_stock_codes_list(Path.cwd() / "config" / "Removecode.txt")
            market_data_ll = calculator.get_stock_codes()  # 使用按照价格和流通股数量筛选的那个表格

            # market_data = list(set(market_data_ll) - set(remove_codes)) + reserved_codes
            market_data = market_data_ll + reserved_codes
            
            # 根据统计时间进行命名
            target_table_name = 'hk_monthly_avg_' + datetime.now().strftime("%Y%m%d")
            
            self.log_message(f"开始处理 {len(market_data)} 支股票...")
            
            # 使用tqdm创建进度条
            for code in tqdm(market_data, desc="处理股票数据", unit="支"):
                # 计算并保存月度均值
                calculator.calculate_and_save_monthly_avg(
                    stock_code =code,
                    target_table = target_table_name
                )

            # self.log_message("月度平均计算完成")
            return True
        except Exception as e:
            self.log_message(f"计算月度平均失败: {str(e)}")
            raise

    def update_float_share(self):
        """更新流通股数量任务"""
        try:
            futuStockFilter = FutuStockFilter(Config.ConfigInfo.db_hk_kline_1d)
            futuStockFilter.run_direct_import()
        except Exception as e:
            self.log_message(f"流通股数量更新失败: {str(e)}")
            raise

    def check_data(self):
        """检查数据任务"""
        try:
            self.log_message("开始检查数据完整性...")
            checker = StockTableChecker(Config.ConfigInfo.db_hk_kline_1d)
            checker.run_check()
            self.log_message("数据检查完成")
            return True
        except Exception as e:
            self.log_message(f"数据检查失败: {str(e)}")
            raise

    def import_stock_data(self):
        """导入股票数据任务"""
        try:
            # 导入必要的模块
            from base.StockDataImporter import StockDataImporter
            from base.MySQLHelper import MySQLHelper
            from pathlib import Path
            
            # 数据库配置
            db_config = {
                'host': 'localhost',
                'user': 'root',
                'password': 'bzskmysql',
                'database': 'hk_kline_1d'
            }
            
            # 列映射配置
            COLUMN_MAPPING = {
                '证券代码': 'stock_code',
                '中文简称': 'stock_name_chn',
                '英文简称': 'stock_name_en',
            }
            
            # 设置数据目录
            current_dir = Path(__file__).parent if "__file__" in locals() else Path.cwd()
            DATA_DIR = current_dir / "data"
            
            # 确保data目录存在
            DATA_DIR.mkdir(exist_ok=True, parents=True)
            
            # 创建导入器
            importer = StockDataImporter(db_config, COLUMN_MAPPING, DATA_DIR)
            
            # 设置上传文件名（使用默认文件名 "港股通标的证券名单.csv"）
            importer.setUploadfile("港股通标的证券名单.csv")
            
            # 执行导入
            if importer.run_import():
                self.log_message("股票数据导入成功!")
                return True
            else:
                self.log_message("股票数据导入失败，请检查日志了解详情")
                return False
                
        except Exception as e:
            self.log_message(f"导入股票数据失败: {str(e)}")
            raise
            
    def closeEvent(self, event):
        """窗口关闭事件"""
        # 确保所有工作线程都已完成
        for thread in self.worker_threads:
            if thread.isRunning():
                thread.terminate()
                thread.wait()
        event.accept()

def main():
    """主函数"""
    # 设置调试环境
    try:
        # 添加 .vscode 目录到 Python 路径
        vscode_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '.vscode')
        if vscode_path not in sys.path:
            sys.path.insert(0, vscode_path)
        
        from debug_helper import setup_debugging, enable_breakpoints
        setup_debugging()
        enable_breakpoints()
        print("Debugging environment initialized")
    except ImportError as e:
        print(f"Debug helper not available: {e}")
    except Exception as e:
        print(f"Error setting up debugging: {e}")
    
    app = QApplication(sys.argv)
    
    # 设置应用程序样式
    app.setStyle('Fusion')
    
    window = MainWindow()
    window.show()
    
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()