import torch
import gc
import sys


def print_xpu_memory_stats():
    """打印 XPU 内存分配统计信息"""
    try:
        if not hasattr(torch, 'xpu') or not torch.xpu.is_available():
            print("XPU 不可用")
            return

        # 清理缓存以获取准确的内存使用情况
        gc.collect()
        torch.xpu.empty_cache()

        # 获取当前内存统计信息
        if hasattr(torch.xpu, 'memory_stats'):
            stats = torch.xpu.memory_stats()
            print("\n===== XPU 内存统计信息 =====")
            print(f"分配的内存: {stats.get('allocated_bytes.all', 0) / (1024 ** 3):.2f} GB")
            print(f"缓存的内存: {stats.get('reserved_bytes.all', 0) / (1024 ** 3):.2f} GB")
            print(f"活跃的内存块: {stats.get('active_bytes.all', 0) / (1024 ** 3):.2f} GB")
            print(f"内存分配次数: {stats.get('allocation.all', 0)}")
        else:
            print("torch.xpu.memory_stats() 不可用")

        # 获取当前设备内存信息
        if hasattr(torch.xpu, 'get_device_properties'):
            device = torch.xpu.current_device()
            props = torch.xpu.get_device_properties(device)
            print("\n===== XPU 设备属性 =====")
            print(f"设备名称: {props.name}")
            print(f"总内存: {props.total_memory / (1024 ** 3):.2f} GB")

            # 获取当前内存使用情况
            if hasattr(torch.xpu, 'memory_allocated') and hasattr(torch.xpu, 'memory_reserved'):
                print("\n===== XPU 当前内存使用 =====")
                print(f"已分配内存: {torch.xpu.memory_allocated() / (1024 ** 3):.2f} GB")
                print(f"已保留内存: {torch.xpu.memory_reserved() / (1024 ** 3):.2f} GB")
                print(f"可用内存: {(props.total_memory - torch.xpu.memory_reserved()) / (1024 ** 3):.2f} GB")
        else:
            print("torch.xpu.get_device_properties() 不可用")

        # 检查内存分配比例设置
        if hasattr(torch.xpu, 'get_memory_fraction'):
            print("\n===== XPU 内存分配比例 =====")
            fraction = torch.xpu.get_memory_fraction()
            print(f"当前内存分配比例: {fraction:.2f}")
        elif hasattr(torch.xpu, 'get_allocator_backend'):
            print("\n===== XPU 分配器后端 =====")
            backend = torch.xpu.get_allocator_backend()
            print(f"当前分配器后端: {backend}")

        # 尝试获取最大内存
        try:
            # 分配测试张量，逐步增加大小直到失败
            max_gb = 0
            step = 1  # 每次增加 1GB

            print("\n===== XPU 最大可分配内存测试 =====")
            print("正在测试最大可分配内存...")

            while True:
                try:
                    size_bytes = int(max_gb * 1024 ** 3 / 4)  # float32 是 4 字节
                    if size_bytes <= 0:
                        max_gb += step
                        continue

                    test_tensor = torch.zeros(size_bytes, dtype=torch.float32, device='xpu')
                    del test_tensor
                    torch.xpu.empty_cache()
                    print(f"成功分配 {max_gb} GB")
                    max_gb += step
                except Exception as e:
                    print(f"在尝试分配 {max_gb} GB 时失败")
                    print(f"最大可分配内存约为: {max_gb - step} GB")
                    break

                # 防止无限循环
                if max_gb > 32:  # 设置上限为 32GB
                    print("达到测试上限 (32GB)")
                    break
        except Exception as e:
            print(f"内存测试失败: {e}")

    except Exception as e:
        print(f"获取 XPU 内存统计信息时出错: {e}")
        import traceback
        print(traceback.format_exc())


# 打印 PyTorch 和系统信息
print(f"PyTorch 版本: {torch.__version__}")
print(f"Python 版本: {sys.version}")

# 检查 IPEX 是否安装
try:
    import intel_extension_for_pytorch as ipex

    print(f"IPEX 版本: {ipex.__version__}")
except ImportError:
    print("IPEX 未安装")

# 打印 XPU 是否可用
if hasattr(torch, 'xpu'):
    print(f"XPU 可用: {torch.xpu.is_available()}")
    if torch.xpu.is_available():
        print(f"XPU 设备数量: {torch.xpu.device_count()}")
        print(f"当前 XPU 设备: {torch.xpu.current_device()}")
else:
    print("XPU 不可用 (torch.xpu 不存在)")

# 打印内存统计信息
print_xpu_memory_stats()

# 测试设置内存分配比例
if hasattr(torch, 'xpu') and torch.xpu.is_available() and hasattr(torch.xpu, 'set_per_process_memory_fraction'):
    print("\n===== 测试设置内存分配比例 =====")
    current_fraction = 0.3  # 默认值

    try:
        # 尝试获取当前值
        if hasattr(torch.xpu, 'get_memory_fraction'):
            current_fraction = torch.xpu.get_memory_fraction()
            print(f"当前内存分配比例: {current_fraction:.2f}")
    except:
        pass

    # 设置新值
    new_fraction = 0.8
    print(f"设置内存分配比例为: {new_fraction:.2f}")
    torch.xpu.set_per_process_memory_fraction(new_fraction)

    # 验证设置是否生效
    try:
        if hasattr(torch.xpu, 'get_memory_fraction'):
            updated_fraction = torch.xpu.get_memory_fraction()
            print(f"更新后的内存分配比例: {updated_fraction:.2f}")
            if abs(updated_fraction - new_fraction) < 0.01:
                print("✓ 内存分配比例设置成功!")
            else:
                print("✗ 内存分配比例设置失败!")
    except:
        print("无法验证内存分配比例设置")