#!/usr/bin/env python3
"""重新分析四川区县预算文件 - v4
修复：1. 只从预算支出表提取 2. 正确识别金额列 3. 关键字合并 4. 金额过滤异常值
"""
import os, json, re, sys, warnings
warnings.filterwarnings('ignore')

import openpyxl
try:
    import xlrd
    HAS_XLRD = True
except:
    HAS_XLRD = False

BASE = '/data/www/files/scQxYjs'
KEYWORDS = ['网络培训','培训','能力提升','视频拍摄','党员教育片','课程制作',
            '直播','视频','课件','系统','信息化','网络','数字','设备','平台','维护','软件']

# 预算表识别关键词（Sheet名含这些才处理）
BUDGET_SHEET_KEYWORDS = ['支出', '预算', '批复', '明细', '经济分类', '三公', '政府采购', '拨款']
# 排除的Sheet名（绩效目标表等）
EXCLUDE_SHEET_KEYWORDS = ['绩效', '目标', '说明']

def is_budget_sheet(sheet_name):
    """判断是否为预算支出表"""
    name = sheet_name.lower() if sheet_name else ''
    for ex in EXCLUDE_SHEET_KEYWORDS:
        if ex in name:
            return False
    for kw in BUDGET_SHEET_KEYWORDS:
        if kw in name:
            return True
    # 默认处理（如果没有明确排除）
    return True

def extract_amount_from_row(row_values, name_col_idx, ncols):
    """从行数据中提取金额，返回(amount, col_idx)
    策略：找name_col右侧的数值列，优先标头含'预算/合计/总计/金额'的列
    """
    best_amount = 0
    best_col = -1
    
    for col_idx in range(name_col_idx + 1, min(name_col_idx + 15, ncols)):
        if col_idx >= len(row_values):
            continue
        val = row_values[col_idx]
        if val is None:
            continue
        # 尝试转为数值
        num = None
        if isinstance(val, (int, float)):
            num = float(val)
        elif isinstance(val, str):
            val_clean = val.strip().replace(',', '').replace('，', '')
            if val_clean and re.match(r'^\d+\.?\d*$', val_clean):
                num = float(val_clean)
        
        if num is not None and num > 0:
            # 合理金额范围：0.01 ~ 99999万元
            if 0.01 <= num <= 99999:
                if num > best_amount:
                    best_amount = num
                    best_col = col_idx
    
    return best_amount, best_col

def find_amount_column(ws, header_row, name_col_idx):
    """通过表头行找到金额列索引"""
    amount_cols = []
    for col_idx in range(name_col_idx + 1, min(name_col_idx + 20, ws.max_column)):
        val = ws.cell(header_row, col_idx).value
        if val and isinstance(val, str):
            v = val.strip()
            if any(kw in v for kw in ['预算', '合计', '总计', '金额', '预算数', '支出数']):
                amount_cols.append(col_idx)
    return amount_cols

def analyze_excel(filepath):
    """分析Excel文件，返回匹配项列表"""
    matches = []
    
    try:
        if filepath.endswith('.xls'):
            if not HAS_XLRD:
                return matches
            wb = xlrd.open_workbook(filepath)
            for sheet_name in wb.sheet_names():
                if not is_budget_sheet(sheet_name):
                    continue
                ws = wb.sheet_by_name(sheet_name)
                nrows = ws.nrows
                ncols = ws.ncols
                
                for row_idx in range(nrows):
                    row_values = ws.row_values(row_idx)
                    for col_idx, cell_val in enumerate(row_values):
                        if not cell_val or not isinstance(cell_val, str):
                            continue
                        text = str(cell_val).strip()
                        if len(text) < 2:
                            continue
                        for kw in KEYWORDS:
                            if kw in text:
                                amount, amt_col = extract_amount_from_row(row_values, col_idx, ncols)
                                matches.append({
                                    '关键字': kw,
                                    '内容': text[:100],
                                    '位置': f'Sheet:{sheet_name},R{row_idx+1},C{col_idx+1}',
                                    '金额': round(amount, 2) if amount else 0,
                                    '文件': filepath
                                })
            return matches
        
        # xlsx
        wb = openpyxl.load_workbook(filepath, data_only=True, read_only=True)
        for sheet_name in wb.sheetnames:
            if not is_budget_sheet(sheet_name):
                continue
            ws = wb[sheet_name]
            nrows = ws.max_row
            ncols = ws.max_column
            if nrows > 5000:
                nrows = 5000
            
            for row_idx in range(1, nrows + 1):
                row_values = []
                for col_idx in range(1, ncols + 1):
                    row_values.append(ws.cell(row_idx, col_idx).value)
                
                for col_idx, cell_val in enumerate(row_values):
                    if not cell_val:
                        continue
                    text = str(cell_val).strip()
                    if len(text) < 2:
                        continue
                    for kw in KEYWORDS:
                        if kw in text:
                            amount, amt_col = extract_amount_from_row(row_values, col_idx, ncols)
                            if amt_col >= 0:
                                pos = f'Sheet:{sheet_name},R{row_idx},C{amt_col+1}'
                            else:
                                pos = f'Sheet:{sheet_name},R{row_idx}'
                            matches.append({
                                '关键字': kw,
                                '内容': text[:100],
                                '位置': pos,
                                '金额': round(amount, 2) if amount else 0,
                                '文件': filepath
                            })
        wb.close()
    except Exception as e:
        sys.stderr.write(f"  Excel错误 {filepath}: {e}\n")
    
    return matches

def analyze_pdf(filepath):
    """分析PDF文件"""
    matches = []
    try:
        import pdfplumber
        with pdfplumber.open(filepath) as pdf:
            for page_num, page in enumerate(pdf.pages[:50]):
                tables = page.extract_tables()
                for table in tables:
                    for row_idx, row in enumerate(table):
                        for col_idx, cell in enumerate(row):
                            if not cell:
                                continue
                            text = str(cell).strip()
                            if len(text) < 2:
                                continue
                            for kw in KEYWORDS:
                                if kw in text:
                                    # 从同行右侧找金额
                                    amount = 0
                                    for c in range(col_idx + 1, min(col_idx + 10, len(row))):
                                        v = row[c]
                                        if v:
                                            v_clean = str(v).strip().replace(',', '')
                                            if re.match(r'^\d+\.?\d*$', v_clean):
                                                num = float(v_clean)
                                                if 0.01 <= num <= 99999:
                                                    amount = num
                                                    break
                                    matches.append({
                                        '关键字': kw,
                                        '内容': text[:100],
                                        '位置': f'PDF:P{page_num+1},R{row_idx+1}',
                                        '金额': round(amount, 2) if amount else 0,
                                        '文件': filepath
                                    })
                
                # 也提取文本
                text = page.extract_text()
                if text:
                    for line_idx, line in enumerate(text.split('\n')):
                        for kw in KEYWORDS:
                            if kw in line:
                                # 提取金额
                                amt_match = re.search(r'(\d+\.?\d*)\s*万元?', line)
                                amount = 0
                                if amt_match:
                                    num = float(amt_match.group(1))
                                    if 0.01 <= num <= 99999:
                                        amount = num
                                matches.append({
                                    '关键字': kw,
                                    '内容': line.strip()[:100],
                                    '位置': f'PDF:P{page_num+1},L{line_idx+1}',
                                    '金额': round(amount, 2) if amount else 0,
                                    '文件': filepath
                                })
    except Exception as e:
        sys.stderr.write(f"  PDF错误 {filepath}: {e}\n")
    
    return matches

def analyze_doc(filepath):
    """分析Word文件"""
    matches = []
    try:
        from docx import Document
        doc = Document(filepath)
        
        # 表格中的内容
        for table_idx, table in enumerate(doc.tables):
            for row_idx, row in enumerate(table.rows):
                cells = [cell.text.strip() for cell in row.cells]
                for col_idx, text in enumerate(cells):
                    if len(text) < 2:
                        continue
                    for kw in KEYWORDS:
                        if kw in text:
                            # 从同行右侧找金额
                            amount = 0
                            for c in range(col_idx + 1, min(col_idx + 10, len(cells))):
                                v = cells[c]
                                if v:
                                    v_clean = v.replace(',', '')
                                    if re.match(r'^\d+\.?\d*$', v_clean):
                                        num = float(v_clean)
                                        if 0.01 <= num <= 99999:
                                            amount = num
                                            break
                            matches.append({
                                '关键字': kw,
                                '内容': text[:100],
                                '位置': f'DOC:T{table_idx+1},R{row_idx+1}',
                                '金额': round(amount, 2) if amount else 0,
                                '文件': filepath
                            })
        
        # 段落中的内容
        for para_idx, para in enumerate(doc.paragraphs):
            text = para.text.strip()
            if len(text) < 2:
                continue
            for kw in KEYWORDS:
                if kw in text:
                    amt_match = re.search(r'(\d+\.?\d*)\s*万元?', text)
                    amount = 0
                    if amt_match:
                        num = float(amt_match.group(1))
                        if 0.01 <= num <= 99999:
                            amount = num
                    matches.append({
                        '关键字': kw,
                        '内容': text[:100],
                        '位置': f'DOC:para{para_idx+1}',
                        '金额': round(amount, 2) if amount else 0,
                        '文件': filepath
                    })
    except Exception as e:
        sys.stderr.write(f"  DOC错误 {filepath}: {e}\n")
    
    return matches

def classify_unit(filepath, filename):
    """判断单位类型：组织部/党校/其他"""
    fn = filename.lower()
    fp = filepath.lower()
    if '党校' in fn or '党校' in fp:
        return '党校'
    if '组织' in fn or '组织' in fp:
        return '组织部'
    if 'zuzhi' in fn or '_org_' in fn or '_dx_' in fn:
        return '组织部' if '_org_' in fn else '党校'
    # 检查上级目录名
    parts = filepath.split('/')
    for p in parts:
        if '党校' in p:
            return '党校'
        if '组织' in p:
            return '组织部'
    return '其他'

def extract_county(filepath, filename, city):
    """提取区县名"""
    # 从文件名提取
    name = filename
    # 去掉扩展名
    for ext in ['.xlsx', '.xls', '.pdf', '.docx', '.doc', '.et']:
        if name.lower().endswith(ext):
            name = name[:-len(ext)]
    
    # 常见区县名匹配
    county_patterns = [
        '龙泉驿', '金堂', '新津', '彭州', '新都', '武侯', '温江', '青羊', 
        '青白江', '郫都', '双流', '锦江', '金牛', '成华', '简阳', '都江堰',
        '邛崃', '大邑', '浦江', '崇州',
        '峨边', '峨眉山', '沐川', '夹江', '井研', '犍为', '金口河', '沙湾',
        '市中区',
        '高县', '长宁', '兴文', '叙州', '屏山', '南溪', '江安', '翠屏',
        '平昌', '通江', '恩阳', '南江',
        '宝兴', '芦山', '天全', '石棉', '汉源', '名山', '荥经', '雨城',
        '顺庆', '高坪', '嘉陵', '营山', '蓬安', '仪陇', '西充', '阆中',
        '理县', '小金', '金川', '壤塘', '马尔康', '汶川', '茂县', '松潘', '九寨沟', '若尔盖', '阿坝', '红原',
        '理塘', '巴塘', '乡城', '得荣', '稻城', '康定', '泸定', '丹巴', '九龙', '雅江', '道孚', '炉霍', '甘孜', '新龙', '德格', '白玉', '石渠', '色达',
        '雷波', '美姑', '甘洛', '越西', '冕宁', '喜德', '昭觉', '金阳', '布拖', '德昌', '会东', '宁南', '会理', '盐源', '木里', '西昌', '普格',
        '东区', '米易', '仁和', '西区', '盐边',
        '丹棱', '青神', '洪雅', '仁寿', '东坡', '彭山',
        '乐至', '雁江', '安岳',
        '梓潼', '北川', '三台', '平武', '游仙', '涪城', '安州', '江油',
        '朝天', '剑阁', '苍溪', '青川', '旺苍', '利州', '昭化',
        '前锋', '华蓥', '邻水', '广安', '武胜', '岳池',
        '纳溪', '合江', '叙永', '古蔺', '龙马潭', '泸县', '江阳',
        '万源', '渠县', '开江', '宣汉', '达川', '通川',
        '船山', '蓬溪', '射洪', '大英', '安居',
        '隆昌', '东兴', '资中',
        '富顺', '荣县', '沿滩', '大安', '贡井', '自流井',
        '绵竹', '什邡', '广汉', '中江', '罗江',
    ]
    
    for p in county_patterns:
        if p in name:
            return p + ('区' if p in ['东坡','利州','船山','安居','前锋','龙马潭','江阳','沿滩','大安','贡井','自流井','游仙','涪城','顺庆','高坪','嘉陵','雨城','名山','沙湾','翠屏','南溪','东区','西区','昭化','朝天','青白江','新都','温江','双流','郫都','锦江','金牛','武侯','青羊','成华','龙泉驿'] else '县' if not p.endswith('市') and not p.endswith('区') and not p.endswith('州') else '')
    
    # 从路径提取
    parts = filepath.split('/')
    for p in parts:
        for cp in county_patterns:
            if cp in p:
                return cp
    
    return name[:20]

def process_all():
    """处理所有文件"""
    results = {}
    file_count = 0
    match_count = 0
    
    # 遍历所有文件
    for root, dirs, files in os.walk(BASE):
        for filename in files:
            if filename.startswith('_') or filename.startswith('.'):
                continue
            if filename.endswith('.html'):
                continue
            
            filepath = os.path.join(root, filename)
            ext = filename.lower().split('.')[-1]
            
            # 跳过非数据文件
            if ext not in ['xlsx', 'xls', 'pdf', 'doc', 'docx', 'et']:
                continue
            
            # 确定市州
            rel_path = os.path.relpath(root, BASE)
            city = rel_path.split('/')[0] if rel_path != '.' else '其他'
            
            if city not in results:
                results[city] = {'组织部': {}, '党校': {}, '其他': {}}
            
            # 分析文件
            matches = []
            if ext in ['xlsx', 'xls', 'et']:
                matches = analyze_excel(filepath)
            elif ext == 'pdf':
                matches = analyze_pdf(filepath)
            elif ext in ['doc', 'docx']:
                matches = analyze_doc(filepath)
            
            if not matches:
                continue
            
            # 分类
            unit_type = classify_unit(filepath, filename)
            county = extract_county(filepath, filename, city)
            
            if county not in results[city][unit_type]:
                results[city][unit_type][county] = []
            
            results[city][unit_type][county].extend(matches)
            match_count += len(matches)
            file_count += 1
            
            if file_count % 10 == 0:
                sys.stderr.write(f"  已处理 {file_count} 文件, {match_count} 匹配项\n")
                sys.stderr.flush()
    
    # 关键字合并 + 金额过滤
    final = {}
    for city, units in results.items():
        final[city] = {}
        for unit_type, counties in units.items():
            final[city][unit_type] = {}
            for county, matches in counties.items():
                # 按关键字合并
                merged = {}
                for m in matches:
                    kw = m['关键字']
                    # 过滤异常金额
                    amt = m.get('金额', 0) or 0
                    if amt > 99999:
                        amt = 0
                    
                    if kw not in merged:
                        merged[kw] = {
                            '关键字': kw,
                            '内容': [],
                            '位置': [],
                            '金额': 0,
                            '文件': m['文件']
                        }
                    merged[kw]['内容'].append(m['内容'])
                    merged[kw]['位置'].append(m['位置'])
                    merged[kw]['金额'] += amt
                
                # 转为列表
                merged_list = []
                for kw, info in merged.items():
                    merged_list.append({
                        '关键字': kw,
                        '内容': ' | '.join(info['内容'][:5]),
                        '位置': ' | '.join(info['位置'][:5]),
                        '金额': round(info['金额'], 2),
                        '文件': info['文件']
                    })
                
                if merged_list:
                    final[city][unit_type][county] = merged_list
    
    # 保存
    out = os.path.join(BASE, '_final_v4.json')
    with open(out, 'w', encoding='utf-8') as f:
        json.dump(final, f, ensure_ascii=False, indent=2)
    
    # 统计
    total_m = 0
    total_a = 0
    for city, units in final.items():
        for ut, counties in units.items():
            for c, matches in counties.items():
                for m in matches:
                    total_m += 1
                    total_a += m['金额']
    
    print(f"\n✅ 完成!")
    print(f"  市州: {len(final)}")
    print(f"  文件: {file_count}")
    print(f"  匹配项(合并后): {total_m}")
    print(f"  总金额: {round(total_a, 2)}万元")
    print(f"  结果文件: {out}")
    
    # 按市州输出
    for city, units in final.items():
        c_m = 0
        c_a = 0
        for ut, counties in units.items():
            for c, matches in counties.items():
                for m in matches:
                    c_m += 1
                    c_a += m['金额']
        if c_m > 0:
            print(f"  {city}: {c_m}项, {round(c_a,2)}万元")

if __name__ == '__main__':
    process_all()