"""
DeepLabCut Toolbox (deeplabcut.org)
© A. & M. Mathis Labs

Licensed under GNU Lesser General Public License v3.0
"""


import sys
import shutil
import warnings

from dlclive import benchmark_videos
import urllib.request
import argparse
from pathlib import Path
from dlclibrary.dlcmodelzoo.modelzoo_download import (
    download_huggingface_model,
)


def urllib_pbar(count, blockSize, totalSize):
    percent = int(count * blockSize * 100 / totalSize)
    outstr = f"{round(percent)}%"
    sys.stdout.write(outstr)
    sys.stdout.write("\b"*len(outstr))
    sys.stdout.flush()


def main(display:bool=None):
    parser = argparse.ArgumentParser(
        description="Test DLC-Live installation by downloading and evaluating a demo DLC project!")
    parser.add_argument('--nodisplay', action='store_false', help="Run the test without displaying tracking")
    args = parser.parse_args()

    if display is None:
        display = args.nodisplay

    if not display:
        print('Running without displaying video')

    # make temporary directory in $HOME
    print("\nCreating temporary directory...\n")
    tmp_dir = Path().home() / 'dlc-live-tmp'
    tmp_dir.mkdir(mode=0o775,exist_ok=True)

    video_file = str(tmp_dir / 'dog_clip.avi')
    model_dir = tmp_dir / 'DLC_Dog_resnet_50_iteration-0_shuffle-0'

    # download dog test video from github:
    print(f"Downloading Video to {video_file}")
    url_link = "https://github.com/DeepLabCut/DeepLabCut-live/blob/master/check_install/dog_clip.avi?raw=True"
    urllib.request.urlretrieve(url_link, video_file, reporthook=urllib_pbar)

    # download exported dog model from DeepLabCut Model Zoo
    if Path(model_dir / 'snapshot-75000.pb').exists():
        print('Model already downloaded, using cached version')
    else:
        print("Downloading full_dog model from the DeepLabCut Model Zoo...")
        download_huggingface_model("full_dog", model_dir)

    # assert these things exist so we can give informative error messages
    assert Path(video_file).exists()
    assert Path(model_dir / 'snapshot-75000.pb').exists()

    # run benchmark videos
    print("\n Running inference...\n")
    # model_dir = "DLC_Dog_resnet_50_iteration-0_shuffle-0"
    # print(video_file)
    benchmark_videos(str(model_dir), video_file, display=display, resize=0.5, pcutoff=0.25)

    # deleting temporary files
    print("\n Deleting temporary files...\n")
    try:
        shutil.rmtree(tmp_dir)
    except PermissionError:
        warnings.warn(f'Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!')

    print("\nDone!\n")


if __name__ == "__main__":


    display = args.nodisplay


    main(display=args.nodisplay)