import cv2
import re
import os
import json
import shutil
import sys
import subprocess
from matplotlib.figure import Figure
import math
import scipy.io as sio
from pathlib import Path
from zebrazoom.code.vars import getGlobalVariables
globalVariables = getGlobalVariables()

from zebrazoom.mainZZ import mainZZ
from zebrazoom.getTailExtremityFirstFrame import getTailExtremityFirstFrame
import zebrazoom.code.popUpAlgoFollow as popUpAlgoFollow
import zebrazoom.code.util as util

try:
  from PyQt6.QtCore import Qt
  from PyQt6.QtWidgets import QCheckBox, QFileDialog, QVBoxLayout
except ImportError:
  from PyQt5.QtCore import Qt
  from PyQt5.QtWidgets import QCheckBox, QFileDialog, QVBoxLayout


LARGE_FONT= ("Verdana", 12)

def chooseVideoToAnalyze(self, justExtractParams, noValidationVideo, chooseFrames, testMode):
    self.videoName, _ = QFileDialog.getOpenFileName(self.window, 'Select file', os.path.expanduser("~"))
    if not self.videoName:
      return
    self.folderName = ''
    self.headEmbedded = 0
    self.sbatchMode = 0

    self.justExtractParams = int(justExtractParams)
    self.noValidationVideo = int(noValidationVideo)
    self.testMode         = testMode
    self.findMultipleROIs = 0

    if chooseFrames:
      def beginningAndEndChosen():
        if "firstFrame" in self.configFile:
          self.firstFrame = self.configFile["firstFrame"]
        if "lastFrame" in self.configFile:
          self.lastFrame = self.configFile["lastFrame"]
        self.configFile.clear()
        self.backgroundExtractionForceUseAllVideoFrames = int(backgroundExtractionForceUseAllVideoFramesCheckbox.isChecked())
        self.show_frame("ConfigFilePromp")
      layout = QVBoxLayout()
      backgroundExtractionForceUseAllVideoFramesCheckbox = QCheckBox("Use all frames to calculate background")
      backgroundExtractionForceUseAllVideoFramesCheckbox.setChecked(True)
      layout.addWidget(backgroundExtractionForceUseAllVideoFramesCheckbox, alignment=Qt.AlignmentFlag.AlignCenter)
      util.chooseBeginningPage(self, self.videoName, "Choose where the analysis of your video should start.", "Ok, I want the tracking to start at this frame!",
                               lambda: util.chooseEndPage(self, self.videoName, "Choose where the analysis of your video should end.", "Ok, I want the tracking to end at this frame!", beginningAndEndChosen),
                               additionalLayout=layout)
    else:
      self.show_frame("ConfigFilePromp")

def chooseFolderToAnalyze(self, justExtractParams, noValidationVideo, sbatchMode):
    self.folderName =  QFileDialog.getExistingDirectory(self.window, 'Select folder', os.path.expanduser("~"))
    if not self.folderName:
      return
    self.headEmbedded = 0
    self.justExtractParams = int(justExtractParams)
    self.noValidationVideo = int(noValidationVideo)
    self.sbatchMode        = int(sbatchMode)
    self.testMode = False
    self.findMultipleROIs = 0
    self.show_frame("ConfigFilePromp")

def chooseFolderForTailExtremityHE(self):
    self.folderName =  QFileDialog.getExistingDirectory(self.window, 'Select folder', os.path.expanduser("~"))
    if not self.folderName:
      return
    self.sbatchMode = 0
    self.headEmbedded = 1
    self.justExtractParams = 0
    self.noValidationVideo = 0
    self.testMode = False
    self.findMultipleROIs = 0
    self.show_frame("ConfigFilePromp")

def chooseFolderForMultipleROIs(self):
    self.folderName =  QFileDialog.getExistingDirectory(self.window, 'Select folder', os.path.expanduser("~"))
    if not self.folderName:
      return
    self.sbatchMode = 0
    self.headEmbedded = 0
    self.justExtractParams = 0
    self.noValidationVideo = 0
    self.testMode = False
    self.findMultipleROIs = 1
    self.show_frame("ConfigFilePromp")

def chooseConfigFile(self):

  cur_dir_path = os.path.dirname(os.path.realpath(__file__))
  path = Path(cur_dir_path)
  path = path.parent.parent
  path = os.path.join(path, 'configuration')

  self.configFileName, _ = QFileDialog.getOpenFileName(self.window, 'Select file', path, "JSON (*.json)")
  if not self.configFileName:
    return
  if len(self.folderName) or globalVariables["mac"] or globalVariables["lin"]:
    self.show_frame("Patience")
  else:
    self.launchZebraZoom()

def findAllFilesRecursivelyInDirectories(folderName):

  allVideos = []

  for name in os.listdir(folderName):
    if os.path.isdir(os.path.join(folderName, name)):
      allVideos2 = findAllFilesRecursivelyInDirectories(os.path.join(folderName, name))
      allVideos = allVideos + allVideos2
    else:
      if len(name) > 3:
        ext = os.path.splitext(name)[1]
        if len(ext):
          if (ext in ['.264', '.3g2', '.3gp', '.3gp2', '.3gpp', '.3gpp2', '.3mm', '.3p2', '.60d', '.787', '.89', '.aaf', '.aec', '.aep', '.aepx', '.aet', '.aetx', '.ajp', '.ale', '.am', '.amc', '.amv', '.amx', '.anim', '.aqt', '.arcut', '.arf', '.asf', '.asx', '.avb', '.avc', '.avd', '.avi', '.avp', '.avs', '.avs', '.avv', '.axm', '.bdm', '.bdmv', '.bdt2', '.bdt3', '.bik', '.bix', '.bmk', '.bnp', '.box', '.bs4', '.bsf', '.bvr', '.byu', '.camproj', '.camrec', '.camv', '.ced', '.cel', '.cine', '.cip', '.clpi', '.cmmp', '.cmmtpl', '.cmproj', '.cmrec', '.cpi', '.cst', '.cvc', '.cx3', '.d2v', '.d3v', '.dat', '.dav', '.dce', '.dck', '.dcr', '.dcr', '.ddat', '.dif', '.dir', '.divx', '.dlx', '.dmb', '.dmsd', '.dmsd3d', '.dmsm', '.dmsm3d', '.dmss', '.dmx', '.dnc', '.dpa', '.dpg', '.dream', '.dsy', '.dv', '.dv-avi', '.dv4', '.dvdmedia', '.dvr', '.dvr-ms', '.dvx', '.dxr', '.dzm', '.dzp', '.dzt', '.edl', '.evo', '.eye', '.ezt', '.f4p', '.f4v', '.fbr', '.fbr', '.fbz', '.fcp', '.fcproject', '.ffd', '.flc', '.flh', '.fli', '.flv', '.flx', '.gfp', '.gl', '.gom', '.grasp', '.gts', '.gvi', '.gvp', '.h264', '.hdmov', '.hkm', '.ifo', '.imovieproj', '.imovieproject', '.ircp', '.irf', '.ism', '.ismc', '.ismv', '.iva', '.ivf', '.ivr', '.ivs', '.izz', '.izzy', '.jss', '.jts', '.jtv', '.k3g', '.kmv', '.ktn', '.lrec', '.lsf', '.lsx', '.m15', '.m1pg', '.m1v', '.m21', '.m21', '.m2a', '.m2p', '.m2t', '.m2ts', '.m2v', '.m4e', '.m4u', '.m4v', '.m75', '.mani', '.meta', '.mgv', '.mj2', '.mjp', '.mjpg', '.mk3d', '.mkv', '.mmv', '.mnv', '.mob', '.mod', '.modd', '.moff', '.moi', '.moov', '.mov', '.movie', '.mp21', '.mp21', '.mp2v', '.mp4', '.mp4v', '.mpe', '.mpeg', '.mpeg1', '.mpeg4', '.mpf', '.mpg', '.mpg2', '.mpgindex', '.mpl', '.mpl', '.mpls', '.mpsub', '.mpv', '.mpv2', '.mqv', '.msdvd', '.mse', '.msh', '.mswmm', '.mts', '.mtv', '.mvb', '.mvc', '.mvd', '.mve', '.mvex', '.mvp', '.mvp', '.mvy', '.mxf', '.mxv', '.mys', '.ncor', '.nsv', '.nut', '.nuv', '.nvc', '.ogm', '.ogv', '.ogx', '.osp', '.otrkey', '.pac', '.par', '.pds', '.pgi', '.photoshow', '.piv', '.pjs', '.playlist', '.plproj', '.pmf', '.pmv', '.pns', '.ppj', '.prel', '.pro', '.prproj', '.prtl', '.psb', '.psh', '.pssd', '.pva', '.pvr', '.pxv', '.qt', '.qtch', '.qtindex', '.qtl', '.qtm', '.qtz', '.r3d', '.rcd', '.rcproject', '.rdb', '.rec', '.rm', '.rmd', '.rmd', '.rmp', '.rms', '.rmv', '.rmvb', '.roq', '.rp', '.rsx', '.rts', '.rts', '.rum', '.rv', '.rvid', '.rvl', '.sbk', '.sbt', '.scc', '.scm', '.scm', '.scn', '.screenflow', '.sec', '.sedprj', '.seq', '.sfd', '.sfvidcap', '.siv', '.smi', '.smi', '.smil', '.smk', '.sml', '.smv', '.spl', '.sqz', '.srt', '.ssf', '.ssm', '.stl', '.str', '.stx', '.svi', '.swf', '.swi', '.swt', '.tda3mt', '.tdx', '.thp', '.tivo', '.tix', '.tod', '.tp', '.tp0', '.tpd', '.tpr', '.trp', '.ts', '.tsp', '.ttxt', '.tvs', '.usf', '.usm', '.vc1', '.vcpf', '.vcr', '.vcv', '.vdo', '.vdr', '.vdx', '.veg','.vem', '.vep', '.vf', '.vft', '.vfw', '.vfz', '.vgz', '.vid', '.video', '.viewlet', '.viv', '.vivo', '.vlab', '.vob', '.vp3', '.vp6', '.vp7', '.vpj', '.vro', '.vs4', '.vse', '.vsp', '.w32', '.wcp', '.webm', '.wlmp', '.wm', '.wmd', '.wmmp', '.wmv', '.wmx', '.wot', '.wp3', '.wpl', '.wtv', '.wve', '.wvx', '.xej', '.xel', '.xesc', '.xfl', '.xlmv', '.xmv', '.xvid', '.y4m', '.yog', '.yuv', '.zeg', '.zm1', '.zm2', '.zm3', '.zmv']):
            allVideos.append(os.path.join(folderName, name))

  return allVideos


def launchZebraZoom(self):
  if self.testMode:
    with open(self.configFileName) as f:
      self.configFile = json.load(f)
    self.videoToCreateConfigFileFor = self.videoName
    self.testConfig(addToHistory=False)
    self.headEmbedded      = 0
    self.justExtractParams = 0
    self.noValidationVideo = 0
    self.testMode         = False
    self.findMultipleROIs  = 0
    self.configFileName = None
    self.videoName = None
    return

  last = 0
  allVideos = []

  if self.sbatchMode:
    commandsFile = open("commands.txt", "w", newline='\n')
    nbVideosToLaunch = 0

  if len(self.folderName):

    allVideos = findAllFilesRecursivelyInDirectories(self.folderName)

  else:
    allVideos = [self.videoName]

  print("allVideos:", allVideos)

  for idx, text in enumerate(allVideos):

    path        = os.path.split(text)[0]
    nameWithExt = os.path.split(text)[1]
    name        = os.path.splitext(nameWithExt)[0]
    videoExt    = os.path.splitext(nameWithExt)[1][1:]

    if self.headEmbedded == 0:
      if len(allVideos) == 1:
        tabParams = ["mainZZ", path, name, videoExt, self.configFileName, "freqAlgoPosFollow", 100, "popUpAlgoFollow", 1, "outputFolder", self.ZZoutputLocation]
      else:
        tabParams = ["mainZZ", path, name, videoExt, self.configFileName, "freqAlgoPosFollow", 100, "outputFolder", self.ZZoutputLocation]
      if hasattr(self, "backgroundExtractionForceUseAllVideoFrames"):
        tabParams.extend(["backgroundExtractionForceUseAllVideoFrames", self.backgroundExtractionForceUseAllVideoFrames])
      if hasattr(self, "firstFrame"):
        tabParams.extend(["firstFrame", self.firstFrame])
      if hasattr(self, "lastFrame"):
        tabParams.extend(["lastFrame", self.lastFrame])
      if self.justExtractParams == 1:
        tabParams = tabParams + ["reloadWellPositions", 1, "reloadBackground", 1, "debugPauseBetweenTrackAndParamExtract", "justExtractParamFromPreviousTrackData"]
      if self.noValidationVideo == 1:
          tabParams = tabParams + ["createValidationVideo", 0]
      if self.findMultipleROIs == 1:
        tabParams = tabParams + ["exitAfterWellsDetection", 1, "saveWellPositionsToBeReloadedNoMatterWhat", 1]
      try:
        if self.sbatchMode:
          commandsFile.write('python -m zebrazoom ' + ' '.join(tabParams[1:4]).replace('\\', '/').replace('//lexport/iss01.', '/network/lustre/iss01/') + ' configFile.json\n')
          nbVideosToLaunch = nbVideosToLaunch + 1
        else:
          mainZZ(path, name, videoExt, self.configFileName, tabParams)
      except ValueError:
        print("moving on to the next video for ROIs identification")
      except NameError:
        self.show_frame("Error")
        return
    else:
      tabParams = ["outputFolder", self.ZZoutputLocation]
      if hasattr(self, "backgroundExtractionForceUseAllVideoFrames"):
        tabParams.extend(["backgroundExtractionForceUseAllVideoFrames", self.backgroundExtractionForceUseAllVideoFrames])
      if hasattr(self, "firstFrame"):
        tabParams.extend(["firstFrame", self.firstFrame])
      if hasattr(self, "lastFrame"):
        tabParams.extend(["lastFrame", self.lastFrame])
      getTailExtremityFirstFrame(path, name, videoExt, self.configFileName, tabParams)

  self.headEmbedded      = 0
  self.justExtractParams = 0
  self.noValidationVideo = 0
  self.testMode         = False
  self.findMultipleROIs  = 0
  if hasattr(self, "backgroundExtractionForceUseAllVideoFrames"):
    del self.backgroundExtractionForceUseAllVideoFrames
  if hasattr(self, "firstFrame"):
    del self.firstFrame
  if hasattr(self, "lastFrame"):
    del self.lastFrame

  if self.sbatchMode:

    commandsFile.close()

    with open(self.configFileName) as f:
      jsonFile = json.load(f)
    nbWells = jsonFile["nbWells"]
    if nbWells > 24:
      nbWells = 24

    launchFile = open("launchZZ.sh", "w", newline='\n')
    linesToWrite = ['#!/bin/sh',
                    '#SBATCH --ntasks=1',
                    '#SBATCH --cpus-per-task='+str(nbWells),
                    '#SBATCH --array=1-'+str(nbVideosToLaunch),
                    '#SBATCH --mem=16G',
                    '#SBATCH --time=23:00:00',
                    '#SBATCH --partition=normal',
                    '#SBATCH --job-name="ZebraZoom-protocole"',
                    '',
                    'module load python/3.8',
                    'source activate zebrazoom',
                    '',
                    'date',
                    '',
                    'export CMD_FILE_PATH=./commands.txt',
                    '',
                    'export CMD=$(sed -n ${SLURM_ARRAY_TASK_ID}p ${CMD_FILE_PATH})',
                    '',
                    'echo $CMD',
                    'eval $CMD',
                    '',
                    'date']
    linesToWrite = [line + '\n' for line in linesToWrite]
    launchFile.writelines(linesToWrite)
    launchFile.close()

    shutil.copy(self.configFileName, 'configFile.json')

    self.show_frame("ZZoutroSbatch")

  else:

    self.show_frame("ZZoutro")


def openConfigurationFileFolder(self, homeDirectory):
  dir_path = os.path.join(homeDirectory,'configuration')
  if sys.platform == "win32":
    os.startfile(dir_path)
  else:
    opener ="open" if sys.platform == "darwin" else "xdg-open"
    subprocess.call([opener, dir_path])


def openZZOutputFolder(self, homeDirectory):
  if len(self.ZZoutputLocation):
    dir_path = self.ZZoutputLocation
  else:
    dir_path = os.path.join(homeDirectory,'ZZoutput')
  if sys.platform == "win32":
    os.startfile(dir_path)
  else:
    opener ="open" if sys.platform == "darwin" else "xdg-open"
    subprocess.call([opener, dir_path])
