# celestial_map.py

import networkx as nx
import matplotlib.pyplot as plt

class CelestialMap:
    """
    基于“差序格局”思想，构建一个描绘产业链上下游“引力关系”的图模型。

    “差序格局”是中国社会学家费孝通提出的概念，指社会关系像水波纹一样，
    由自己这个中心推出去，一层层地变淡。

    在这个模型中：
    - 恒星 (Star): 代表产业链的核心或最上游的公司/板块，是引力中心。
    - 行星 (Planet): 代表与核心紧密相关的公司/板块，围绕恒星运转。
    - 卫星 (Satellite): 代表更下游或关联度更弱的公司/板块，围绕行星运转。

    引力 (Gravity): 代表实体间的影响力或情绪传导强度，距离越远，引力越弱。
    """
    def __init__(self):
        self.graph = nx.DiGraph()

    def add_entity(self, name, entity_type, **kwargs):
        """
        添加一个天体实体。

        :param name: 实体名称 (e.g., 'NVIDIA')
        :param entity_type: 实体类型 ('star', 'planet', 'satellite')
        :param kwargs: 其他属性 (e.g., market_cap, sector)
        """
        self.graph.add_node(name, entity_type=entity_type, **kwargs)

    def establish_influence(self, source, target, gravity_strength=1.0, relationship_type='supply_chain'):
        """
        在两个实体间建立影响关系（引力）。

        差序格局的体现：引力强度(gravity_strength)默认为1.0，
        当从恒星 -> 行星 -> 卫星传递时，可以通过计算路径权重来体现引力的衰减。

        :param source: 施加影响的实体名称
        :param target: 受到影响的实体名称
        :param gravity_strength: 基础引力强度 (1.0 / 距离)
        :param relationship_type: 关系类型 (e.g., 'supply_chain', 'competitor', 'capital')
        """
        if source not in self.graph:
            raise ValueError(f"Source entity '{source}' not found in the map.")
        if target not in self.graph:
            raise ValueError(f"Target entity '{target}' not found in the map.")

        # 使用 1/gravity_strength 作为边的权重，模拟距离
        # 强引力 = 低权重 = 短距离
        distance = 1.0 / gravity_strength
        self.graph.add_edge(source, target, weight=distance, relationship=relationship_type)

    def get_influence_path(self, start_node, end_node):
        """
        计算两个实体间的最短（最强）影响路径。

        :param start_node: 路径起点
        :param end_node: 路径终点
        :return: 路径节点列表和总路径长度（总距离）
        """
        try:
            path = nx.shortest_path(self.graph, source=start_node, target=end_node, weight='weight')
            length = nx.shortest_path_length(self.graph, source=start_node, target=end_node, weight='weight')
            return path, length
        except nx.NetworkXNoPath:
            return None, float('inf')

    def get_downstream_entities(self, start_node, max_depth=None):
        """
        获取一个实体的所有下游受影响实体。

        :param start_node: 起点实体
        :param max_depth: 搜索的最大深度（层级）
        :return: 下游实体列表
        """
        if start_node not in self.graph:
            return []
        return list(nx.dfs_preorder_nodes(self.graph, source=start_node, depth_limit=max_depth))

    def visualize(self, layout='spring'):
        """
        将星图可视化。
        """
        pos = None
        if layout == 'spring':
            pos = nx.spring_layout(self.graph, k=0.9)
        elif layout == 'kamada_kawai':
            pos = nx.kamada_kawai_layout(self.graph)
        elif layout == 'circular':
            pos = nx.circular_layout(self.graph)
        else:
            pos = nx.spring_layout(self.graph)


        node_colors = []
        for node in self.graph.nodes(data=True):
            if node[1]['entity_type'] == 'star':
                node_colors.append('yellow')
            elif node[1]['entity_type'] == 'planet':
                node_colors.append('skyblue')
            elif node[1]['entity_type'] == 'satellite':
                node_colors.append('lightgreen')
            else:
                node_colors.append('gray')

        edge_labels = nx.get_edge_attributes(self.graph, 'relationship')

        plt.figure(figsize=(14, 10))
        nx.draw(self.graph, pos, with_labels=True, node_size=3000, node_color=node_colors, font_size=10, font_weight='bold', arrowsize=20)
        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels, font_color='red')
        plt.title("Celestial Map of Industry Influence (差序格局)")
        plt.show()

if __name__ == '__main__':
    # --- 示例：构建一个AI芯片产业链的“差序格局”图 ---

    # 1. 初始化星图
    industry_map = CelestialMap()

    # 2. 添加实体 (恒星、行星、卫星)
    industry_map.add_entity('NVIDIA', 'star', sector='Semiconductors')
    industry_map.add_entity('TSMC', 'planet', sector='Manufacturing')
    industry_map.add_entity('ASML', 'planet', sector='Equipment')
    industry_map.add_entity('Supermicro', 'planet', sector='Servers')
    industry_map.add_entity('Apple', 'satellite', sector='Consumer Electronics')
    industry_map.add_entity('Tesla', 'satellite', sector='Automotive')
    industry_map.add_entity('OpenAI', 'satellite', sector='AI Research')

    # 3. 建立影响关系 (引力)
    # 恒星 -> 行星 (强引力)
    industry_map.establish_influence('NVIDIA', 'TSMC', gravity_strength=0.9, relationship_type='chip_design')
    industry_map.establish_influence('NVIDIA', 'Supermicro', gravity_strength=0.8, relationship_type='hardware_demand')
    industry_map.establish_influence('ASML', 'TSMC', gravity_strength=0.9, relationship_type='euv_machines')

    # 行星 -> 卫星 (中等引力)
    industry_map.establish_influence('TSMC', 'Apple', gravity_strength=0.7, relationship_type='chip_supply')
    industry_map.establish_influence('Supermicro', 'OpenAI', gravity_strength=0.6, relationship_type='server_supply')
    industry_map.establish_influence('NVIDIA', 'Tesla', gravity_strength=0.5, relationship_type='ai_chip_supply') # 直接影响

    # 4. 分析与查询
    print("--- 情绪传导路径分析 ---")
    path, total_distance = industry_map.get_influence_path('ASML', 'Apple')
    if path:
        print(f"从 'ASML' 到 'Apple' 的最强影响路径: {' -> '.join(path)}")
        # 总距离越大，代表关系越疏远，影响力衰减越多
        print(f"总影响距离 (1/gravity): {total_distance:.2f}")

    path, total_distance = industry_map.get_influence_path('NVIDIA', 'OpenAI')
    if path:
        print(f"从 'NVIDIA' 到 'OpenAI' 的最强影响路径: {' -> '.join(path)}")
        print(f"总影响距离 (1/gravity): {total_distance:.2f}")


    print("\n--- 下游影响范围查询 ---")
    downstream = industry_map.get_downstream_entities('NVIDIA', max_depth=2)
    print(f"'NVIDIA' 在两层关系内的所有下游影响实体: {downstream}")

    # 5. 可视化
    industry_map.visualize(layout='kamada_kawai')