CostPrediction/scripts/calculate_features.py

"""
计算并填充数据库中的派生/计算特征字段，包括：
1. 火箭炮：火力密度、射程比、各项评分
2. 巡飞弹：长宽比、重量射程比等 + 估算缺失的原始字段
   (engine_power_kw, engine_thrust_n, min/max_altitude_m,
    guidance_accuracy_m, datalink_range_km)

用法：
    python -m scripts.calculate_features
"""

import sys
import os
import math

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.database.db_connection import get_db_connection, init_db
from src.logger import setup_logger

logger = setup_logger(__name__)


def clamp(value, low=1, high=10):
    return max(low, min(high, value))


def safe_float(val, default=0.0):
    if val is None:
        return default
    return float(val)


# ────────────────────────────────────────────────
#  火箭炮
# ────────────────────────────────────────────────

def calculate_rocket_features(conn):
    cursor = conn.cursor()

    cursor.execute("""
        SELECT rap.id, rap.equipment_id,
               rap.rate_of_fire, rap.rocket_weight_kg, rap.rocket_length_m,
               rap.rocket_diameter_mm, rap.min_range_km, rap.max_range_km,
               rap.speed_kmh, rap.power_hp, rap.combat_weight_kg,
               rap.travel_range_km, rap.firing_angle_horizontal,
               rap.firing_angle_vertical,
               cp.length_m, cp.width_m, cp.weight_kg
        FROM rocket_artillery_params rap
        LEFT JOIN common_params cp ON rap.equipment_id = cp.equipment_id
    """)

    rows = cursor.fetchall()

    # ── 第一遍：收集原始值，计算 min/max（用于归一化）──
    all_vals = {
        'spd': [], 'power_ratio': [], 'tr': [],
        'angle_cov': [], 'deploy_factor': [], 'rof': [],
        'cw': []
    }
    for row in rows:
        spd = safe_float(row['speed_kmh'])
        pwr = safe_float(row['power_hp'])
        cw = safe_float(row['combat_weight_kg'])
        tr = safe_float(row['travel_range_km'])
        ah = safe_float(row['firing_angle_horizontal'])
        av = safe_float(row['firing_angle_vertical'])
        rof = safe_float(row['rate_of_fire'])

        all_vals['spd'].append(spd)
        all_vals['power_ratio'].append(pwr / max(1, cw))
        all_vals['tr'].append(tr)
        all_vals['angle_cov'].append((ah / 360) * (av / 90))
        all_vals['deploy_factor'].append(1.0 / max(1, cw))  # 越小越难部署
        all_vals['rof'].append(rof)
        all_vals['cw'].append(cw)

    def norm10(v, vals):
        """Min-max 归一化到 1-10"""
        mn, mx = min(vals), max(vals)
        span = mx - mn
        if span == 0:
            return 5.0
        return 1.0 + (v - mn) / span * 9.0

    # ── 第二遍：逐条计算并更新 ──
    updates = []

    for row in rows:
        rof = safe_float(row['rate_of_fire'])
        rw = safe_float(row['rocket_weight_kg'])
        rl = safe_float(row['rocket_length_m'])
        rd = safe_float(row['rocket_diameter_mm'])
        min_r = safe_float(row['min_range_km'])
        max_r = safe_float(row['max_range_km'])
        spd = safe_float(row['speed_kmh'])
        pwr = safe_float(row['power_hp'])
        cw = safe_float(row['combat_weight_kg'])
        tr = safe_float(row['travel_range_km'])
        ah = safe_float(row['firing_angle_horizontal'])
        av = safe_float(row['firing_angle_vertical'])

        # 火力密度 (kg/min)
        fire_density = round(rw * rof, 2)

        # 射程比
        range_ratio = round(max_r / max(1, min_r), 4)

        # 火箭弹功重比
        if rd > 0 and rl > 0:
            rocket_power_ratio = round((rw * rl * rof) / (rd * 10), 4)
        else:
            rocket_power_ratio = 0.0

        # 平台效率 (m/kg)
        platform_efficiency = round(max_r / max(1, cw) * 1000, 4)

        power_ratio = pwr / max(1, cw)
        angle_cov = (ah / 360) * (av / 90)
        deploy_factor = 1.0 / max(1, cw)

        # 机动性评分 (1-10)：综合速度、功重比、行程，min-max 归一化
        s_speed = norm10(spd, all_vals['spd'])
        s_power = norm10(power_ratio, all_vals['power_ratio'])
        s_range = norm10(tr, all_vals['tr'])
        mobility_score = int(round((s_speed + s_power + s_range) / 3))

        # 战备状态评分 (1-10)：射界覆盖、部署便捷度（轻量化）、射速
        s_angle = norm10(angle_cov, all_vals['angle_cov'])
        s_deploy = norm10(deploy_factor, all_vals['deploy_factor'])
        s_rof = norm10(rof, all_vals['rof'])
        combat_readiness_score = int(round((s_angle + s_deploy + s_rof) / 3))

        # 部署评分 (1-10)：轻量化 + 速度
        s_weight = 11.0 - norm10(cw, all_vals['cw'])  # 越轻越高
        s_speed2 = norm10(spd, all_vals['spd'])
        deployment_score = int(round((s_weight + s_speed2) / 2))

        # 地形适应性评分 (1-10)：功重比 + 速度
        terrain_score = int(round((s_power + s_speed2) / 2))

        updates.append((
            fire_density, range_ratio, mobility_score,
            combat_readiness_score, deployment_score, terrain_score,
            rocket_power_ratio, platform_efficiency,
            row['id']
        ))

    if updates:
        cursor.executemany("""
            UPDATE rocket_artillery_params
            SET fire_density = ?, range_ratio = ?, mobility_score = ?,
                combat_readiness_score = ?, deployment_score = ?,
                terrain_adaptability_score = ?, rocket_power_ratio = ?,
                platform_efficiency = ?
            WHERE id = ?
        """, updates)
        logger.info(f"火箭炮：更新了 {len(updates)} 条记录")

    return len(updates)


# ────────────────────────────────────────────────
#  巡飞弹
# ────────────────────────────────────────────────

def estimate_guidance_accuracy(guidance_system):
    """根据制导系统文字估算制导精度 (m)"""
    gs = (guidance_system or '').lower()
    # 统计制导模式数量作为精度指标
    modes = [m.strip() for m in guidance_system.split('/') if m.strip()] if guidance_system else []
    mode_count = len(modes)

    # 基础精度：模式越多越精确
    base_accuracy = {
        1: 50, 2: 30, 3: 20, 4: 15, 5: 10, 6: 5, 7: 3, 8: 2, 9: 1.5, 10: 1,
    }.get(mode_count, max(1, 60 - mode_count * 6))

    # 根据制导类型微调
    if any(k in gs for k in ['卫星', 'satellite', 'gps']):
        base_accuracy = max(3, base_accuracy * 0.7)
    if any(k in gs for k in ['激光', 'laser']):
        base_accuracy = max(1, base_accuracy * 0.4)
    if any(k in gs for k in ['红外', 'infrared', 'ir', '光电']):
        base_accuracy = max(2, base_accuracy * 0.6)
    if any(k in gs for k in ['雷达', 'radar']):
        base_accuracy = max(5, base_accuracy * 0.8)
    if any(k in gs for k in ['ai', '辅助']):
        base_accuracy = base_accuracy * 0.8

    return round(base_accuracy, 1)


def calculate_guidance_score(guidance_system, estimated_accuracy):
    """计算制导系统评分 (1-10)"""
    gs = (guidance_system or '').lower()
    modes = [m.strip() for m in guidance_system.split('/') if m.strip()] if guidance_system else []

    # 制导模式数量 (1-7 映射到 1-6 分)
    mode_score = clamp(len(modes) / 2, 0.5, 6)

    # 精度分 (精度越高分越高)
    if estimated_accuracy <= 3:
        accuracy_score = 5
    elif estimated_accuracy <= 5:
        accuracy_score = 4.5
    elif estimated_accuracy <= 10:
        accuracy_score = 4
    elif estimated_accuracy <= 20:
        accuracy_score = 3
    elif estimated_accuracy <= 30:
        accuracy_score = 2
    elif estimated_accuracy <= 50:
        accuracy_score = 1
    else:
        accuracy_score = 0.5

    # 高端制导方式加成
    high_end_bonus = 0
    if any(k in gs for k in ['卫星', 'satellite']):
        high_end_bonus += 0.5

    return int(round(clamp(mode_score + accuracy_score + high_end_bonus, 1, 10)))


def calculate_warhead_score(warhead_type, warhead_weight_kg):
    """计算战斗部威力评分 (1-10)"""
    wt_lower = (warhead_type or '').lower()

    # 重量得分: 5kg→2, 10kg→4, 20kg→6, 30kg→8, 50kg→10
    weight_score = clamp(warhead_weight_kg * 0.2, 1, 10)

    # 类型加成
    type_kw_scores = {
        '模块化': 3, 'modular': 3,
        '高爆': 2, 'explosive': 2, 'he': 2,
        '云爆': 3, 'thermobaric': 3,
        '破片杀伤': 1.5, 'fragmentation': 1.5,
        '穿甲': 2.5, 'armor': 2.5, 'penetrat': 2.5,
        '动能': 2, 'kinetic': 2,
        '双用': 2, 'dual': 2,
        '破甲': 2, 'heat': 2,
        '子母': 2, 'cluster': 2, 'submunition': 2
    }
    type_bonus = 0
    for kw, val in type_kw_scores.items():
        if kw in wt_lower:
            type_bonus = max(type_bonus, val)

    return int(round(clamp(weight_score + type_bonus, 1, 10)))


def calculate_loitering_features(conn):
    """计算巡飞弹的派生特征 + 估算缺失的原始字段"""
    cursor = conn.cursor()

    cursor.execute("""
        SELECT lmp.id, lmp.equipment_id,
               lmp.wingspan_m, lmp.warhead_weight_kg, lmp.max_speed_ms,
               lmp.max_range_km, lmp.guidance_accuracy_m, lmp.datalink_range_km,
               lmp.engine_power_kw, lmp.engine_thrust_n,
               lmp.min_altitude_m, lmp.max_altitude_m,
               lmp.guidance_system, lmp.warhead_type,
               lmp.cruise_speed_kmh, lmp.endurance_min,
               lmp.ceiling_altitude_m, lmp.max_payload_kg,
               lmp.combat_radius_km,
               cp.length_m, cp.width_m, cp.weight_kg
        FROM loitering_munition_params lmp
        LEFT JOIN common_params cp ON lmp.equipment_id = cp.equipment_id
    """)

    rows = cursor.fetchall()
    updates = []

    for row in rows:
        length = safe_float(row['length_m'])
        width = safe_float(row['width_m'])
        weight = safe_float(row['weight_kg'])
        warhead_weight = safe_float(row['warhead_weight_kg'])
        max_speed = safe_float(row['max_speed_ms'])
        cruise_speed = safe_float(row['cruise_speed_kmh'])
        max_range = safe_float(row['max_range_km'])
        endurance = safe_float(row['endurance_min'])
        ceiling = safe_float(row['ceiling_altitude_m'])
        payload = safe_float(row['max_payload_kg'])
        combat_radius = safe_float(row['combat_radius_km'])
        guidance_system = row['guidance_system'] or ''
        warhead_type = row['warhead_type'] or ''

        # ── 1. 长宽比 ──
        length_width_ratio = round(length / max(0.1, width), 4)

        # ── 2. 重量射程比 (kg/km) ──
        weight_range_ratio = round(weight / max(1, max_range), 4)

        # ── 3. 速度重量比 ──
        speed_weight_ratio = round(max_speed / max(1, weight), 4)

        # ── 4. 估算缺失的原始字段 ──

        # 制导精度 (m)：当前全为 0，根据制导系统文字估算
        estimated_accuracy = estimate_guidance_accuracy(guidance_system)

        # 数据链距离 (km)：用战斗半径或最大射程的 80% 估算
        base_ref = combat_radius if combat_radius > 0 else max_range
        estimated_datalink = round(base_ref * 0.85, 1)

        # 发动机功率 (kw)：基于重量和速度估算
        # 小型无人机/巡飞弹功率密度约 0.05-0.3 kW/kg
        total_weight = weight + payload
        power_density = 0.08 + (max_speed / 200) * 0.12  # 越快功率密度越高
        power_density = min(power_density, 0.28)  # 上限避免重型弹药估算过高
        estimated_engine_kw = round(total_weight * power_density, 1)

        # 发动机推力 (N)：基于总重估算，推力 ≈ 重量 * 推重比
        # 巡飞弹推重比通常 0.2-0.5，但重型滑翔弹药不具备大推力发动机
        thrust_ratio = 0.15 + (max_speed / 150) * 0.20
        thrust_ratio = min(thrust_ratio, 0.3)  # 上限
        estimated_thrust_n = round(total_weight * 9.8 * thrust_ratio, 1)

        # 最小作战高度 (m)：巡飞弹通常 50-500m
        if cruise_speed > 0:
            estimated_min_alt = round(max(30, min(300, cruise_speed * 0.8)), 0)
        else:
            estimated_min_alt = round(max(30, min(300, max_speed * 2.5)), 0)

        # 最大作战高度 (m)：优先用 ceiling_altitude_m，否则估算
        if ceiling > 0:
            estimated_max_alt = ceiling
        else:
            estimated_max_alt = round(max(500, min(10000, endurance * 15)), 0)

        # ── 5. 制导系统评分 (1-10) ──
        guidance_system_score = calculate_guidance_score(guidance_system, estimated_accuracy)

        # ── 6. 战斗部威力评分 (1-10) ──
        warhead_power_score = calculate_warhead_score(warhead_type, warhead_weight)

        updates.append((
            length_width_ratio, weight_range_ratio, speed_weight_ratio,
            guidance_system_score, warhead_power_score,
            estimated_accuracy, estimated_datalink,
            estimated_engine_kw, estimated_thrust_n,
            estimated_min_alt, estimated_max_alt,
            row['id']
        ))

    if updates:
        cursor.executemany("""
            UPDATE loitering_munition_params
            SET length_width_ratio = ?, weight_range_ratio = ?,
                speed_weight_ratio = ?, guidance_system_score = ?,
                warhead_power_score = ?,
                guidance_accuracy_m = ?, datalink_range_km = ?,
                engine_power_kw = ?, engine_thrust_n = ?,
                min_altitude_m = ?, max_altitude_m = ?
            WHERE id = ?
        """, updates)
        logger.info(f"巡飞弹：更新了 {len(updates)} 条记录")

    return len(updates)


def main():
    init_db()
    conn = None
    try:
        with get_db_connection() as conn:
            rc = calculate_rocket_features(conn)
            lc = calculate_loitering_features(conn)
            conn.commit()
            print(f"\n完成！")
            print(f"  火箭炮：{rc} 条记录已更新")
            print(f"  巡飞弹：{lc} 条记录已更新")
    except Exception as e:
        if conn:
            conn.rollback()
        logger.error(f"特征计算失败: {e}")
        raise


if __name__ == '__main__':
    main()