# -*- coding: utf-8 -*-
"""
模拟数据生成器（当前为演示数据，后期可替换为真实数据源）
生成内容：
  1. stocks       股票基础信息（60 只，覆盖 20+ 行业）
  2. stock_daily  180 个交易日 OHLCV 行情（随机游走 + 趋势/题材分化）
  3. market_index 上证/深证/创业板 三大指数序列
  4. news         财经新闻（280+ 条，含情感标签/关联个股，作为 RAG 语料）
  5. institutions 机构实体（公募/券商/保险资管/外资/私募）
  6. inst_ratings 机构评级记录
  7. fund_holdings 基金季度持仓
  8. 向量索引：stock_news_v1（新闻正文）+ stock_profiles_v1（公司概况）

用法：python seed_data.py [--skip-vector]
"""
import argparse
import datetime as dt
import json
import math
import random
import sys

from config import (CHROMA_NEWS_COLLECTION, CHROMA_PROFILE_COLLECTION, LOG_DIR)
from database import init_db, executemany, query_one, wipe_all
from rag import vector_store as vs

random.seed(42)

# ===================================================================== 股票池
# (code, name, industry, board, base_price, float_shares(亿股), trend(0多/1平/2空), vol, biz)
STOCKS = [
    ("600519", "贵州茅台", "白酒", "主板", 1680, 12.56, 1, 0.020, "高端白酒龙头"),
    ("000858", "五粮液", "白酒", "主板", 128, 38.8, 1, 0.021, "浓香型白酒头部企业"),
    ("600809", "山西汾酒", "白酒", "主板", 198, 12.2, 0, 0.024, "清香型白酒代表"),
    ("000568", "泸州老窖", "白酒", "主板", 118, 14.7, 1, 0.022, "国窖1573高端白酒"),
    ("300750", "宁德时代", "动力电池", "创业板", 248, 43.9, 0, 0.026, "全球动力电池龙头"),
    ("002594", "比亚迪", "新能源汽车", "主板", 235, 11.6, 0, 0.025, "新能源整车+电池一体化"),
    ("601012", "隆基绿能", "光伏", "主板", 16.8, 75.8, 2, 0.028, "单晶硅片与组件龙头"),
    ("600438", "通威股份", "光伏", "主板", 20.5, 45.0, 2, 0.027, "硅料+电池片双龙头"),
    ("300274", "阳光电源", "光伏储能", "创业板", 78, 14.9, 0, 0.026, "光伏逆变器与储能龙头"),
    ("002460", "赣锋锂业", "锂电材料", "主板", 36, 20.2, 2, 0.030, "锂盐龙头"),
    ("688981", "中芯国际", "半导体", "科创板", 88, 26.5, 0, 0.025, "晶圆代工龙头"),
    ("688012", "中微公司", "半导体设备", "科创板", 178, 6.2, 0, 0.029, "刻蚀设备龙头"),
    ("002371", "北方华创", "半导体设备", "主板", 352, 5.3, 0, 0.028, "半导体设备平台型龙头"),
    ("603986", "兆易创新", "半导体", "主板", 118, 6.7, 1, 0.027, "存储芯片+MCU"),
    ("603501", "韦尔股份", "半导体", "主板", 96, 12.2, 1, 0.028, "CIS图像传感器龙头"),
    ("688041", "海光信息", "芯片设计", "科创板", 128, 23.2, 0, 0.028, "国产CPU/DCU"),
    ("600276", "恒瑞医药", "创新药", "主板", 46, 63.7, 1, 0.022, "创新药龙头"),
    ("603259", "药明康德", "CXO", "主板", 56, 29.6, 1, 0.023, "医药研发外包龙头"),
    ("300760", "迈瑞医疗", "医疗器械", "创业板", 268, 12.1, 1, 0.021, "医疗器械平台龙头"),
    ("600085", "同仁堂", "中药", "主板", 42, 13.7, 1, 0.019, "老字号中药"),
    ("000538", "云南白药", "中药", "主板", 52, 17.8, 1, 0.020, "中药+健康消费品"),
    ("300015", "爱尔眼科", "医疗服务", "创业板", 13.5, 93.2, 1, 0.026, "连锁眼科医疗"),
    ("600036", "招商银行", "银行", "主板", 36, 206.3, 1, 0.018, "零售银行龙头"),
    ("601398", "工商银行", "银行", "主板", 6.1, 2696.0, 1, 0.015, "国有大行"),
    ("601166", "兴业银行", "银行", "主板", 19.8, 207.7, 1, 0.017, "股份行"),
    ("600000", "浦发银行", "银行", "主板", 8.9, 293.5, 2, 0.016, "股份行"),
    ("600030", "中信证券", "券商", "主板", 28, 148.2, 1, 0.022, "券商龙头"),
    ("601688", "华泰证券", "券商", "主板", 16.5, 90.8, 1, 0.021, "互联网券商"),
    ("300059", "东方财富", "互联网金融", "创业板", 13.8, 158.0, 0, 0.027, "互联网财富管理"),
    ("600999", "招商证券", "券商", "主板", 15.2, 86.9, 1, 0.020, "综合券商"),
    ("601633", "长城汽车", "汽车", "主板", 24, 85.4, 1, 0.024, "SUV与新能源整车"),
    ("000625", "长安汽车", "汽车", "主板", 14.5, 99.2, 1, 0.023, "自主品牌整车"),
    ("601238", "广汽集团", "汽车", "主板", 8.6, 104.7, 2, 0.022, "整车集团"),
    ("600104", "上汽集团", "汽车", "主板", 14.2, 115.7, 2, 0.020, "整车集团"),
    ("002475", "立讯精密", "消费电子", "主板", 32, 72.1, 0, 0.026, "消费电子+汽车连接器"),
    ("002241", "歌尔股份", "消费电子", "主板", 24, 34.2, 1, 0.026, "声学与VR整机"),
    ("000725", "京东方A", "面板", "主板", 4.2, 376.0, 1, 0.024, "显示面板龙头"),
    ("000100", "TCL科技", "面板", "主板", 4.8, 187.7, 1, 0.024, "面板+光伏材料"),
    ("000333", "美的集团", "家电", "主板", 68, 68.3, 1, 0.019, "白电龙头"),
    ("000651", "格力电器", "家电", "主板", 40, 56.3, 1, 0.020, "空调龙头"),
    ("600690", "海尔智家", "家电", "主板", 28, 94.5, 1, 0.020, "白电+智慧家庭"),
    ("600309", "万华化学", "化工", "主板", 78, 31.4, 1, 0.022, "MDI与新材料龙头"),
    ("601899", "紫金矿业", "有色", "主板", 18, 263.0, 0, 0.023, "黄金铜矿龙头"),
    ("002648", "卫星化学", "化工", "主板", 16.5, 33.7, 1, 0.024, "轻烃一体化化工"),
    ("600031", "三一重工", "工程机械", "主板", 16.8, 84.9, 1, 0.023, "挖掘机龙头"),
    ("000425", "徐工机械", "工程机械", "主板", 7.2, 118.3, 1, 0.022, "工程机械"),
    ("600760", "中航沈飞", "军工", "主板", 52, 27.6, 1, 0.024, "战斗机主机厂"),
    ("603288", "海天味业", "食品饮料", "主板", 42, 55.6, 1, 0.020, "调味品龙头"),
    ("600887", "伊利股份", "食品饮料", "主板", 27, 63.7, 1, 0.019, "乳制品龙头"),
    ("000895", "双汇发展", "食品饮料", "主板", 25, 34.6, 2, 0.019, "肉制品龙头"),
    ("002415", "海康威视", "安防", "主板", 32, 92.3, 1, 0.022, "智能物联龙头"),
    ("002230", "科大讯飞", "人工智能", "主板", 46, 23.1, 0, 0.028, "AI语音与大模型"),
    ("688111", "金山办公", "软件", "科创板", 268, 4.6, 0, 0.027, "办公软件WPS"),
    ("300033", "同花顺", "金融科技", "创业板", 158, 5.4, 0, 0.028, "金融信息服务"),
    ("600570", "恒生电子", "软件", "主板", 26, 19.0, 1, 0.023, "金融IT龙头"),
    ("601888", "中国中免", "免税", "主板", 68, 20.7, 1, 0.023, "免税龙头"),
    ("600900", "长江电力", "电力", "主板", 28, 244.7, 1, 0.015, "水电龙头"),
    ("601088", "中国神华", "煤炭", "主板", 42, 198.7, 1, 0.017, "煤电一体化"),
    ("000002", "万科A", "房地产", "主板", 6.8, 119.3, 2, 0.026, "地产开发"),
    ("600048", "保利发展", "房地产", "主板", 8.4, 119.7, 2, 0.025, "地产开发"),
    ("600941", "中国移动", "通信运营", "主板", 102, 95.7, 1, 0.016, "通信运营龙头"),
    ("000063", "中兴通讯", "通信设备", "主板", 28, 47.8, 1, 0.024, "通信设备"),
    ("002714", "牧原股份", "养殖", "主板", 42, 54.6, 1, 0.026, "生猪养殖龙头"),
    ("300498", "温氏股份", "养殖", "创业板", 18, 66.5, 1, 0.025, "生猪+禽养殖"),
    ("601318", "中国平安", "保险", "主板", 48, 181.7, 1, 0.019, "综合金融"),
    ("601628", "中国人寿", "保险", "主板", 30, 282.6, 1, 0.018, "寿险龙头"),
    ("600111", "北方稀土", "稀土", "主板", 21, 36.2, 1, 0.027, "稀土龙头"),
    ("601600", "中国铝业", "有色", "主板", 7.5, 171.6, 1, 0.022, "电解铝龙头"),
]

# ===================================================================== 机构
INSTITUTIONS = [
    # (name, type, desc, 重点行业)
    ("易方达基金", "公募基金", "国内头部公募基金，管理规模超1.6万亿，深耕消费与科技赛道。", ["白酒", "消费电子", "医药"]),
    ("华夏基金", "公募基金", "老牌公募，主动权益与ETF双轮驱动。", ["半导体", "新能源", "医药"]),
    ("嘉实基金", "公募基金", "全牌照公募，均衡配置大盘蓝筹与成长。", ["银行", "食品饮料", "家电"]),
    ("广发基金", "公募基金", "权益投资见长，聚焦科技成长。", ["半导体", "人工智能", "光伏"]),
    ("南方基金", "公募基金", "老十家公募之一，稳健价值风格。", ["银行", "券商", "白酒"]),
    ("富国基金", "公募基金", "主动权益明星基金较多，偏好行业景气。", ["医药", "半导体", "机械"]),
    ("汇添富基金", "公募基金", "价值成长均衡，重仓消费与医药。", ["白酒", "医药", "消费电子"]),
    ("中欧基金", "公募基金", "权益投资口碑公募，聚焦成长赛道。", ["新能源", "人工智能", "软件"]),
    ("中信证券研究部", "券商研究", "国内综合实力领先的券商研究所，覆盖全行业。", ["券商", "半导体", "军工"]),
    ("华泰证券研究所", "券商研究", "科技+金融双主线研究见长。", ["消费电子", "金融科技", "半导体"]),
    ("国泰君安研究所", "券商研究", "老牌研究所，宏观策略与大金融研究扎实。", ["银行", "券商", "保险"]),
    ("招商证券研究所", "券商研究", "以行业景气跟踪见长，深度覆盖成长板块。", ["汽车", "消费电子", "医药"]),
    ("广发证券研究所", "券商研究", "策略+行业结合，偏成长挖掘。", ["人工智能", "软件", "新能源"]),
    ("兴业证券研究所", "券商研究", "产业链研究见长，覆盖周期成长。", ["化工", "有色", "银行"]),
    ("东吴证券研究所", "券商研究", "新经济研究活跃，覆盖TMT。", ["半导体", "软件", "通信"]),
    ("浙商证券研究所", "券商研究", "后起之秀，新能源与制造研究领先。", ["光伏", "新能源", "机械"]),
    ("中国人寿资管", "保险资管", "险资巨头，偏好高股息蓝筹与稳定现金流。", ["银行", "电力", "保险"]),
    ("平安资管", "保险资管", "综合金融背景，长期资金代表。", ["银行", "食品饮料", "保险"]),
    ("泰康资产", "保险资管", "保险资管头部，注重绝对收益。", ["医药", "家电", "白酒"]),
    ("高盛", "外资机构", "全球投行，A股研究覆盖核心资产。", ["白酒", "新能源", "消费电子"]),
    ("摩根士丹利", "外资机构", "全球投行，偏好大盘龙头与全球化公司。", ["消费电子", "新能源", "汽车"]),
    ("瑞银证券", "外资机构", "外资券商，深入研究A股核心蓝筹。", ["银行", "白酒", "医药"]),
    ("贝莱德", "外资机构", "全球最大资管，长期配置中国核心资产。", ["银行", "食品饮料", "电力"]),
    ("高毅资产", "私募游资", "头部私募，逆向投资与深度研究。", ["医药", "消费", "科技"]),
    ("淡水泉投资", "私募游资", "老牌私募，擅长困境反转与逆向布局。", ["化工", "地产", "养殖"]),
    ("景林资产", "私募游资", "价值投资私募，重仓优质成长。", ["白酒", "互联网", "家电"]),
    ("幻方量化", "私募游资", "头部量化私募，捕捉市场波动机会。", ["金融", "周期", "TMT"]),
]

# ===================================================================== 新闻
NEWS_SOURCES = ["证券时报", "上海证券报", "中国证券报", "财联社", "界面新闻",
                "每日经济新闻", "澎湃新闻", "21世纪经济报道", "第一财经"]

NEWS_TEMPLATES = [
    # (category, positive, title_tpl, content_tpls)
    ("业绩", True, "{name}发布业绩预告：前三季度净利润同比增长{pct}%",
     ["公司公告显示，受益于{industry}行业景气度提升，{name}核心业务收入实现较快增长，业绩超市场一致预期。",
      "多家券商点评认为，{name}盈利质量改善明显，费用管控有效，全年业绩有望延续高增。",
      "业内人士表示，行业需求回暖背景下，{name}作为{industry}领域头部公司，市占率有望进一步提升。"]),
    ("业绩", True, "{name}半年度净利大增{pct}%，创上市以来新高",
     ["{name}中报披露，报告期内实现营业收入同比增长{num}%，净利润同比增长{pct}%，均超市场预期。",
      "公司称，新产能释放叠加产品结构优化，带动毛利率显著提升。",
      "多家机构预计，随着行业景气延续，{name}未来业绩增长确定性较强。"]),
    ("业绩", False, "{name}业绩不及预期：第三季度营收同比下滑{pct}%",
     ["{name}三季报显示，受行业需求疲软影响，公司营业收入同比下滑{pct}%，净利润降幅扩大。",
      "公司解释称，原材料成本上行及产品价格承压是主要原因。",
      "市场人士认为，短期{name}基本面仍面临压力，需观察行业拐点信号。"]),
    ("行业", True, "政策加码！{industry}行业迎来{num}亿产业基金支持",
     ["相关部门发文明确，将设立{num}亿元产业投资基金，重点支持{industry}产业链关键环节技术攻关与产能建设。",
      "分析人士指出，政策红利有望带动{industry}板块整体估值修复，相关龙头企业将直接受益。",
      "板块内多只个股盘中异动，资金关注度明显提升。"]),
    ("行业", False, "原材料价格波动，{industry}行业盈利承压",
     ["近期上游原材料价格波动加大，{industry}行业部分企业毛利率受到侵蚀。",
      "业内调研显示，中小企业已出现减产观望情绪，头部公司凭借成本优势影响相对可控。",
      "机构提示，短期需关注库存去化进度与价格企稳信号。"]),
    ("公司", True, "{name}拟回购{num}亿元股份，彰显发展信心",
     ["{name}公告，拟以自有资金{num}亿元回购公司股份，用于员工持股计划或股权激励，回购价格不超过{price}元/股。",
      "公司表示，回购基于对未来发展前景的信心及对公司价值的认可。",
      "分析人士称，回购计划落地有望对股价形成支撑，彰显管理层信心。"]),
    ("公司", True, "{name}中标{num}亿元重大项目，订单持续饱满",
     ["{name}公告，近日中标{industry}领域重大工程项目，中标金额合计约{num}亿元。",
      "公司称，该项目是公司在核心客户与重点市场的重要突破，有利于巩固行业地位。",
      "机构预计，在手订单充足将支撑{name}未来收入增长确定性。"]),
    ("公司", False, "{name}股东拟减持不超过{pct}%股份",
     ["{name}公告，持股{num}%的股东计划在未来6个月内减持不超过公司总股本{pct}%的股份。",
      "公司称，减持系股东自身资金安排，不影响公司正常经营。",
      "市场对此反应谨慎，分析人士提醒关注减持节奏对股价的短期压制。"]),
    ("机构观点", True, "{inst}上调{name}评级至「买入」，目标价{price}元",
     ["{inst}发布研报认为，{name}受益于{industry}行业景气回升，业绩进入加速释放期，将评级由「增持」上调至「买入」。",
      "研报给出目标价{price}元，较当前股价存在一定上行空间。",
      "研报强调，{name}核心竞争力稳固，估值具备吸引力，建议积极配置。"]),
    ("机构观点", False, "{inst}下调{name}评级至「中性」，提示估值风险",
     ["{inst}研报指出，{name}短期涨幅较大，当前估值已透支部分预期，将评级由「买入」下调至「中性」。",
      "研报认为，行业景气虽有支撑，但股价上行空间收窄，建议等待更好的介入时点。",
      "市场人士表示，机构评级下调或引发短期情绪扰动。"]),
    ("市场", True, "沪指放量上涨{pct}%，两市成交额突破{num}万亿",
     ["A股市场情绪回暖，沪指放量上行{pct}%，深成指、创业板指同步走强。",
      "盘面上，{industry}等板块领涨，赚钱效应明显，两市成交额突破{num}万亿元。",
      "分析人士认为，市场风险偏好回升，中期趋势向好，可关注业绩确定性方向。"]),
    ("市场", False, "大盘缩量回调{pct}%，市场观望情绪升温",
     ["A股缩量调整，沪指收跌{pct}%，两市成交额较前期明显萎缩。",
      "盘面上热点轮动加快，缺乏持续性主线，资金观望情绪浓厚。",
      "机构提示，短期指数或以震荡为主，建议控制仓位、关注结构机会。"]),
]

FILLERS = [
    "相关消息发布后，市场反应总体平稳。",
    "多位市场人士对此进行了讨论。",
    "后续进展值得持续跟踪。",
    "公司方面暂未就此事进一步置评。",
    "受此影响，相关产业链公司受到市场关注。",
    "整体来看，基本面对股价中期走势具有决定性影响。",
]


def _gen_trading_dates(n=180):
    """生成最近 n 个交易日（跳过周末），终止于最近工作日"""
    dates = []
    d = dt.date.today()
    # 回退到最近的非周末
    while d.weekday() >= 5:
        d -= dt.timedelta(days=1)
    while len(dates) < n:
        if d.weekday() < 5:
            dates.append(d.isoformat())
        d -= dt.timedelta(days=1)
    return list(reversed(dates))


# ===================================================================== 生成
def gen_stocks():
    rows = []
    for code, name, industry, board, base, float_shares, trend, vol, biz in STOCKS:
        total = round(float_shares * random.uniform(1.0, 1.6), 2)
        pe = round(random.uniform(15, 60), 1)
        pb = round(random.uniform(1.5, 8), 2)
        desc = (f"{name}是{industry}领域{biz}。"
                f"公司主营产品广泛应用于核心客户，行业地位稳固，近年来持续加大研发投入，"
                f"积极拓展新增长曲线。当前总股本约{total}亿股，流通市值位居行业前列。")
        rows.append((code, name, industry, board, total, float_shares,
                     f"200{random.randint(0, 9)}-{random.randint(1, 12):02d}-{random.randint(1, 28):02d}",
                     pe, pb, 0.0, desc))
    return rows


def gen_daily(dates):
    """生成个股日线 + 指数序列"""
    daily = []
    index = {}
    sh, sz, cy = 3245.0, 10580.0, 2120.0
    price = {}

    for code, name, industry, board, base, float_shares, trend, vol, biz in STOCKS:
        p = base
        drift = {0: 0.0011, 1: 0.00025, 2: -0.00085}[trend]
        # 最近30天加速（制造趋势分化，让荐股有区分度）
        recent_drift = {0: 0.0045, 1: 0.0001, 2: -0.0045}[trend]
        # 基准成交量（万股）= 流通盘 × 0.4%~1.5% 日换手（贴近真实市场）
        base_vol = float_shares * 10000 * random.uniform(0.004, 0.015)
        for i, d in enumerate(dates):
            phase = max(0, i - (len(dates) - 30))
            dr = drift + (recent_drift if phase > 0 else 0)
            r = random.gauss(dr, vol)
            if random.random() < 0.02:          # 偶发跳空
                r += random.gauss(0, vol * 1.6)
            prev = p
            p = max(0.8, p * (1 + r))
            open_p = prev * (1 + random.gauss(0, vol * 0.5))
            high = max(open_p, p) * (1 + abs(random.gauss(0, vol * 0.35)))
            low = min(open_p, p) * (1 - abs(random.gauss(0, vol * 0.35)))
            volume = base_vol * (1 + 1.5 * abs(r) / vol) * random.uniform(0.6, 1.4)
            amount = volume * (open_p + p) / 2          # 万元
            chg = (p - prev) / prev * 100
            daily.append((code, d, round(open_p, 2), round(high, 2), round(low, 2),
                          round(p, 2), round(volume, 0), round(amount, 0), round(chg, 2)))
            price[code] = p

        # 指数序列（独立随机游走）
    for i, d in enumerate(dates):
        if i == 0:
            index[d] = {"sh": round(sh, 2), "sz": round(sz, 2), "cy": round(cy, 2)}
            continue
        sh_r = sum(random.gauss(0.0004, 0.008) for _ in range(6)) / 6
        sz_r = sh_r + random.gauss(0, 0.004)
        cy_r = sh_r + random.gauss(0, 0.006)
        sh *= (1 + sh_r); sz *= (1 + sz_r); cy *= (1 + cy_r)
        index[d] = {"sh": round(sh, 2), "sz": round(sz, 2), "cy": round(cy, 2)}
    return daily, index, price


def gen_news(dates, price):
    """生成新闻（关联最近90天，偏近分布）"""
    news = []
    recent_dates = dates[-95:]
    stock_map = {s[0]: s for s in STOCKS}
    for _ in range(300):
        tpl = random.choice(NEWS_TEMPLATES)
        category, positive, title_tpl, contents = tpl
        s = random.choice(STOCKS)
        code, name, industry = s[0], s[1], s[2]
        pct = random.randint(5, 85) if positive else random.randint(5, 60)
        num = random.choice([5, 10, 20, 30, 50, 80, 100, 120, 150, 200])
        price_t = round(price.get(code, 10) * random.uniform(1.02, 1.35), 2)
        inst = random.choice(INSTITUTIONS)[0]
        title = title_tpl.format(name=name, pct=pct, num=num, price=price_t,
                                 inst=inst, industry=industry)
        content = " ".join(c.format(name=name, pct=pct, num=num, price=price_t,
                                    inst=inst, industry=industry) for c in contents)
        content += " " + " ".join(random.sample(FILLERS, random.randint(1, 3)))
        date = random.choice(recent_dates)
        source = random.choice(NEWS_SOURCES)
        sentiment = round(random.uniform(0.25, 0.85) if positive else random.uniform(-0.85, -0.25), 2)
        # 关联股票：主股 + 同行业0~2只
        related = [code]
        same = [x[0] for x in STOCKS if x[2] == industry and x[0] != code]
        random.shuffle(same)
        related += same[:random.randint(0, 2)]
        news.append({
            "title": title, "content": content, "source": source, "category": category,
            "publish_date": date, "related": ",".join(related),
            "sentiment": sentiment, "is_positive": 1 if positive else 0,
        })
    news.sort(key=lambda x: x["publish_date"])
    return news


def gen_institutions():
    rows = []
    for name, typ, desc, _focus in INSTITUTIONS:
        rows.append((name, typ, desc))
    return rows


def gen_ratings(price):
    """机构评级：每家机构关注重点行业内的股票"""
    rows = []
    for name, typ, desc, focus in INSTITUTIONS:
        pool = [s for s in STOCKS if s[2] in focus]
        if not pool:
            pool = STOCKS
        picks = random.sample(pool, min(8, len(pool)))
        for s in picks:
            code = s[0]
            r = random.choices(["买入", "增持", "中性", "减持"], weights=[5, 4, 2, 1])[0]
            up = {"买入": random.uniform(0.15, 0.35), "增持": random.uniform(0.05, 0.18),
                  "中性": random.uniform(-0.05, 0.06), "减持": random.uniform(-0.18, -0.08)}[r]
            tp = round(price.get(code, 10) * (1 + up), 2)
            # 用近期交易日作为评级日
            days = _recent_days(45)
            rows.append((0, name, code, r, tp, random.choice(days),
                         random.choice(["中性", "增持", "买入", "增持", "买入"])))
    return rows


def _recent_days(n):
    dates = _gen_trading_dates(n)
    return dates


GLOBAL_INDICES = [
    ("dji", "道琼斯", 34000), ("nasdaq", "纳斯达克", 12800), ("sp500", "标普500", 4400),
    ("hsi", "恒生指数", 17500), ("nikkei", "日经225", 33000), ("kospi", "韩国KOSPI", 2500),
    ("dax", "德国DAX", 16000), ("cac", "法国CAC40", 7000), ("ftse", "英国FTSE100", 7500),
]


def _gen_global(dates):
    """生成全球主要指数模拟序列（随机游走，chg 基于前一交易日）"""
    vals = {k: v for k, _, v in GLOBAL_INDICES}
    prev = dict(vals)
    out = {}
    for d in dates:
        row = {}
        for k, label, _v in GLOBAL_INDICES:
            vals[k] *= (1 + random.gauss(0.0002, 0.009))
            row[k] = {"label": label, "value": round(vals[k], 2),
                      "chg": round((vals[k] - prev[k]) / prev[k] * 100, 2)}
        for k in prev:
            prev[k] = vals[k]
        out[d] = row
    return out


def gen_holdings(price, dates):
    """基金季度持仓：2025Q4 / 2026Q1 / 2026Q2"""
    rows = []
    quarters = ["2025Q4", "2026Q1", "2026Q2"]
    funds = [i for i in INSTITUTIONS if i[1] == "公募基金"] + \
            [i for i in INSTITUTIONS if i[1] == "保险资管"] + \
            [i for i in INSTITUTIONS if i[1] in ("外资机构", "私募游资")]
    for name, typ, desc, focus in funds:
        pool = [s for s in STOCKS if s[2] in focus] or STOCKS
        picks = random.sample(pool, min(random.randint(6, 12), len(pool)))
        for s in picks:
            code = s[0]
            last_hold = random.uniform(2000, 40000)   # 万股
            prev_hold = None
            for qi, q in enumerate(quarters):
                if qi == 0:
                    hold = last_hold * random.uniform(0.5, 1.1)
                else:
                    hold = last_hold
                chg = 0.0
                if qi == 2:  # 最新季度给出增减方向，与趋势挂钩
                    trend = s[6]
                    delta = {0: random.uniform(5, 30), 1: random.uniform(-12, 12), 2: random.uniform(-25, -3)}[trend]
                    chg = hold * delta / 100
                    hold = max(500, hold + chg)
                value = hold * price.get(code, 10)
                rows.append((0, name, code, q, round(hold, 0), round(value, 0),
                             round(chg, 0), round(chg / max(hold - chg, 1) * 100, 2)))
                last_hold = hold
    return rows


# ===================================================================== 入库
def build_vectors(news, stocks):
    """构建 Chroma 向量索引：新闻 + 公司概况"""
    print(">>> 构建新闻向量索引 ...")
    vs.delete_collection(CHROMA_NEWS_COLLECTION)
    ids, docs, metas = [], [], []
    for n in news:
        for code in n["related"].split(","):
            ids.append(f"news-{n['title']}-{code}")
            docs.append(f"{n['title']}\n{n['content']}")
            metas.append({"code": code, "title": n["title"], "date": n["publish_date"],
                          "category": n["category"], "sentiment": n["sentiment"],
                          "news_id": 0})
    for i in range(0, len(ids), 16):
        vs.add_documents(ids[i:i + 16], docs[i:i + 16], metas[i:i + 16], CHROMA_NEWS_COLLECTION)
        print(f"    news {min(i+16, len(ids))}/{len(ids)}")

    print(">>> 构建公司概况向量索引 ...")
    vs.delete_collection(CHROMA_PROFILE_COLLECTION)
    ids, docs, metas = [], [], []
    for s in stocks:
        ids.append(f"profile-{s[0]}")
        docs.append(f"{s[1]}（{s[0]}），所属行业：{s[2]}。{s[10]}")
        metas.append({"code": s[0], "name": s[1], "industry": s[2]})
    vs.add_documents(ids, docs, metas, CHROMA_PROFILE_COLLECTION)
    print(f"    profiles {len(ids)}")
    print(f"    新闻索引条数: {vs.collection_count(CHROMA_NEWS_COLLECTION)}")
    print(f"    概况索引条数: {vs.collection_count(CHROMA_PROFILE_COLLECTION)}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--skip-vector", action="store_true", help="跳过向量索引重建")
    parser.add_argument("--no-strategies", action="store_true", help="跳过量化策略回测")
    args = parser.parse_args()

    print(">>> 初始化数据库 ...")
    init_db()
    wipe_all()

    dates = _gen_trading_dates()
    print(f">>> 生成 {len(STOCKS)} 只股票日线行情（{dates[0]} ~ {dates[-1]}）...")
    stocks = gen_stocks()
    executemany(
        "INSERT INTO stocks(code,name,industry,board,total_shares,float_shares,list_date,pe,pb,market_cap,description) "
        "VALUES(?,?,?,?,?,?,?,?,?,?,?)",
        [(s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], 0.0, s[10]) for s in stocks])
    daily, index, price = gen_daily(dates)
    executemany(
        "INSERT OR REPLACE INTO stock_daily(code,date,open,high,low,close,volume,amount,change_pct) "
        "VALUES(?,?,?,?,?,?,?,?,?)", daily)
    executemany("INSERT OR REPLACE INTO market_index(date,sh,sz,cy) VALUES(?,?,?,?)",
                [(d, v["sh"], v["sz"], v["cy"]) for d, v in index.items()])
    # 全球市场指数（模拟）
    gm = _gen_global(dates)
    executemany("INSERT OR REPLACE INTO global_markets(date,data) VALUES(?,?)",
                [(d, json.dumps(v, ensure_ascii=False)) for d, v in gm.items()])
    # 回填市值
    for code, name, industry, board, base, fs, trend, vol, biz in STOCKS:
        from database import execute as ex
        ex("UPDATE stocks SET market_cap=ROUND((SELECT close FROM stock_daily WHERE code=? ORDER BY date DESC LIMIT 1)*total_shares,2) WHERE code=?", (code, code))

    print(">>> 生成财经新闻 ...")
    news = gen_news(dates, price)
    executemany(
        "INSERT INTO news(title,content,source,category,publish_date,related_stocks,sentiment,is_positive) "
        "VALUES(?,?,?,?,?,?,?,?)",
        [(n["title"], n["content"], n["source"], n["category"], n["publish_date"],
          n["related"], n["sentiment"], n["is_positive"]) for n in news])

    print(">>> 生成机构与评级 ...")
    insts = gen_institutions()
    executemany("INSERT INTO institutions(name,type,description) VALUES(?,?,?)", insts)
    inst_map = {}
    for i, (name, typ, desc) in enumerate(insts, start=1):
        inst_map[name] = i
    ratings = gen_ratings(price)
    executemany(
        "INSERT INTO inst_ratings(inst_id,inst_name,stock_code,rating,target_price,rating_date,prev_rating) "
        "VALUES(?,?,?,?,?,?,?)",
        [(inst_map[r[1]], r[1], r[2], r[3], r[4], r[5], r[6]) for r in ratings])

    print(">>> 生成基金持仓 ...")
    holdings = gen_holdings(price, dates)
    executemany(
        "INSERT INTO fund_holdings(inst_id,inst_name,stock_code,quarter,hold_shares,hold_value,change_shares,change_pct) "
        "VALUES(?,?,?,?,?,?,?,?)",
        [(inst_map.get(h[1], 0), h[1], h[2], h[3], h[4], h[5], h[6], h[7]) for h in holdings])

    print(">>> 向量索引构建 ...")
    if not args.skip_vector:
        build_vectors(news, stocks)
    else:
        print("    (跳过)")

    print(">>> 量化策略全市场回测 ...")
    if not args.no_strategies:
        from engine.strategies import build_all, STRATEGIES
        cnt = build_all()
        print(f"    回测记录 {cnt} 条（{len(STRATEGIES)} 策略 × 全市场）")
    else:
        print("    (跳过)")

    from database import table_count
    print("=" * 50)
    print("数据库统计：")
    for t in ("stocks", "stock_daily", "news", "institutions", "inst_ratings",
              "fund_holdings", "watchlist", "analysis_cache", "market_index"):
        print(f"  {t:16s} {table_count(t)} 条")
    print("✅ 数据生成完成")


if __name__ == "__main__":
    main()