from pathlib import Path
import numpy as np
import json
import cv2
import zebrazoom.videoFormatConversion.zzVideoReading as zzVideoReading
from zebrazoom.code.GUI.getCoordinates import findWellLeft, findWellRight, findHeadCenter, findBodyExtremity
from zebrazoom.code.GUI.automaticallyFindOptimalParameters import automaticallyFindOptimalParameters
import math
from zebrazoom.code.findWells import findWells
from zebrazoom.code.getHyperparameters import getHyperparametersSimple
from zebrazoom.code.getBackground import getBackground
from zebrazoom.code.getImage.getForegroundImage import getForegroundImage
import cvui
import pickle
from zebrazoom.mainZZ import mainZZ
import json
import os

try:
  from PyQt6.QtCore import Qt, QTimer
  from PyQt6.QtWidgets import QApplication, QFileDialog, QGridLayout, QLabel, QVBoxLayout
except ImportError:
  from PyQt5.QtCore import Qt, QTimer
  from PyQt5.QtWidgets import QApplication, QFileDialog, QGridLayout, QLabel, QVBoxLayout

import zebrazoom.code.util as util


def getMainArguments(self):
  s            = self.videoToCreateConfigFileFor
  arr          = s.split("/")
  nameWithExt  = arr.pop()
  pathToVideo  = '/'.join(arr) + '/'
  nameWithExtL = nameWithExt.split(".")
  videoExt     = nameWithExtL.pop()
  videoName    = '.'.join(nameWithExtL)
  configFile   = self.configFile
  argv         = []
  return [pathToVideo, videoName, videoExt, configFile, argv]

def chooseVideoToCreateConfigFileFor(self, controller, reloadConfigFile, freelySwimAutomaticParameters=False, boutDetectionsOnly=False):

  if int(reloadConfigFile):

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

    configFileName, _ =  QFileDialog.getOpenFileName(self.window, "Select configuration file", pathconf, "All files(*)")
    with open(configFileName) as f:
      self.configFile = json.load(f)

  self.videoToCreateConfigFileFor, _ = QFileDialog.getOpenFileName(self.window, "Select video to create config file for", os.path.expanduser("~"), "All files(*)")

  if boutDetectionsOnly:
    controller.calculateBackgroundFreelySwim(controller, 0, False, False, True)
  else:
    if freelySwimAutomaticParameters:
      controller.calculateBackgroundFreelySwim(controller, 0, False, True)
    else:
      controller.show_frame("ChooseGeneralExperiment")

def chooseGeneralExperimentFirstStep(self, controller, freeZebra, headEmbZebra, drosophilia, rodent, other, fastScreen):
  if int(fastScreen):
    controller.show_frame("HomegeneousWellsLayout")
  else:
    self.configFile["extractAdvanceZebraParameters"] = 0
    if int(freeZebra):
      controller.show_frame("FreelySwimmingExperiment")
    else:
      chooseGeneralExperiment(self, controller, freeZebra, headEmbZebra, drosophilia, rodent, other, 0)

def chooseGeneralExperiment(self, controller, freeZebra, headEmbZebra, drosophilia, rodent, other, freeZebra2):
  self.configFile["extractAdvanceZebraParameters"] = 0
  if int(freeZebra):
    self.organism = 'zebrafish'
    self.configFile["headEmbeded"] = 0
    controller.show_frame("WellOrganisation")
  elif int(freeZebra2):
    self.organism = 'zebrafishNew'
    self.configFile["headEmbeded"] = 0
    controller.show_frame("WellOrganisation")
  elif int(headEmbZebra):
    self.organism = 'headembeddedzebrafish'
    self.configFile["headEmbeded"] = 1
    chooseBeginningAndEndOfVideo(self, controller)
  else:
    self.organism = 'drosoorrodent'
    self.configFile["headEmbeded"] = 0
    self.configFile["freeSwimmingTailTrackingMethod"] = "none"
    controller.show_frame("WellOrganisation")

def wellOrganisation(self, controller, circular, rectangular, roi, other, multipleROIs, groupSameSizeAndShapeEquallySpacedWells):
  if multipleROIs:
    controller.show_frame("NbRegionsOfInterest")
  else:
    if groupSameSizeAndShapeEquallySpacedWells:
      self.shape = 'groupSameSizeAndShapeEquallySpacedWells'
      self.configFile["groupOfMultipleSameSizeAndShapeEquallySpacedWells"] = 1
      controller.show_frame("CircularOrRectangularWells")
    else:
      if rectangular:
        self.shape = 'rectangular'
        self.configFile["wellsAreRectangles"] = 1
        controller.show_frame("CircularOrRectangularWells")
      else:
        if circular and self.organism != 'drosoorrodent': # should remove the self.organism != 'drosoorrodent' at some point
          self.shape = 'circular'
          controller.show_frame("CircularOrRectangularWells")
        else:
          self.shape = 'other'
          if roi:
            cap = zzVideoReading.VideoCapture(controller.videoToCreateConfigFileFor)
            cap.set(1, 10)
            ret, frame = cap.read()
            rect = util.getRectangle(frame, "Click on the top left and bottom right of the region of interest")
            self.configFile["oneWellManuallyChosenTopLeft"], self.configFile["oneWellManuallyChosenBottomRight"] = rect
            self.configFile["nbWells"] = 1
            chooseBeginningAndEndOfVideo(self, controller)
          else:
            self.configFile["noWellDetection"] = 1
            self.configFile["nbWells"] = 1
            chooseBeginningAndEndOfVideo(self, controller)


def regionsOfInterest(self, controller, nbwells):

  self.configFile["multipleROIsDefinedDuringExecution"] = 1
  self.configFile["nbWells"] = int(nbwells)

  chooseBeginningAndEndOfVideo(self, controller)


def rectangularWells(self, controller, nbwells, nbRowsOfWells, nbWellsPerRows):

  [pathToVideo, videoName, videoExt, configFile, argv] = getMainArguments(self)

  configFile["adjustRectangularWellsDetect"] = 1

  app = QApplication.instance()
  with app.busyCursor():
    try:
      mainZZ(pathToVideo, videoName, videoExt, configFile, argv)
    except ValueError:
      newhyperparameters = pickle.load(open('newhyperparameters', 'rb'))
      for index in newhyperparameters:
        configFile[index] = newhyperparameters[index]
    except NameError:
      print("Configuration file parameters changes discarded.")

  configFile["adjustRectangularWellsDetect"] = 0

  self.configFile = configFile

  chooseBeginningAndEndOfVideo(self, controller)


def homegeneousWellsLayout(self, controller, nbwells, nbRowsOfWells, nbWellsPerRows):

  self.configFile = {"nbAnimalsPerWell": 1, "nbWells": 8, "nbRowsOfWells": 2, "nbWellsPerRows": 4, "groupOfMultipleSameSizeAndShapeEquallySpacedWells": 1, "postProcessMultipleTrajectories": 1,   "postProcessRemoveLowProbabilityDetection" : 1, "postProcessLowProbabilityDetectionPercentOfMaximum" : 0.2, "trackingPointSizeDisplay": 4, "extractAdvanceZebraParameters": 0,  "validationVideoPlotHeading": 0, "trackTail": 0, "freqAlgoPosFollow": 100, "fasterMultiprocessing": 1, "copyOriginalVideoToOutputFolderForValidation": 0, "backgroundSubtractorKNN": 1}

  self.configFile["nbWells"]          = int(nbwells)

  if len(nbRowsOfWells):
    self.configFile["nbRowsOfWells"]  = int(nbRowsOfWells)
  else:
    self.configFile["nbRowsOfWells"]  = 1

  if len(nbWellsPerRows):
    self.configFile["nbWellsPerRows"]  = int(nbWellsPerRows)
  else:
    self.configFile["nbWellsPerRows"]  = 4

  controller.show_frame("FinishConfig")


def morePreciseFastScreen(self, controller, nbwells, nbRowsOfWells, nbWellsPerRows):

  # The gaussian image filtering should be added here in the future
  self.configFile = {"minPixelDiffForBackExtract": 20, "backgroundPreProcessParameters": [[3]], "backgroundPreProcessMethod": ["erodeThenMin"], "trackingPointSizeDisplay": 1, "nbAnimalsPerWell": 1, "extractAdvanceZebraParameters": 0, "trackTail": 0, "nbWells": 1, "noWellDetection": 1, "backgroundExtractionForceUseAllVideoFrames": 1, "headSize": 2, "createValidationVideo": 0, "lastFrame": 1000}

  self.configFile["nbAnimalsPerWell"] = int(nbwells)

  # self.configFile["nbWells"]          = int(nbwells)

  if len(nbRowsOfWells):
    self.configFile["nbRowsOfWells"]  = int(nbRowsOfWells)
  else:
    self.configFile["nbRowsOfWells"]  = 1

  if len(nbWellsPerRows):
    self.configFile["nbWellsPerRows"]  = int(nbWellsPerRows)
  else:
    self.configFile["nbWellsPerRows"]  = 4

  self.calculateBackgroundFreelySwim(controller, 0, True)

  # controller.show_frame("FinishConfig")


def circularOrRectangularWells(self, controller, nbwells, nbRowsOfWells, nbWellsPerRows):
  self.configFile["nbWells"]        = int(nbwells)

  if len(nbRowsOfWells):
    self.configFile["nbRowsOfWells"]  = int(nbRowsOfWells)
  else:
    self.configFile["nbRowsOfWells"]  = 1

  if len(nbWellsPerRows):
    self.configFile["nbWellsPerRows"]  = int(nbWellsPerRows)
  else:
    self.configFile["nbWellsPerRows"]  = 4

  if self.shape == 'circular':
    controller.show_frame("ChooseCircularWellsLeft")
  else:
    if self.shape == 'groupSameSizeAndShapeEquallySpacedWells':
      chooseBeginningAndEndOfVideo(self, controller)
    else:
      rectangularWells(self, controller, nbwells, nbRowsOfWells, nbWellsPerRows)


def chooseCircularWellsLeft(self, controller):
  cap = zzVideoReading.VideoCapture(self.videoToCreateConfigFileFor)
  ret, frame = cap.read()
  [x, y] = findWellLeft(frame)
  self.wellLeftBorderX = x
  self.wellLeftBorderY = y
  controller.show_frame("ChooseCircularWellsRight")

def chooseCircularWellsRight(self, controller):
  cap = zzVideoReading.VideoCapture(self.videoToCreateConfigFileFor)
  ret, frame = cap.read()
  [xRight, yRight] = findWellRight(frame)
  xLeft = self.wellLeftBorderX
  yLeft = self.wellLeftBorderY
  dist = math.sqrt((xLeft - xRight)**2 + (yLeft - yRight)**2)
  self.configFile["minWellDistanceForWellDetection"] = int(dist)
  self.configFile["wellOutputVideoDiameter"]         = int(dist + dist * 0.2)
  chooseBeginningAndEndOfVideo(self, controller)

def chooseBeginningAndEndOfVideo(self, controller):
  if util.chooseBeginning(controller, controller.videoToCreateConfigFileFor, "Choose where the analysis of your video should start.", "Ok, I want the tracking to start at this frame!", allowWholeVideo=True):
    util.chooseEnd(controller, controller.videoToCreateConfigFileFor, "Choose where the analysis of your video should end.", "Ok, I want the tracking to end at this frame!")

  if int(self.configFile["headEmbeded"]) == 1:
    controller.show_frame("HeadEmbeded")
  else:
    if self.organism == 'zebrafishNew':
      controller.show_frame("NumberOfAnimals2")
    elif self.organism == 'zebrafish':
      controller.show_frame("NumberOfAnimals")
    else:
      controller.show_frame("NumberOfAnimalsCenterOfMass")

def getImageForMultipleAnimalGUI(l, nx, ny, max_l, videoToCreateConfigFileFor, background, wellPositions, hyperparameters):

  [frame, a1, a2] = getForegroundImage(videoToCreateConfigFileFor, background, l, 0, [], hyperparameters)

  lengthX = nx * 2
  lengthY = ny

  frame2 = frame
  ret,thresh2 = cv2.threshold(frame2,hyperparameters["thresholdForBlobImg"],255,cv2.THRESH_BINARY)
  kernel  = np.ones((hyperparameters["erodeSize"],hyperparameters["erodeSize"]), np.uint8)
  thresh2 = cv2.dilate(thresh2, kernel, iterations=hyperparameters["dilateIter"])

  thresh = thresh2
  thresh2 = cv2.cvtColor(thresh2, cv2.COLOR_GRAY2RGB)
  frame   = cv2.cvtColor(frame,   cv2.COLOR_GRAY2RGB)
  thresh[0, :]                = 255
  thresh[:, 0]                = 255
  thresh[len(thresh)-1, :]    = 255
  thresh[:, len(thresh[0])-1] = 255
  areaList = []
  contours, hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
  for contour in contours:
    area = cv2.contourArea(contour)
    if area < (len(thresh) * len(thresh[0]))/2:
      areaList.append(area)
    if (area > hyperparameters["minArea"]) and (area < hyperparameters["maxArea"]):
      M = cv2.moments(contour)
      if M['m00']:
        x = int(M['m10']/M['m00'])
        y = int(M['m01']/M['m00'])
        cv2.circle(thresh2, (x,y), 3, (0,0,255), -1)
      else:
        x = 0
        y = 0

  frame = cv2.line(frame, (len(frame[0])-5, 0), (len(frame[0])-5, len(frame)), (255, 0, 0), 5)

  frame = np.concatenate((frame, thresh2), axis=1)

  if len(areaList):
    maxToReturn = int((max(areaList)+2)*2)
  else:
    maxToReturn = (len(thresh) * len(thresh[0]))/5

  return [frame, maxToReturn]

def _createWidget(layout, values, key, minn, maxx, name, updateFrame):
  slider = util.SliderWithSpinbox(values[key], minn, maxx, name=name)

  def valueChanged():
    values[key] = slider.value()
    updateFrame()
  slider.valueChanged.connect(valueChanged)

  if name != "Frame number":
    elements = layout.count() - 2  # frame, frameSlider
    row = elements // 2 + 2
    col = elements % 2
    layout.addWidget(slider, row, col, Qt.AlignmentFlag.AlignLeft if col else Qt.AlignmentFlag.AlignRight)
  else:
    layout.addWidget(slider, 1, 0, 1, 2, Qt.AlignmentFlag.AlignCenter)

  return slider

def identifyMultipleHead(self, controller, nbanimals):
  self.configFile["videoName"] = "configFilePrep"

  tempConfig = self.configFile

  app = QApplication.instance()
  with app.busyCursor():
    # Getting hyperparameters, wellPositions, and background
    hyperparameters = getHyperparametersSimple(tempConfig)
    wellPositions = findWells(self.videoToCreateConfigFileFor, hyperparameters)
    background    = getBackground(self.videoToCreateConfigFileFor, hyperparameters)

  cur_dir_path = os.path.dirname(os.path.realpath(__file__))
  cur_dir_path = Path(cur_dir_path)

  def imagesGenerator():
    images = (cv2.imread(os.path.join(cur_dir_path, 'no1.png')),
              cv2.imread(os.path.join(cur_dir_path, 'no2.png')),
              cv2.imread(os.path.join(cur_dir_path, 'ok1.png')))
    while True:
      yield from images
  images = imagesGenerator()

  label = QLabel()
  label.setMinimumSize(1, 1)
  layout = QVBoxLayout()
  layout.addWidget(label, alignment=Qt.AlignmentFlag.AlignCenter)
  timer = QTimer()
  timer.setInterval(2000)
  timer.timeout.connect(lambda: util.setPixmapFromCv(next(images), label))
  timer.start()
  util.pageOrDialog(layout, title="", dialog=False, buttons=(("Ok, I understand!", lambda: timer.stop()),), labelInfo=(next(images), label))

  # Manual parameters adjustements
  cap        = zzVideoReading.VideoCapture(self.videoToCreateConfigFileFor)
  nx         = int(cap.get(3))
  ny         = int(cap.get(4))
  max_l      = int(cap.get(7))

  hyperparameters["minArea"] = 5
  hyperparameters["maxArea"] = 800

  frameNum = {"frameNum": hyperparameters.get("firstFrame", 1)}
  minPixelDiffForBackExtract = [hyperparameters["minPixelDiffForBackExtract"]]
  thresholdForBlobImg        = [hyperparameters["thresholdForBlobImg"]]
  dilateIter                 = [hyperparameters["dilateIter"]]
  minArea                    = [hyperparameters["minArea"]]
  maxArea                    = [hyperparameters["maxArea"]]
  firstFrame = hyperparameters["firstFrame"] if "firstFrame" in hyperparameters else 1
  lastFrame  = hyperparameters["lastFrame"]-1 if "lastFrame" in hyperparameters else max_l - 10

  frame, maxAreaBlobs = getImageForMultipleAnimalGUI(frameNum["frameNum"], nx, ny, max_l, self.videoToCreateConfigFileFor, background, wellPositions, hyperparameters)

  label = QLabel()
  label.setMinimumSize(1, 1)
  layout = QGridLayout()
  layout.addWidget(label, 0, 0, 1, 2, Qt.AlignmentFlag.AlignCenter)
  layout.setRowStretch(0, 1)

  def updateFrame():
    frame, maxAreaBlobs = getImageForMultipleAnimalGUI(frameNum["frameNum"], nx, ny, max_l, self.videoToCreateConfigFileFor, background, wellPositions, hyperparameters)
    util.setPixmapFromCv(frame, label)
    minAreaWidget.setMaximum(maxAreaBlobs)
    maxAreaWidget.setMaximum(maxAreaBlobs)

  _createWidget(layout, frameNum, "frameNum", firstFrame, lastFrame, "Frame number", updateFrame)
  _createWidget(layout, hyperparameters, "minPixelDiffForBackExtract", 0, 255, "Threshold left image", updateFrame)
  _createWidget(layout, hyperparameters, "thresholdForBlobImg", 0, 255, "Threshold right image", updateFrame)
  _createWidget(layout, hyperparameters, "dilateIter", 0, 15, "Area dilatation", updateFrame)
  minAreaWidget = _createWidget(layout, hyperparameters, "minArea", 0, maxAreaBlobs, "Minimum area", updateFrame)
  maxAreaWidget = _createWidget(layout, hyperparameters, "maxArea", 0, maxAreaBlobs, "Maximum area", updateFrame)

  util.pageOrDialog(layout, title="Adjust Parameters: As much as possible, you must see red points on and only on animals on the right image.", dialog=False, buttons=(("Ok, done!", None),), labelInfo=(frame, label))

  del self.configFile["videoName"]

  self.configFile["minPixelDiffForBackExtract"] = int(hyperparameters["minPixelDiffForBackExtract"])
  self.configFile["thresholdForBlobImg"]        = int(hyperparameters["thresholdForBlobImg"])
  self.configFile["dilateIter"]                 = int(hyperparameters["dilateIter"])
  self.configFile["minArea"]                    = int(hyperparameters["minArea"])
  self.configFile["maxArea"]                    = int(hyperparameters["maxArea"])
  self.configFile["headSize"]        = math.sqrt((int(hyperparameters["minArea"]) + int(hyperparameters["maxArea"])) / 2)


def numberOfAnimals(self, controller, nbanimals, yes, noo, forceBlobMethodForHeadTracking, yesBouts, nooBouts, recommendedMethod, alternativeMethod, yesBends, nooBends, adjustBackgroundExtractionBasedOnNumberOfBlackPixels):

  self.configFile["noBoutsDetection"] = 1
  self.configFile["noChecksForBoutSelectionInExtractParams"] = 1
  self.configFile["trackingPointSizeDisplay"] = 4
  self.configFile["validationVideoPlotHeading"] = 0

  if int(yesBends):
    self.configFile["extractAdvanceZebraParameters"] = 1

  nbanimals = int(nbanimals)
  if nbanimals == self.configFile["nbWells"]:
    self.configFile["nbAnimalsPerWell"] = 1
  else:
    self.configFile["nbAnimalsPerWell"] = int(nbanimals / self.configFile["nbWells"])
    if noo:
      self.configFile["multipleHeadTrackingIterativelyRelaxAreaCriteria"] = 0
    else:
      self.configFile["multipleHeadTrackingIterativelyRelaxAreaCriteria"] = 1

  self.forceBlobMethodForHeadTracking = int(forceBlobMethodForHeadTracking)
  if self.forceBlobMethodForHeadTracking:
    self.configFile["forceBlobMethodForHeadTracking"] = self.forceBlobMethodForHeadTracking

  if self.organism == 'zebrafish':
    controller.show_frame("IdentifyHeadCenter")
  elif self.organism == 'zebrafishNew':
    detectBouts = 0
    if int(yesBouts):
      detectBouts = 1
    method = 0
    if int(alternativeMethod):
      method = 1
    automaticallyFindOptimalParameters(self, controller, True, detectBouts, method, int(noo), int(adjustBackgroundExtractionBasedOnNumberOfBlackPixels))
  elif self.organism == 'drosoorrodent' and int(recommendedMethod):
    automaticallyFindOptimalParameters(self, controller, True, 0, 0, int(noo), int(adjustBackgroundExtractionBasedOnNumberOfBlackPixels))
  else:
    identifyMultipleHead(self, controller, nbanimals)
    controller.show_frame("FinishConfig")

def chooseHeadCenter(self, controller):
  cap = zzVideoReading.VideoCapture(self.videoToCreateConfigFileFor)
  ret, frame = cap.read()
  [x, y] = findHeadCenter(frame)
  self.headCenterX = x
  self.headCenterY = y
  controller.show_frame("IdentifyBodyExtremity")

def chooseBodyExtremity(self, controller):
  cap = zzVideoReading.VideoCapture(self.videoToCreateConfigFileFor)
  ret, frame = cap.read()
  [extX, extY] = findBodyExtremity(frame)
  headCenterX = self.headCenterX
  headCenterY = self.headCenterY
  dist = math.sqrt((extX - headCenterX)**2 + (extY - headCenterY)**2)

  if self.organism == 'zebrafish':
    minArea = int(dist * (dist * 0.1))
    maxArea = int(dist * (dist * 0.4))
    self.configFile["minArea"]     = minArea
    self.configFile["maxArea"]     = maxArea
    self.configFile["minAreaBody"] = minArea
    self.configFile["maxAreaBody"] = maxArea
    self.configFile["headSize"]    = int(dist * 0.3)
    self.configFile["minTailSize"] = int(dist * 0.05)
    self.configFile["maxTailSize"] = int(dist * 2)
    self.configFile["paramGaussianBlur"] = int(int(dist*(31/52))/2)*2 + 1
  else:
    minArea = int(((2 * dist) * (2 * dist)) * 0.2)
    maxArea = int(((2 * dist) * (2 * dist)) * 1.5)
    self.configFile["minArea"]     = minArea
    self.configFile["maxArea"]     = maxArea
    self.configFile["minAreaBody"] = minArea
    self.configFile["maxAreaBody"] = maxArea
    self.configFile["headSize"]    = int(dist * 2)

  self.configFile["extractBackWhiteBackground"] = 1

  self.configFile["noBoutsDetection"] = 1

  if int(self.configFile["nbAnimalsPerWell"]) > 1 or self.forceBlobMethodForHeadTracking:
    identifyMultipleHead(self, controller, int(self.configFile["nbAnimalsPerWell"]))

  controller.show_frame("GoToAdvanceSettings")


def goToAdvanceSettings(self, controller, yes, no):

  if int(no):
    controller.show_frame("FinishConfig")
  else:
    self.configFile["noBoutsDetection"] = 0
    self.calculateBackgroundFreelySwim(controller, 0)

def finishConfig(self, controller, reference):

  # Ideally would like to remove these four lines below, once the problem with wrong 'firstFrame' and 'lastFrame' being saved in the configuration file is solved
  if "lastFrame" in self.configFile:
    del self.configFile["lastFrame"]
  if "firstFrame" in self.configFile:
    del self.configFile["firstFrame"]

  with open(reference, 'w') as outfile:
    json.dump(self.configFile, outfile)
  self.configFile = {}
  self.videoToCreateConfigFileFor = ''
  self.wellLeftBorderX = 0
  self.wellLeftBorderY = 0
  self.headCenterX = 0
  self.headCenterY = 0
  self.organism = ''

  controller.show_frame("StartPage")
