__version__ = "5.0"

from meshroom.core import desc
from meshroom.core.utils import EXR_STORAGE_DATA_TYPE, VERBOSE_LEVEL

# List of supported video extensions (provided by OpenImageIO)
videoExts = [".avi", ".mov", ".mp4", ".m4a", ".m4v", ".3gp", ".3g2", ".mj2", ".m4v", ".mpg"]

class KeyframeSelectionNodeSize(desc.DynamicNodeSize):
    def computeSize(self, node):
        inputPathsSize = super(KeyframeSelectionNodeSize, self).computeSize(node)
        s = 0
        finalSize = 0
        defaultParam = self._param

        # Compute the size for each entry in the list of input paths
        for input in node.attribute("inputPaths").value:
            self._param = input.fullName
            s = s + super(KeyframeSelectionNodeSize, self).computeSize(node)

        # Retrieve the maximum number of keyframes for the smart selection
        # (which is high by default)
        maxFramesSmart = node.attribute("selectionMethod.smartSelection.maxNbOutFrames").value

        # If the smart selection is enabled and the number of input frames is available
        # (s is not equal to the number of input paths), set the size as the minimum between the
        # number of input frames and maximum number of output keyframes. If the number of input
        # frames is not available, set the size to the maximum number of output keyframes.
        smartSelectionOn = node.attribute("selectionMethod.useSmartSelection").value
        if smartSelectionOn:
            if s != inputPathsSize:
                finalSize = min(s, maxFramesSmart)
            else:
                finalSize = maxFramesSmart

        # If the smart selection is not enabled, the maximum number of output keyframes
        # for the regular mode can be used if and only if it has been set, in the same way as
        # for the smart selection. If the maximum number of frames has not been set, then the
        # size is either the minimum between the maximum number of output keyframes for the smart
        # selection and the number of input frames if it is available, or the maximum number of
        # output keyframes for the smart selection if the number of input frames is not available.
        else:
            maxFrames = node.attribute("selectionMethod.regularSelection.maxNbOutFrames").value
            if maxFrames > 0 and s != inputPathsSize:
                finalSize = min(s, maxFrames)
            elif maxFrames > 0 and s == inputPathsSize:
                finalSize = maxFrames
            elif maxFrames <= 0 and s != inputPathsSize:
                finalSize = min(s, maxFramesSmart)
            else:
                finalSize = maxFramesSmart

        # Reset the param used to compute size to the default one: if the size is computed again,
        # this will prevent having an inputPathsSize that is erroneous
        self._param = defaultParam
        return finalSize


class KeyframeSelection(desc.AVCommandLineNode):
    commandLine = "aliceVision_keyframeSelection {allParams}"
    size = KeyframeSelectionNodeSize("inputPaths")

    category = "Utils"
    documentation = """
Allows to extract keyframes from a video and insert metadata.
It can extract frames from a synchronized multi-cameras rig.

You can extract frames at regular interval by configuring only the min/maxFrameStep.
"""

    inputs = [
        desc.ListAttribute(
            elementDesc=desc.File(
                name="inputPath",
                label="Input Path",
                description="Input path.",
                value="",
            ),
            name="inputPaths",
            label="Input Paths",
            description="Input video files, image sequence directories or SfMData file.",
        ),
        desc.ListAttribute(
            elementDesc=desc.File(
                name="brand",
                label="Brand",
                description="Camera brand.",
                value="",
            ),
            name="brands",
            label="Brands",
            description="Camera brands.",
        ),
        desc.ListAttribute(
            elementDesc=desc.File(
                name="model",
                label="Model",
                description="Camera model.",
                value="",
            ),
            name="models",
            label="Models",
            description="Camera models.",
        ),
        desc.ListAttribute(
            elementDesc=desc.FloatParam(
                name="mmFocal",
                label="Focal",
                description="Focal in mm (will be used if not 0).",
                value=0.0,
                range=(0.0, 500.0, 1.0),
            ),
            name="mmFocals",
            label="Focals",
            description="Focals in mm (will be used if not 0).",
        ),
        desc.File(
            name="sensorDbPath",
            label="Sensor Database",
            description="Camera sensor width database path.",
            value="${ALICEVISION_SENSOR_DB}",
        ),
        desc.ListAttribute(
            elementDesc=desc.File(
                name="masks",
                label="Masks Path",
                description="Directory containing masks to apply to the frames.",
                value="",
            ),
            name="maskPaths",
            label="Masks",
            description="Masks (e.g. segmentation masks) used to exclude some parts of the frames"
                        "from the score computations\nfor the smart keyframe selection.",
            enabled=lambda node: node.selectionMethod.useSmartSelection.value,
        ),
        desc.GroupAttribute(
            name="selectionMethod",
            label="Keyframe Selection Method",
            description="Selection the regular or smart method for the keyframe selection.\n"
                        "- With the regular method, keyframes are selected regularly over the "
                        "sequence with respect to the set parameters.\n"
                        "- With the smart method, keyframes are selected based on their sharpness "
                        "and optical flow scores.",
            group=None,  # skip group from command line
            groupDesc=[
                desc.BoolParam(
                    name="useSmartSelection",
                    label="Use Smart Keyframe Selection",
                    description="Use the smart keyframe selection.",
                    value=True,
                ),
                desc.GroupAttribute(
                    name="regularSelection",
                    label="Regular Keyframe Selection",
                    description="Parameters for the regular keyframe selection.\n"
                                "Keyframes are selected regularly over the sequence with respect "
                                "to the set parameters.",
                    group=None,  # skip group from command line
                    enabled=lambda node: node.selectionMethod.useSmartSelection.value is False,
                    groupDesc=[
                        desc.IntParam(
                            name="minFrameStep",
                            label="Min Frame Step",
                            description="Minimum number of frames between two keyframes.",
                            value=12,
                            range=(1, 1000, 1),
                            enabled=lambda node: node.selectionMethod.regularSelection.enabled,
                        ),
                        desc.IntParam(
                            name="maxFrameStep",
                            label="Max Frame Step",
                            description="Maximum number of frames between two keyframes.\n"
                                        "Ignored if equal to 0.",
                            value=0,
                            range=(0, 1000, 1),
                            enabled=lambda node: node.selectionMethod.regularSelection.enabled,
                        ),
                        desc.IntParam(
                            name="maxNbOutFrames",
                            label="Max Nb Output Frames",
                            description="Maximum number of output frames (0 = no limit).\n"
                                        "'minFrameStep' and 'maxFrameStep' will always be "
                                        "respected, so combining them with this parameter\n"
                                        "might cause the selection to stop before reaching "
                                        "the end of the input sequence(s).",
                            value=0,
                            range=(0, 10000, 1),
                            enabled=lambda node: node.selectionMethod.regularSelection.enabled,
                        ),
                    ],
                ),
                desc.GroupAttribute(
                    name="smartSelection",
                    label="Smart Keyframe Selection",
                    description="Parameters for the smart keyframe selection.\n"
                                "Keyframes are selected based on their sharpness and optical "
                                "flow scores.",
                    group=None,  # skip group from command line
                    enabled=lambda node: node.selectionMethod.useSmartSelection.value,
                    groupDesc=[
                        desc.FloatParam(
                            name="pxDisplacement",
                            label="Pixel Displacement",
                            description="The percentage of pixels in the frame that need to have "
                                        "moved since the last keyframe to be considered for the "
                                        "selection.",
                            value=10.0,
                            range=(0.0, 100.0, 1.0),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                        ),
                        desc.IntParam(
                            name="minNbOutFrames",
                            label="Min Nb Output Frames",
                            description="Minimum number of frames selected to be keyframes.",
                            value=40,
                            range=(1, 100, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                        ),
                        desc.IntParam(
                            name="maxNbOutFrames",
                            label="Max Nb Output Frames",
                            description="Maximum number of frames selected to be keyframes.",
                            value=2000,
                            range=(1, 10000, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                        ),
                        desc.IntParam(
                            name="rescaledWidthSharpness",
                            label="Rescaled Frame's Width For Sharpness",
                            description="Width, in pixels, of the frame used for the sharpness "
                                        "score computation after a rescale.\n"
                                        "Aspect ratio will be preserved. No rescale will be "
                                        "performed if equal to 0.",
                            value=720,
                            range=(0, 4000, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                            advanced=True,
                        ),
                        desc.IntParam(
                            name="rescaledWidthFlow",
                            label="Rescaled Frame's Width For Motion",
                            description="Width, in pixels, of the frame used for the motion score "
                                        "computation after a rescale.\n"
                                        "Aspect ratio will be preserved. No rescale will be "
                                        "performed if equal to 0.",
                            value=720,
                            range=(0, 4000, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                            advanced=True,
                        ),
                        desc.IntParam(
                            name="sharpnessWindowSize",
                            label="Sharpness Window Size",
                            description="The size, in pixels, of the sliding window used to "
                                        "evaluate a frame's sharpness.",
                            value=200,
                            range=(1, 10000, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                            advanced=True,
                        ),
                        desc.IntParam(
                            name="flowCellSize",
                            label="Optical Flow Cell Size",
                            description="The size, in pixels, of the cells within a frame in "
                                        "which the optical flow scores is evaluated.",
                            value=90,
                            range=(10, 2000, 1),
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                            advanced=True,
                        ),
                        desc.IntParam(
                            name="minBlockSize",
                            label="Multi-Threading Minimum Block Size",
                            description="The minimum number of frames to process for a thread "
                                        "to be spawned.\n"
                                        "If using all the available threads implies processing "
                                        "less than this value in every thread, then less threads "
                                        "should be spawned,\n"
                                        "and each will process at least 'minBlockSize' "
                                        "(except maybe for the very last thread, that might "
                                        "process less).",
                            value=10,
                            range=(1, 1000, 1),
                            invalidate=False,
                            enabled=lambda node: node.selectionMethod.smartSelection.enabled,
                            advanced=True,
                        ),
                    ],
                ),
            ],
        ),
        desc.BoolParam(
            name="renameKeyframes",
            label="Rename Output Keyframes",
            description="Instead of using the selected keyframes' index as their name, name them "
                        "as consecutive output frames.\n"
                        "If the selected keyframes are at index [15, 294, 825], they will be "
                        "written as [00000.exr, 00001.exr, 00002.exr] with this\n"
                        "option enabled instead of [00015.exr, 00294.exr, 00825.exr].",
            value=False,
            enabled=lambda node: node.outputExtension.value != "none",
        ),
        desc.ChoiceParam(
            name="outputExtension",
            label="Keyframes File Extension",
            description="File extension of the written keyframes.\n"
                        "If 'none' is selected, no keyframe will be written on disk.\n"
                        "For input videos, 'none' should not be used since the written keyframes "
                        "are used to generate the output SfMData file.",
            value="none",
            values=["none", "exr", "jpg", "png"],
            validValue=lambda node: not (any(ext in input.value.lower()
                                             for ext in videoExts for
                                             input in node.inputPaths.value)
                                             and node.outputExtension.value == "none"),
            errorMessage="A video input has been provided. The output extension should be "
                         "different from 'none'.",
        ),
        desc.ChoiceParam(
            name="storageDataType",
            label="EXR Storage Data Type",
            description="Storage image data type for keyframes written to EXR files:\n"
                        " - float: Use full floating point (32 bits per channel).\n"
                        " - half: Use half float (16 bits per channel).\n"
                        " - halfFinite: Use half float, but clamp values to avoid non-finite "
                        "values.\n"
                        " - auto: Use half float if all values can fit, else use full float.",
            values=EXR_STORAGE_DATA_TYPE,
            value="float",
            enabled=lambda node: node.outputExtension.value == "exr",
            advanced=True,
        ),
        desc.GroupAttribute(
            name="debugOptions",
            label="Debug Options",
            description="Debug options for the Smart keyframe selection method.",
            group=None,  # skip group from command line
            enabled=lambda node: node.selectionMethod.useSmartSelection.value,
            advanced=True,
            groupDesc=[
                desc.GroupAttribute(
                    name="debugScores",
                    label="Export Scores",
                    description="Export the computed sharpness and optical flow scores to a file.",
                    group=None,  # skip group from command line
                    enabled=lambda node: node.debugOptions.enabled,
                    groupDesc=[
                        desc.BoolParam(
                            name="exportScores",
                            label="Export Scores To CSV",
                            description="Export the computed sharpness and optical flow scores to "
                                        "a CSV file.",
                            value=False,
                        ),
                        desc.StringParam(
                            name="csvFilename",
                            label="CSV Filename",
                            description="Name of the CSV file to export. It will be written in "
                                        "the node's output folder.",
                            value="scores.csv",
                            enabled=lambda node: node.debugOptions.debugScores.exportScores.value,
                        ),
                        desc.BoolParam(
                            name="exportSelectedFrames",
                            label="Export Selected Frames",
                            description="Add a column in the CSV file containing 1s for frames "
                                        "that were selected and 0s for those that were not.",
                            value=False,
                            enabled=lambda node: node.debugOptions.debugScores.exportScores.value,
                        ),
                    ],
                ),
                desc.GroupAttribute(
                    name="opticalFlowVisualisation",
                    label="Optical Flow Visualisation",
                    description="Visualise the motion vectors for each input frame in HSV.",
                    group=None,  # skip group from command line
                    enabled=lambda node: node.debugOptions.enabled,
                    groupDesc=[
                        desc.BoolParam(
                            name="exportFlowVisualisation",
                            label="Visualise Optical Flow",
                            description="Export each frame's optical flow HSV visualisation as "
                                        "PNG images.",
                            value=False,
                            enabled=lambda node: node.debugOptions.opticalFlowVisualisation.enabled,
                        ),
                        desc.BoolParam(
                            name="flowVisualisationOnly",
                            label="Only Visualise Optical Flow",
                            description="Export each frame's optical flow HSV visualisation as "
                                        "PNG images, but do not perform any score computation or "
                                        "frame selection.\n"
                                        "If this option is selected, all the other options will "
                                        "be ignored.",
                            value=False,
                            enabled=lambda node: node.debugOptions.opticalFlowVisualisation.enabled,
                        ),
                    ],
                ),
                desc.BoolParam(
                    name="skipSharpnessComputation",
                    label="Skip Sharpness Computation",
                    description="Skip the sharpness score computation.\n"
                                "A fixed score of 1.0 will be applied by default to all the "
                                "frames.",
                    value=False,
                    enabled=lambda node: node.debugOptions.enabled,
                ),
                desc.BoolParam(
                    name="skipSelection",
                    label="Skip Frame Selection",
                    description="Compute the sharpness and optical flow scores, but do not "
                                "proceed to the frame selection.",
                    value=False,
                    enabled=lambda node: node.debugOptions.enabled,
                ),
            ],
        ),
        desc.ChoiceParam(
            name="verboseLevel",
            label="Verbose Level",
            description="Verbosity level (fatal, error, warning, info, debug, trace).",
            values=VERBOSE_LEVEL,
            value="info",
        ),
    ]

    outputs = [
        desc.File(
            name="outputFolder",
            label="Folder",
            description="Output keyframes folder for extracted frames.",
            value="{nodeCacheFolder}",
        ),
        desc.File(
            name="outputSfMDataKeyframes",
            label="Keyframes SfMData",
            description="Output SfMData file containing all the selected keyframes.",
            value="{nodeCacheFolder}/keyframes.sfm",
        ),
        desc.File(
            name="outputSfMDataFrames",
            label="Frames SfMData",
            description="Output SfMData file containing all the frames that were not selected "
                        "as keyframes.\n"
                        "If the input contains videos, this file will not be written since all "
                        "the frames that were not selected do not actually exist on disk.",
            value="{nodeCacheFolder}/frames.sfm",
        ),
    ]
