# Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
#
import os
from enum import Enum, IntEnum
import carb


class CaptureMovieType(IntEnum):
    SEQUENCE = 0
    SUNSTUDY = 1
    PLAYLIST = 2


class CaptureRangeType(IntEnum):
    FRAMES = 0
    SECONDS = 1


class CaptureRenderPreset(IntEnum):
    PATH_TRACE = 0
    RAY_TRACE = 1
    IRAY = 2


class CaptureDebugMaterialType(IntEnum):
    SHADED = 0
    WHITE = 1


EXR_COMPRESSION_METHODS = {"zip", "zips", "dwaa", "dwab", "piz", "rle", "b44", "b44a"}
MP4_ENCODING_PRESETS = {
    "PRESET_DEFAULT",
    "PRESET_HP",
    "PRESET_HQ",
    "PRESET_BD",
    "PRESET_LOW_LATENCY_DEFAULT",
    "PRESET_LOW_LATENCY_HQ",
    "PRESET_LOW_LATENCY_HP",
    "PRESET_LOSSLESS_DEFAULT",
    "PRESET_LOSSLESS_HP"
}
MP4_ENCODING_PROFILES = {
    "H264_PROFILE_BASELINE",
    "H264_PROFILE_MAIN",
    "H264_PROFILE_HIGH",
    "H264_PROFILE_HIGH_444",
    "H264_PROFILE_STEREO",
    "H264_PROFILE_SVC_TEMPORAL_SCALABILITY",
    "H264_PROFILE_PROGRESSIVE_HIGH",
    "H264_PROFILE_CONSTRAINED_HIGH",
    "HEVC_PROFILE_MAIN",
    "HEVC_PROFILE_MAIN10",
    "HEVC_PROFILE_FREXT"
}
MP4_ENCODING_RC_MODES = {
    "RC_CONSTQP",
    "RC_VBR",
    "RC_CBR",
    "RC_CBR_LOWDELAY_HQ",
    "RC_CBR_HQ",
    "RC_VBR_HQ"
}


class CaptureOptions:
    """ All Capture options that will be used when capturing.

        Note: When adding an attribute make sure it is exposed via the constructor.
        Not doing this will cause erorrs when serializing and deserializing this object.
    """
    INVALID_ANIMATION_FPS = -1

    def __init__(
        self,
        camera="camera",
        range_type=CaptureRangeType.FRAMES,
        capture_every_nth_frames=-1,
        fps=24,
        start_frame=1,
        end_frame=40,
        start_time=0,
        end_time=10,
        res_width=1920,
        res_height=1080,
        render_preset=CaptureRenderPreset.PATH_TRACE,
        debug_material_type=CaptureDebugMaterialType.SHADED,
        spp_per_iteration=1,
        path_trace_spp=1,
        ptmb_subframes_per_frame=1,
        ptmb_fso=0.0,
        ptmb_fsc=1.0,
        output_folder="",
        file_name="Capture",
        file_name_num_pattern=".####",
        file_type=".tga",
        save_alpha=False,
        hdr_output=False,
        show_pathtracing_single_frame_progress=False,
        preroll_frames=0,
        overwrite_existing_frames=False,
        movie_type=CaptureMovieType.SEQUENCE,
        sunstudy_start_time=0.0,
        sunstudy_current_time=0.0,
        sunstudy_end_time=0.0,
        sunstudy_movie_length_in_seconds=0,
        sunstudy_player=None,
        real_time_settle_latency_frames=0,
        renumber_negative_frame_number_from_0=False,
        render_product="",
        exr_compression_method="zips",
        mp4_encoding_bitrate=16777216,
        mp4_encoding_iframe_interval=60,
        mp4_encoding_preset="PRESET_DEFAULT",
        mp4_encoding_profile="H264_PROFILE_HIGH",
        mp4_encoding_rc_mode="RC_VBR",
        mp4_encoding_rc_target_quality=0,
        mp4_encoding_video_full_range_flag=False,
        app_level_capture=False,
        animation_fps=INVALID_ANIMATION_FPS
    ):
        self._camera = camera
        self._range_type = range_type
        self._capture_every_nth_frames = capture_every_nth_frames
        self._fps = fps
        self._start_frame = start_frame
        self._end_frame = end_frame
        self._start_time = start_time
        self._end_time = end_time
        self._res_width = res_width
        self._res_height = res_height
        self._render_preset = render_preset
        self._debug_material_type = debug_material_type
        self._spp_per_iteration = spp_per_iteration
        self._path_trace_spp = path_trace_spp
        self._ptmb_subframes_per_frame = ptmb_subframes_per_frame
        self._ptmb_fso = ptmb_fso
        self._ptmb_fsc = ptmb_fsc
        self._output_folder = output_folder
        self._file_name = file_name
        self._file_name_num_pattern = file_name_num_pattern
        self._file_type = file_type
        self._save_alpha = save_alpha
        self._hdr_output = hdr_output
        self._show_pathtracing_single_frame_progress = show_pathtracing_single_frame_progress
        self._preroll_frames = preroll_frames
        self._overwrite_existing_frames = overwrite_existing_frames
        self._movie_type = movie_type
        self._sunstudy_start_time = sunstudy_start_time
        self._sunstudy_current_time = sunstudy_current_time
        self._sunstudy_end_time = sunstudy_end_time
        self._sunstudy_movie_length_in_seconds = sunstudy_movie_length_in_seconds
        self._sunstudy_player = sunstudy_player
        self._real_time_settle_latency_frames = real_time_settle_latency_frames
        self._renumber_negative_frame_number_from_0 = renumber_negative_frame_number_from_0
        self._render_product = render_product
        self.exr_compression_method = exr_compression_method
        self.mp4_encoding_bitrate = mp4_encoding_bitrate
        self.mp4_encoding_iframe_interval = mp4_encoding_iframe_interval
        self.mp4_encoding_preset = mp4_encoding_preset
        self.mp4_encoding_profile = mp4_encoding_profile
        self.mp4_encoding_rc_mode = mp4_encoding_rc_mode
        self.mp4_encoding_rc_target_quality = mp4_encoding_rc_target_quality
        self.mp4_encoding_video_full_range_flag = mp4_encoding_video_full_range_flag
        self._app_level_capture = app_level_capture
        self.animation_fps = animation_fps

    def to_dict(self):
        data = vars(self)
        return {key.lstrip("_"): value for key, value in data.items()}

    @classmethod
    def from_dict(cls, options):
        return cls(**options)

    @property
    def camera(self):
        return self._camera

    @camera.setter
    def camera(self, value):
        self._camera = value

    @property
    def range_type(self):
        return self._range_type

    @range_type.setter
    def range_type(self, value):
        self._range_type = value

    @property
    def capture_every_Nth_frames(self):
        return self._capture_every_nth_frames

    @capture_every_Nth_frames.setter
    def capture_every_Nth_frames(self, value):
        self._capture_every_nth_frames = value

    @property
    def fps(self):
        return self._fps

    @fps.setter
    def fps(self, value):
        self._fps = value

    @property
    def start_frame(self):
        return self._start_frame

    @start_frame.setter
    def start_frame(self, value):
        self._start_frame = value

    @property
    def end_frame(self):
        return self._end_frame

    @end_frame.setter
    def end_frame(self, value):
        self._end_frame = value

    @property
    def start_time(self):
        return self._start_time

    @start_time.setter
    def start_time(self, value):
        self._start_time = value

    @property
    def end_time(self):
        return self._end_time

    @end_time.setter
    def end_time(self, value):
        self._end_time = value

    @property
    def res_width(self):
        return self._res_width

    @res_width.setter
    def res_width(self, value):
        self._res_width = value

    @property
    def res_height(self):
        return self._res_height

    @res_height.setter
    def res_height(self, value):
        self._res_height = value

    @property
    def render_preset(self):
        return self._render_preset

    @render_preset.setter
    def render_preset(self, value):
        self._render_preset = value

    @property
    def debug_material_type(self):
        return self._debug_material_type

    @debug_material_type.setter
    def debug_material_type(self, value):
        self._debug_material_type = value

    @property
    def spp_per_iteration(self):
        return self._spp_per_iteration

    @spp_per_iteration.setter
    def spp_per_iteration(self, value):
        self._spp_per_iteration = value

    @property
    def path_trace_spp(self):
        return self._path_trace_spp

    @path_trace_spp.setter
    def path_trace_spp(self, value):
        self._path_trace_spp = value

    @property
    def ptmb_subframes_per_frame(self):
        return self._ptmb_subframes_per_frame

    @ptmb_subframes_per_frame.setter
    def ptmb_subframes_per_frame(self, value):
        self._ptmb_subframes_per_frame = value

    @property
    def ptmb_fso(self):
        return self._ptmb_fso

    @ptmb_fso.setter
    def ptmb_fso(self, value):
        self._ptmb_fso = value

    @property
    def ptmb_fsc(self):
        return self._ptmb_fsc

    @ptmb_fsc.setter
    def ptmb_fsc(self, value):
        self._ptmb_fsc = value

    @property
    def output_folder(self):
        return self._output_folder

    @output_folder.setter
    def output_folder(self, value):
        self._output_folder = value

    @property
    def file_name(self):
        return self._file_name

    @file_name.setter
    def file_name(self, value):
        self._file_name = value

    @property
    def file_name_num_pattern(self):
        return self._file_name_num_pattern

    @file_name_num_pattern.setter
    def file_name_num_pattern(self, value):
        self._file_name_num_pattern = value

    @property
    def file_type(self):
        return self._file_type

    @file_type.setter
    def file_type(self, value):
        self._file_type = value

    @property
    def save_alpha(self):
        return self._save_alpha

    @save_alpha.setter
    def save_alpha(self, value):
        self._save_alpha = value

    @property
    def hdr_output(self):
        return self._hdr_output

    @hdr_output.setter
    def hdr_output(self, value):
        self._hdr_output = value

    @property
    def show_pathtracing_single_frame_progress(self):
        return self._show_pathtracing_single_frame_progress

    @show_pathtracing_single_frame_progress.setter
    def show_pathtracing_single_frame_progress(self, value):
        self._show_pathtracing_single_frame_progress = value

    @property
    def preroll_frames(self):
        return self._preroll_frames

    @preroll_frames.setter
    def preroll_frames(self, value):
        self._preroll_frames = value

    @property
    def overwrite_existing_frames(self):
        return self._overwrite_existing_frames

    @overwrite_existing_frames.setter
    def overwrite_existing_frames(self, value):
        self._overwrite_existing_frames = value

    @property
    def movie_type(self):
        return self._movie_type

    @movie_type.setter
    def movie_type(self, value):
        self._movie_type = value

    @property
    def sunstudy_start_time(self):
        return self._sunstudy_start_time

    @sunstudy_start_time.setter
    def sunstudy_start_time(self, value):
        self._sunstudy_start_time = value

    @property
    def sunstudy_current_time(self):
        return self._sunstudy_current_time

    @sunstudy_current_time.setter
    def sunstudy_current_time(self, value):
        self._sunstudy_current_time = value

    @property
    def sunstudy_end_time(self):
        return self._sunstudy_end_time

    @sunstudy_end_time.setter
    def sunstudy_end_time(self, value):
        self._sunstudy_end_time = value

    @property
    def sunstudy_movie_length_in_seconds(self):
        return self._sunstudy_movie_length_in_seconds

    @sunstudy_movie_length_in_seconds.setter
    def sunstudy_movie_length_in_seconds(self, value):
        self._sunstudy_movie_length_in_seconds = value

    @property
    def sunstudy_player(self):
        return self._sunstudy_player

    @sunstudy_player.setter
    def sunstudy_player(self, value):
        self._sunstudy_player = value

    @property
    def real_time_settle_latency_frames(self):
        return self._real_time_settle_latency_frames

    @real_time_settle_latency_frames.setter
    def real_time_settle_latency_frames(self, value):
        self._real_time_settle_latency_frames = value

    @property
    def renumber_negative_frame_number_from_0(self):
        return self._renumber_negative_frame_number_from_0

    @renumber_negative_frame_number_from_0.setter
    def renumber_negative_frame_number_from_0(self, value):
        self._renumber_negative_frame_number_from_0 = value

    @property
    def render_product(self):
        return self._render_product

    @render_product.setter
    def render_product(self, value):
        self._render_product = value

    @property
    def exr_compression_method(self):
        return self._exr_compression_method

    @exr_compression_method.setter
    def exr_compression_method(self, value):
        global EXR_COMPRESSION_METHODS
        if value in EXR_COMPRESSION_METHODS:
            self._exr_compression_method = value
        else:
            self._exr_compression_method = "zips"
            carb.log_warn(f"Can't set unsupported compression method {value} for exr format, set to default zip16. Supported values are: {EXR_COMPRESSION_METHODS}")

    @property
    def mp4_encoding_bitrate(self):
        return self._mp4_encoding_bitrate

    @mp4_encoding_bitrate.setter
    def mp4_encoding_bitrate(self, value):
        self._mp4_encoding_bitrate = value

    @property
    def mp4_encoding_iframe_interval(self):
        return self._mp4_encoding_iframe_interval

    @mp4_encoding_iframe_interval.setter
    def mp4_encoding_iframe_interval(self, value):
        self._mp4_encoding_iframe_interval = value

    @property
    def mp4_encoding_preset(self):
        return self._mp4_encoding_preset

    @mp4_encoding_preset.setter
    def mp4_encoding_preset(self, value):
        global MP4_ENCODING_PRESETS
        if value in MP4_ENCODING_PRESETS:
            self._mp4_encoding_preset = value
        else:
            self._mp4_encoding_preset = "PRESET_DEFAULT"
            carb.log_warn(f"Can't set unsupported mp4 encoding preset {value}, set to default {self._mp4_encoding_preset}. Supported values are: {MP4_ENCODING_PRESETS}")

    @property
    def mp4_encoding_profile(self):
        return self._mp4_encoding_profile

    @mp4_encoding_profile.setter
    def mp4_encoding_profile(self, value):
        global MP4_ENCODING_PROFILES
        if value in MP4_ENCODING_PROFILES:
            self._mp4_encoding_profile = value
        else:
            self._mp4_encoding_profile = "H264_PROFILE_HIGH"
            carb.log_warn(f"Can't set unsupported mp4 encoding profile {value}, set to default {self._mp4_encoding_profile}. Supported values are: {MP4_ENCODING_PROFILES}")

    @property
    def mp4_encoding_rc_mode(self):
        return self._mp4_encoding_rc_mode

    @mp4_encoding_rc_mode.setter
    def mp4_encoding_rc_mode(self, value):
        global MP4_ENCODING_RC_MODES
        if value in MP4_ENCODING_RC_MODES:
            self._mp4_encoding_rc_mode = value
        else:
            self._mp4_encoding_rc_mode = "RC_VBR"
            carb.log_warn(f"Can't set unsupported mp4 encoding rate control mode {value}, set to default {self._mp4_encoding_rc_mode}. Supported values are: {MP4_ENCODING_RC_MODES}")

    @property
    def mp4_encoding_rc_target_quality(self):
        return self._mp4_encoding_rc_target_quality

    @mp4_encoding_rc_target_quality.setter
    def mp4_encoding_rc_target_quality(self, value):
        if 0 <= value and value <= 51:
            self._mp4_encoding_rc_target_quality = value
        else:
            self._mp4_encoding_rc_target_quality = 0
            carb.log_warn(f"Can't set unsupported mp4 encoding rate control target quality {value}, set to default {self._mp4_encoding_rc_target_quality}. Supported range is [0, 51]")

    @property
    def mp4_encoding_video_full_range_flag(self):
        return self._mp4_encoding_video_full_range_flag

    @mp4_encoding_video_full_range_flag.setter
    def mp4_encoding_video_full_range_flag(self, value):
        self._mp4_encoding_video_full_range_flag = value

    @property
    def app_level_capture(self):
        return self._app_level_capture

    @app_level_capture.setter
    def app_level_capture(self, value):
        self._app_level_capture = value

    @property
    def animation_fps(self):
        return self._animation_fps

    @animation_fps.setter
    def animation_fps(self, value):
        self._animation_fps = value

    def is_video(self):
        return self.file_type == ".mp4"

    def is_capturing_nth_frames(self):
        return self._capture_every_nth_frames > 0

    def is_capturing_pathtracing_single_frame(self):
        return (
            self.is_video() is False
            and self.is_capturing_nth_frames() is False
            and (self.render_preset == CaptureRenderPreset.PATH_TRACE or self.render_preset == CaptureRenderPreset.IRAY)
        )

    def is_capturing_frame(self):
        return self._range_type == CaptureRangeType.FRAMES

    def get_full_path(self):
        return os.path.join(self._output_folder, self._file_name + self._file_type)
