cosmo/backend/scripts/migrate_interstellar_data.py

#!/usr/bin/env python3
"""
迁移 static_data 中的 interstellar 数据到 star_systems 和 celestial_bodies 表
包含自动中文名翻译功能
"""
import asyncio
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from sqlalchemy import select, func, update
from sqlalchemy.dialects.postgresql import insert
from app.database import AsyncSessionLocal
from app.models.db.static_data import StaticData
from app.models.db.star_system import StarSystem
from app.models.db.celestial_body import CelestialBody


# 恒星名称中文翻译字典（常见恒星）
STAR_NAME_ZH = {
    'Proxima Cen': '比邻星',
    "Barnard's star": '巴纳德星',
    'eps Eri': '天苑四',
    'Lalande 21185': '莱兰21185',
    '61 Cyg A': '天鹅座61 A',
    '61 Cyg B': '天鹅座61 B',
    'tau Cet': '天仓五',
    'Kapteyn': '开普敦星',
    'Lacaille 9352': '拉卡伊9352',
    'Ross 128': '罗斯128',
    'Wolf 359': '狼359',
    'Sirius': '天狼星',
    'Alpha Centauri': '南门二',
    'TRAPPIST-1': 'TRAPPIST-1',
    'Kepler-442': '开普勒-442',
    'Kepler-452': '开普勒-452',
    'Gliese 581': '格利泽581',
    'Gliese 667C': '格利泽667C',
    'HD 40307': 'HD 40307',
}

# 常见恒星系后缀翻译
SYSTEM_SUFFIX_ZH = {
    'System': '系统',
    'system': '系统',
}


def translate_star_name(english_name: str) -> str:
    """
    翻译恒星名称为中文
    优先使用字典，否则保留英文名
    """
    # 直接匹配
    if english_name in STAR_NAME_ZH:
        return STAR_NAME_ZH[english_name]

    # 移除常见后缀尝试匹配
    base_name = english_name.replace(' A', '').replace(' B', '').replace(' C', '').strip()
    if base_name in STAR_NAME_ZH:
        suffix = english_name.replace(base_name, '').strip()
        return STAR_NAME_ZH[base_name] + suffix

    # Kepler/TRAPPIST 等编号星
    if english_name.startswith('Kepler-'):
        return f'开普勒-{english_name.split("-")[1]}'
    if english_name.startswith('TRAPPIST-'):
        return f'TRAPPIST-{english_name.split("-")[1]}'
    if english_name.startswith('Gliese '):
        return f'格利泽{english_name.split(" ")[1]}'
    if english_name.startswith('GJ '):
        return f'GJ {english_name.split(" ")[1]}'
    if english_name.startswith('HD '):
        return f'HD {english_name.split(" ")[1]}'
    if english_name.startswith('HIP '):
        return f'HIP {english_name.split(" ")[1]}'

    # 默认返回英文名
    return english_name


def translate_system_name(english_name: str) -> str:
    """翻译恒星系名称"""
    if ' System' in english_name:
        star_name = english_name.replace(' System', '').strip()
        star_name_zh = translate_star_name(star_name)
        return f'{star_name_zh}系统'
    return translate_star_name(english_name)


def translate_planet_name(english_name: str) -> str:
    """
    翻译系外行星名称
    格式：恒星名 + 行星字母
    """
    # 分离恒星名和行星字母
    parts = english_name.rsplit(' ', 1)
    if len(parts) == 2:
        star_name, planet_letter = parts
        star_name_zh = translate_star_name(star_name)
        return f'{star_name_zh} {planet_letter}'
    return english_name


async def deduplicate_planets(planets: list) -> list:
    """
    去除重复的行星记录
    保留字段最完整的记录
    """
    if not planets:
        return []

    planet_map = {}
    for planet in planets:
        name = planet.get('name', '')
        if not name:
            continue

        if name not in planet_map:
            planet_map[name] = planet
        else:
            # 比较字段完整度
            existing = planet_map[name]
            existing_fields = sum(1 for v in existing.values() if v is not None and v != '')
            current_fields = sum(1 for v in planet.values() if v is not None and v != '')

            if current_fields > existing_fields:
                planet_map[name] = planet

    return list(planet_map.values())


async def migrate_star_systems():
    """迁移恒星系统数据"""
    async with AsyncSessionLocal() as session:
        print("=" * 60)
        print("开始迁移系外恒星系数据...")
        print("=" * 60)

        # 读取所有 interstellar 数据
        result = await session.execute(
            select(StaticData)
            .where(StaticData.category == 'interstellar')
            .order_by(StaticData.name)
        )
        interstellar_data = result.scalars().all()

        print(f"\n📊 共找到 {len(interstellar_data)} 个恒星系统")

        migrated_systems = 0
        migrated_planets = 0
        skipped_systems = 0

        for star_data in interstellar_data:
            try:
                data = star_data.data
                star_name = star_data.name

                # 翻译中文名
                star_name_zh = translate_star_name(star_name)
                system_name = f"{star_name} System"
                system_name_zh = translate_system_name(system_name)

                # 创建恒星系统记录
                system = StarSystem(
                    name=system_name,
                    name_zh=system_name_zh,
                    host_star_name=star_name,
                    distance_pc=data.get('distance_pc'),
                    distance_ly=data.get('distance_ly'),
                    ra=data.get('ra'),
                    dec=data.get('dec'),
                    position_x=data.get('position', {}).get('x') if 'position' in data else None,
                    position_y=data.get('position', {}).get('y') if 'position' in data else None,
                    position_z=data.get('position', {}).get('z') if 'position' in data else None,
                    spectral_type=data.get('spectral_type'),
                    radius_solar=data.get('radius_solar'),
                    mass_solar=data.get('mass_solar'),
                    temperature_k=data.get('temperature_k'),
                    magnitude=data.get('magnitude'),
                    color=data.get('color', '#FFFFFF'),
                    planet_count=0,  # 将在迁移行星后更新
                    description=f"距离地球 {data.get('distance_ly', 0):.2f} 光年的恒星系统。"
                )

                session.add(system)
                await session.flush()  # 获取 system.id

                print(f"\n✅ 恒星系: {system_name} ({system_name_zh})")
                print(f"   距离: {data.get('distance_pc', 0):.2f} pc (~{data.get('distance_ly', 0):.2f} ly)")

                # 处理行星数据
                planets = data.get('planets', [])
                if planets:
                    # 去重
                    unique_planets = await deduplicate_planets(planets)
                    print(f"   行星: {len(planets)} 条记录 → {len(unique_planets)} 颗独立行星（去重 {len(planets) - len(unique_planets)} 条）")

                    # 迁移行星
                    for planet_data in unique_planets:
                        planet_name = planet_data.get('name', '')
                        if not planet_name:
                            continue

                        planet_name_zh = translate_planet_name(planet_name)

                        # 创建系外行星记录
                        planet = CelestialBody(
                            id=f"exo-{system.id}-{planet_name.replace(' ', '-')}",  # 生成唯一ID
                            name=planet_name,
                            name_zh=planet_name_zh,
                            type='planet',
                            system_id=system.id,
                            description=f"{system_name_zh}的系外行星。",
                            extra_data={
                                'semi_major_axis_au': planet_data.get('semi_major_axis_au'),
                                'period_days': planet_data.get('period_days'),
                                'eccentricity': planet_data.get('eccentricity'),
                                'radius_earth': planet_data.get('radius_earth'),
                                'mass_earth': planet_data.get('mass_earth'),
                                'temperature_k': planet_data.get('temperature_k'),
                            }
                        )

                        session.add(planet)
                        migrated_planets += 1
                        print(f"      • {planet_name} ({planet_name_zh})")

                    # 更新恒星系的行星数量
                    system.planet_count = len(unique_planets)

                migrated_systems += 1

                # 每100个系统提交一次
                if migrated_systems % 100 == 0:
                    await session.commit()
                    print(f"\n💾 已提交 {migrated_systems} 个恒星系统...")

            except Exception as e:
                print(f"\n❌ 错误：迁移 {star_name} 失败 - {str(e)[:200]}")
                skipped_systems += 1
                # 简单回滚，继续下一个
                try:
                    await session.rollback()
                except:
                    pass
                continue

        # 最终提交
        await session.commit()

        print("\n" + "=" * 60)
        print("迁移完成！")
        print("=" * 60)
        print(f"✅ 成功迁移恒星系: {migrated_systems}")
        print(f"✅ 成功迁移行星: {migrated_planets}")
        print(f"⚠️  跳过的恒星系: {skipped_systems}")
        print(f"📊 平均每个恒星系: {migrated_planets / migrated_systems:.1f} 颗行星")


async def update_solar_system_count():
    """更新太阳系的天体数量"""
    async with AsyncSessionLocal() as session:
        result = await session.execute(
            select(func.count(CelestialBody.id))
            .where(CelestialBody.system_id == 1)
        )
        count = result.scalar()

        await session.execute(
            update(StarSystem)
            .where(StarSystem.id == 1)
            .values(planet_count=count - 1)  # 减去太阳本身
        )
        await session.commit()

        print(f"\n✅ 更新太阳系天体数量: {count} (不含太阳: {count - 1})")


async def verify_migration():
    """验证迁移结果"""
    async with AsyncSessionLocal() as session:
        print("\n" + "=" * 60)
        print("验证迁移结果...")
        print("=" * 60)

        # 统计恒星系
        result = await session.execute(select(func.count(StarSystem.id)))
        system_count = result.scalar()
        print(f"\n📊 恒星系统总数: {system_count}")

        # 统计各系统的行星数量
        result = await session.execute(
            select(StarSystem.name, StarSystem.name_zh, StarSystem.planet_count)
            .order_by(StarSystem.planet_count.desc())
            .limit(10)
        )
        print("\n🏆 行星最多的恒星系（前10）:")
        for name, name_zh, count in result:
            print(f"   {name} ({name_zh}): {count} 颗行星")

        # 统计天体类型分布
        result = await session.execute(
            select(CelestialBody.type, CelestialBody.system_id, func.count(CelestialBody.id))
            .group_by(CelestialBody.type, CelestialBody.system_id)
            .order_by(CelestialBody.system_id, CelestialBody.type)
        )
        print("\n📈 天体类型分布:")
        for type_, system_id, count in result:
            system_name = "太阳系" if system_id == 1 else f"系外恒星系"
            print(f"   {system_name} - {type_}: {count}")


async def main():
    """主函数"""
    print("\n" + "=" * 60)
    print("Cosmo 系外恒星系数据迁移工具")
    print("=" * 60)

    try:
        # 执行迁移
        await migrate_star_systems()

        # 更新太阳系统计
        await update_solar_system_count()

        # 验证结果
        await verify_migration()

        print("\n✅ 所有操作完成！")

    except Exception as e:
        print(f"\n❌ 迁移失败: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())