#!/usr/bin/env python3
"""[第1步] 补 adj_factor(复权因子)列到四池日线。
- 有 hfq_close+raw_close 的: factor = hfq_close/raw_close (已验证与 tushare 一致)
- 缺 raw 的(66只): 从 tushare 拉全史 factor
- 统一列顺序; 保留现有 qfq/hfq/raw 列(第2步再消减)
DRY=1 只扫描不写盘。
"""
import os, glob
import pandas as pd

BASE = os.path.expanduser('~/zt_app/backtest_zt_full')
DRY = os.environ.get('DRY') == '1'

STD = ['ts_code','trade_date',
       'qfq_open','qfq_high','qfq_low','qfq_close',
       'hfq_open','hfq_high','hfq_low','hfq_close',
       'raw_open','raw_high','raw_low','raw_close',
       'adj_factor','pre_close','change','pct_chg','vol','amount']

files = sorted(glob.glob(f'{BASE}/daily_*/*.csv'))
need_ts, already, ok = [], 0, 0
for f in files:
    with open(f) as fh:
        cols = fh.readline().strip().split(',')
    if 'adj_factor' in cols:
        already += 1
    elif 'hfq_close' in cols and 'raw_close' in cols:
        ok += 1
    else:
        need_ts.append(f)

print(f'扫描: 共{len(files)} | 已有factor {already} | 可反推 {ok} | 需tushare {len(need_ts)}')

if not DRY:
    written = 0
    for f in files:
        if f in need_ts:
            continue
        if os.path.basename(f) in [os.path.basename(n) for n in need_ts]:
            continue
        try:
            d = pd.read_csv(f, dtype={'trade_date': str})
            if 'adj_factor' in d.columns:
                continue
            d['adj_factor'] = (d['hfq_close'] / d['raw_close']).astype(float).round(6)
            keep = [c for c in STD if c in d.columns] + [c for c in d.columns if c not in STD]
            d = d[keep]
            d.to_csv(f, index=False)
            written += 1
            if written % 500 == 0:
                print(f'  反推已写 {written}', flush=True)
        except Exception as e:
            print(f'  反推失败 {f}: {e}', flush=True)
    print(f'反推写入 {written} 只')

    if need_ts:
        import tushare as ts
        ts.set_token('edf6739fe1a4de0d747600cc753a8b4bf335cf27ef0f5aea2d2aa64c')
        pro = ts.pro_api()
        nts = 0
        for f in need_ts:
            code = os.path.basename(f)[:-4]
            try:
                af = pro.adj_factor(ts_code=code, start_date='20210101', end_date='20260911')
                if af is None or len(af) == 0:
                    continue
                af = af[['trade_date', 'adj_factor']].sort_values('trade_date')
                d = pd.read_csv(f, dtype={'trade_date': str})
                d = d.merge(af, on='trade_date', how='left')
                keep = [c for c in STD if c in d.columns] + [c for c in d.columns if c not in STD]
                d = d[keep]
                d.to_csv(f, index=False)
                nts += 1
            except Exception as e:
                print(f'  tushare失败 {code}: {e}', flush=True)
        print(f'tushare 补写 {nts} 只')
print('完成')