#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-A】　通知音のダウンロード

#@markdown ■■■■■■■■■■■■■■■■■■

from enum import Enum

class SoundType(Enum):
    START = 0
    SUCCESS = 1
    FAIL = 2

# ----------------------------------

! wget --no-check-certificate -c "https://soundeffect-lab.info/sound/anime/mp3/sceneswitch1.mp3"
! wget --no-check-certificate -c "https://soundeffect-lab.info/sound/anime/mp3/incorrect1.mp3"
! wget --no-check-certificate -c "https://soundeffect-lab.info/sound/anime/mp3/switch1.mp3"

from IPython.display import Audio, display
import torch

def play_sound(sound_type: SoundType, autoplay=True):
    try:
        if sound_type == SoundType.START:
            file_name = "/content/switch1.mp3"
        elif sound_type == SoundType.SUCCESS:
            file_name = "/content/sceneswitch1.mp3"
        else:
            file_name = "/content/incorrect1.mp3"

        display(Audio(file_name, autoplay=autoplay, normalize=False))
    except:
        print("■■■■■■■■■■■■■■■")
        print("■　効果音が再生できませんでした")
        print("■■■■■■■■■■■■■■■")

# ----------------------------------

exec_dict = {
    '①-A': {'exec': False, 'premise': []},
    '①-B': {'exec': False, 'premise': ['①-A', ]},
    '①-C': {'exec': False, 'premise': ['①-A', '①-B', ]},
    '①-D': {'exec': False, 'premise': ['①-A', '①-B', '①-C', ]},
    '①-E': {'exec': False, 'premise': ['①-A', '①-B', '①-C', '①-D', ]},
    '①-F': {'exec': False, 'premise': ['①-A', '①-B', '①-C', '①-D', '①-E', ]},
    '②-A': {'exec': False, 'premise': ['①-A', '①-B', '①-C', '①-D', '①-E', '①-F', ]},
    '②-B': {'exec': False, 'premise': ['①-A', '①-B', '①-C', '①-D', '①-E', '①-F', '②-A' ]},
}

NVIDIA_VISIBLE_DEVICES=all

class IpyExit(SystemExit):
    def __init__(self):
        play_sound(SoundType.FAIL)
        pass

def check_exec_dict(cell_key: str):
    if not torch.cuda.is_available():
        print("■■■■■■■■■■■■■■■")
        print("■　** ERROR **")
        print("■　GPUが有効になってません。")
        print("■　準備編を参考にしながらランタイムをGPUに変更してください。")
        print("■■■■■■■■■■■■■■■")
        raise IpyExit

    if not exec_dict:
        print("■■■■■■■■■■■■■■■")
        print("■　** ERROR **")
        print("■　セル【①-A】が実行されていない可能性があります。")
        print("■　目次から戻って実行してください。")
        print("■■■■■■■■■■■■■■■")
        raise IpyExit

    for p in exec_dict[cell_key]['premise']:
        # 事前に実行されるべきセルがすべて実行されているか
        if not exec_dict[p]['exec']:
            print("■■■■■■■■■■■■■■■")
            print("■　** ERROR **")
            print(f"■　セル【{p}】が実行されていない可能性があります。")
            print("■　目次から戻って実行してください。")
            print("■■■■■■■■■■■■■■■")
            raise IpyExit

    # 全部実行されていたらOK
    play_sound(SoundType.START)
    return True

def finish_cell(cell_key: str):
    play_sound(SoundType.SUCCESS)

    print("■■■■■■■■■■■■■■■")
    print("■　** OK **")
    print(f"■　セル【{cell_key}】の実行に成功しました。")
    print("■■■■■■■■■■■■■■■")

    exec_dict[cell_key]['exec'] = True

# ----------------------------------

# セル実行可否
ckey = "①-A"
check_exec_dict(ckey)

# ----------------------------------

# セル終了

finish_cell(ckey)

print("")
print("")
print("■ セル実行開始")
play_sound(SoundType.START, autoplay=False)

print("")
print("")
print("■ セル実行失敗")
play_sound(SoundType.FAIL, autoplay=False)

print("")
print("")
print("■ セル実行成功")
play_sound(SoundType.SUCCESS, autoplay=False)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-B】　GPUランタイムの確認

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "①-B"
check_exec_dict(ckey)

# ----------------------------------
! nvidia-smi

import subprocess
try:
    subprocess.check_output("nvidia-smi", shell=True)
except:
    print("■■■■■■■■■■■■■■■")
    print("■　** ERROR **")
    print("■　nvidia-smiコマンドの実行に失敗しました。")
    print("■　① まだGPUを使ってない場合、ランタイムの設定を行っていない場合")
    print("■　　準備編を参考にしながら、ランタイムをGPUに設定してください")
    print("■　　※ランタイムを変更すると環境がリセットされるため、【①-A】からやり直してください")
    print("■　② GPUの利用上限に達して、CPUでしか使えない場合")
    print("■　　無料域の場合、24時間以上置いていただくと再度使用できるようになります")
    print("■■■■■■■■■■■■■■■")

# ----------------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-C】　Googleドライブとの連携

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "①-C"
check_exec_dict(ckey)

# ----------------------------------

from google.colab import drive
import os

# Googleドライブマウント
drive.mount('/gdrive')

# 起点ディレクトリ
base_dir_path = "/gdrive/My Drive/autotrace"

if os.path.exists(base_dir_path):
    print("■■■■■■■■■■■■■■■")
    print("■　** OK **")
    print("■　autotraceフォルダ内との連携が成功しました。")
    print("■■■■■■■■■■■■■■■")
else:
    print("■■■■■■■■■■■■■■■")
    print("■　** ERROR **")
    print("■　autotraceフォルダがGoogleドライブの直下に見つかりませんでした。")
    print("■　準備編を参考にしながら、autotraceフォルダを作成してください。")
    print("■　※スペルミスだった場合、ファイル名などの変更は行わず、新しいフォルダもしくはファイルで試してください。")
    print("■　　Googleドライブとの連携後の更新はうまく認識されません。")
    print("■　※少しタイムラグがあるので、修正した直後は認識されない場合があります。")
    print("■　　10秒くらい待ってからセルを再実行してください。")
    print("■■■■■■■■■■■■■■■")
    raise IpyExit

# ----------------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-D】　Google SDK との連携

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "①-D"
check_exec_dict(ckey)

# ----------------------------------

# Googleドライブアクセスライブラリ
!pip install -U -q PyDrive

from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from google.colab import auth
from oauth2client.client import GoogleCredentials

try:
    auth.authenticate_user()
    gauth = GoogleAuth()
    gauth.credentials = GoogleCredentials.get_application_default()
    drive = GoogleDrive(gauth)

    print("■■■■■■■■■■■■■■■")
    print("■　** OK **")
    print("■　Google SDK との連携に成功しました。")
    print("■■■■■■■■■■■■■■■")

except Exception as e:
    print(e.message)
    print("■■■■■■■■■■■■■■■")
    print("■　** Error **")
    print("■　Google SDK との連携に失敗しました。")
    print("■　このノートブックを作者に共有してください。")
    print("■■■■■■■■■■■■■■■")
    raise IpyExit

# ----------------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-E】　コードセットアップ

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "①-E"
check_exec_dict(ckey)

# ----------------------------

# https://www.usagi1975.com/2019152355/
! apt-key adv --keyserver keyserver.ubuntu.com --recv aa9795093fe2f255acb3500d85304cfd1408787d
! echo "deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu bionic main " > /etc/apt/sources.list.d/python.list
! echo "deb-src http://ppa.launchpad.net/deadsnakes/ppa/ubuntu bionic main" >> /etc/apt/sources.list.d/python.list
! apt update

# Python3.9にアップデート
!sudo apt-get -y update
!sudo apt-get -y install python3.9
!sudo apt-get -y install python3.9-dev
!sudo apt-get -y install python3-pip
!sudo apt-get -y install python3.9-distutils
!python3.9 -m pip install --upgrade setuptools
!python3.9 -m pip install --upgrade pip
!python3.9 -m pip install --upgrade distlib

! sudo apt install libosmesa6-dev
! sudo apt-get install build-essential libgl1-mesa-dev libglu1-mesa-dev freeglut3-dev gcc libyaml-dev

import shutil

if os.path.exists('./mmd-auto-trace-3'):
    shutil.rmtree('./mmd-auto-trace-3')

# MMD自動トレースv3 バージョンタグ
version_tag = "develop"

! git clone  --recursive --depth 1 -b "$version_tag" "https://github.com/miu200521358/mmd-auto-trace-3.git"

! python3.9 -m pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113
! python3.9 -m pip install cython
! python3.9 -m pip install six
! python3.9 -m pip install terminaltables
! python3.9 -m pip install scipy
! python3.9 -m pip install opencv-python
! python3.9 -m pip install matplotlib
! python3.9 -m pip install visdom
! python3.9 -m pip install tqdm
! python3.9 -m pip install tensorboardx
! python3.9 -m pip install easydict
! python3.9 -m pip install pyyaml
! python3.9 -m pip install halpecocotools
! python3.9 -m pip install munkres
! python3.9 -m pip install timm==0.1.20
! python3.9 -m pip install natsort
! python3.9 -m pip install opendr
! python3.9 -m pip install pycocotools 
! python3.9 -m pip install cython-bbox
! python3.9 -m pip install --upgrade -r "./mmd-auto-trace-3/requirements.txt"
! cd /content/mmd-auto-trace-3/src/AlphaPose && python3.9 setup.py build develop

# ----------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【①-F】　データ配置

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "①-F"
check_exec_dict(ckey)

# ----------------------------

# 3d-multi-pose-pose.zip
downloaded = drive.CreateFile({'id': '1LtUFDGOSLW52qQwg59azrJxOfXTCtFDm'})
downloaded.GetContentFile('/content/mmd-auto-trace-3/data/3d-multi-pose-pose.zip')
! cd mmd-auto-trace-3/data && unzip -o ./3d-multi-pose-pose.zip

# alphapose.zip
downloaded = drive.CreateFile({'id': '16MITkRde9HVyr-6CVC55PkEYlTokwdG_'})
downloaded.GetContentFile('/content/mmd-auto-trace-3/data/alphapose.zip')
! cd mmd-auto-trace-3/data && unzip -o ./alphapose.zip

# posetriplet.zip
downloaded = drive.CreateFile({'id': '1Z_L0bq78sA66piDbcP3OEMPv0pImdrHU'})
downloaded.GetContentFile('/content/mmd-auto-trace-3/data/posetriplet.zip')
! cd mmd-auto-trace-3/data && unzip -o ./posetriplet.zip

# ----------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【②-A】　トレース元動画配置

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "②-A"
check_exec_dict(ckey)

# ----------------------------

input_video_name = "input.mp4"  #@param {type: "string"}

import os
input_video_path = os.path.join(base_dir_path, input_video_name)

if not os.path.exists(input_video_path):
    print("■■■■■■■■■■■■■■■")
    print("■　** ERROR **")
    print("■　autotrace 直下に 指定動画 が見つかりませんでした。")
    print("■　・マイドライブ > 「autotrace」フォルダ > 指定画像 の順番でファイルが配置されているか再確認してください")
    print("■　・配置した直後だと、認識が間に合わない場合がありますので、10秒ほど待ってから再度実行してください")
    print("■　ファイルパス: " + input_video_path)
    print("■■■■■■■■■■■■■■■")
    raise IpyExit
else:
    ! cp "$input_video_path" .

    input_video_path = f"/content/{input_video_name}"

    print("■■■■■■■■■■■■■■■")
    print("■　** OK **")
    print("■　トレース元動画の配置が完了しました。")
    print("■　ファイルパス: " + input_video_path)
    print("■■■■■■■■■■■■■■■")


# ----------------------------

# セル終了
finish_cell(ckey)

#@markdown ■■■■■■■■■■■■■■■■■■

#@markdown 【②-B】　MMD自動トレースv3 実行

#@markdown ■■■■■■■■■■■■■■■■■■

# セル実行可否
ckey = "②-B"
check_exec_dict(ckey)

# ----------------------------

# import torch
# torch.multiprocessing.set_sharing_strategy("file_descriptor")
# torch.multiprocessing.set_start_method('spawn')

import datetime
process_img_dir_name = f"{os.path.basename(input_video_path).replace('.', '_')}_{datetime.datetime.now():%Y%m%d_%H%M%S}"
process_img_dir_path = os.path.join('/content/output', process_img_dir_name)
os.makedirs(process_img_dir_path, exist_ok=True)

! cd /content/mmd-auto-trace-3/src && python3.9 executor.py --video-file "$input_video_path" --parent-dir "$process_img_dir_path" --process prepare,alphapose,multipose,posetriplet,mix,motion --verbose 20 --log-mode 0 --lang ja

import os
if os.path.exists("/content/mmd-auto-trace-3/log/quit.log"):
    raise IpyExit

# 結果をGoogleドライブにコピー
output_path = os.path.join(base_dir_path, process_img_dir_name)
os.makedirs(output_path, exist_ok=True)

import shutil

# alphapose.mp4
shutil.copy(os.path.join(process_img_dir_path, "02_alphapose", "alphapose.mp4"), os.path.join(output_path, "alphapose.mp4"))

# mix(json)
shutil.copytree(os.path.join(process_img_dir_path, "05_mix"), os.path.join(output_path, "01_json"))

# viz_pos.html
shutil.copy(os.path.join("/content/mmd-auto-trace-3/data/viz_pos.html"), os.path.join(output_path, "01_json", "viz_pos.html"))

# motion(vmd)
shutil.copytree(os.path.join(process_img_dir_path, "06_motion"), os.path.join(output_path, "02_motion"))

os.makedirs(os.path.join(output_path, "trace_model"))

# pmx
shutil.copy(os.path.join("/content/mmd-auto-trace-3/data/pmx/trace_model.pmx"), os.path.join(output_path, "trace_model", "trace_model.pmx"))

# pmx(tex)
shutil.copytree("/content/mmd-auto-trace-3/data/pmx/tex", os.path.join(output_path, "trace_model/tex"))

# readme
shutil.copy(os.path.join("/content/mmd-auto-trace-3/data/readme_ja.txt"), os.path.join(output_path, "readme.txt"))

# ----------------------------

# セル終了
finish_cell(ckey)