# encoding=utf-8
__author__ = "Aaron Randreth"
__copyright__ = "Copyright 2015+, Consortium MonPaGe"
__license__ = "Creative Commons 4.0 By-Nc-Sa"
__maintainer__ = "Roland Trouville"
__email__ = "contact.monpage@gmail.com"
__status__ = "Production"

import math
from typing import Any, Callable, Optional, Sequence, Union

import numpy as np
import pyqtgraph as pg
from pyqtgraph import LinearRegionItem

# pg.setConfigOptions(foreground="k", background="d")
from PySide6.QtWidgets import QGraphicsProxyWidget, QVBoxLayout, QWidget

import cotation.acoustic.struct.praat.spectrogram as display
# from cotation.acoustic.cotation_acoustic_window import CotationAcousticWindow
from cotation.acoustic.display.audio_cursor import AudioCursor
from cotation.acoustic.display.double_axis_plot_item import DoubleAxisPlotItem
from cotation.acoustic.display.zoom_toolbar import ZoomToolbar
from cotation.acoustic.struct.pm.parselmouth_sound import ParselmouthSound
from cotation.acoustic.struct.pm.parselmouth_spectrogram import ParselmouthSpectrogram


class QtGraphDisplay(pg.GraphicsLayoutWidget):
	"""
	A widget for visualizing and interacting with acoustic signals and spectrograms using PyQtGraph.

	This class provides synchronized visual representations of a waveform (acoustic signal) and its
	corresponding spectrogram, allowing for interactive selection, zooming, and playback control.

	Attributes:
	    sound (ParselmouthSound): The sound data containing waveform and timestamp information.
	    spect (ParselmouthSpectrogram): The spectrogram data associated with the sound.
	    selection_size (float): The initial size of the selection region in seconds.
	    selection_region_acoustic_signal (pg.LinearRegionItem): Region selector for the waveform.
	    selection_region_spectrogram (pg.LinearRegionItem): Region selector for the spectrogram.
	    soundwave_plot (pg.PlotItem): Plot item for the waveform.
	    spectrogram_plot (DoubleAxisPlotItem): Plot item for the spectrogram with dual y-axes.
	    play_fill (pg.LinearRegionItem): Visual element indicating playback progress.
	    was_moved (bool): Flag indicating if the selection was modified by the user.
	    _selection_changed_callback (Optional[Callable[[], Any]]): Callback triggered on selection change.
	    is_selection_movable (bool): Whether the selection can be moved.
	    is_selection_manually_resizable (bool): Whether the selection can be resized manually.
	    selection_max_length (float): Maximum allowable selection duration.
	    selection_min_length (float): Minimum allowable selection duration.
	    zoom (ZoomToolbar): Interactive zoom toolbar.

	"""

	sound: ParselmouthSound
	spect: ParselmouthSpectrogram

	selection_size: float
	selection_region_acoustic_signal: pg.LinearRegionItem
	selection_region_spectrogram: pg.LinearRegionItem

	soundwave_plot: pg.PlotItem
	spectrogram_plot: DoubleAxisPlotItem

	play_fill: pg.LinearRegionItem

	# interval_change_finished = Signal()
	was_moved: bool

	_selection_changed_callback: Optional[Callable[[], Any]]


	is_selection_movable: bool = True
	is_selection_manually_resizable: bool = True
	selection_min_length: float = 4
	selection_max_length: float = 4
	fixed_size_allowed_margin: float = 1.0

	zoom: ZoomToolbar



	def __init__(
		self,
		sound: ParselmouthSound,
		spectrogram: ParselmouthSpectrogram,
		selection_size: float,
		color_map: str = "",
		not_moved_color: str = "#FF000044",
		was_moved_color: str = "#00005544",
		fixed_size_allowed_margin:float=1.0
	):
		super().__init__()
		self.sound = sound
		self.spect = spectrogram
		self.fixed_size_allowed_margin = fixed_size_allowed_margin
		self.selection_size = selection_size
		self.selection_min_length = self.selection_size
		self.selection_max_length = self.selection_min_length * self.fixed_size_allowed_margin

		self.was_moved = False
		self.not_moved_color = not_moved_color
		self.was_moved_color = was_moved_color

		self.init_ui()

	def init_ui(self):
		"""
		Initialize the user interface components for acoustic signal visualization.

		- Create soundwave and spectrogram plots with auto-ranging viewboxes.
		- Generate and add a spectrogram display item.
		- Set axis labels for frequency and amplitude.
		- Initialize selection regions on both plots with defined limits and sizes.
		- Connect region change signals to update handlers.
		- Add an audio cursor to visualize playback position.
		- Setup zoom functionality and finalize view ranges.
		"""
		self.soundwave_plot = pg.PlotItem()
		self.spectrogram_plot = DoubleAxisPlotItem()

		if viewbox := self.soundwave_plot.getViewBox():
			viewbox.enableAutoRange()

		if viewbox := self.spectrogram_plot.getViewBox():
			viewbox.enableAutoRange()

		spect = display.Spectrogram(
			np.array(self.spect.frequencies),
			np.array(self.spect.timestamps),
			np.array(self.spect.data_matrix),
			zoom_blur=False,
		)

		xlim = (0, max(self.sound.timestamps))
		self.add_aligned_plots([self.soundwave_plot, self.spectrogram_plot], *xlim)

		self.soundwave_plot.plot(x=self.sound.timestamps, y=self.sound.amplitudes[0])
		self.spectrogram_plot.init_axes()
		self.spectrogram_plot.update_views()
		self.spectrogram_plot.add_item_main_axis(spect)

		self.spectrogram_plot.setLabel("left", "Fréquence", units="Hz")
		self.soundwave_plot.setLabel("left", "Amplitude")

		# Initialisation
		self.selection_region_acoustic_signal = self.__set_plot_selection_min_max(
			self.soundwave_plot, 0, max(self.sound.timestamps)
		)
		self.selection_region_acoustic_signal.setRegion((0, self.selection_size))
		self.selection_region_spectrogram = self.__set_plot_selection_min_max(
			self.spectrogram_plot, 0, max(self.sound.timestamps)
		)
		self.selection_region_spectrogram.setMovable(False)
		self.selection_region_spectrogram.setRegion((0, self.selection_size))

		self.selection_region_acoustic_signal.sigRegionChanged.connect(
			lambda: self.on_region_selection_change(
				self.selection_region_acoustic_signal
			)
		)
		self.selection_region_spectrogram.sigRegionChanged.connect(
			lambda: self.on_region_selection_change(self.selection_region_spectrogram)
		)

		# self.selection_region_spectrogram.sigRegionChangeFinished.connect(self.on_interval_selection_changed)
		# # self.selection_region_acoustic_signal.sigRegionChanged.connect(lambda: self.__change)

		self.play_fill = AudioCursor(self.selection_region_acoustic_signal)
		self.soundwave_plot.addItem(self.play_fill)
		self.setup_zoom()

		if viewbox := self.soundwave_plot.getViewBox():
			viewbox.autoRange()

		if viewbox := self.spectrogram_plot.getViewBox():
			viewbox.autoRange()

	def set_selection_fixed_size(self, new_size_in_sec: float):
		"""
		Sets the selection region to a fixed size. If CotationAcousticWindow.FIXED_SIZE_ALLOWED_MARGIN is set to more
		than one, the fixed size allow for a margin.
		The resizable is computed if there is a margin, and the max is now the fixed size multiplied by the margin

		Parameters:
			new_size_in_sec (float): The fixed size of the selection region in seconds.

		Returns:
			None
		"""
		self.set_selection_parameters(True, self.fixed_size_allowed_margin>1,
			new_size_in_sec, new_size_in_sec*self.fixed_size_allowed_margin
		)

	def set_selection_parameters(
		self,
		is_movable: bool,
		is_resizable: bool,
		min_length: float,
		max_length: float,
	) -> None:
		"""
		Configures the selection region's movability, resizability, and size constraints.

		If resizing is disabled, sets both minimum and maximum lengths to the midpoint between them.

		Parameters:
			is_movable (bool): Whether the selection region can be moved.
			is_resizable (bool): Whether the selection region can be resized.
			min_length (float): Minimum allowed length of the selection region.
			max_length (float): Maximum allowed length of the selection region (can be None).

		Returns:
			None
		"""
		# self.selection_region_acoustic_signal.se
		self.selection_region_acoustic_signal.setMovable(is_movable)
		self.selection_region_spectrogram.setMovable(is_movable)
		self.is_selection_movable = is_movable
		self.selection_min_length = min_length
		if max_length is None:
			self.selection_max_length = min_length
		else:
			self.selection_max_length = max_length
		if not is_resizable:
			tmp = (min_length + max_length) / 2.0
			self.selection_min_length = tmp
			self.selection_max_length = tmp

	def on_region_selection_change(self, sender: LinearRegionItem):
		"""
		Handles changes in the selection region triggered by a LinearRegionItem sender.

		Applies constraints to the selection interval, synchronizes the selection regions between
		spectrogram and acoustic signal displays, and marks the selection as moved.

		Parameters:
			sender (LinearRegionItem): The region item that triggered the change.

		Returns:
			None
		"""
		self.selection_region_spectrogram.blockSignals(True)
		self.selection_region_acoustic_signal.blockSignals(True)

		actual_selection = sender.getRegion()
		target_selection = self.calculate_interval_start_end_with_constraints(
			actual_selection
		)
		self.selection_region_acoustic_signal.setRegion(target_selection)
		self.selection_region_spectrogram.setRegion(target_selection)

		# A verifier
		self.set_was_moved()

		self.selection_region_spectrogram.blockSignals(False)
		self.selection_region_acoustic_signal.blockSignals(False)

	def calculate_interval_start_end_with_constraints(
		self, actual_selection: tuple
	) -> tuple:
		"""
		Calculates a constrained interval start and end based on min and max selection lengths.

		If the actual selection duration exceeds the max length, it is trimmed centered around the midpoint.
		If it is less than the min length, it is expanded centered around the midpoint.
		Otherwise, returns the original selection.

		Parameters:
			actual_selection (tuple): A tuple (start, end) representing the current selection interval.

		Returns:
			tuple: A tuple (start, end) adjusted according to the constraints.
		"""
		duration = actual_selection[1] - actual_selection[0]
		mid_point = (actual_selection[0] + actual_selection[1]) / 2.0
		if duration > self.selection_max_length:
			return (
				mid_point - self.selection_max_length / 2.0,
				mid_point + self.selection_max_length / 2.0,
			)
			# return (actual_selection[0], actual_selection[0]+self.selection_max_length)
		if duration < self.selection_min_length:
			return (
				mid_point - self.selection_min_length / 2.0,
				mid_point + self.selection_min_length / 2.0,
			)
			# return (actual_selection[0], actual_selection[0] + self.selection_min_length)
		return actual_selection

	def __set_plot_selection_min_max(
		self, plot: pg.PlotItem, min_bounds: int, max_bounds: int
	) -> pg.LinearRegionItem:
		"""
		Creates and adds a LinearRegionItem to a plot with given minimum and maximum bounds.

		Raises a ValueError if min_bounds is greater than max_bounds.

		Parameters:
			plot (pg.PlotItem): The plot item to add the region to.
			min_bounds (int): The minimum bound of the selection region.
			max_bounds (int): The maximum bound of the selection region.

		Returns:
			pg.LinearRegionItem: The created region item.
		"""
		if min_bounds > max_bounds:
			raise ValueError(f"{min_bounds} > {max_bounds}, not good")

		region = pg.LinearRegionItem(
			bounds=(min_bounds, max_bounds),
			swapMode="sort",
			pen=pg.mkPen(color="#000055AA", width=3),
		)
		region.setBrush(pg.mkBrush(color=self.not_moved_color))
		region.setZValue(100)
		plot.addItem(region)

		return region

	def add_aligned_plots(self, plot_list: list[pg.PlotItem], xmin: float, xmax: float):
		"""
		Adds a list of plots aligned vertically with synchronized x-axis limits and linked view boxes.

		Sets the bottom axis label on the first plot and ensures all plots share the same horizontal limits
		and disable vertical mouse interaction.

		Parameters:
			plot_list (list[pg.PlotItem]): A list of PlotItem instances to be added and aligned.
			xmin (float): The minimum x-axis value limit.
			xmax (float): The maximum x-axis value limit.

		Returns:
			None
		"""
		if len(plot_list) == 0:
			return

		plot_list[0].setLabel("bottom", "Temps", "s")

		for p in plot_list:
			p.getAxis("left").setWidth(60)  # Aligns the 0
			p.getAxis("right").setWidth(60)
			if viewbox := p.getViewBox():
				viewbox.setLimits(xMin=xmin, xMax=xmax)  # Set a common limit

				viewbox.setXLink(plot_list[0])
				viewbox.setMouseEnabled(x=True, y=False)

			self.addItem(p)
			self.nextRow()  # pyright: ignore [reportAttributeAccessIssue]

	def setup_zoom(self):
		"""
		Initializes and configures the zoom toolbar linked to the selection region.

		Links the zoom functionality to both the soundwave and spectrogram plots,
		wraps the toolbar widget for integration into the graphics scene, and adds it to the layout.

		Returns:
			None
		"""
		self.zoom = ZoomToolbar(self.selection_region_acoustic_signal)
		self.zoom.link_viewbox(self.soundwave_plot)
		self.zoom.link_viewbox(self.spectrogram_plot)

		toolbar_widget = QWidget()
		toolbar_layout = QVBoxLayout()
		toolbar_layout.addWidget(self.zoom)
		toolbar_widget.setLayout(toolbar_layout)

		# Standard Qt widgets (QWidget and its derivatives) are not natively
		# designed to be added to QGraphicsScenes. QGraphicsProxyWidget wraps a
		# standard Qt widget, making it behave like a QGraphicsItem that can be
		# managed by the QGraphicsScene.
		proxy_widget = QGraphicsProxyWidget()
		proxy_widget.setWidget(toolbar_widget)

		self.addItem(proxy_widget)
		self.nextRow()  # pyright: ignore [reportAttributeAccessIssue]

	def set_was_moved(self):
		"""
		Marks the selection as moved if not already set, and updates the brush color
		of the acoustic signal and spectrogram selection regions to indicate this state.

		Returns:
			None
		"""
		if self.was_moved:
			return

		self.was_moved = True
		self.selection_region_acoustic_signal.setBrush(
			pg.mkBrush(color=self.was_moved_color)
		)
		self.selection_region_spectrogram.setBrush(
			pg.mkBrush(color=self.was_moved_color)
		)

	def add_secondary_plotlines(
		self, plotlines: tuple[list[pg.PlotDataItem], float, float]
	):
		"""
		Adds secondary plot lines to the spectrogram plot and configures the secondary axis range.

		Parameters:
			plotlines (tuple[list[pg.PlotDataItem], float, float]):
				A tuple where the first element is a list of plot data items to add,
				and the next two elements specify the range for the secondary axis.

		Returns:
			None
		"""
		if len(plotlines[0]) == 0:
			return

		self.soundwave_plot.showAxis("right")
		self.spectrogram_plot.set_secondary_axis_range(*plotlines[1:])

		for c in plotlines[0]:
			self.spectrogram_plot.add_item_secondary_axis(c)

	def get_selection_bounds(self) -> tuple[int, int]:
		"""
		Retrieves the current selection bounds from the acoustic signal selection region.

		Returns:
			tuple[int, int]: The start and end positions of the selection region.
		"""
		return self.selection_region_acoustic_signal.getRegion()  # pyright: ignore [reportReturnType]

	def set_selection_bounds(self, start: int, end: int):
		"""
		Sets the selection bounds of the acoustic signal selection region.

		Parameters:
			start (int): The start position of the selection.
			end (int): The end position of the selection.

		Returns:
			None
		"""
		self.selection_region_acoustic_signal.setRegion((start, end))

	def get_index_selection_range(self) -> tuple[int, int]:
		"""
		Returns the index range corresponding to the current selection bounds by finding the nearest
		indices in the sound timestamps array.

		Returns:
			tuple[int, int]: A tuple containing the lower and upper index bounds of the selection.
		"""
		lbound, ubound = self.get_selection_bounds()

		return (
			QtGraphDisplay.find_idx_of_nearest(self.sound.timestamps, lbound),
			QtGraphDisplay.find_idx_of_nearest(self.sound.timestamps, ubound),
		)

	# def __change(self, callback: Optional[Callable[[None], None]] = None):
	# 	print("__change")
	# 	self.selection_region_spectrogram.setRegion(self.get_selection_bounds())
	#
	# 	if callback is None:
	# 		return
	#
	# 	callback()

	# def on_selection_change(self, callback: Optional[Callable[[None], None]] = None):
	# 	self.selection_region.sigRegionChanged.connect(lambda: self.__change(callback))

	def start_play_animation(self):
		"""
		Starts the playback animation and disables moving of the spectrogram selection region.

		Returns:
			None
		"""
		self.play_fill.start()
		self.selection_region_spectrogram.setMovable(False)

	def stop_play_animation(self):
		"""
		Stops the playback animation and re-enables moving of the spectrogram selection region.

		Returns:
			None
		"""
		self.play_fill.stop()
		self.selection_region_spectrogram.setMovable(True)

	@staticmethod
	def find_idx_of_nearest(array: Sequence[Union[float, int]], value: float) -> int:
		"""
		Finds the index of the closest element to `value` inside of `array`.

		Parameters:
		- array: list of floats sorted in ascending order.
		- value: the value to find.

		Returns:
		- int: the index
		"""

		idx: int = int(np.searchsorted(array, value, side="left"))

		is_last_el = idx == len(array)
		left_is_closest = math.fabs(value - array[idx - 1]) < math.fabs(
			value - array[idx]
		)

		if idx > 0 and (is_last_el or left_is_closest):
			return idx - 1
		return idx

	def set_selection_changed_callback(self, callback: Callable[[], Any]) -> None:
		"""
		Sets the callback function to be invoked when the selection changes.

		Parameters:
			callback (Callable[[], Any]): A function with no arguments to handle selection change events.

		Returns:
			None
		"""
		self._selection_changed_callback = callback
