from PyQt5.QtWidgets import QApplication, QComboBox, QFrame, QCheckBox, QVBoxLayout, QWidget, QLabel, QHBoxLayout, QPushButton
from PyQt5.QtCore import Qt, QSize
from PyQt5.QtGui import QIcon
from celldetective.gui.gui_utils import center_window, ListWidget, DistanceChoice
from superqt.fonticon import icon
from fonticon_mdi6 import MDI6
import numpy as np
import json
import os
from glob import glob
import pandas as pd
from celldetective.gui.viewers import CellSizeViewer, CellEdgeVisualizer
from celldetective.gui import Styles

class ConfigNeighborhoods(QWidget, Styles):
	
	"""
	Widget to configure neighborhood measurements.

	"""

	def __init__(self, neighborhood_type='distance_threshold',neighborhood_parameter_name='threshold distance', parent_window=None, *args, **kwargs):
		
		super().__init__(*args, **kwargs)
		self.parent_window = parent_window
		self.attr_parent = self.parent_window.parent_window
		self.setWindowIcon(QIcon(os.sep.join(['celldetective','icons','logo.png'])))

		self.neighborhood_type = neighborhood_type
		self.neighborhood_parameter_name = neighborhood_parameter_name

		self.setWindowTitle('Configure neighborhoods')
		self.neigh_instructions = self.attr_parent.exp_dir + os.sep.join(["configs","neighborhood_instructions.json"])
		self.clear_previous = False
		self.not_status_reference = False
		self.not_status_neighbor = False	

		self.screen_height = self.attr_parent.screen_height
		self.setMinimumWidth(750)
		self.setMinimumHeight(int(0.5*self.screen_height))
		self.setMaximumHeight(int(0.95*self.screen_height))
		
		self.generate_main_layout()
		self.load_previous_neighborhood_instructions()
		center_window(self)
		self.setAttribute(Qt.WA_DeleteOnClose)

	def generate_main_layout(self):

		main_layout = QVBoxLayout(self)
		main_layout.setContentsMargins(30,30,30,30)

		populations_layout = QHBoxLayout()

		# Reference population
		self.reference_population_frame = QFrame()
		self.reference_population_frame.setFrameStyle(QFrame.StyledPanel | QFrame.Raised)
		self.populate_reference_frame()
		populations_layout.addWidget(self.reference_population_frame, 50)

		# Neighbor population
		self.neigh_population_frame = QFrame()
		self.neigh_population_frame.setFrameStyle(QFrame.StyledPanel | QFrame.Raised)
		self.populate_neighbor_frame()
		populations_layout.addWidget(self.neigh_population_frame, 50)
		main_layout.addLayout(populations_layout)

		# Measurements

		self.measurement_frame = QFrame()
		self.measurement_frame.setFrameStyle(QFrame.StyledPanel | QFrame.Raised)
		self.populate_measurement_frame()
		# if self.neighborhood_type=='distance_threshold':
		# 	self.populate_radii_frame()
		# elif self.neighborhood_type=='mask_contact':
		# 	self.populate_contact_frame()
		main_layout.addWidget(self.measurement_frame)

		self.clear_previous_btn = QCheckBox('clear previous neighborhoods')
		self.clear_previous_btn.setToolTip('Clear all previous neighborhood measurements.')
		self.clear_previous_btn.setIcon(icon(MDI6.broom, color='black'))

		main_layout.addWidget(self.clear_previous_btn, alignment=Qt.AlignRight)

		main_layout.addWidget(QLabel(''))
		self.submit_btn = QPushButton('Set')
		self.submit_btn.setStyleSheet(self.button_style_sheet)
		self.submit_btn.clicked.connect(self.write_instructions)
		main_layout.addWidget(self.submit_btn)

		self.adjustSize()
		QApplication.processEvents()

	def populate_measurement_frame(self):
		
		"""
		Add widgets and layout in the radii frame.
		"""

		grid = QVBoxLayout(self.measurement_frame)

		self.dist_lbl = QLabel(f"NEIGHBORHOOD {self.neighborhood_parameter_name.upper()}")
		self.dist_lbl.setStyleSheet(self.block_title)
		grid.addWidget(self.dist_lbl, alignment=Qt.AlignCenter)

		self.ContentsMeasurements = QFrame()
		layout = QVBoxLayout(self.ContentsMeasurements)
		layout.setContentsMargins(0,0,0,0)

		list_header_layout = QHBoxLayout()
		meas = self.neighborhood_parameter_name.split(' ')[-1]
		lbl = QLabel(f'{meas} [px]:')
		lbl.setToolTip('')
		list_header_layout.addWidget(lbl, 85)

		self.delete_measurement_btn = QPushButton("")
		self.delete_measurement_btn.setStyleSheet(self.button_select_all)
		self.delete_measurement_btn.setIcon(icon(MDI6.trash_can,color="black"))
		self.delete_measurement_btn.setToolTip("Remove measurement.")
		self.delete_measurement_btn.setIconSize(QSize(20, 20))
		list_header_layout.addWidget(self.delete_measurement_btn, 5)

		self.add_measurement_btn = QPushButton("")
		self.add_measurement_btn.setStyleSheet(self.button_select_all)
		self.add_measurement_btn.setIcon(icon(MDI6.plus,color="black"))
		self.add_measurement_btn.setToolTip("Add a neighborhood measurement.")
		self.add_measurement_btn.setIconSize(QSize(20, 20))	
		list_header_layout.addWidget(self.add_measurement_btn, 5)

		self.viewer_btn = QPushButton()
		self.viewer_btn.setStyleSheet(self.button_select_all)
		self.viewer_btn.setIcon(icon(MDI6.image_check, color="black"))
		self.viewer_btn.setToolTip("View stack and set measurement.")
		self.viewer_btn.setIconSize(QSize(20, 20))
		if self.neighborhood_type=='distance_threshold':
			self.viewer_btn.clicked.connect(self.view_current_stack_with_circle)
		elif self.neighborhood_type=='mask_contact':
			self.viewer_btn.clicked.connect(self.view_current_stack_with_edge)			
		list_header_layout.addWidget(self.viewer_btn, 5)

		layout.addLayout(list_header_layout)

		self.measurements_list = ListWidget(DistanceChoice, initial_features=["60"], dtype=int)
		self.measurements_list.setToolTip('Neighborhoods to compute.')
		layout.addWidget(self.measurements_list)

		self.delete_measurement_btn.clicked.connect(self.measurements_list.removeSel)
		self.add_measurement_btn.clicked.connect(self.measurements_list.addItem)

		grid.addWidget(self.ContentsMeasurements)

	def view_current_stack_with_circle(self):
		
		self.parent_window.parent_window.locate_image()
		if self.parent_window.parent_window.current_stack is not None:
			self.viewer = CellSizeViewer(
										  initial_diameter = 100,
										  parent_list_widget = self.measurements_list.list_widget,
										  set_radius_in_list = True,
										  stack_path=self.parent_window.parent_window.current_stack,
										  window_title=f'Position {self.parent_window.parent_window.position_list.currentText()}',
										  frame_slider = True,
										  contrast_slider = True,
										  channel_cb = True,
										  diameter_slider_range = (0,300),
										  channel_names = self.parent_window.parent_window.exp_channels,
										  n_channels = self.parent_window.parent_window.nbr_channels,
										  PxToUm = 1,
										 )
			self.viewer.show()

	def view_current_stack_with_edge(self):
		
		self.attr_parent.locate_image()
		if self.attr_parent.current_stack is not None:
			self.viewer = CellEdgeVisualizer(
										  cell_type='effectors',
										  edge_range=(1,30),
										  invert=True,
										  initial_edge=3,
										  parent_list_widget = self.measurements_list.list_widget,
										  stack_path=self.attr_parent.current_stack,
										  window_title=f'Position {self.attr_parent.position_list.currentText()}',
										  frame_slider = True,
										  contrast_slider = True,
										  channel_cb = True,
										  channel_names = self.attr_parent.exp_channels,
										  n_channels = self.attr_parent.nbr_channels,
										  PxToUm = 1,
										 )
			self.viewer.show()


	def populate_reference_frame(self):
		
		"""
		Add widgets and layout in the reference population frame.
		"""

		grid = QVBoxLayout(self.reference_population_frame)
		grid.setSpacing(15)
		self.ref_lbl = QLabel("REFERENCE")
		self.ref_lbl.setStyleSheet(self.block_title)
		self.ref_lbl.setToolTip('Reference population settings.')
		grid.addWidget(self.ref_lbl, 30, alignment=Qt.AlignCenter)

		self.ContentsReference = QFrame()
		layout = QVBoxLayout(self.ContentsReference)
		layout.setContentsMargins(15,15,15,15)

		population_layout = QHBoxLayout()
		population_layout.addWidget(QLabel('population: '),30)
		self.reference_population_cb = QComboBox()
		self.reference_population_cb.addItems(['targets','effectors'])
		self.reference_population_cb.setToolTip('Select a reference population.')
		population_layout.addWidget(self.reference_population_cb,70)
		layout.addLayout(population_layout)

		status_layout = QHBoxLayout()

		#status_layout.addWidget(QLabel('status: '), 30)

		status_sublayout = QHBoxLayout()
		self.reference_population_status_cb = QComboBox()
		self.reference_population_status_cb.setToolTip('Status of the reference population.')
		self.reference_population_status_cb.hide()
		status_sublayout.addWidget(self.reference_population_status_cb,95)

		self.reference_switch_status_btn = QPushButton("")
		self.reference_switch_status_btn.setStyleSheet(self.button_select_all)
		self.reference_switch_status_btn.setIcon(icon(MDI6.invert_colors,color="black"))
		self.reference_switch_status_btn.setIconSize(QSize(20, 20))
		self.reference_switch_status_btn.clicked.connect(self.switch_not_reference)
		self.reference_switch_status_btn.setToolTip('Invert status values.')
		self.reference_switch_status_btn.hide()

		status_sublayout.addWidget(self.reference_switch_status_btn, 5)

		status_layout.addLayout(status_sublayout, 70)
		layout.addLayout(status_layout)

		event_layout = QHBoxLayout()
		event_layout.addWidget(QLabel('event time: '),30)
		self.event_time_cb = QComboBox()
		self.event_time_cb.setToolTip('Compute average neighborhood metrics before and after this event time.')
		event_layout.addWidget(self.event_time_cb,70)
		layout.addLayout(event_layout)
		
		self.fill_cbs_of_reference_population()
		self.reference_population_cb.currentIndexChanged.connect(self.fill_cbs_of_reference_population)

		grid.addWidget(self.ContentsReference, 70)

	def populate_neighbor_frame(self):

		"""
		Add widgets and layout in the neighbor population frame.
		"""

		grid = QVBoxLayout(self.neigh_population_frame)
		grid.setSpacing(15)
		self.ref_lbl = QLabel("NEIGHBOR")
		self.ref_lbl.setStyleSheet(self.block_title)
		self.ref_lbl.setToolTip('Neighbor population settings.')
		grid.addWidget(self.ref_lbl, 30, alignment=Qt.AlignCenter)

		self.ContentsNeigh = QFrame()
		layout = QVBoxLayout(self.ContentsNeigh)
		layout.setContentsMargins(15,15,15,15)

		population_layout = QHBoxLayout()
		population_layout.addWidget(QLabel('population: '),30)
		self.neighbor_population_cb = QComboBox()
		self.neighbor_population_cb.addItems(['targets','effectors'])
		self.neighbor_population_cb.setToolTip('Select a neighbor population.')
		population_layout.addWidget(self.neighbor_population_cb,70)
		layout.addLayout(population_layout)

		status_layout = QHBoxLayout()

		status_layout.addWidget(QLabel('status: '), 30)
		status_sublayout = QHBoxLayout()

		self.neighbor_population_status_cb = QComboBox()
		self.neighbor_population_status_cb.setToolTip('Status of the neighbor population.')
		status_sublayout.addWidget(self.neighbor_population_status_cb,95)

		self.neighbor_switch_status_btn = QPushButton("")
		self.neighbor_switch_status_btn.setStyleSheet(self.button_select_all)
		self.neighbor_switch_status_btn.setIcon(icon(MDI6.invert_colors,color="black"))
		self.neighbor_switch_status_btn.setToolTip("Invert status values.")
		self.neighbor_switch_status_btn.setIconSize(QSize(20, 20))
		self.neighbor_switch_status_btn.clicked.connect(self.switch_not_neigh)
		status_sublayout.addWidget(self.neighbor_switch_status_btn, 5)
		status_layout.addLayout(status_sublayout, 70)
		layout.addLayout(status_layout)

		self.cumulated_presence_btn = QCheckBox('cumulated presence')
		self.cumulated_presence_btn.setToolTip("Compute the cumulated presence time of each neighbor around a reference cell.")
		self.cumulated_presence_btn.setIcon(icon(MDI6.timer_outline, color='black'))

		layout.addWidget(self.cumulated_presence_btn)

		# self.symmetrize_btn = QCheckBox('symmetrize')
		# self.symmetrize_btn.setToolTip("Write the neighborhood of the neighbor cells with respect to the reference cells.")
		# layout.addWidget(self.symmetrize_btn)
		
		self.fill_cbs_of_neighbor_population()
		self.neighbor_population_cb.currentIndexChanged.connect(self.fill_cbs_of_neighbor_population)

		grid.addWidget(self.ContentsNeigh, 70)

	def fill_cbs_of_neighbor_population(self):

		population = self.neighbor_population_cb.currentText()
		class_cols, status_cols, group_cols, time_cols = self.locate_population_specific_columns(population)
		self.neighbor_population_status_cb.clear()
		self.neighbor_population_status_cb.addItems(['--','class', 'status']+class_cols+status_cols+group_cols)

	def fill_cbs_of_reference_population(self):

		population = self.reference_population_cb.currentText()
		class_cols, status_cols, group_cols, time_cols = self.locate_population_specific_columns(population)
		self.reference_population_status_cb.clear()
		self.reference_population_status_cb.addItems(['--','class', 'status']+class_cols+status_cols+group_cols)
		self.event_time_cb.addItems(['--', 't0']+time_cols)

	def switch_not_reference(self):
		
		self.not_status_reference = not self.not_status_reference
		if self.not_status_reference:
			self.reference_switch_status_btn.setIcon(icon(MDI6.invert_colors,color=self.celldetective_blue))
			self.reference_switch_status_btn.setIconSize(QSize(20, 20))
		else:
			self.reference_switch_status_btn.setIcon(icon(MDI6.invert_colors,color="black"))
			self.reference_switch_status_btn.setIconSize(QSize(20, 20))

	def switch_not_neigh(self):
		
		self.not_status_neighbor = not self.not_status_neighbor
		if self.not_status_neighbor:
			self.neighbor_switch_status_btn.setIcon(icon(MDI6.invert_colors,color=self.celldetective_blue))
			self.neighbor_switch_status_btn.setIconSize(QSize(20, 20))
		else:
			self.neighbor_switch_status_btn.setIcon(icon(MDI6.invert_colors,color="black"))
			self.neighbor_switch_status_btn.setIconSize(QSize(20, 20))


	def locate_population_specific_columns(self, population):

		# Look for all classes and times
		tables = glob(self.attr_parent.exp_dir+os.sep.join(['W*','*','output','tables',f'trajectories_{population}.csv']))
		self.all_columns = []
		for tab in tables:
			cols = pd.read_csv(tab, nrows=1).columns.tolist()
			self.all_columns.extend(cols)
		self.all_columns = np.unique(self.all_columns)

		class_idx = np.array([s.startswith('class_') for s in self.all_columns])
		status_idx = np.array([s.startswith('status_') for s in self.all_columns])
		group_idx = np.array([s.startswith('group_') for s in self.all_columns])		
		time_idx = np.array([s.startswith('t_') for s in self.all_columns])

		if len(class_idx)>0:
			class_columns = list(self.all_columns[class_idx])
			for c in ['class_id', 'class_color']:
				if c in class_columns:
					class_columns.remove(c)
		else:
			class_columns = []

		if len(status_idx)>0:
			status_columns = list(self.all_columns[status_idx])
		else:
			status_columns = []

		if len(group_idx)>0:
			group_columns = list(self.all_columns[group_idx])
		else:
			group_columns = []

		if len(time_idx)>0:
			time_columns = list(self.all_columns[time_idx])
		else:
			time_columns = []

		return class_columns, status_columns, group_columns, time_columns

	def write_instructions(self):

		"""
		Write the selected options in a json file for later reading by the software.
		"""

		print('Writing instructions...')
		
		neighborhood_options = {}
		pop = [self.reference_population_cb.currentText(), self.neighbor_population_cb.currentText()]
		neighborhood_options.update({'population': pop})
		
		status_options = [self.reference_population_status_cb.currentText(), self.neighbor_population_status_cb.currentText()]
		for k in range(2):
			if status_options[k]=='--':
				status_options[k] = None
		if pop[0]!=pop[1]:
			mode = 'two-pop'
		else:
			mode = 'self'

		# TO ADAPT
		distances = self.measurements_list.getItems()
		neighborhood_options.update({'neighborhood_type': self.neighborhood_type})
		neighborhood_options.update({'distance': distances})	
		neighborhood_options.update({'clear_neigh': self.clear_previous_btn.isChecked()})
		event_time_col = self.event_time_cb.currentText()
		if event_time_col=='--':
			event_time_col = None
		neighborhood_options.update({'event_time_col': event_time_col})		

		neighborhood_kwargs = {'mode': mode, 'status': status_options, 'not_status_option': [self.not_status_reference, self.not_status_neighbor], 
							  'compute_cum_sum': self.cumulated_presence_btn.isChecked(), 'attention_weight': True, 'symmetrize': False,
							  'include_dead_weight': True}

		neighborhood_options.update({'neighborhood_kwargs': neighborhood_kwargs})

		print('Neighborhood instructions: ', neighborhood_options)
		file_name = self.neigh_instructions
		with open(file_name, 'w') as f:
			json.dump(neighborhood_options, f, indent=4)


		self.parent_window.protocols.append(neighborhood_options)
		correction_description = ""
		for index, (key, value) in enumerate(neighborhood_options.items()):
			if index > 0:
				correction_description += ", "
			correction_description += str(key) + " : " + str(value)
		self.parent_window.protocol_list.addItem(correction_description)
		
		print('Done.')
		self.close()

	def load_previous_neighborhood_instructions(self):

		"""
		Read the measurmeent options from a previously written json file and format properly for the UI.
		"""

		print('Reading instructions..')
		if os.path.exists(self.neigh_instructions):
			with open(self.neigh_instructions, 'r') as f:
				
				neigh_instructions = json.load(f)
				print(f"Loading the instructions: {neigh_instructions}...")

				if 'neighborhood_type' not in neigh_instructions:
					neigh_instructions.update({'neighborhood_type': self.neighborhood_type})

				if self.neighborhood_type==neigh_instructions['neighborhood_type']:

					if 'distance' in neigh_instructions:
						distances = neigh_instructions['distance']
						distances = [str(d) for d in distances]
						self.measurements_list.list_widget.clear()
						self.measurements_list.list_widget.addItems(distances)

					if 'population' in neigh_instructions:

						pop = neigh_instructions['population']
						idx0 = self.reference_population_cb.findText(pop[0])
						self.reference_population_cb.setCurrentIndex(idx0)
						idx1 = self.neighbor_population_cb.findText(pop[1])
						self.neighbor_population_cb.setCurrentIndex(idx1)

					if 'clear_neigh' in neigh_instructions:
						clear_neigh = neigh_instructions['clear_neigh']
						self.clear_previous_btn.setChecked(clear_neigh)

					if 'event_time_col' in neigh_instructions:
						event_time_col = neigh_instructions['event_time_col']
						if event_time_col is None:
							event_time_col = '--'
						idx = self.event_time_cb.findText(event_time_col)
						self.event_time_cb.setCurrentIndex(idx)

					if 'neighborhood_kwargs' in neigh_instructions:
						neighborhood_kwargs = neigh_instructions['neighborhood_kwargs']
						if 'compute_cum_sum' in neighborhood_kwargs:
							self.cumulated_presence_btn.setChecked(neighborhood_kwargs['compute_cum_sum'])
						if 'status' in neighborhood_kwargs:
							status_options = neighborhood_kwargs['status']
							status_options = ['--' if s is None else s for s in status_options]
							idx0 = self.reference_population_status_cb.findText(status_options[0])
							self.reference_population_status_cb.setCurrentIndex(idx0)
							idx1 = self.neighbor_population_status_cb.findText(status_options[1])
							self.neighbor_population_status_cb.setCurrentIndex(idx1)
						if 'not_status_option' in neighborhood_kwargs:
							not_status_option = neighborhood_kwargs['not_status_option']
							if not_status_option[0]:
								self.reference_switch_status_btn.click()
							if not_status_option[1]:
								self.neighbor_switch_status_btn.click()