utils/data_generator.py

# data_generator.py
import torch
import numpy as np
from torch_geometric.data import HeteroData
from datetime import datetime, timedelta
import random
from typing import Dict, List, Tuple
from collections import defaultdict


class IPEcosystemGenerator:
    """香港知识产权生态网络数据生成器"""

    def __init__(self, seed=42):
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)

        # 香港真实产业分布
        self.industries = [
            '金融科技', '生物医药', '人工智能', '半导体',
            '新能源', '电子商务', '物流科技', '智能制造'
        ]

        # 香港地区分布
        self.districts = [
            '中环', '湾仔', '尖沙咀', '观塘', '荃湾',
            '科学园', '数码港', '将军澳工业邨'
        ]

        # IPC分类 (国际专利分类)
        self.ipc_classes = {
            '金融科技': ['G06Q', 'G06F', 'H04L'],
            '生物医药': ['A61K', 'C07D', 'A61P', 'C12N'],
            '人工智能': ['G06N', 'G06K', 'G06T'],
            '半导体': ['H01L', 'H01S', 'G11C'],
            '新能源': ['H01M', 'H02J', 'F03D'],
            '电子商务': ['G06Q', 'H04W', 'G06F'],
            '物流科技': ['G01S', 'B65G', 'G06Q'],
            '智能制造': ['B25J', 'G05B', 'B23Q']
        }

    def generate(self,
                 n_companies: int = 500,
                 n_patents: int = 3000,
                 n_trademarks: int = 1500,
                 n_persons: int = 2000,
                 n_institutions: int = 50,
                 time_span_years: int = 10) -> HeteroData:
        """
        生成异构图数据

        参数:
            n_companies: 企业数量
            n_patents: 专利数量
            n_trademarks: 商标数量
            n_persons: 人员数量（发明人/代理人）
            n_institutions: 机构数量
            time_span_years: 时间跨度（年）
        """

        data = HeteroData()

        # ========== 生成节点 ==========
        print("🏢 生成企业节点...")
        companies = self._generate_companies(n_companies)
        data['company'].x = companies['features']
        data['company'].company_id = companies['ids']
        data['company'].industry = companies['industry']

        print("📜 生成专利节点...")
        patents = self._generate_patents(n_patents, time_span_years)
        data['patent'].x = patents['features']
        data['patent'].patent_id = patents['ids']
        data['patent'].year = patents['year']

        print("™️  生成商标节点...")
        trademarks = self._generate_trademarks(n_trademarks, time_span_years)
        data['trademark'].x = trademarks['features']
        data['trademark'].trademark_id = trademarks['ids']

        print("👤 生成人员节点...")
        persons = self._generate_persons(n_persons)
        data['person'].x = persons['features']
        data['person'].person_id = persons['ids']

        print("🏛️  生成机构节点...")
        institutions = self._generate_institutions(n_institutions)
        data['institution'].x = institutions['features']
        data['institution'].institution_id = institutions['ids']

        # ========== 生成边关系 ==========
        print("\n🔗 生成关系边...")

        # 1. 企业 → 专利 (owns)
        company_patent_edges = self._generate_ownership_edges(
            n_companies, n_patents, companies['industry'], patents['ipc']
        )
        data['company', 'owns', 'patent'].edge_index = company_patent_edges['edge_index']
        data['company', 'owns', 'patent'].edge_attr = company_patent_edges['edge_attr']

        # 2. 企业 → 商标 (holds)
        company_trademark_edges = self._generate_trademark_edges(
            n_companies, n_trademarks, avg_per_company=3
        )
        data['company', 'holds', 'trademark'].edge_index = company_trademark_edges

        # 3. 人员 → 专利 (invents)
        person_patent_edges = self._generate_invention_edges(
            n_persons, n_patents, avg_inventors_per_patent=2.5
        )
        data['person', 'invents', 'patent'].edge_index = person_patent_edges['edge_index']
        data['person', 'invents', 'patent'].edge_attr = person_patent_edges['edge_attr']

        # 4. 企业 → 人员 (employs)
        company_person_edges = self._generate_employment_edges(
            n_companies, n_persons, company_patent_edges['edge_index'],
            person_patent_edges['edge_index']
        )
        data['company', 'employs', 'person'].edge_index = company_person_edges

        # 5. 专利 → 专利 (cites)
        patent_citation_edges = self._generate_citation_edges(
            n_patents, patents['year'], avg_citations=5
        )
        data['patent', 'cites', 'patent'].edge_index = patent_citation_edges['edge_index']
        data['patent', 'cites', 'patent'].edge_attr = patent_citation_edges['edge_attr']

        # 6. 企业 → 企业 (cooperates)
        company_cooperation_edges = self._generate_cooperation_edges(
            n_companies, companies['industry'], companies['district']
        )
        data['company', 'cooperates', 'company'].edge_index = company_cooperation_edges['edge_index']
        data['company', 'cooperates', 'company'].edge_attr = company_cooperation_edges['edge_attr']

        # 7. 企业 → 机构 (collaborates)
        company_institution_edges = self._generate_collaboration_edges(
            n_companies, n_institutions, companies['industry']
        )
        data['company', 'collaborates', 'institution'].edge_index = company_institution_edges

        # 8. 机构 → 专利 (produces)
        institution_patent_edges = self._generate_institution_patent_edges(
            n_institutions, n_patents
        )
        data['institution', 'produces', 'patent'].edge_index = institution_patent_edges

        # ========== 添加反向边 ==========
        data['patent', 'owned_by', 'company'].edge_index = data['company', 'owns', 'patent'].edge_index.flip(0)
        data['trademark', 'held_by', 'company'].edge_index = data['company', 'holds', 'trademark'].edge_index.flip(0)
        data['patent', 'invented_by', 'person'].edge_index = data['person', 'invents', 'patent'].edge_index.flip(0)
        data['person', 'employed_by', 'company'].edge_index = data['company', 'employs', 'person'].edge_index.flip(0)
        data['institution', 'collaborated_by', 'company'].edge_index = data[
            'company', 'collaborates', 'institution'].edge_index.flip(0)
        data['patent', 'produced_by', 'institution'].edge_index = data[
            'institution', 'produces', 'patent'].edge_index.flip(0)

        print("\n✅ 数据生成完成!")
        self._print_statistics(data)

        return data

    def _generate_companies(self, n: int) -> Dict:
        """生成企业节点特征"""
        features = []
        industries = []
        districts = []

        for i in range(n):
            # 企业规模分布 (幂律分布)
            size = np.random.pareto(2) * 50 + 10  # 10-500人
            size = min(size, 500)

            # 成立年限 (0-30年)
            age = np.random.exponential(8)
            age = min(age, 30)

            # 研发投入比例 (0-30%)
            rd_ratio = np.random.beta(2, 5) * 0.3

            # 产业选择 (有聚类效应)
            industry_idx = np.random.choice(len(self.industries),
                                            p=self._industry_distribution())
            industry = self.industries[industry_idx]
            industries.append(industry_idx)

            # 地区选择 (产业相关)
            district_idx = self._select_district(industry)
            districts.append(district_idx)

            # 国际化程度 (0-1)
            international = np.random.beta(2, 5)

            # 创新能力评分 (0-100)
            innovation_score = rd_ratio * 200 + np.random.normal(50, 15)
            innovation_score = np.clip(innovation_score, 0, 100)

            # 年营收 (对数正态分布，单位：百万港币)
            revenue = np.random.lognormal(3, 1.5)

            feature = [
                size / 500,  # 归一化
                age / 30,
                rd_ratio / 0.3,
                international,
                innovation_score / 100,
                np.log1p(revenue) / 10,
                industry_idx / len(self.industries),
                district_idx / len(self.districts)
            ]
            features.append(feature)

        return {
            'features': torch.tensor(features, dtype=torch.float),
            'ids': [f'HK-CO-{i:05d}' for i in range(n)],
            'industry': torch.tensor(industries, dtype=torch.long),
            'district': torch.tensor(districts, dtype=torch.long)
        }

    def _generate_patents(self, n: int, time_span: int) -> Dict:
        """生成专利节点特征"""
        features = []
        years = []
        ipcs = []

        end_year = 2025
        start_year = end_year - time_span

        for i in range(n):
            # 申请年份 (近年增长趋势)
            year_prob = np.linspace(0.5, 2.0, time_span)
            year_prob = year_prob / year_prob.sum()
            year = np.random.choice(range(start_year, end_year), p=year_prob)
            years.append(year)

            # IPC分类
            industry = random.choice(self.industries)
            ipc = random.choice(self.ipc_classes[industry])
            ipcs.append(ipc)

            # 权利要求数量 (1-30)
            claims = np.random.poisson(10) + 1
            claims = min(claims, 30)

            # 发明人数量 (1-10)
            inventors = np.random.poisson(2) + 1
            inventors = min(inventors, 10)

            # 引用数量 (随时间增长)
            age = end_year - year
            citations = np.random.poisson(age * 0.8)

            # 技术领域宽度 (1-5个IPC小类)
            tech_breadth = np.random.poisson(2) + 1
            tech_breadth = min(tech_breadth, 5)

            # 专利价值评分 (0-100)
            value_score = (claims * 2 + citations * 5 + tech_breadth * 10) / 2
            value_score = min(value_score, 100)

            # 是否获得授权 (时间越长概率越高)
            grant_prob = 0.3 + min(age * 0.1, 0.6)
            is_granted = 1.0 if random.random() < grant_prob else 0.0

            feature = [
                (year - start_year) / time_span,
                claims / 30,
                inventors / 10,
                citations / 50,
                tech_breadth / 5,
                value_score / 100,
                is_granted,
                hash(ipc) % 100 / 100  # IPC编码
            ]
            features.append(feature)

        return {
            'features': torch.tensor(features, dtype=torch.float),
            'ids': [f'HK-PT-{i:06d}' for i in range(n)],
            'year': torch.tensor(years, dtype=torch.long),
            'ipc': ipcs
        }

    def _generate_trademarks(self, n: int, time_span: int) -> Dict:
        """生成商标节点特征"""
        features = []

        end_year = 2025
        start_year = end_year - time_span

        for i in range(n):
            # 注册年份
            year = np.random.choice(range(start_year, end_year))

            # 商标类别 (1-45类)
            nice_class = np.random.randint(1, 46)

            # 续展次数 (0-3)
            renewals = min(np.random.poisson(0.5), 3)

            # 商标类型: 文字/图形/组合
            tm_type = np.random.choice([0, 1, 2], p=[0.4, 0.3, 0.3])

            # 知名度评分 (0-100)
            fame_score = np.random.beta(2, 8) * 100

            # 争议记录 (0-10)
            disputes = np.random.poisson(0.3)
            disputes = min(disputes, 10)

            feature = [
                (year - start_year) / time_span,
                nice_class / 45,
                renewals / 3,
                tm_type / 2,
                fame_score / 100,
                disputes / 10
            ]
            features.append(feature)

        return {
            'features': torch.tensor(features, dtype=torch.float),
            'ids': [f'HK-TM-{i:06d}' for i in range(n)]
        }

    def _generate_persons(self, n: int) -> Dict:
        """生成人员节点特征"""
        features = []

        for i in range(n):
            # 学历 (0: 本科, 1: 硕士, 2: 博士)
            education = np.random.choice([0, 1, 2], p=[0.3, 0.5, 0.2])

            # 工作年限 (0-40年)
            experience = np.random.gamma(3, 3)
            experience = min(experience, 40)

            # 专利发明数量
            patent_count = np.random.negative_binomial(2, 0.3)
            patent_count = min(patent_count, 50)

            # 技术领域数量 (1-5)
            tech_fields = np.random.poisson(2) + 1
            tech_fields = min(tech_fields, 5)

            # H指数 (科研影响力)
            h_index = np.random.poisson(education * 5 + 2)
            h_index = min(h_index, 50)

            # 跨界合作能力 (0-1)
            collaboration = np.random.beta(3, 3)

            feature = [
                education / 2,
                experience / 40,
                patent_count / 50,
                tech_fields / 5,
                h_index / 50,
                collaboration
            ]
            features.append(feature)

        return {
            'features': torch.tensor(features, dtype=torch.float),
            'ids': [f'HK-PS-{i:05d}' for i in range(n)]
        }

    def _generate_institutions(self, n: int) -> Dict:
        """生成机构节点特征"""
        features = []

        institution_types = ['大学', '研究所', '孵化器', '政府实验室']

        for i in range(n):
            # 机构类型
            inst_type = np.random.choice(len(institution_types))

            # 建立年限
            age = np.random.exponential(15)
            age = min(age, 100)

            # 研究人员数量
            researchers = np.random.lognormal(4, 1)
            researchers = min(researchers, 1000)

            # 年度专利产出
            patent_output = np.random.poisson(20)

            # 国际排名 (归一化后的倒数)
            ranking = np.random.pareto(2) + 1
            ranking_score = 1 / ranking

            # 产学研合作数量
            industry_collab = np.random.poisson(15)

            feature = [
                inst_type / len(institution_types),
                age / 100,
                np.log1p(researchers) / np.log1p(1000),
                patent_output / 100,
                ranking_score,
                industry_collab / 50
            ]
            features.append(feature)

        return {
            'features': torch.tensor(features, dtype=torch.float),
            'ids': [f'HK-IN-{i:04d}' for i in range(n)]
        }

    def _generate_ownership_edges(self, n_companies: int, n_patents: int,
                                  company_industry: torch.Tensor,
                                  patent_ipc: List[str]) -> Dict:
        """生成企业-专利所有权边"""
        edges = []
        edge_attrs = []

        # 每个专利1-3个所有者
        for patent_idx in range(n_patents):
            n_owners = np.random.choice([1, 2, 3], p=[0.7, 0.25, 0.05])

            # 选择同产业或相关产业的企业
            patent_ipc_code = patent_ipc[patent_idx]

            # 找到相关产业的企业
            candidate_companies = []
            for industry_idx, industry in enumerate(self.industries):
                if patent_ipc_code in self.ipc_classes[industry]:
                    candidate_companies.extend(
                        (company_industry == industry_idx).nonzero(as_tuple=True)[0].tolist()
                    )

            if not candidate_companies:
                candidate_companies = list(range(n_companies))

            owners = np.random.choice(candidate_companies,
                                      size=min(n_owners, len(candidate_companies)),
                                      replace=False)

            for company_idx in owners:
                edges.append([company_idx, patent_idx])

                # 边属性: 所有权份额
                ownership_share = 1.0 / n_owners
                # 重要性评分
                importance = np.random.beta(3, 2)

                edge_attrs.append([ownership_share, importance])

        return {
            'edge_index': torch.tensor(edges, dtype=torch.long).t().contiguous(),
            'edge_attr': torch.tensor(edge_attrs, dtype=torch.float)
        }

    def _generate_trademark_edges(self, n_companies: int, n_trademarks: int,
                                  avg_per_company: float) -> torch.Tensor:
        """生成企业-商标持有边"""
        edges = []

        # 每个企业持有的商标数量 (泊松分布)
        for company_idx in range(n_companies):
            n_tm = np.random.poisson(avg_per_company)
            if n_tm > 0:
                trademarks = np.random.choice(n_trademarks,
                                              size=min(n_tm, n_trademarks),
                                              replace=False)
                for tm_idx in trademarks:
                    edges.append([company_idx, tm_idx])

        return torch.tensor(edges, dtype=torch.long).t().contiguous()

    def _generate_invention_edges(self, n_persons: int, n_patents: int,
                                  avg_inventors_per_patent: float) -> Dict:
        """生成人员-专利发明边"""
        edges = []
        edge_attrs = []

        for patent_idx in range(n_patents):
            n_inventors = max(1, int(np.random.poisson(avg_inventors_per_patent)))
            n_inventors = min(n_inventors, 10)

            inventors = np.random.choice(n_persons, size=n_inventors, replace=False)

            for rank, person_idx in enumerate(inventors):
                edges.append([person_idx, patent_idx])

                # 边属性: 发明人排序 (第一发明人=1.0)
                inventor_rank = 1.0 - (rank / n_inventors)
                # 贡献度
                contribution = np.random.beta(5, 2) if rank == 0 else np.random.beta(2, 5)

                edge_attrs.append([inventor_rank, contribution])

        return {
            'edge_index': torch.tensor(edges, dtype=torch.long).t().contiguous(),
            'edge_attr': torch.tensor(edge_attrs, dtype=torch.float)
        }

    def _generate_employment_edges(self, n_companies: int, n_persons: int,
                                   company_patent_edges: torch.Tensor,
                                   person_patent_edges: torch.Tensor) -> torch.Tensor:
        """生成企业-人员雇佣边 (基于专利合作推断)"""
        edges = set()

        # 基于共同专利推断雇佣关系
        company_patents = defaultdict(set)
        person_patents = defaultdict(set)

        for company, patent in company_patent_edges.t().tolist():
            company_patents[company].add(patent)

        for person, patent in person_patent_edges.t().tolist():
            person_patents[person].add(patent)

        for person_idx in range(n_persons):
            person_patent_set = person_patents[person_idx]
            if not person_patent_set:
                continue

            # 找到拥有相同专利的企业
            candidate_companies = []
            for company_idx in range(n_companies):
                overlap = len(company_patents[company_idx] & person_patent_set)
                if overlap > 0:
                    candidate_companies.append((company_idx, overlap))

            if candidate_companies:
                # 选择专利重叠最多的企业
                candidate_companies.sort(key=lambda x: x[1], reverse=True)
                employer = candidate_companies[0][0]
                edges.add((employer, person_idx))
            else:
                # 随机分配雇主
                employer = np.random.randint(0, n_companies)
                edges.add((employer, person_idx))

        return torch.tensor(list(edges), dtype=torch.long).t().contiguous()

    def _generate_citation_edges(self, n_patents: int, patent_years: torch.Tensor,
                                 avg_citations: float) -> Dict:
        """生成专利引用边 (只能引用更早的专利)"""
        edges = []
        edge_attrs = []

        for patent_idx in range(n_patents):
            current_year = patent_years[patent_idx].item()

            # 可引用的专利 (更早的专利)
            earlier_patents = (patent_years < current_year).nonzero(as_tuple=True)[0]

            if len(earlier_patents) > 0:
                n_citations = np.random.poisson(avg_citations)
                n_citations = min(n_citations, len(earlier_patents))

                if n_citations > 0:
                    # 偏向引用更近期的专利
                    weights = 1.0 / (current_year - patent_years[earlier_patents].float() + 1)
                    weights = weights / weights.sum()

                    cited_patents = np.random.choice(
                        earlier_patents.numpy(),
                        size=n_citations,
                        replace=False,
                        p=weights.numpy()
                    )

                    for cited_idx in cited_patents:
                        edges.append([patent_idx, cited_idx])

                        # 引用重要性
                        citation_importance = np.random.beta(3, 2)
                        # 时间差
                        time_diff = (current_year - patent_years[cited_idx].item()) / 10.0

                        edge_attrs.append([citation_importance, time_diff])

        return {
            'edge_index': torch.tensor(edges, dtype=torch.long).t().contiguous(),
            'edge_attr': torch.tensor(edge_attrs, dtype=torch.float)
        }

    def _generate_cooperation_edges(self, n_companies: int,
                                    company_industry: torch.Tensor,
                                    company_district: torch.Tensor) -> Dict:
        """生成企业合作边"""
        edges = []
        edge_attrs = []

        for company_idx in range(n_companies):
            # 每个企业0-5个合作伙伴
            n_partners = np.random.poisson(2)
            n_partners = min(n_partners, 5)

            if n_partners > 0:
                # 优先选择同产业或同地区的企业
                same_industry = (company_industry == company_industry[company_idx]).nonzero(as_tuple=True)[0]
                same_district = (company_district == company_district[company_idx]).nonzero(as_tuple=True)[0]

                # 去掉自己
                same_industry = same_industry[same_industry != company_idx]
                same_district = same_district[same_district != company_idx]

                # 70% 同产业, 30% 跨产业
                candidates = []
                if len(same_industry) > 0 and random.random() < 0.7:
                    candidates = same_industry.tolist()
                else:
                    candidates = list(range(n_companies))
                    candidates.remove(company_idx)

                if candidates:
                    partners = np.random.choice(candidates,
                                                size=min(n_partners, len(candidates)),
                                                replace=False)

                    for partner_idx in partners:
                        # 避免重复边
                        if company_idx < partner_idx:
                            edges.append([company_idx, partner_idx])

                            # 合作强度
                            strength = np.random.beta(3, 3)
                            # 合作时长 (年)
                            duration = np.random.exponential(2)
                            duration = min(duration, 10) / 10

                            edge_attrs.append([strength, duration])

        return {
            'edge_index': torch.tensor(edges, dtype=torch.long).t().contiguous(),
            'edge_attr': torch.tensor(edge_attrs, dtype=torch.float)
        }

    def _generate_collaboration_edges(self, n_companies: int, n_institutions: int,
                                      company_industry: torch.Tensor) -> torch.Tensor:
        """生成企业-机构合作边"""
        edges = []

        for company_idx in range(n_companies):
            # 30% 的企业与机构有合作
            if random.random() < 0.3:
                n_collabs = np.random.poisson(1) + 1
                n_collabs = min(n_collabs, 3)

                institutions = np.random.choice(n_institutions, size=n_collabs, replace=False)
                for inst_idx in institutions:
                    edges.append([company_idx, inst_idx])

        return torch.tensor(edges, dtype=torch.long).t().contiguous()

    def _generate_institution_patent_edges(self, n_institutions: int,
                                           n_patents: int) -> torch.Tensor:
        """生成机构-专利产出边"""
        edges = []

        # 每个机构产出多个专利
        for inst_idx in range(n_institutions):
            n_patents_produced = np.random.poisson(20)
            n_patents_produced = min(n_patents_produced, n_patents // 2)

            if n_patents_produced > 0:
                patents = np.random.choice(n_patents, size=n_patents_produced, replace=False)
                for patent_idx in patents:
                    edges.append([inst_idx, patent_idx])

        return torch.tensor(edges, dtype=torch.long).t().contiguous()

    def _industry_distribution(self) -> np.ndarray:
        """产业分布 (香港实际情况: 金融科技和AI占比较大)"""
        weights = np.array([0.25, 0.15, 0.20, 0.10, 0.08, 0.10, 0.07, 0.05])
        return weights / weights.sum()

    def _select_district(self, industry: str) -> int:
        """根据产业选择地区"""
        district_map = {
            '金融科技': ['中环', '湾仔', '尖沙咀'],
            '生物医药': ['科学园', '将军澳工业邨'],
            '人工智能': ['科学园', '数码港'],
            '半导体': ['科学园', '将军澳工业邨'],
            '新能源': ['科学园', '观塘'],
            '电子商务': ['数码港', '荃湾'],
            '物流科技': ['观塘', '荃湾'],
            '智能制造': ['将军澳工业邨', '观塘']
        }

        preferred = district_map.get(industry, self.districts)
        district = random.choice(preferred)
        return self.districts.index(district)

    def _print_statistics(self, data: HeteroData):
        """打印数据统计信息"""
        print("\n" + "=" * 60)
        print("📊 数据集统计")
        print("=" * 60)

        print(f"\n【节点数量】")
        for node_type in data.node_types:
            print(f"  {node_type:15s}: {data[node_type].num_nodes:6d} 节点")

        print(f"\n【边数量】")
        for edge_type in data.edge_types:
            src, rel, dst = edge_type
            num_edges = data[edge_type].num_edges
            print(f"  {src:12s} --[{rel:12s}]--> {dst:12s}: {num_edges:6d} 条边")

        print(f"\n【特征维度】")
        for node_type in data.node_types:
            print(f"  {node_type:15s}: {data[node_type].num_features} 维")

        print("\n" + "=" * 60)


    @staticmethod
    def load_data(dataset_size='medium', data_dir='data'):
        """
        加载已保存的数据集

        参数:
            dataset_size: 数据集规模 ('test', 'medium', 'large')
            data_dir: 数据目录路径

        返回:
            HeteroData: 加载的异构图数据
        """
        import os

        filename_map = {
            'test': 'hk_ip_test.pt',
            'medium': 'hk_ip_medium.pt',
            'large': 'hk_ip_large.pt'
        }

        if dataset_size not in filename_map:
            raise ValueError(f"不支持的数据集规模: {dataset_size}. 支持的选项: {list(filename_map.keys())}")

        filepath = os.path.join(data_dir, filename_map[dataset_size])

        if not os.path.exists(filepath):
            print(f"❌ 数据文件不存在: {filepath}")
            print(f"🔄 正在生成 {dataset_size} 规模的数据集...")

            # 自动生成缺失的数据集
            generator = IPEcosystemGenerator(seed=42)

            size_configs = {
                'test': {
                    'n_companies': 100,
                    'n_patents': 500,
                    'n_trademarks': 200,
                    'n_persons': 300,
                    'n_institutions': 20,
                    'time_span_years': 5
                },
                'medium': {
                    'n_companies': 500,
                    'n_patents': 3000,
                    'n_trademarks': 1500,
                    'n_persons': 2000,
                    'n_institutions': 50,
                    'time_span_years': 10
                },
                'large': {
                    'n_companies': 2000,
                    'n_patents': 15000,
                    'n_trademarks': 8000,
                    'n_persons': 10000,
                    'n_institutions': 200,
                    'time_span_years': 15
                }
            }

            config = size_configs[dataset_size]
            data = generator.generate(**config)

            # 确保目录存在
            os.makedirs(data_dir, exist_ok=True)
            torch.save(data, filepath)
            print(f"✅ 数据已生成并保存到: {filepath}")

            return data

        try:
            print(f"📂 加载数据集: {filepath}")
            data = torch.load(filepath, weights_only=False)
            print(f"✅ 成功加载 {dataset_size} 规模数据集")

            # 打印数据集信息
            print(f"\n📊 数据集概览:")
            for node_type in data.node_types:
                print(f"  {node_type:15s}: {data[node_type].num_nodes:6d} 节点")

            total_edges = sum(data[edge_type].num_edges for edge_type in data.edge_types)
            print(f"  {'总边数':15s}: {total_edges:6d} 条边")

            return data

        except Exception as e:
            print(f"❌ 加载数据失败: {e}")
            print(f"🔄 尝试重新生成数据...")
            # 删除损坏的文件
            if os.path.exists(filepath):
                os.remove(filepath)
            # 递归调用重新生成
            return IPEcosystemGenerator.load_data(dataset_size, data_dir)


    @staticmethod
    def list_available_datasets(data_dir='data'):
        """
        列出可用的数据集

        参数:
            data_dir: 数据目录路径
        """
        import os

        if not os.path.exists(data_dir):
            print(f"📁 数据目录不存在: {data_dir}")
            return

        filename_map = {
            'hk_ip_test.pt': 'test (小规模)',
            'hk_ip_medium.pt': 'medium (中等规模)',
            'hk_ip_large.pt': 'large (大规模)'
        }

        print(f"\n📋 可用数据集 (目录: {data_dir}):")
        print("-" * 50)

        found_any = False
        for filename, description in filename_map.items():
            filepath = os.path.join(data_dir, filename)
            if os.path.exists(filepath):
                file_size = os.path.getsize(filepath) / (1024 * 1024)  # MB
                print(f"✅ {description:20s} - {filename:20s} ({file_size:.1f} MB)")
                found_any = True
            else:
                print(f"❌ {description:20s} - {filename:20s} (不存在)")

        if not found_any:
            print("🚫 未找到任何数据集文件")
            print("💡 运行 data_generator.py 来生成数据集")

        print("-" * 50)


    @staticmethod
    def get_dataset_info(dataset_size='medium', data_dir='data'):
        """
        获取数据集详细信息而不加载数据

        参数:
            dataset_size: 数据集规模
            data_dir: 数据目录路径
        """
        filename_map = {
            'test': 'hk_ip_test.pt',
            'medium': 'hk_ip_medium.pt',
            'large': 'hk_ip_large.pt'
        }

        filepath = os.path.join(data_dir, filename_map[dataset_size])

        if not os.path.exists(filepath):
            print(f"❌ 数据文件不存在: {filepath}")
            return None

        try:
            # 只加载metadata而不加载完整数据
            data = torch.load(filepath, map_location='cpu')

            info = {
                'dataset_size': dataset_size,
                'filepath': filepath,
                'file_size_mb': os.path.getsize(filepath) / (1024 * 1024),
                'node_types': data.node_types,
                'edge_types': data.edge_types,
                'node_counts': {ntype: data[ntype].num_nodes for ntype in data.node_types},
                'edge_counts': {etype: data[etype].num_edges for etype in data.edge_types},
                'total_nodes': sum(data[ntype].num_nodes for ntype in data.node_types),
                'total_edges': sum(data[etype].num_edges for etype in data.edge_types)
            }

            return info

        except Exception as e:
            print(f"❌ 读取数据集信息失败: {e}")
            return None


# ========== 使用示例 (更新) ==========
if __name__ == "__main__":
    # 生成数据集
    generator = IPEcosystemGenerator(seed=42)

    # 确保data目录存在
    import os

    os.makedirs('../data', exist_ok=True)

    # 检查现有数据集
    print("🔍 检查现有数据集...")
    IPEcosystemGenerator.list_available_datasets()

    # 小规模测试
    print("\n🧪 生成测试数据集...")
    test_data = generator.generate(
        n_companies=100,
        n_patents=500,
        n_trademarks=200,
        n_persons=300,
        n_institutions=20,
        time_span_years=5
    )

    # 中等规模
    print("\n🏗️  生成中等规模数据集...")
    medium_data = generator.generate(
        n_companies=500,
        n_patents=3000,
        n_trademarks=1500,
        n_persons=2000,
        n_institutions=50,
        time_span_years=10
    )

    # 大规模
    print("\n🚀 生成大规模数据集...")
    large_data = generator.generate(
        n_companies=2000,
        n_patents=15000,
        n_trademarks=8000,
        n_persons=10000,
        n_institutions=200,
        time_span_years=15
    )

    # 保存数据
    torch.save(test_data, '../data/hk_ip_test.pt')
    torch.save(medium_data, '../data/hk_ip_medium.pt')
    torch.save(large_data, '../data/hk_ip_large.pt')

    print("\n✅ 数据已保存到 data/ 目录")

    # 显示最终状态
    print("\n" + "=" * 60)
    IPEcosystemGenerator.list_available_datasets()

    # 测试加载功能
    print("\n🧪 测试数据加载功能...")
    loaded_data = IPEcosystemGenerator.load_data('medium')
    print("✅ 数据加载测试成功!")