{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 量化投资数据批量导入\n",
    "\n",
    "## 目标\n",
    "将 Tushare 的日线行情数据及公司基本面数据批量导入 Docker PostgreSQL (192.168.27.15:12345)\n",
    "\n",
    "## 数据库结构概览\n",
    "\n",
    "| 表名 | 说明 | Tushare 接口 |\n",
    "|------|------|-------------|\n",
    "| stock_basic | 股票基本信息 | stock_basic |\n",
    "| trade_cal | 交易日历 | trade_cal |\n",
    "| daily | 日线行情 | daily |\n",
    "| daily_basic | 每日指标(估值/基本面) | daily_basic |\n",
    "| adj_factor | 复权因子 | adj_factor |\n",
    "| income | 利润表 | income |\n",
    "| balancesheet | 资产负债表 | balancesheet |\n",
    "| cashflow | 现金流量表 | cashflow |\n",
    "| fina_indicator | 财务指标 | fina_indicator |\n",
    "| moneyflow | 个股资金流向 | moneyflow |\n",
    "| index_daily | 指数日线行情 | index_daily |\n",
    "\n",
    "## 使用步骤\n",
    "1. 复制 `.env.example` 为 `.env`，填入真实密码和 Token\n",
    "2. 逐 Cell 运行本 Notebook（敏感信息通过环境变量读取）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 0: 检查环境 & 连接测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "os.chdir(r\"t:\\jupyter\\notebook\\quantitative_data\")\n",
    "print(f\"工作目录: {os.getcwd()}\")\n",
    "print(f\"Python 版本: {sys.version}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检查依赖包\n",
    "!pip list | findstr -i \"tushare pandas psycopg2 sqlalchemy\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如果需要安装依赖，取消注释下面这行\n",
    "# !pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入核心模块\n",
    "from importer import (\n",
    "    get_pg_connection,\n",
    "    get_ts_pro,\n",
    "    init_database,\n",
    "    import_stock_basic,\n",
    "    import_trade_cal,\n",
    "    import_daily_batch,\n",
    "    import_daily_by_year,\n",
    "    import_daily_basic,\n",
    "    import_daily_basic_by_date,\n",
    "    import_adj_factor,\n",
    "    import_adj_factor_batch,\n",
    "    import_financial_statements,\n",
    "    import_index_daily,\n",
    "    get_all_stock_codes,\n",
    "    get_stock_codes_from_db,\n",
    "    full_import,\n",
    "    batch_insert,\n",
    "    logger,\n",
    ")\n",
    "from config import DB_CONFIG, TUSHARE_TOKEN, START_DATE, END_DATE\n",
    "\n",
    "print(\"模块导入成功!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试数据库连接\n",
    "try:\n",
    "    conn = get_pg_connection()\n",
    "    cursor = conn.cursor()\n",
    "    cursor.execute(\"SELECT version()\")\n",
    "    version = cursor.fetchone()[0]\n",
    "    print(f\"✓ PostgreSQL 连接成功!\")\n",
    "    print(f\"  服务器版本: {version}\")\n",
    "    print(f\"  连接信息: {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}\")\n",
    "    cursor.close()\n",
    "    conn.close()\n",
    "except Exception as e:\n",
    "    print(f\"✗ 连接失败: {e}\")\n",
    "    print(\"请检查:\")\n",
    "    print(\"  1. Docker 容器是否已启动: docker ps | findstr postgres\")\n",
    "    print(\"  2. 环境变量 (.env) 中的连接参数是否正确\")\n",
    "    print(\"  3. 防火墙是否开放 12345 端口\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试 Tushare API 连接\n",
    "try:\n",
    "    pro = get_ts_pro()\n",
    "    # 简单测试: 获取一只股票信息\n",
    "    df = pro.stock_basic(ts_code=\"000001.SZ\", fields=\"ts_code,name,industry\")\n",
    "    print(f\"✓ Tushare API 连接成功!\")\n",
    "    print(f\"  测试查询: {df.iloc[0].to_dict()}\")\n",
    "except Exception as e:\n",
    "    print(f\"✗ Tushare API 连接失败: {e}\")\n",
    "    print(\"请检查环境变量 TUSHARE_TOKEN 是否正确（或 .env 文件）\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 1: 初始化数据库 Schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 执行 DDL，创建所有表结构\n",
    "init_database()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 2: 导入股票基本信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入全量 A 股股票基本信息（含上市和退市）\n",
    "import_stock_basic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证：查看导入结果\n",
    "conn = get_pg_connection()\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"SELECT COUNT(*) FROM stock_basic\")\n",
    "print(f\"stock_basic 总记录数: {cursor.fetchone()[0]}\")\n",
    "cursor.execute(\"SELECT list_status, COUNT(*) FROM stock_basic GROUP BY list_status\")\n",
    "for row in cursor.fetchall():\n",
    "    print(f\"  状态 '{row[0]}': {row[1]} 只\")\n",
    "cursor.close()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 3: 导入交易日历"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入交易日历 (默认从环境变量 QUANT_START_DATE ~ QUANT_END_DATE)\n",
    "import_trade_cal()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 或者指定日期范围\n",
    "# import_trade_cal(start_date=\"2020-01-01\", end_date=\"2025-12-31\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证\n",
    "conn = get_pg_connection()\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"\"\"\n",
    "    SELECT exchange, MIN(cal_date) AS first_date, MAX(cal_date) AS last_date, COUNT(*) AS total\n",
    "    FROM trade_cal\n",
    "    GROUP BY exchange\n",
    "\"\"\")\n",
    "for row in cursor.fetchall():\n",
    "    print(f\"  {row[0]}: {row[1]} ~ {row[2]}, 共 {row[3]} 条\")\n",
    "cursor.close()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 4: 导入日线行情 (核心表，最耗时)\n",
    "\n",
    "### 4.1 获取股票列表"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取所有需要导入的股票代码\n",
    "stock_list = get_stock_codes_from_db()\n",
    "print(f\"共 {len(stock_list)} 只股票需要导入日线行情\")\n",
    "print(f\"前 10 只: {stock_list[:10]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 按年批量导入 (推荐 - 断点续传友好)\n",
    "\n",
    "数据量估算: 约5000只股票 × 250交易日/年 × 16年 ≈ 2000万条记录"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 按年份逐批导入日线行情\n",
    "# 如果中断，可以修改年份范围从断点继续\n",
    "import_daily_by_year(\n",
    "    stock_list,\n",
    "    start_year=2010,\n",
    "    end_year=2025,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 单只股票导入 (补充/重试)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入单只股票日线行情 (用于补充导入或测试)\n",
    "conn = get_pg_connection()\n",
    "from importer import import_daily_for_stock\n",
    "\n",
    "n = import_daily_for_stock(\"000001.SZ\", \"2020-01-01\", \"2020-12-31\", conn)\n",
    "print(f\"导入 000001.SZ 2020年数据: {n} 条\")\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证日线数据\n",
    "conn = get_pg_connection()\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"\"\"\n",
    "    SELECT \n",
    "        COUNT(*) AS total_records,\n",
    "        COUNT(DISTINCT ts_code) AS stock_count,\n",
    "        MIN(trade_date) AS first_date,\n",
    "        MAX(trade_date) AS last_date\n",
    "    FROM daily\n",
    "\"\"\")\n",
    "for row in cursor.fetchall():\n",
    "    print(f\"  总记录数: {row[0]:,}\")\n",
    "    print(f\"  股票数量: {row[1]}\")\n",
    "    print(f\"  日期范围: {row[2]} ~ {row[3]}\")\n",
    "cursor.close()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 5: 导入每日指标 (估值数据)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 按交易日导入每日指标 (PE/PB/PS/总市值/流通市值等)\n",
    "import_daily_basic_by_date(\n",
    "    start_date=\"2010-01-01\",\n",
    "    end_date=\"2025-12-31\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证\n",
    "conn = get_pg_connection()\n",
    "cursor = conn.cursor()\n",
    "cursor.execute(\"SELECT COUNT(*), MIN(trade_date), MAX(trade_date) FROM daily_basic\")\n",
    "row = cursor.fetchone()\n",
    "print(f\"  daily_basic: {row[0]:,} 条, {row[1]} ~ {row[2]}\")\n",
    "cursor.close()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 6: 导入复权因子"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入复权因子 (用于前复权/后复权价格计算)\n",
    "import_adj_factor_batch(\n",
    "    stock_list,\n",
    "    start_date=\"2010-01-01\",\n",
    "    end_date=\"2025-12-31\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 7: 导入财务数据 (三大报表 + 财务指标)\n",
    "\n",
    "⚠ 此步骤耗时较长，约需数小时（取决于股票数量）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入利润表、资产负债表、现金流量表、财务指标\n",
    "import_financial_statements(\n",
    "    stock_list,\n",
    "    start_date=\"2010-01-01\",\n",
    "    end_date=\"2025-12-31\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 验证财务数据\n",
    "conn = get_pg_connection()\n",
    "cursor = conn.cursor()\n",
    "for table in [\"income\", \"balancesheet\", \"cashflow\", \"fina_indicator\"]:\n",
    "    cursor.execute(f\"SELECT COUNT(*) FROM {table}\")\n",
    "    count = cursor.fetchone()[0]\n",
    "    print(f\"  {table}: {count:,} 条\")\n",
    "cursor.close()\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Step 8: 导入指数日线行情"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入主要指数日线行情\n",
    "import_index_daily(\n",
    "    index_codes=[\n",
    "        \"000001.SH\",  # 上证指数\n",
    "        \"399001.SZ\",  # 深证成指\n",
    "        \"000300.SH\",  # 沪深300\n",
    "        \"000905.SH\",  # 中证500\n",
    "        \"399006.SZ\",  # 创业板指\n",
    "        \"000688.SH\",  # 科创50\n",
    "        \"000016.SH\",  # 上证50\n",
    "        \"399005.SZ\",  # 中小100\n",
    "        \"000852.SH\",  # 中证1000\n",
    "    ],\n",
    "    start_date=\"2010-01-01\",\n",
    "    end_date=\"2025-12-31\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 一键全量导入 (可选)\n",
    "\n",
    "如果不想逐步执行，可以运行下面这个 Cell 一键完成所有导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 一键全量导入 (需数小时~数十小时，请谨慎)\n",
    "# full_import(\n",
    "#     start_date=\"2010-01-01\",\n",
    "#     end_date=\"2025-12-31\",\n",
    "#     import_financials=True,  # 设为 False 跳过财务数据加快速度\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## 数据验证与查询示例"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import psycopg2\n",
    "\n",
    "conn = get_pg_connection()\n",
    "\n",
    "# 各表统计\n",
    "tables = [\"stock_basic\", \"daily\", \"daily_basic\", \"adj_factor\",\n",
    "          \"income\", \"balancesheet\", \"cashflow\", \"fina_indicator\",\n",
    "          \"trade_cal\", \"index_daily\"]\n",
    "\n",
    "print(f\"{'表名':<20} {'记录数':>12} {'最早日期':>12} {'最晚日期':>12}\")\n",
    "print(\"-\" * 60)\n",
    "for table in tables:\n",
    "    try:\n",
    "        count_sql = f\"SELECT COUNT(*) FROM {table}\"\n",
    "        count = pd.read_sql(count_sql, conn).iloc[0, 0]\n",
    "        \n",
    "        # 尝试获取日期范围\n",
    "        date_col = None\n",
    "        if table == \"daily\":\n",
    "            date_col = \"trade_date\"\n",
    "        elif table == \"daily_basic\":\n",
    "            date_col = \"trade_date\"\n",
    "        elif table in [\"income\", \"balancesheet\", \"cashflow\"]:\n",
    "            date_col = \"end_date\"\n",
    "        elif table == \"fina_indicator\":\n",
    "            date_col = \"end_date\"\n",
    "        elif table == \"trade_cal\":\n",
    "            date_col = \"cal_date\"\n",
    "        elif table == \"index_daily\":\n",
    "            date_col = \"trade_date\"\n",
    "        elif table == \"adj_factor\":\n",
    "            date_col = \"trade_date\"\n",
    "            \n",
    "        if date_col:\n",
    "            date_sql = f\"SELECT MIN({date_col}), MAX({date_col}) FROM {table}\"\n",
    "            min_d, max_d = pd.read_sql(date_sql, conn).iloc[0]\n",
    "            print(f\"{table:<20} {count:>12,} {str(min_d)[:10]:>12} {str(max_d)[:10]:>12}\")\n",
    "        else:\n",
    "            print(f\"{table:<20} {count:>12,}\")\n",
    "    except Exception as e:\n",
    "        print(f\"{table:<20} {'错误':>12}: {str(e)[:40]}\")\n",
    "\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 示例查询 1: 查询某股票最近10个交易日数据\n",
    "query1 = \"\"\"\n",
    "SELECT trade_date, open, high, low, close, vol, amount, pct_chg\n",
    "FROM daily\n",
    "WHERE ts_code = '000001.SZ'\n",
    "ORDER BY trade_date DESC\n",
    "LIMIT 10\n",
    "\"\"\"\n",
    "conn = get_pg_connection()\n",
    "df1 = pd.read_sql(query1, conn)\n",
    "print(\"平安银行(000001.SZ) 最近10个交易日:\")\n",
    "display(df1)\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 示例查询 2: 日线行情 + 估值指标联合查询 (使用视图)\n",
    "query2 = \"\"\"\n",
    "SELECT *\n",
    "FROM v_daily_with_valuation\n",
    "WHERE ts_code = '000001.SZ'\n",
    "  AND trade_date >= '2024-01-01'\n",
    "ORDER BY trade_date DESC\n",
    "LIMIT 10\n",
    "\"\"\"\n",
    "conn = get_pg_connection()\n",
    "df2 = pd.read_sql(query2, conn)\n",
    "print(\"平安银行 - 日线+估值:\")\n",
    "display(df2[['trade_date', 'close', 'pct_chg', 'pe', 'pe_ttm', 'pb', 'total_mv']])\n",
    "conn.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 示例查询 3: 最新财务指标 Top 20 (按 ROE 排序)\n",
    "query3 = \"\"\"\n",
    "SELECT *\n",
    "FROM v_latest_financials\n",
    "WHERE roe IS NOT NULL\n",
    "  AND roe > 0\n",
    "ORDER BY roe DESC\n",
    "LIMIT 20\n",
    "\"\"\"\n",
    "conn = get_pg_connection()\n",
    "df3 = pd.read_sql(query3, conn)\n",
    "print(\"ROE Top 20:\")\n",
    "display(df3[['ts_code', 'name', 'industry', 'roe', 'roa', 'eps', 'debt_to_assets']])\n",
    "conn.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}