# Setting up env variables for cleaner command line commands.
import os

%env KEY=nvidia_tlt
%env NUM_GPUS=1
%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/gazenet
%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/gazenet/data

# Set this path if you don't run the notebook from the samples directory.
# %env NOTEBOOK_ROOT=~/tao-samples/gazenet

# Please define this local project directory that needs to be mapped to the TAO docker session.
# The dataset is expected to be present in $LOCAL_PROJECT_DIR/gazenet/data, while the results for the steps
# in this notebook will be stored at $LOCAL_PROJECT_DIR/gazenet
# !PLEASE MAKE SURE TO UPDATE THIS PATH!.
%env LOCAL_PROJECT_DIR=FIXME

# $SAMPLES_DIR is the path to the sample notebook folder and the dependency folder
# $SAMPLES_DIR/deps should exist for dependency installation
%env SAMPLES_DIR=FIXME

os.environ["LOCAL_DATA_DIR"] = os.path.join(
    os.getenv("LOCAL_PROJECT_DIR", os.getcwd()),
    "gazenet/data"
)
os.environ["LOCAL_EXPERIMENT_DIR"] = os.path.join(
    os.getenv("LOCAL_PROJECT_DIR", os.getcwd()),
    "gazenet"
)

# The sample spec files are present in the same path as the downloaded samples.
os.environ["LOCAL_SPECS_DIR"] = os.path.join(
    os.getenv("NOTEBOOK_ROOT", os.getcwd()),
    "specs"
)
%env SPECS_DIR=/workspace/tao-experiments/gazenet/specs

# Showing list of specification files.
!ls -rlt $LOCAL_SPECS_DIR

# Mapping up the local directories to the TAO docker.
import json
mounts_file = os.path.expanduser("~/.tao_mounts.json")

# Define the dictionary with the mapped drives
drive_map = {
    "Mounts": [
        # Mapping the data directory
        {
            "source": os.environ["LOCAL_PROJECT_DIR"],
            "destination": "/workspace/tao-experiments"
        },
        # Mapping the specs directory.
        {
            "source": os.environ["LOCAL_SPECS_DIR"],
            "destination": os.environ["SPECS_DIR"]
        },
    ]
}

# Writing the mounts file.
with open(mounts_file, "w") as mfile:
    json.dump(drive_map, mfile, indent=4)

!cat ~/.tao_mounts.json

# Install requirement
!pip3 install -r $SAMPLES_DIR/deps/requirements-pip.txt

# Skip this cell if the TAO launcher was already installed.
!pip3 install nvidia-pyindex
!pip3 install nvidia-tao

# View the version of the TAO launcher
!tao info

# check if the label file is presented
!if [ ! -f $SAMPLES_DIR/gazenet/sample_labels/data_factory.zip ]; then echo 'Label file not found, please check your sample path.'; else echo 'Found label file.';fi

# Check the dataset is present
!mkdir -p $LOCAL_DATA_DIR
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze ]; then echo 'Data folder not found, please download.'; else echo 'Found Data folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/data_factory ]; then echo 'Label folder not found, please download.'; else echo 'Found Labels folder.';fi

# Sample json label.
!sed -n 1,201p $LOCAL_DATA_DIR/MPIIFaceGaze/data_factory/day03/p01/p01_day03.json

!python3 mpiifacegaze_convert.py --data_path $LOCAL_DATA_DIR/MPIIFaceGaze \
                                 --json_label_root_path $LOCAL_DATA_DIR/MPIIFaceGaze

# Check the generated data is present
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/p01-day03 ]; then echo 'Generated data folder not found, please regenerated.'; else echo 'Found generated data folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/p01-day03/Config ]; then echo 'Config folder not found, please regenerated.'; else echo 'Found Config folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/p01-day03/Data ]; then echo 'Data folder not found, please regenerated.'; else echo 'Found Data folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/p01-day03/json_datafactory_v2 ]; then echo 'Labels folder not found, please regenerated.'; else echo 'Found Labels folder.';fi

# Check the inference data is present
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/inference-set ]; then echo 'Inference data folder not found, please regenerated.'; else echo 'Found inference data folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/inference-set/Config ]; then echo 'Config folder not found, please regenerated.'; else echo 'Found Config folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/inference-set/Data ]; then echo 'Data folder not found, please regenerated.'; else echo 'Found Data folder.';fi
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/inference-set/json_datafactory_v2 ]; then echo 'Labels folder not found, please regenerated.'; else echo 'Found Labels folder.';fi

# Installing NGC CLI on the local machine.
## Download and install
%env CLI=ngccli_cat_linux.zip
!mkdir -p $LOCAL_PROJECT_DIR/ngccli

# Remove any previously existing CLI installations
!rm -rf $LOCAL_PROJECT_DIR/ngccli/*
!wget "https://ngc.nvidia.com/downloads/$CLI" -P $LOCAL_PROJECT_DIR/ngccli
!unzip -u "$LOCAL_PROJECT_DIR/ngccli/$CLI" -d $LOCAL_PROJECT_DIR/ngccli/
!rm $LOCAL_PROJECT_DIR/ngccli/*.zip 
os.environ["PATH"]="{}/ngccli:{}".format(os.getenv("LOCAL_PROJECT_DIR", ""), os.getenv("PATH", ""))

# List models available in the model registry.
!ngc registry model list nvidia/tao/gazenet:*

# Create the target destination to download the model.
!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrain_models/

# Download the pretrained model from NGC
!ngc registry model download-version nvidia/tao/gazenet:trainable_v1.0 \
    --dest $LOCAL_EXPERIMENT_DIR/pretrain_models/

!ls -rlt $LOCAL_EXPERIMENT_DIR/pretrain_models/gazenet_vtrainable_v1.0

# Check the dataset is present
!if [ ! -f $LOCAL_EXPERIMENT_DIR/pretrain_models/gazenet_vtrainable_v1.0/model.tlt ]; then echo 'Pretrain model file not found, please download.'; else echo 'Found Pretrain model file.';fi

!tao gazenet dataset_convert -folder-suffix pipeline \
                             -norm_folder_name Norm_Data \
                             -sets p01-day03 \
                             -data_root_path $DATA_DOWNLOAD_DIR/MPIIFaceGaze/sample-dataset

!ls -rl $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset

# check the tfrecords are presented
!if [ ! -d $LOCAL_DATA_DIR/MPIIFaceGaze/sample-dataset/p01-day03/Ground_Truth_DataFactory_pipeline ]; then echo 'Tfrecords folder not found, please generate.'; else echo 'Found Tfrecords folder.';fi

!cat $LOCAL_SPECS_DIR/gazenet_tlt_pretrain.yaml

!tao gazenet train -e $SPECS_DIR/gazenet_tlt_pretrain.yaml \
                   -r $USER_EXPERIMENT_DIR/experiment_result/exp1 \
                   -k $KEY

!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_result/exp1

!tao gazenet evaluate -type kpi_testing \
                      -m $USER_EXPERIMENT_DIR/experiment_result/exp1 \
                      -e $SPECS_DIR/gazenet_tlt_pretrain.yaml \
                      -k $KEY

!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_result/exp1/KPI_TMP

!tao gazenet inference -e $SPECS_DIR/gazenet_tlt_pretrain.yaml \
                       -i $DATA_DOWNLOAD_DIR/MPIIFaceGaze/sample-dataset/inference-set \
                       -m $USER_EXPERIMENT_DIR/experiment_result/exp1/model.tlt \
                       -o $USER_EXPERIMENT_DIR/experiment_result/exp1 \
                       -k $KEY

!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_result/exp1/result.txt

import sys
import cv2
import numpy as np
import os
import json
import IPython.display
import PIL.Image
from utils_gazeviz import load_cam_intrinsics,\
        get_landmarks_dict, visualize_frame

# load data
data_root_path = os.path.join(os.environ['LOCAL_DATA_DIR'],
                              'MPIIFaceGaze/sample-dataset/inference-set')
print(data_root_path)
# load calibration
config_path = os.path.join(data_root_path, 'Config')
calib = {}
camera_mat, distortion_coeffs = load_cam_intrinsics(config_path)
distortion_coeffs = distortion_coeffs[0:5]
calib['cam'] = camera_mat
calib['dist'] = distortion_coeffs

# load json files
json_file_folder = os.path.join(data_root_path, 'json_datafactory_v2')
landmarks_dict = get_landmarks_dict(json_file_folder)
assert len(landmarks_dict.keys()) > 0

# visualize each frame in the result file
num_viz_frames = 5
result_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'],
                           "experiment_result/exp1/result.txt")

with open(result_path, 'r') as reader:
    lines = reader.readlines()

num_lines = len(lines)
num_viz_frames = min(num_viz_frames, num_lines)
for k in range(0, num_viz_frames):
    content = lines[k]
    line_info = content.split(' ')
    old_frame_path = line_info[0]
    sub_path = old_frame_path.split(os.environ['DATA_DOWNLOAD_DIR'])[-1]
    frame_path = os.environ['LOCAL_DATA_DIR'] + sub_path
    cam_coord = np.array(line_info[1:4], dtype=np.float32)
    frame_name = frame_path.split('/')[-1]
    landmarks_2D = landmarks_dict[frame_name]
    display_frame, le_px, le_por, re_px, re_por = visualize_frame(frame_path, landmarks_2D, cam_coord, calib)
    # Visualize selected landmarks
    cv2.arrowedLine(display_frame, tuple(le_px), tuple(le_por), (0, 255, 0), thickness=2, tipLength=0.05)
    cv2.arrowedLine(display_frame, tuple(re_px), tuple(re_por), (0, 255, 0), thickness=2, tipLength=0.05)
    IPython.display.display(PIL.Image.fromarray(display_frame))

!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_final
# Removing a pre-existing copy of the etlt if there has been any.
import os
output_file=os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'],
                         "experiment_dir_final/gazenet_onnx.etlt")
if os.path.exists(output_file):
    os.system("rm {}".format(output_file))

!tao gazenet export -m $USER_EXPERIMENT_DIR/experiment_result/exp1/model.tlt \
                    -o $USER_EXPERIMENT_DIR/experiment_dir_final/gazenet_onnx.etlt \
                    -t tfonnx \
                    -k $KEY

# check the Deployed file is presented
!if [ ! -f $LOCAL_EXPERIMENT_DIR/experiment_dir_final/gazenet_onnx.etlt ]; then echo 'Deployed file not found, please generate.'; else echo 'Found Deployed file folder.';fi