from PyQt5.QtWidgets import * 
from PyQt5 import QtCore, QtGui
from PyQt5.QtGui import * 
from PyQt5.QtCore import * 
from PyQt5.QtCore import pyqtSlot as Slot
import os
import yaml
from pathlib import Path
import argparse
import re
import glob

from dmriprep_ui.protocoltab import ProtocolTab
from dmriprep_ui.executiontab import ExecutionTab
from dmriprep_ui.protocol import Protocol
from dmriprep_ui.quickview import QuickView

class Widgets(QWidget):
    
  def __init__(self, parent, protocol_template, args):
    super(QWidget, self).__init__(parent)
    self.protocol_template = protocol_template
    self.unsaved_changes = True

    # User preferences
    user_directory = os.path.expanduser("~/.niral-dti")
    if not os.path.exists(user_directory + "/user_preferences.yml"):
      self.CreateUserPreferencesFile()
    preferences_yml = yaml.safe_load(open(user_directory + "/user_preferences.yml", 'r'))
    
    # Load yml file used to initialize modules
    def_protocol_path = Path(__file__).parent.joinpath('data/init_protocol.yml')
    protocol_yml = yaml.safe_load(open(def_protocol_path.__str__(),'r'))
    
    # instanciation of classes
    
    self.protocol_tab = ProtocolTab(protocol_template, protocol_yml, preferences_yml) #elements related to the protocol tab
    self.execution_tab = ExecutionTab(protocol_template, args) #elements related to the execution tab
    init_module_id = self.protocol_tab.selector.modules_list_widget.count()
    self.protocol = Protocol(protocol_template, protocol_yml, preferences_yml, init_module_id) #functions related to the manipulations of the protocol

    # Initialize tabs
    self.tabs = QTabWidget()    
    self.tabs.addTab(self.protocol_tab.tab, "Protocol")
    self.tabs.addTab(self.execution_tab.tab, "Execution")       

    # Add tabs to widget
    layout = QVBoxLayout(self)
    layout.addWidget(self.tabs)
    self.setLayout(layout)

    # Create red and green icons : could be a function
    pixmap = QPixmap(100, 100)
    pixmap.fill(QColor("red"))
    self.red_icon = QIcon(pixmap)
    pixmap.fill(QColor("green"))
    self.green_icon = QIcon(pixmap)

    # Communication between QuickView and Dmriprep
    QuickView.signal_quickview.execute_exclude_gradients.connect(self.ExecuteExcludeGradientsFromQuickView)
    # Communicatin between modules_details classes and Protocol class
    self.protocol_tab.details.baselineaverage.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.brainmask.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.dtiestimate.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.eddymotion.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.exclude.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.interlacecheck.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.slicecheck.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.susceptibility.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.susceptibility.communicate.call_update_user_preferences_file.connect(self.UpdateUserPreferencesFile)
    self.protocol_tab.details.utilheader.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.utilmerge.communicate.send_params.connect(self.UpdateModuleParams)
    self.protocol_tab.details.qcreport.communicate.send_params.connect(self.UpdateModuleParams)
    # Communication from module selector subtab
    self.protocol_tab.selector.communicate.module_added_to_protocol_drop.connect(self.protocol.ModuleAddedToProtocolDrop)
    self.protocol_tab.selector.communicate.call_execute_exclude_gradients_popup.connect(self.protocol_tab.exclude_popup.exec_)
    self.protocol_tab.selector.communicate.call_execute_merge_images_popup.connect(self.protocol_tab.merge_popup.exec_)
    # Communication from module details subtab
    self.protocol_tab.details.communicate.set_update_module_details.connect(self.SetUpdateModuleDetails)
    self.protocol_tab.details.communicate.call_OrderModulesInProtocolList.connect(self.GetDicProtocolKeys)
    self.protocol_tab.details.communicate.get_selected_item_data.connect(self.GetDicProtocolSelectedItem)
    # Communication from general subtab
    self.protocol_tab.general.communicate.call_UnsavedModif.connect(self.UnsavedModif)    
    # Communication from execution tab
    self.execution_tab.communicate.call_DeleteTempDirectory.connect(self.DeleteTempDirectory)
    self.execution_tab.communicate.check_manual_exclude.connect(self.CheckManualExcludeInProtocol)
    self.execution_tab.communicate.call_write_protocol_yml.connect(self.protocol.WriteProtocolYML)
    self.execution_tab.communicate.check_unsaved_changes.connect(self.CheckUnsavedChangesBeforeCompute)
    self.execution_tab.communicate.call_save_protocol.connect(self.protocol_tab.SaveProtocol)
    # Communication from protocol tab
    self.protocol_tab.communicate.call_module_added_to_protocol.connect(self.protocol.ModuleAddedToProtocol)
    self.protocol_tab.communicate.call_unsaved_modif.connect(self.UnsavedModif)
    self.protocol_tab.communicate.call_order_modules_in_protocol_list.connect(self.GetDicProtocolKeys)
    self.protocol_tab.communicate.set_loading_protocol.connect(self.protocol.SetLoadingProtocol)
    self.protocol_tab.communicate.set_module_id.connect(self.protocol.SetModuleId)
    self.protocol_tab.communicate.get_module_id.connect(self.GetModuleId)
    self.protocol_tab.communicate.call_set_default_protocol.connect(self.protocol.SettingDefaultProtocol)
    self.protocol_tab.communicate.call_read_protocol.connect(self.protocol.ReadProtocol)
    self.protocol_tab.communicate.call_remove_module_from_dic_protocol.connect(self.protocol.RemoveModuleFromDicProtocol)
    self.protocol_tab.communicate.set_susceptibility_param_eddymotion.connect(self.protocol.SetSusceptibilityParamEddyMotion)
    self.protocol_tab.communicate.set_dic_protocol.connect(self.protocol.SetDicProtocol)
    self.protocol_tab.communicate.call_no_unsaved_modif.connect(self.NoUnsavedModif)
    self.protocol_tab.communicate.call_write_protocol_yml.connect(self.protocol.WriteProtocolYML)
    self.protocol_tab.communicate.module_added_to_protocol_loading.connect(self.protocol.ModuleAddedToProtocolLoading)
    self.protocol_tab.communicate.call_update_user_preferences_file.connect(self.UpdateUserPreferencesFile)
    # Communication from protocol class
    self.protocol.communicate.call_clear_protocol.connect(self.protocol_tab.ClearProtocol)
    self.protocol.communicate.set_general_param_from_read_protocol.connect(self.SetGeneralParam)
    self.protocol.communicate.add_module_to_protocol_widget.connect(self.protocol_tab.AddModuleToProtocolWidget)
    self.protocol.communicate.get_protocol_io.connect(self.GetProtocolIO)
    self.protocol.communicate.call_unsaved_modif.connect(self.UnsavedModif)
    self.protocol.communicate.get_list_widget_modules.connect(self.protocol_tab.GetListWidgetModules)
    self.protocol.communicate.get_selected_modules.connect(self.protocol_tab.selector.GetListWidgetSelectedModules)
    self.protocol.communicate.set_module_data.connect(self.protocol_tab.SetModuleData)
    self.protocol.communicate.call_execute_exclude_gradients_popup.connect(self.protocol_tab.exclude_popup.exec_)
    self.protocol.communicate.call_execute_merge_images_popup.connect(self.protocol_tab.merge_popup.exec_)
    self.protocol.communicate.enable_modules_list_widget.connect(self.protocol_tab.selector.EnableModulesListWidget)
    self.protocol.communicate.enable_exclude_gradients_module.connect(self.protocol_tab.selector.EnableExcludeGradientsModule)
    self.protocol.communicate.enable_merge_images_module.connect(self.protocol_tab.selector.EnableMergeImagesModule)
    self.protocol.communicate.update_selector_data.connect(self.protocol_tab.selector.UpdateSelectorData)


    self.tabs.tabBarClicked.connect(self.MultiInput)
    
    # Command Line arguments
    if args.protocol:
      #arg_protocol = yaml.safe_load(open(args.protocol,'r'))
      self.protocol.ReadProtocol(args.protocol)
    if args.output_directory:
      self.execution_tab.output_line.setText(args.output_directory)
    if args.image:
      if len(args.image) == 1:
        self.execution_tab.input_filename = args.image[0]
        if args.quickview:
          self.protocol.dic_protocol[0] = self.protocol.exclude_default_params
          self.execution_tab.multi_input = 0
          self.execution_tab.Compute([True, True]) #launch QuickView window
      else:
        self.execution_tab.input_filename1 = args.image[0]
        self.execution_tab.input_filename2 = args.image[1]


    self.unsaved_changes = False    

  def UnsavedModif(self):
    if self.unsaved_changes == False:
      self.unsaved_changes = True
      self.tabs.setTabIcon(0, self.red_icon)
      self.tabs.setTabToolTip(0, "Unsaved modifications have been made to protocol")

  def NoUnsavedModif(self):
    self.unsaved_changes = False
    self.tabs.setTabIcon(0, self.green_icon)
    self.tabs.setTabToolTip(0, "Protocol saved")

  def CreateUserPreferencesFile(self):
    user_directory = os.path.expanduser("~/.niral-dti")
    list_of_dir = glob.glob(user_directory + "/*")
    if user_directory + "/user_preferences.yml" in list_of_dir:
      list_of_dir.remove(user_directory + "/user_preferences.yml")
    latest_dir = max(list_of_dir, key = os.path.getctime)
    default = {"fslConfigurationFilePath": latest_dir + "/parameters/fsl/fsl_regb02b0.cnf",
      "showExcludeGradientsPopup": True, 
      "showMergeImagesPopup": True}
    with open(user_directory + "/user_preferences.yml", 'w')as filename:
      yaml.dump(default, filename)

  def UpdateUserPreferencesFile(self, preferences):
    user_directory = os.path.expanduser("~/.niral-dti")
    with open(user_directory + "/user_preferences.yml", 'w')as filename:
      yaml.dump(preferences, filename)
    
  def DeleteTempDirectory(self):
    if os.path.exists("temp_dmriprep_ui/protocol.yml"):
      os.remove("temp_dmriprep_ui/protocol.yml")
    if os.path.exists("temp_dmriprep_ui"):
      os.rmdir("temp_dmriprep_ui")

  def ExecuteExcludeGradientsFromQuickView(self, list_gradients):
    key = list(self.protocol.dic_protocol.keys())[0]
    self.protocol.dic_protocol[key][2]["protocol"]['gradientsToExclude'] = list_gradients

    if not os.path.exists("temp_dmriprep_ui"):
      os.mkdir("temp_dmriprep_ui")
      print("Directory 'temp_dmriprep_ui' created")
    else:
      print("Directory 'temp_dmriprep_ui' already exists")
    self.protocol.WriteProtocolYML("temp_dmriprep_ui/protocol.yml")

    self.execution_tab.StartComputation("temp_dmriprep_ui/protocol.yml")

  def SetUpdateModuleDetails(self, bool_update_module_details):
    self.protocol.update_module_details = bool_update_module_details
  
  def GetDicProtocolSelectedItem(self, data):
    key = data[0]
    index = data[1]
    selected_module = self.protocol.dic_protocol[key]
    self.protocol_tab.details.communicate.SendDicProtocolSelectedItem([selected_module, index])

  def MultiInput(self):
    multi_input = 0
    for index in self.protocol.dic_protocol:
      for module_yml in self.protocol_tab.details.modules_yml_list:
        if self.protocol.dic_protocol[index][1] in module_yml["name"]:
          if "multi_input" in module_yml["process_attributes"]:
            multi_input = 1
            if self.protocol.dic_protocol[index][0] == "Merge Images":
              multi_input = 2

    self.execution_tab.communicate.NumberOfInputs(multi_input)

  def CheckManualExcludeInProtocol(self):
    computation_details = [False, False]
    key = list(self.protocol.dic_protocol.keys())[0]
    if self.protocol.dic_protocol[key][1] == "MANUAL_Exclude":
      computation_details[0] = True
      if self.protocol_tab.details.exclude.selectionMethod.currentText() == "QuickView":
        computation_details[1] = True
    self.execution_tab.communicate.CallCompute(computation_details)


  def CheckUnsavedChangesBeforeCompute(self):
    if self.unsaved_changes == True:
      self.execution_tab.unsaved_changes_popup.exec_()
    filename = self.protocol_tab.protocol_filename
    self.execution_tab.StartComputation(filename)

  def GetModuleId(self):
    identifier = self.protocol.module_id
    self.protocol_tab.DefaultProtocol(identifier)
  
  def SetModuleId(self, new_id):
    self.protocol.module_id = new_id

  def SetGeneralParam(self, protocol):
    self.execution_tab.output_line.setText(protocol["output_directory"])
    if protocol["no_output_image"] == False:
      self.protocol_tab.general.no_output_image_false.setChecked(True)
    else:
      self.protocol_tab.general.no_output_image_true.setChecked(True)
    image_type = ["output_format"]
    for ite in self.protocol_template["options"]["io"]["output_format"]["candidates"]:
      if ite["value"] == image_type:
        self.protocol_tab.general.output_format.setCurrentText(ite["caption"])
        self.protocol_tab.general.GetOutputImageTypeIt(ite["caption"])
    self.protocol_tab.general.b0.setValue(protocol["baseline_threshold"])

  def GetProtocolIO(self, data):
    data[0]["io"] = {
        'baseline_threshold': self.protocol_tab.general.b0.value(),
        'no_output_image': self.protocol_tab.general.no_output_image_true.isChecked(),
        'output_directory': '',
        'output_format': self.protocol_template["options"]["io"]["output_format"]["candidates"][self.protocol_tab.general.outputimagetype_it]["value"]
      }
    self.protocol.Writing(data)

  def GetIndexAndKey(self):
    index = self.protocol_tab.protocol_list_widget.currentRow() #in protocol_list_widget
    key = self.protocol_tab.protocol_list_widget.item(index).data(QtCore.Qt.UserRole) #module_id, in dic_protocol (dictionary)
    return index, key

  def UpdateModuleParams(self, params):
    self.UnsavedModif()
    index, key = self.GetIndexAndKey()
    if index >= 0 and self.protocol.update_module_details == True:
      self.protocol.dic_protocol[key] = params

  def GetDicProtocolKeys(self):
    list_new_keys = []
    for module_index in range(len(self.protocol.dic_protocol)): 
      list_new_keys.append(self.protocol_tab.protocol_list_widget.item(module_index).data(QtCore.Qt.UserRole))
    self.protocol.OrderModulesInProtocolList(list_new_keys)