# ImagePlaneLoader.py
# A streamlined method of importing and optimising reference footage.
# ---------------------------------------------------------------------------------------
# Created by Andy Latham in collaboration with Ilse Zamarripa
# andylatham82@hotmail.com
# March 2022
# Provided with the permission of Steamroller Animation - https://steamrollerstudios.com/
# ---------------------------------------------------------------------------------------

# USAGE:
# import ImagePlaneLoader as ipl
# reload(ipl)
# r = ipl.ImagePlaneLoader()

# OR FOR MAYA 2022:
# import importlib
# import ImagePlaneLoader as ipl
# importlib.reload(ipl)
# r = ipl.ImagePlaneLoader()

# If a camera is selected, the tool will load the image plane into that camera.
# If no camera is selected, a "reference" camera will be created.

version = 2.5

from argparse import FileType
import maya.cmds as cmds
import os
# import subprocess
from subprocess import Popen, PIPE
import maya.OpenMayaUI as omUI
from PySide6 import QtWidgets
import shiboken6
import time
import maya.mel as mel

class ImagePlaneLoader():
  def __init__(self):
    # Show version number.
    print("ImagePlaneLoader v{}".format(version))

    self.fastFormat = False
    self.freePlane = False

    self.CreateUI()

    # Run the loader.
    # self.ImagePlaneLoader()


  def ImagePlaneLoader(self, state):
    # Get user selection and 
    selection = ""
    try:
      selection = cmds.ls(sl = True)[-1]
    except:
      pass

    # Open file dialog box and save the chosen location.
    videoFilePath = cmds.fileDialog2(dialogStyle = 1, fileMode = 1)

    # End script if no path has been selected.
    if videoFilePath == None:
      return
    
    # Format the path without unicode.
    videoFilePath = [str(r) for r in videoFilePath][0]

    # Check if user has selected a video or image sequence.
    # If it's a video, convert it. If not, then just use it.
    importPath = ""
    extensionsToCheck = (".mov", ".mp4", ".avi", ".gif")
    if videoFilePath.endswith(extensionsToCheck):
      # Get video file information.
      videoFileName = os.path.basename(videoFilePath)
      videoRawName, videoExtension = os.path.splitext(videoFileName)
      videoFolderPath = videoFilePath.replace(videoFileName, "")
      
      # Determine if current Maya scene is part of a project.
      project = False
      projectPath = cmds.workspace(q = True, active = True)
      filePath = cmds.file(q = True, sn = True)
      if projectPath in filePath:
        project = True

      if videoExtension == ".gif":
        fileType = "png"
      else:
        fileType = "jpg"

      # Decide whether to store image sequence in project "images" folder or next to video file.
      imagePath = ""
      if project:
        imageFolderPath = projectPath + "/images/" + videoRawName + "/"
        imagePath = imageFolderPath + videoRawName + "_%04d."+ fileType
        importPath = imageFolderPath + videoRawName + "_0001."+ fileType
      else:
        imageFolderPath = videoFolderPath + videoRawName + "/"
        imagePath = imageFolderPath + videoRawName + "_%04d."+ fileType
        importPath = imageFolderPath + videoRawName + "_0001."+ fileType

      # Create image sequence folder if it doesn't exist.
      if not os.path.exists(imageFolderPath):
        os.makedirs(imageFolderPath)

      # Define location of ffmpeg.
      ffmpegLocation = os.path.join(os.path.dirname(__file__), 'ffmpeg.exe')

      # Get frame rate of Maya scene.
      fps = mel.eval('float $fps = `currentTimeUnitToFPS`')

      # Convert video to image sequence.
      if self.fastFormat:
        # p = Popen([ffmpegLocation, '-i', videoFilePath, '-vf', 'scale=320:-1,format=gray', '-qscale:v', '2', '-filter:v', 'fps=' + str(fps), imagePath], stdin=PIPE, stdout=PIPE, stderr=PIPE)
        p = Popen([ffmpegLocation, '-i', videoFilePath, '-vf', 'scale=320:-1,format=gray,fps=' + str(fps), '-qscale:v', '2', imagePath], stdin=PIPE, stdout=PIPE, stderr=PIPE)
      else:
        p = Popen([ffmpegLocation, '-i', videoFilePath, '-qscale:v', '2', '-filter:v', 'fps=' + str(fps), imagePath], stdin=PIPE, stdout=PIPE, stderr=PIPE)
      p.wait()

    else:
      importPath = videoFilePath
      # Get video file information.
      videoFileName = os.path.basename(videoFilePath)
      videoRawName, videoExtension = os.path.splitext(videoFileName)

    # Check if selection is a camera and create one if not.
    cameraList = cmds.listCameras()
    if selection == "" or selection not in cameraList:
      selection = cmds.camera()
      selection = cmds.rename(selection[0], "reference")

    

    # Create an image plane on the selected camera.
    if self.freePlane:
      # Create free image plane.
      imagePlane = cmds.imagePlane(n = videoRawName + "_ImagePlane", showInAllViews = False, maintainRatio = True, lookThrough = selection)[1]

      # Match plane location to camera.
      cmds.matchTransform(imagePlane, selection, pos = True, rot = True)

      # Create group.
      if cmds.objExists(selection + "_ImagePlaneGroup"):
        imagePlaneGroup = selection + "_ImagePlaneGroup"
      else:
        imagePlaneGroup = cmds.group(n = selection + "_ImagePlaneGroup", empty = True)

      # Match group location to camera.
      cmds.matchTransform(imagePlaneGroup, selection, pos = True, rot = True)

      # Put image plane into group.
      cmds.parent(imagePlane, imagePlaneGroup)

      # Move image plane in front of camera.
      cmds.move(0, 0, -150, imagePlaneGroup, relative = True, objectSpace = True, worldSpaceDistance = True)

      # Parent constrain group to camera.
      # cmds.parent(imagePlaneGroup, selection)
      cmds.parentConstraint(selection, imagePlaneGroup, mo = True, weight = 1.0)
    else:
      # Create fixed image plane.
      imagePlane = cmds.imagePlane(n = videoRawName + "_ImagePlane", camera = selection, showInAllViews = False, maintainRatio = True, lookThrough = selection)[1]

    # Set image plane fit to "best".
    cmds.setAttr(imagePlane + ".fit", 1)

    # Choose to use an image sequence as the file type.
    cmds.setAttr(imagePlane + ".type", 0)
    
    # Set image plane to use a sequence.
    cmds.setAttr(imagePlane + ".useFrameExtension", True)

    # Set image plane to use selected file.
    cmds.setAttr(imagePlane + ".imageName", importPath, type = "string")
    
    size = cmds.imagePlane(imagePlane, q = True, imageSize = True)
    cmds.imagePlane(imagePlane, e = True, width = size[0] / 10)
    cmds.imagePlane(imagePlane, e = True, height = size[1] / 10)



  def ChangeFastFormatSelection(self, state):
    self.fastFormat = state

  def ChangeFreePlaneSelection(self, state):
    self.freePlane = state
  



  def CreateUI(self):
    # Set UI window name.
    name = "ImagePlaneLoader"

    # Close any previously created UI.
    self.CloseOldWindow(name)

    # Default window dimensions.
    width = 160
    height = 10

    # Create a new window.
    global window
    window = self.CreateWindow(name, "Image Plane Loader v" + str(version), width, height, -1)

    cmds.columnLayout(adjustableColumn = True, rowSpacing = -1)

    # Create text.
    messageText = cmds.text(label = "Select a camera and choose\nyour options. A reference camera\nwill be created if no camera is selected", height = 50, font = "boldLabelFont")
    # Get text's Qt Widget.
    messageTextWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( messageText ) ), QtWidgets.QWidget )
    # Edit text's StyleSheet.
    messageTextWidget.setStyleSheet( 'QWidget {color: white; margin-bottom: 0px; margin-left: 10px; margin-right: 10px} ')

    cmds.separator(height = 6, style = "none")

    cmds.rowColumnLayout(adjustableColumn = True, numberOfColumns = 3, columnWidth = [(1, 30), (2, 175), (3, 1)])
    gao = cmds.text(label = "", font = "boldLabelFont")
    fastCheckbox = cmds.checkBox("Fast Mode (video only)", height = 20, annotation = "Fast", changeCommand = lambda a: self.ChangeFastFormatSelection(a))
    fastCheckboxWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( fastCheckbox ) ), QtWidgets.QWidget )
    # Edit button's StyleSheet.
    fastCheckboxWidget.setStyleSheet( 'QWidget {color: rgb(0, 186, 186); font-size: 10px; font-style: italic; font-weight: bold; margin-top: -2px; margin-left: 12px} ')
    gao = cmds.text(label = "", font = "boldLabelFont")
    gao = cmds.text(label = "", font = "boldLabelFont")
    freeCheckbox = cmds.checkBox("Free Image Plane", height = 20, annotation = "Fast", changeCommand = lambda a: self.ChangeFreePlaneSelection(a))
    freeCheckboxWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( freeCheckbox ) ), QtWidgets.QWidget )
    # Edit button's StyleSheet.
    freeCheckboxWidget.setStyleSheet( 'QWidget {color: rgb(0, 186, 186); font-size: 10px; font-style: italic; font-weight: bold; margin-top: -2px; margin-left: 12px} ')
    cmds.setParent("..")
    cmds.separator(height = 6, style = "none")
    
    # Create Animate button.
    animateButton = cmds.button(label = "CHOOSE VIDEO OR\nIMAGE SEQUENCE", height = 40, width = 20, command = lambda a: self.ImagePlaneLoader(a), backgroundColor = [0.173, 0.235, 0.451], highlightColor = [1, 1, 1], annotation = "Save your notes")
    # Get button's Qt Widget.
    animateButtonWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( animateButton ) ), QtWidgets.QWidget )
    # Edit button's StyleSheet.
    animateButtonWidget.setStyleSheet( 'QWidget {color: white; font-size: 12px; font-weight: bold; margin-top: -0px;} ')

    # Initialise the window.
    cmds.window(window, edit = True, width = width, height = height)
    cmds.showWindow(window)



  def CreateMessageWindow(self, state, noteText):
    # Set UI window name.
    name = "MessageWindow"

    # Close any previously created UI.
    self.CloseOldWindow(name)

    # Default window dimensions.
    width = 200
    height = 40

    # Create a new window.
    global window
    window = self.CreateWindow(name, noteText, width, height, -1)

    cmds.columnLayout(adjustableColumn = True, rowSpacing = -1)

    textHeight = (noteText.count("\n") + 1) * 20

    # Create text.
    messageText = cmds.text(label = "Temp text", height = textHeight, font = "boldLabelFont")
    # Get text's Qt Widget.
    messageTextWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( messageText ) ), QtWidgets.QWidget )
    # Edit text's StyleSheet.
    messageTextWidget.setStyleSheet( 'QWidget {color: white;} ')

    cmds.separator(height = 5, style= "none")

    cmds.rowColumnLayout(adjustableColumn = True, numberOfColumns = 1, columnWidth=[(1, 95)] )

    # Create OK button.
    noButton = cmds.button(label = "OK", height = 20, width = 20, command = lambda a, v = name: self.CloseWindow(a, v), backgroundColor = [0.000, 0.486, 0.482], highlightColor = [1, 1, 1], annotation = "Playblast your selected cameras!")
    # Get button's Qt Widget.
    noButtonWidget = shiboken6.wrapInstance( int( omUI.MQtUtil.findControl( noButton ) ), QtWidgets.QWidget )
    # Edit button's StyleSheet.
    noButtonWidget.setStyleSheet( 'QWidget {color: white; font-size: 15px; font-weight: bold; margin-top: -0px} ')

    # Initialise the window.
    cmds.window(window, edit = True, width = width, height = height)
    cmds.showWindow(window)



  def CloseWindow(self, state, name):
    # Close a specific window.
    self.CloseOldWindow(name)



  def CloseOldWindow(self, name):
    # Close any previously created UI, returning its position.
    pos = [0, 0]
    if (cmds.window(name, exists = True)):
      pos = cmds.window(name, query = True, tlc = True )
      cmds.deleteUI(name)
    return pos



  def CreateWindow(self, name, title, width, height, rowSpacing):
    # Create a new window.
    window = cmds.window(name, title = title, iconName = "SAF", widthHeight = [width, height], backgroundColor = [0.525, 0.376, 0.561], maximizeButton = False, minimizeButton = False, sizeable = False)
    cmds.columnLayout(adjustableColumn = True, rowSpacing = rowSpacing)
    #cmds.windowPref(window, w = width, h = height, tlc = pos)
    return window