@echo off
setlocal enabledelayedexpansion

REM 设置环境变量
SET MICROMAMBA_EXE=%~dp0micromamba.exe
SET CUDA_HOME="C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8"
SET PATH=%CUDA_HOME%\bin;%PATH%
SET PROJECT_DIR=%~dp0
SET PYTHONDONTWRITEBYTECODE=1
SET GDOWN_CACHE=cache\gdown
SET TORCH_HOME=cache\torch
SET HF_HOME=cache\huggingface
SET BLENDER_DIR="C:\Program Files\Blender Foundation\Blender 3.6"
SET VS2019_DIR="C:\Program Files (x86)\Microsoft Visual Studio\2019\BuildTools"
SET VS2019_VCVARS=%VS2019_DIR%\VC\Auxiliary\Build\vcvars64.bat

REM 检查micromamba
IF NOT EXIST "%MICROMAMBA_EXE%" (
    echo ERROR: micromamba not found at %MICROMAMBA_EXE%
    echo Please install micromamba from https://mamba.readthedocs.io/en/latest/installation.html
    exit /b 1
)

REM 设置micromamba根目录
SET MAMBA_ROOT_PREFIX=%PROJECT_DIR%\micromamba
mkdir %MAMBA_ROOT_PREFIX%

REM 初始化micromamba
CALL "%MICROMAMBA_EXE%" shell init --root-prefix=%MAMBA_ROOT_PREFIX%

REM 检查必要的环境和依赖
IF NOT EXIST "%CUDA_HOME%" (
    echo ERROR: CUDA 11.8 not found at %CUDA_HOME%
    exit /b 1
)
IF NOT EXIST "%BLENDER_DIR%" (
    echo ERROR: Blender 3.6 not found at %BLENDER_DIR%
    echo Please install Blender 3.6 from https://www.blender.org/download/lts/3-6/
    exit /b 1
)

REM 检查Visual Studio
IF NOT EXIST "%VS2019_VCVARS%" (
    echo ERROR: Visual Studio 2019 not found
    echo Please install Visual Studio 2019 Build Tools with C++ development tools
    exit /b 1
)

REM 检查CUDA版本
nvcc --version | findstr "release 11.8" >nul
IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: CUDA 11.8 not found or version mismatch
    exit /b 1
)

REM 检查COLMAP
where colmap >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: COLMAP not found in PATH
    echo Please install COLMAP and add it to PATH
    exit /b 1
)

REM 创建缓存目录
mkdir cache\gdown 2>nul
mkdir cache\torch 2>nul
mkdir cache\huggingface 2>nul

REM 创建ext目录
mkdir ext 2>nul
cd ext

REM 克隆外部库
git clone --depth 1 https://github.com/CMU-Perceptual-Computing-Lab/openpose
cd openpose
git submodule update --init --recursive --remote
cd ..

git clone https://github.com/hustvl/Matte-Anything
cd Matte-Anything
git clone https://github.com/IDEA-Research/GroundingDINO.git
cd ..

git clone https://github.com/egorzakharov/NeuralHaircut.git --recursive
git clone https://github.com/facebookresearch/pytorch3d
cd pytorch3d
git checkout 2f11ddc5ee7d6bd56f2fb6744a16776fab6536f7
cd ..

git clone https://github.com/camenduru/simple-knn
cd diff_gaussian_rasterization_hair\third_party
git clone https://github.com/g-truc/glm
cd glm
git checkout 5c46b9c07008ae65cb81ab79cd677ecc1934b903
cd ..\..\..

git clone --recursive https://github.com/NVIDIAGameWorks/kaolin
cd kaolin
git checkout v0.15.0
cd ..

git clone https://github.com/SSL92/hyperIQA

REM 创建环境
CALL "%MICROMAMBA_EXE%" create -p %MAMBA_ROOT_PREFIX%\envs\gaussian_splatting_hair python=3.8 pytorch=2.0.0 torchvision pytorch-cuda=11.8 cmake ninja setuptools=58.2.0 -c pytorch -c nvidia -c conda-forge -y
CALL "%MICROMAMBA_EXE%" create -p %MAMBA_ROOT_PREFIX%\envs\matte_anything pytorch=2.0.0 pytorch-cuda=11.8 torchvision tensorboard timm=0.5.4 opencv=4.5.3 mkl=2024.0 setuptools=58.2.0 easydict wget scikit-image gradio=3.46.1 fairscale -c pytorch -c nvidia -c conda-forge -y
CALL "%MICROMAMBA_EXE%" create -p %MAMBA_ROOT_PREFIX%\envs\openpose python=3.8 cmake=3.20 -c conda-forge -y
CALL "%MICROMAMBA_EXE%" create -p %MAMBA_ROOT_PREFIX%\envs\pixie-env python=3.8 pytorch=2.0.0 torchvision pytorch-cuda=11.8 fvcore pytorch3d==0.7.5 kornia matplotlib -c pytorch -c nvidia -c fvcore -c conda-forge -c pytorch3d -c bottler -c iopath -y

REM 安装 gaussian_splatting_hair 环境
CALL "%MICROMAMBA_EXE%" activate -p %MAMBA_ROOT_PREFIX%\envs\gaussian_splatting_hair
pip install -r requirements.txt
cd %PROJECT_DIR%\ext\pytorch3d
pip install -e .
cd %PROJECT_DIR%\ext\NeuralHaircut\npbgpp
pip install -e .
cd %PROJECT_DIR%\ext\simple-knn
pip install -e .
cd %PROJECT_DIR%\ext\diff_gaussian_rasterization_hair
pip install -e .
cd %PROJECT_DIR%\ext\kaolin
pip install -e .
cd %PROJECT_DIR%

REM 下载 Neural Haircut 文件
cd %PROJECT_DIR%\ext\NeuralHaircut
gdown --folder https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8
cd pretrained_models\diffusion_prior
gdown 1_9EOUXHayKiGH5nkrayncln3d6m1uV7f
cd ..\..\PIXIE
gdown 1mPcGu62YPc4MdkT8FFiOCP629xsENHZf
tar -xf pixie_data.tar.gz
del pixie_data.tar.gz

REM 安装 Matte-Anything 环境
CALL activate_matte_anything.bat
cd %PROJECT_DIR%\ext\Matte-Anything

REM 修改pytorch文件以支持Windows
IF EXIST "%CONDA_PREFIX%\Lib\site-packages\torch\include\torch\csrc\jit\runtime\argument_spec.h" (
    echo Patching argument_spec.h...
    powershell -Command "(gc '%CONDA_PREFIX%\Lib\site-packages\torch\include\torch\csrc\jit\runtime\argument_spec.h') -replace 'static constexpr size_t DEPTH_LIMIT = 128;', 'static const size_t DEPTH_LIMIT = 128;' | Out-File -encoding ASCII '%CONDA_PREFIX%\Lib\site-packages\torch\include\torch\csrc\jit\runtime\argument_spec.h'"
)
IF EXIST "%CONDA_PREFIX%\Lib\site-packages\torch\include\pybind11\cast.h" (
    echo Patching cast.h...
    powershell -Command "(gc '%CONDA_PREFIX%\Lib\site-packages\torch\include\pybind11\cast.h') -replace 'explicit operator type&\(\) { return \*\(this->value\); }', 'explicit operator type&() { return *((type*)this->value); }' | Out-File -encoding ASCII '%CONDA_PREFIX%\Lib\site-packages\torch\include\pybind11\cast.h'"
)

REM 安装SAM
pip install git+https://github.com/facebookresearch/segment-anything.git

REM 安装detectron2
pip install 'git+https://github.com/facebookresearch/detectron2.git@v0.6'
pip install -e GroundingDINO
pip install supervision==0.22.0
mkdir pretrained
cd pretrained
curl -LO https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
curl -LO https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
cd ..
gdown 1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW

REM 安装 OpenPose
CALL activate_openpose.bat
cd %PROJECT_DIR%\ext\openpose
gdown 1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV
tar -xf models.tar.gz
del models.tar.gz
mkdir build
cd build
CALL "%VS2019_VCVARS%"
cmake .. -DBUILD_PYTHON=true -DUSE_CUDNN=off -G "Visual Studio 16 2019" -A x64
cmake --build . --config Release

REM 安装 PIXIE 环境
CALL activate_pixie-env.bat
cd %PROJECT_DIR%\ext\PIXIE
pip install pyyaml==5.4.1
pip install git+https://github.com/1adrianb/face-alignment.git@54623537fd9618ca7c15688fd85aba706ad92b59

echo Installation completed!