pixi global list
# 或者简写
pixi g list

# 创建 Python 3.10 环境（不指定 --expose 来测试默认行为）
pixi global install --environment scanpy \
  --expose "py310_scanpy=python" \
  python=3.10 pip numpy pandas scanpy
# 创建 Python 3.11 环境
pixi global install --environment scgpt \
  --expose "py311_scgpt=python" \
  python=3.11 jupyter

--expose <EXPOSE>
Add one or more mapping which describe which executables are exposed. 
The syntax is `exposed_name=executable_name`, so for example `python3.10=python`. 
Alternatively, you can input only an executable_name and `executable_name=executable_name` is assumed

# 将 python 暴露为 py311，将 pip 暴露为 pip311
pixi global install --environment python311 \
  --expose "py311=python" \
  --expose "pip311=pip" \
  --expose "jupyter311=jupyter" \
  python=3.11 pip jupyter

# 给命令起有意义的名字
pixi global install --environment python311 \
  --expose "python-data=python" \
  --expose "pip-data=pip" \
  --expose "notebook=jupyter" \
  python=3.11 pip jupyter matplotlib pandas

# 使用版本号避免冲突
pixi global install --environment python311 \
  --expose "python3.11=python" \
  --expose "pip3.11=pip" \
  --expose "jupyter3.11=jupyter" \
  python=3.11 pip jupyter

# 按照警告提示，先同步环境
pixi global sync

version = 1

[envs.r]
channels = ["https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r", "nvidia"]
dependencies = { r-base = "*", r-essentials = "*", r-pak = "*", pkg-config = "*", freetype = "*", fontconfig = "*", harfbuzz = "*", fribidi = "*", libpng = "*", libjpeg-turbo = "*", libtiff = "*", glpk = "*", libgit2 = "*", pandoc = "*", imagemagick = "*", perl = "*", libarrow = "*", arrow-cpp = "*", r-arrow = "*", r-sf = "*", r-spdep = "*", gdal = "*", proj = "*", geos = "*", udunits2 = "*" }
exposed = { R = "R", Rscript = "Rscript" }

[envs.scanpy]
channels = ["https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free", "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r", "nvidia"]
dependencies = { python = "3.10.*", pip = "*", numpy = "*", pandas = "*", scanpy = "*" }
exposed = { py310_scanpy = "python", pip310_scanpy = "pip" }

# 添加 scanpy 包
pixi g add --environment scanpy scanpy

[tasks]
# 一键安装命令
# 这里depends-on的意思就是运行此命令会先自动运行这些依赖的命令
setup-all = { depends-on = ["install-torch", "install-flash-attn", "check-env"] }

# 🚀 [迁移核心]：在此处统一修改系统 CUDA 路径
# 理由：编译需要系统 NVCC，运行 JAX 需要系统驱动库。pixi环境中的 CUDA 像是零件包，无法满足需求。
_system_cuda = "export CUDA_HOME=/usr/local/cuda-12.4 && export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH"

# 1. 安装支持 CUDA 12.1 的 PyTorch
install-torch = "pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://mirror.nju.edu.cn/pytorch/whl/cu121"

# 2. 现场编译 Flash Attention (会自动引用 _setup_cuda 中的路径)
# 💡 迁移提示：修改 _setup_cuda 任务中的路径即可全局生效
install-flash-attn = { cmd = "pip install flash-attn==1.0.4 --no-build-isolation --no-cache-dir", depends-on = ["_system_cuda"] }

# 3. 环境全项验证
check-env = { cmd = "python -c 'import jax; import torch; print(f\"PyTorch GPU: {torch.cuda.is_available()}\"); print(f\"JAX Devices: {jax.devices()}\")'", depends-on = ["_system_cuda"] }

# 4. 使用python脚本全项验证
test-scgpt = { cmd = "python test_scgpt.py", depends-on = ["_system_cuda"] }

[workspace]
name = "scGPT"
version = "0.1.0"
authors = ["kleikli <x1246050927@gmail.com>"]
channels = [
    "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge",
    "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda",
    "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch",
    "nvidia"
]
platforms = ["linux-64"]

[pypi-options]
index-url = "https://pypi.tuna.tsinghua.edu.cn/simple"

[dependencies]
python = "3.11.*"
pip = ">=25.2,<26"
gcc_linux-64 = "12.*"
gxx_linux-64 = "12.*"
ninja = "*"
pyarrow = "<15.0"
jupyter = ">=1.1.1,<2"
ipykernel = ">=6.0,<7"

[pypi-dependencies]
torchvision = "==0.18.0"
torchaudio = "==2.3.0"
packaging = ">=25.0, <26"
ipykernel = ">=6.30.1, <7"
jax = "==0.4.13"
jaxlib = { url = "https://storage.googleapis.com/jax-releases/cuda12/jaxlib-0.4.13%2Bcuda12.cudnn89-cp311-cp311-manylinux2014_x86_64.whl" }
scgpt = ">=0.2.4, <0.3"
scvi-tools = "==0.20.3"
anndata = "==0.9.2"
scanpy = ">=1.10.0"
scipy = "<1.13.0"
wandb = ">=0.22.2, <0.23"
gseapy = ">=1.1.10, <2"
faiss-cpu = ">=1.12.0, <2"
gpustat = ">=1.1.1, <2"

[tasks]
# 一键安装命令
setup-all = { depends-on = ["install-torch", "install-flash-attn", "check-env"] }

# 🚀 [迁移核心]：在此处统一修改系统 CUDA 路径
# 理由：编译需要系统 NVCC，运行 JAX 需要系统驱动库。pixi环境中的 CUDA 像是零件包，无法满足需求。
_system_cuda = "export CUDA_HOME=/usr/local/cuda-12.4 && export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH"

# 1. 安装支持 CUDA 12.1 的 PyTorch
install-torch = "pip install torch==2.3.0 torchvision==0.18.0 torchaudio==2.3.0 --index-url https://mirror.nju.edu.cn/pytorch/whl/cu121"

# 2. 现场编译 Flash Attention (会自动引用 _setup_cuda 中的路径)
# 💡 迁移提示：修改 _setup_cuda 任务中的路径即可全局生效
install-flash-attn = { cmd = "pip install flash-attn==1.0.4 --no-build-isolation --no-cache-dir", depends-on = ["_system_cuda"] }

# 3. 环境全项验证
check-env = { cmd = "python -c 'import jax; import torch; print(f\"PyTorch GPU: {torch.cuda.is_available()}\"); print(f\"JAX Devices: {jax.devices()}\")'", depends-on = ["_system_cuda"] }

# 4. 使用python脚本全项验证
test-scgpt = { cmd = "python test_scgpt.py", depends-on = ["_system_cuda"] 

import os
import sys
import warnings
import torch
import numpy as np

# ==========================================
# 1. 警告与日志过滤 (保持终端整洁)
# ==========================================
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" 
try:
    import torchtext
    torchtext.disable_torchtext_deprecation_warning()
except (ImportError, AttributeError):
    pass

warnings.filterwarnings("ignore", message=".*The 'frozen' attribute with value True.*")
warnings.filterwarnings("ignore", message=".*enable_nested_tensor is True, but self.use_nested_tensor is False.*")
warnings.filterwarnings("ignore", category=UserWarning, module="lightning_fabric")

# 针对 RTX 4090 的优化设置
if torch.cuda.is_available():
    torch.set_float32_matmul_precision('high')

def test_environment():
    print(f"{'='*20} 1. 环境基础检查 {'='*20}")
    print(f"Python 路径: {sys.executable}")
    print(f"PyTorch 版本: {torch.__version__}")
    print(f"CUDA 是否可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"当前 GPU: {torch.cuda.get_device_name(0)}")

    # ==========================================
    # 2. JAX GPU 加速检测
    # ==========================================
    print(f"\n{'='*20} 2. JAX 后端检测 {'='*20}")
    try:
        import jax
        devices = jax.devices()
        print(f"✅ JAX 加载成功! 版本: {jax.__version__}")
        print(f"✅ JAX 识别设备: {devices}")
        if "gpu" not in str(devices[0]).lower():
            print("⚠️ 警告: JAX 未识别到 GPU，请检查 jaxlib 版本。")
    except Exception as e:
        print(f"❌ JAX 检测失败: {e}")

    # ==========================================
    # 3. 单细胞核心算法库加载
    # ==========================================
    print(f"\n{'='*20} 3. 核心算法库加载 {'='*20}")
    try:
        import scvi
        import scanpy as sc
        import scgpt
        print(f"✅ scvi-tools 版本: {scvi.__version__}")
        print(f"✅ scanpy 版本: {sc.__version__}")
        print(f"✅ scGPT 模块路径: {scgpt.__path__[0]}")
    except Exception as e:
        print(f"❌ 库加载失败: {e}")

    # ==========================================
    # 4. Flash Attention 与 scGPT 模型验证
    # ==========================================
    print(f"\n{'='*20} 4. Flash Attention 算子验证 {'='*20}")
    try:
        import flash_attn
        print(f"✅ Flash Attention 库已安装 (版本: {getattr(flash_attn, '__version__', '1.0.4')})")
        
        from scgpt.model import TransformerModel
        
        # 构造模拟词表 (ntoken=100)
        fake_vocab = {"<pad>": 0, "<mask>": 1, "<cls>": 2}
        for i in range(3, 100):
            fake_vocab[f"gene_{i}"] = i
            
        # 初始化模型
        model = TransformerModel(
            ntoken=100, 
            d_model=128, 
            nhead=4, 
            d_hid=128, 
            nlayers=2,
            vocab=fake_vocab,
            pad_token="<pad>",
            use_fast_transformer=True, 
            fast_transformer_backend="flash"
        ).cuda()

        # RTX 4090 性能优化：使用半精度进行前向传播测试
        model.half()
        
        # 准备模拟数据 (Batch=2, Seq_len=10)
        src = torch.randint(1, 100, (2, 10)).cuda()
        values = torch.randn(2, 10).cuda().half()
        # 关键修复：构造 src_key_padding_mask
        # 0 代表该位置是真实的基因，1 代表是填充位
        src_key_padding_mask = torch.zeros(2, 10, dtype=torch.bool).cuda()

        with torch.no_grad():
            # 执行前向传播
            output = model(
                src, 
                values, 
                src_key_padding_mask=src_key_padding_mask
            )
        
        print("✅ scGPT 模型初始化成功并已载入 GPU")
        print(f"✅ 前向传播测试成功！")
        print(f"✅ 细胞嵌入 (Cell Embedding) 形状: {output['cell_emb'].shape}")
        print("\n🚀 结论：您的 scGPT 环境已达满血状态，Flash Attention 运行正常！")

    except Exception as e:
        print(f"❌ 模型验证环节出错: {e}")
        print("💡 建议：检查 scgpt 源码版本是否与参数匹配。")

    print(f"\n{'='*50}")
    print("🌟 所有测试已完成！")
    print(f"{'='*50}")

if __name__ == "__main__":
    test_environment()