"""This module offers a range of distance functions for crystals."""

import functools
import time
import warnings
from collections.abc import Callable
from multiprocessing import Pool, cpu_count
from typing import Literal, overload

import amd
import numpy as np
from ElMD import elmd
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.core import Structure
from scipy.spatial.distance import squareform

from .constants import (
	CONTINUOUS_UNNORMALIZED_DISTANCES,
	DIST_WO_EMB,
	ELMD_AMD_COEF_AMD,
	ELMD_AMD_COEF_ELMD,
	TYPE_EMB_ALL,
	TYPE_EMB_AMD,
	TYPE_EMB_COMP,
	TYPE_EMB_ELMD,
	TYPE_EMB_MAGPIE,
	TYPE_EMB_PDD,
	TYPE_EMB_WYCKOFF,
)
from .crystal import Crystal, _to_crystal, _to_crystal_list


class _EmbeddingError:
	"""Sentinel wrapping a failed embedding computation."""

	def __init__(self, exc: Exception) -> None:
		self.exc = exc


def _normalize(d: float) -> float:
	"""Normalize d to [0, 1] via d / (1 + d).

	Args:
		d (float): Distance value.

	Returns:
		float: Normalized distance.
	"""
	return d / (1 + d)


def _resolve_elmd_amd_coefs(args_dist: dict) -> tuple[float, float]:
	"""Read and normalize elmd+amd blending coefficients from args_dist.

	Args:
		args_dist (dict): Distance arguments dict, optionally containing "coefs".

	Returns:
		tuple[float, float]: Normalized (coef_elmd, coef_amd).
	"""
	if "coefs" in args_dist:
		coefs = args_dist["coefs"]
		coef_elmd, coef_amd = coefs["elmd"], coefs["amd"]
		total = coef_elmd + coef_amd
		return coef_elmd / total, coef_amd / total
	return ELMD_AMD_COEF_ELMD, ELMD_AMD_COEF_AMD


def _insert_nan_rows_cols(
	d_mtx: np.ndarray,
	error_indices_1: list,
	error_indices_2: list | None,
) -> np.ndarray:
	"""Insert NaN rows and columns for failed embeddings.

	Args:
		d_mtx (np.ndarray): Distance matrix computed from valid embeddings only.
		error_indices_1 (list): Row indices where embeddings failed.
		error_indices_2 (list | None): Column indices where embeddings failed.
			If None, the same indices as error_indices_1 are used (symmetric case).

	Returns:
		np.ndarray: Distance matrix with NaN rows/columns inserted.
	"""
	for i in error_indices_1:
		d_mtx = np.insert(d_mtx, i, float("nan"), axis=0)
	col_indices = error_indices_2 if error_indices_2 is not None else error_indices_1
	for j in col_indices:
		d_mtx = np.insert(d_mtx, j, float("nan"), axis=1)
	return d_mtx


def _blend_elmd_amd_distance(
	d_elmd_unnorm: float | np.ndarray,
	d_amd_unnorm: float | np.ndarray,
	args_dist: dict,
) -> float | np.ndarray:
	"""Blend normalized elmd and amd distances.

	Args:
		d_elmd_unnorm (float | np.ndarray): Unnormalized ElMD distance(s).
		d_amd_unnorm (float | np.ndarray): Unnormalized AMD distance(s).
		args_dist (dict): Distance arguments, optionally containing "coefs".

	Returns:
		float | np.ndarray: Blended distance(s).
	"""
	coef_elmd, coef_amd = _resolve_elmd_amd_coefs(args_dist)
	return coef_elmd * _normalize(d_elmd_unnorm) + coef_amd * _normalize(d_amd_unnorm)


def _d_smat(
	emb_1: Crystal, emb_2: Crystal, matcher: StructureMatcher | None = None, **kwargs
) -> float:
	"""Compute the binary distance based on pymatgen's StructureMatcher.

	Args:
		emb_1 (Crystal): Embedding 1.
		emb_2 (Crystal): Embedding 2.
		matcher (StructureMatcher | None): Pre-initialized StructureMatcher object.
		**kwargs: Additional keyword arguments for StructureMatcher, e.g., ltol, stol,
			angle_tol. If matcher is provided, these kwargs will be ignored.

	Returns:
		float: StructureMatcher-based distance (0.0 or 1.0).

	Note:
		d_smat does not allow pre-computation of embeddings.
	"""
	if matcher is None:
		matcher = StructureMatcher(**kwargs)
	return 0.0 if matcher.fit(emb_1, emb_2) else 1.0


def _d_comp(emb_1: TYPE_EMB_COMP, emb_2: TYPE_EMB_COMP, **kwargs) -> float:
	"""Compute the binary distance based on the match of compositions.

	Args:
		emb_1 (TYPE_EMB_COMP): Embedding 1.
		emb_2 (TYPE_EMB_COMP): Embedding 2.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		float: Composition distance (0.0 or 1.0).
	"""
	return 0.0 if emb_1 == emb_2 else 1.0


def _d_wyckoff(emb_1: TYPE_EMB_WYCKOFF, emb_2: TYPE_EMB_WYCKOFF, **kwargs) -> float:
	"""Compute the binary distance based on the match of Wyckoff representations.

	Args:
		emb_1 (TYPE_EMB_WYCKOFF): Embedding 1.
		emb_2 (TYPE_EMB_WYCKOFF): Embedding 2.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		float: Wyckoff distance (0.0 or 1.0).
	"""
	if isinstance(emb_1, _EmbeddingError) or isinstance(emb_2, _EmbeddingError):
		return float("nan")
	return 0.0 if emb_1 == emb_2 else 1.0


def _d_magpie(emb_1: TYPE_EMB_MAGPIE, emb_2: TYPE_EMB_MAGPIE, **kwargs) -> float:
	"""Compute the continuous distance using compositional Magpie fingerprints.

	Args:
		emb_1 (TYPE_EMB_MAGPIE): Embedding 1.
		emb_2 (TYPE_EMB_MAGPIE): Embedding 2.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		float: Magpie distance.

	References:
		- Ward et al., (2016). A general-purpose machine learning framework for
			predicting properties of inorganic materials. npj Computational Materials,
			2(1), 1-7. https://doi.org/10.1038/npjcompumats.2016.28
	"""
	return np.sqrt(np.sum((np.array(emb_1) - np.array(emb_2)) ** 2)).item()


def _d_pdd(emb_1: TYPE_EMB_PDD, emb_2: TYPE_EMB_PDD, **kwargs) -> float:
	"""Compute the continuous distance using pointwise distance distribution (PDD).

	Args:
		emb_1 (TYPE_EMB_PDD): Embedding 1.
		emb_2 (TYPE_EMB_PDD): Embedding 2.
		**kwargs: Additional arguments for amd.PDD_cdist.

	Returns:
		float: PDD distance.

	References:
		- Widdowson et al., (2022). Resolving the data ambiguity for periodic
			crystals. Advances in Neural Information Processing Systems, 35,
			24625--24638. https://openreview.net/forum?id=4wrB7Mo9_OQ
	"""
	if isinstance(emb_1, _EmbeddingError) or isinstance(emb_2, _EmbeddingError):
		return float("nan")
	return amd.PDD_cdist([emb_1], [emb_2], **kwargs)[0][0].item()


def _d_amd(emb_1: TYPE_EMB_AMD, emb_2: TYPE_EMB_AMD, **kwargs) -> float:
	"""Compute the continuous distance using average minimum distance (AMD).

	Args:
		emb_1 (TYPE_EMB_AMD): Embedding 1.
		emb_2 (TYPE_EMB_AMD): Embedding 2.
		**kwargs: Additional arguments for amd.AMD_cdist.

	Returns:
		float: AMD distance.

	References:
		- Widdson et al., (2022). Average Minimum Distances of periodic point sets -
			foundational invariants for mapping periodic crystals. MATCH
			Communications in Mathematical and in Computer Chemistry, 87(3), 529-559,
			https://doi.org/10.46793/match.87-3.529W
	"""
	if isinstance(emb_1, _EmbeddingError) or isinstance(emb_2, _EmbeddingError):
		return float("nan")
	return amd.AMD_cdist([emb_1], [emb_2], **kwargs)[0][0].item()


def _d_elmd(emb_1: TYPE_EMB_ELMD, emb_2: TYPE_EMB_ELMD, **kwargs) -> float:
	"""Compute continuous Element Movers Distance (ElMD).

	Args:
		emb_1 (TYPE_EMB_ELMD): Embedding 1.
		emb_2 (TYPE_EMB_ELMD): Embedding 2.
		**kwargs: Additional arguments for ElMD.elmd.

	Returns:
		float: ElMD distance.

	References:
		- Hargreaves et al., (2020). The Earth Mover's Distance as a Metric for the
			Space of Inorganic Compositions. Chemistry of Materials, 32(24),
			10610-10620.
	"""
	if "return_assignments" in kwargs and kwargs["return_assignments"] is True:
		warnings.warn(
			"return_assignments=True is not supported. Automatically set to False.",
			UserWarning,
			stacklevel=2,
		)
		kwargs["return_assignments"] = False
	return elmd(emb_1, emb_2, **kwargs)


def _row_func_magpie(emb: np.ndarray, embs_2: np.ndarray) -> np.ndarray:
	"""Compute one row of the d_magpie distance matrix.

	Args:
		emb (np.ndarray): Single MAGPIE feature vector.
		embs_2 (np.ndarray): 2D array of target MAGPIE feature vectors.

	Returns:
		np.ndarray: Euclidean distances from emb to each row of embs_2.
	"""
	return np.sqrt(np.sum((emb[np.newaxis, :] - embs_2) ** 2, axis=1))


def _set_n_processes(n_processes: int | None = None) -> int:
	"""Set the number of processes for multiprocessing.

	Args:
		n_processes (int | None): Number of processes. If None, use
			max(cpu_count() - 1, 1). Default is None.

	Returns:
		int: Number of processes to use.
	"""
	if n_processes is None:
		return max(cpu_count() - 1, 1)
	else:
		return max(min(n_processes, cpu_count() - 1), 1)


def _distance_matrix_template(
	embs_1: list,
	embs_2: list | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	d_func: Callable | None = None,
	initializer: Callable | None = None,
	initargs: tuple = (),
	row_func: Callable | None = None,
) -> np.ndarray:
	"""Template to compute the distance matrix between two sets of embeddings.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list): Embeddings
		embs_2 (list | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		d_func (Callable | None): Distance function.
		initializer (Callable | None): Initializer function for worker processes.
		initargs (tuple): Arguments for the initializer function.
		row_func (Callable | None): Row function ``(emb_i, embs_2_arr) -> np.ndarray``
			computing all distances from one source embedding to every target.
			When provided, uses row-wise computation instead of pairwise d_func.

	Returns:
		np.ndarray: Distance matrix.

	Raises:
		ValueError: If both d_func and row_func are None.
	"""
	if d_func is None and row_func is None:
		raise ValueError("Either d_func or row_func must be provided.")
	if row_func is not None:
		embs_2_local = embs_2 if embs_2 is not None else embs_1
		embs_2_arr = np.array(embs_2_local)
		if multiprocessing:
			embs_1_arr = np.array(embs_1)
			with Pool(
				processes=_set_n_processes(n_processes),
				initializer=_init_row_worker,
				initargs=(row_func, embs_2_arr),
			) as pool:
				rows = pool.map(_row_worker, embs_1_arr)
		else:
			rows = [row_func(np.array(emb_i), embs_2_arr) for emb_i in embs_1]
		return np.array(rows)
	given_two_sets = embs_2 is not None
	if given_two_sets:
		if multiprocessing:
			with Pool(
				processes=_set_n_processes(n_processes),
				initializer=initializer,
				initargs=initargs,
			) as pool:
				results = pool.starmap(
					d_func,
					((emb_i, emb_j) for emb_i in embs_1 for emb_j in embs_2),
				)
				d_mtx = np.array(results).reshape(len(embs_1), len(embs_2))
		else:
			d_mtx = np.zeros((len(embs_1), len(embs_2)))
			for i, emb_i in enumerate(embs_1):
				for j, emb_j in enumerate(embs_2):
					d_mtx[i, j] = d_func(emb_i, emb_j)
	else:
		if multiprocessing:
			with Pool(
				processes=_set_n_processes(n_processes),
				initializer=initializer,
				initargs=initargs,
			) as pool:
				results = pool.starmap(
					d_func,
					(
						(embs_1[i], embs_1[j])
						for i in range(len(embs_1))
						for j in range(i + 1)
					),
				)
				d_mtx = np.zeros((len(embs_1), len(embs_1)))
				for i in range(len(embs_1)):
					for j in range(i + 1):
						d_mtx[i, j] = results[i * (i + 1) // 2 + j]
						d_mtx[j, i] = d_mtx[i, j]
		else:
			d_mtx = np.zeros((len(embs_1), len(embs_1)))
			for i, emb_i in enumerate(embs_1):
				for j, emb_j in enumerate(embs_1[: i + 1]):
					d_mtx[i, j] = d_func(emb_i, emb_j)
					d_mtx[j, i] = d_mtx[i, j]
	return d_mtx


_global_matcher: StructureMatcher | None = None


def _init_worker_distance_matrix_d_smat(kwargs_dict: dict):
	"""Initialize global StructureMatcher for worker processes.

	Args:
		kwargs_dict (dict): Additional keyword arguments for StructureMatcher.
	"""
	global _global_matcher
	_global_matcher = StructureMatcher(**kwargs_dict)


def _worker_distance_matrix_d_smat(emb_i, emb_j):
	"""Worker function for computing d_smat in multiprocessing.

	Args:
		emb_i: Embedding i.
		emb_j: Embedding j.

	Returns:
		float: d_smat distance between emb_i and emb_j.
	"""
	global _global_matcher
	return _d_smat(emb_i, emb_j, matcher=_global_matcher)


def _distance_matrix_d_smat(
	embs_1: list[Crystal],
	embs_2: list[Crystal] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
):
	"""Compute the distance matrix between two sets of embeddings based on d_smat.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[Crystal]): Embeddings
		embs_2 (list[Crystal] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Additional keyword arguments for StructureMatcher.

	Returns:
		np.ndarray: d_smat distance matrix.
	"""
	return _distance_matrix_template(
		embs_1,
		embs_2,
		multiprocessing,
		n_processes,
		_worker_distance_matrix_d_smat,
		_init_worker_distance_matrix_d_smat,
		(kwargs,),
	)


def _distance_matrix_d_comp(
	embs_1: list[TYPE_EMB_COMP],
	embs_2: list[TYPE_EMB_COMP] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_comp.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_COMP]): Embeddings
		embs_2 (list[TYPE_EMB_COMP] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		np.ndarray: d_comp distance matrix.
	"""
	return _distance_matrix_template(
		embs_1, embs_2, multiprocessing, n_processes, _d_comp
	)


def _distance_matrix_d_wyckoff(
	embs_1: list[TYPE_EMB_WYCKOFF],
	embs_2: list[TYPE_EMB_WYCKOFF] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_wyckoff.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_WYCKOFF]): Embeddings
		embs_2 (list[TYPE_EMB_WYCKOFF] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		np.ndarray: d_wyckoff distance matrix.
	"""
	return _distance_matrix_template(
		embs_1, embs_2, multiprocessing, n_processes, _d_wyckoff
	)


_global_row_func: Callable | None = None
_global_embs_2_arr: np.ndarray | None = None


def _init_row_worker(row_func: Callable, embs_2_arr: np.ndarray):
	"""Initialize generic row-worker processes.

	Args:
		row_func (Callable): Row function ``(emb_i, embs_2) -> np.ndarray``.
		embs_2_arr (np.ndarray): 2D array of all target embeddings.
	"""
	global _global_row_func, _global_embs_2_arr
	_global_row_func = row_func
	_global_embs_2_arr = embs_2_arr


def _row_worker(emb_i: np.ndarray) -> np.ndarray:
	"""Generic row-worker: compute one row of the distance matrix.

	Args:
		emb_i (np.ndarray): Single source embedding.

	Returns:
		np.ndarray: Distance row against all target embeddings.
	"""
	global _global_row_func, _global_embs_2_arr
	return _global_row_func(emb_i, _global_embs_2_arr)


def _distance_matrix_d_magpie(
	embs_1: list[TYPE_EMB_MAGPIE],
	embs_2: list[TYPE_EMB_MAGPIE] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_magpie.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_MAGPIE]): Embeddings
		embs_2 (list[TYPE_EMB_MAGPIE] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Ignored; present for uniform interface.

	Returns:
		np.ndarray: d_magpie distance matrix.
	"""
	return _distance_matrix_template(
		embs_1, embs_2, multiprocessing, n_processes, row_func=_row_func_magpie
	)


def _distance_matrix_d_pdd(
	embs_1: list[TYPE_EMB_PDD],
	embs_2: list[TYPE_EMB_PDD] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_pdd.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_PDD]): Embeddings
		embs_2 (list[TYPE_EMB_PDD] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Additional arguments for amd.PDD_pdist and amd.PDD_cdist.

	Returns:
		np.ndarray: d_pdd distance matrix.
	"""
	if multiprocessing and "n_jobs" not in kwargs:
		kwargs["n_jobs"] = _set_n_processes(n_processes)

	given_two_sets = embs_2 is not None
	valids_1 = [emb for emb in embs_1 if not isinstance(emb, _EmbeddingError)]
	error_indices_1 = [
		i for i, emb in enumerate(embs_1) if isinstance(emb, _EmbeddingError)
	]
	if not given_two_sets:
		d_mtx = squareform(amd.PDD_pdist(valids_1, **kwargs))
		d_mtx = _insert_nan_rows_cols(d_mtx, error_indices_1, None)
	else:
		valids_2 = [emb for emb in embs_2 if not isinstance(emb, _EmbeddingError)]
		error_indices_2 = [
			i for i, emb in enumerate(embs_2) if isinstance(emb, _EmbeddingError)
		]
		d_mtx = amd.PDD_cdist(valids_1, valids_2, **kwargs)
		d_mtx = _insert_nan_rows_cols(d_mtx, error_indices_1, error_indices_2)
	return d_mtx


def _distance_matrix_d_amd(
	embs_1: list[TYPE_EMB_AMD],
	embs_2: list[TYPE_EMB_AMD] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_amd.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_AMD]): Embeddings
		embs_2 (list[TYPE_EMB_AMD] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Multiprocessing is
			not implemented for AMD; if True, a warning is issued and the computation
			proceeds without it. Default is False.
		n_processes (int | None): Ignored (multiprocessing not supported). Default is
			None.
		**kwargs: Additional arguments for amd.AMD_pdist and amd.AMD_cdist.

	Returns:
		np.ndarray: d_amd distance matrix.
	"""
	if multiprocessing:
		warnings.warn(
			"Multiprocessing is not implemented for _distance_matrix_d_amd. "
			"Proceeding without multiprocessing.",
			stacklevel=2,
		)

	given_two_sets = embs_2 is not None
	valids_1 = [emb for emb in embs_1 if not isinstance(emb, _EmbeddingError)]
	error_indices_1 = [
		i for i, emb in enumerate(embs_1) if isinstance(emb, _EmbeddingError)
	]
	if not given_two_sets:
		d_mtx = squareform(amd.AMD_pdist(valids_1, **kwargs))
		d_mtx = _insert_nan_rows_cols(d_mtx, error_indices_1, None)
	else:
		valids_2 = [emb for emb in embs_2 if not isinstance(emb, _EmbeddingError)]
		error_indices_2 = [
			i for i, emb in enumerate(embs_2) if isinstance(emb, _EmbeddingError)
		]
		d_mtx = amd.AMD_cdist(valids_1, valids_2, **kwargs)
		d_mtx = _insert_nan_rows_cols(d_mtx, error_indices_1, error_indices_2)
	return d_mtx


def _distance_matrix_d_elmd(
	embs_1: list[TYPE_EMB_ELMD],
	embs_2: list[TYPE_EMB_ELMD] | None = None,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	**kwargs,
) -> np.ndarray:
	"""Compute the distance matrix between two sets of embeddings based on d_elmd.

	If embs_2 is None, compute the distance matrix within embs_1.

	Args:
		embs_1 (list[TYPE_EMB_ELMD]): Embeddings
		embs_2 (list[TYPE_EMB_ELMD] | None): Embeddings or None. Default is None.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Additional arguments for ElMD.elmd.

	Returns:
		np.ndarray: d_elmd distance matrix.
	"""
	return _distance_matrix_template(
		embs_1,
		embs_2,
		multiprocessing,
		n_processes,
		functools.partial(_d_elmd, **kwargs),
	)


def _compute_embedding_worker(xtal: Crystal, distance: str, **kwargs) -> TYPE_EMB_ALL:
	"""Worker function for computing embedding in multiprocessing.

	Compute the embedding for a crystal based on the specified distance metric.

	Args:
		xtal (Crystal): A Crystal object.
		distance (str): The distance metric to use.
		**kwargs: Additional arguments for embedding methods if needed.

	Returns:
		TYPE_EMB_ALL | _EmbeddingError: The embedding, or an _EmbeddingError
			wrapping the exception if embedding computation failed.
	"""
	try:
		return xtal.get_embedding(distance, **kwargs)
	except Exception as e:
		return _EmbeddingError(e)


def _compute_embeddings(
	distance: str,
	xtals: Crystal | list[Crystal],
	multiprocessing: bool,
	n_processes: int | None = None,
	**kwargs,
) -> TYPE_EMB_ALL | list[TYPE_EMB_ALL]:
	"""Compute embedding(s) for given crystal(s) based on the specified distance metric.

	Args:
		distance (str): The distance metric to use.
		xtals (Crystal | list[Crystal]): A Crystal object or a list of Crystal objects.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		**kwargs: Additional arguments for embedding methods if needed.

	Returns:
		TYPE_EMB_ALL | list[TYPE_EMB_ALL]: A single embedding or a list of embeddings.
	"""
	single = False
	if isinstance(xtals, Crystal):
		xtals = [xtals]
		single = True

	if multiprocessing:
		_compute_embedding_worker_partial = functools.partial(
			_compute_embedding_worker, distance=distance, **kwargs
		)

		with Pool(processes=_set_n_processes(n_processes)) as pool:
			embs = pool.map(
				_compute_embedding_worker_partial,
				xtals,
			)
	else:
		embs = [_compute_embedding_worker(xtal, distance, **kwargs) for xtal in xtals]
	if single:
		return embs[0]
	else:
		return embs


# Registry: metric name → pairwise distance function
_DIST_FUNCS: dict[str, Callable] = {
	"smat": _d_smat,
	"comp": _d_comp,
	"wyckoff": _d_wyckoff,
	"magpie": _d_magpie,
	"pdd": _d_pdd,
	"amd": _d_amd,
	"elmd": _d_elmd,
}

# Registry: metric name → distance matrix function
_DIST_MTX_FUNCS: dict[str, Callable] = {
	"smat": _distance_matrix_d_smat,
	"comp": _distance_matrix_d_comp,
	"wyckoff": _distance_matrix_d_wyckoff,
	"magpie": _distance_matrix_d_magpie,
	"pdd": _distance_matrix_d_pdd,
	"amd": _distance_matrix_d_amd,
	"elmd": _distance_matrix_d_elmd,
}


@overload
def distance(
	distance: str,
	xtal_1: Structure | Crystal | TYPE_EMB_ALL,
	xtal_2: Structure | Crystal | TYPE_EMB_ALL,
	normalize: bool = ...,
	verbose: Literal[False] = ...,
	**kwargs,
) -> float: ...


@overload
def distance(
	distance: str,
	xtal_1: Structure | Crystal | TYPE_EMB_ALL,
	xtal_2: Structure | Crystal | TYPE_EMB_ALL,
	normalize: bool = ...,
	verbose: Literal[True] = ...,
	**kwargs,
) -> tuple[float, TYPE_EMB_ALL, TYPE_EMB_ALL]: ...


def distance(
	distance: str,
	xtal_1: Structure | Crystal | TYPE_EMB_ALL,
	xtal_2: Structure | Crystal | TYPE_EMB_ALL,
	normalize: bool = True,
	verbose: bool = False,
	**kwargs,
) -> float | tuple[float, TYPE_EMB_ALL, TYPE_EMB_ALL]:
	"""Compute the distance between two crystals.

	Args:
		distance (str): The distance metric to use. Currently supported metrics are
			listed in SUPPORTED_DISTANCES in constants.py. For more detailed information
			about each distance metric, please refer to the `tutorial notebook`_.
		xtal_1 (Structure | Crystal | TYPE_EMB_ALL): pymatgen Structure or
			Crystal or an embedding.
		xtal_2 (Structure | Crystal | TYPE_EMB_ALL): pymatgen Structure or
			Crystal or an embedding.
		normalize (bool): Whether to normalize the distance d to [0, 1] by using d'
			= d / (1 + d). This argument is only considered when d is a continuous
			distance that is not normalized to [0, 1]. Such distances are listed in
			CONTINUOUS_UNNORMALIZED_DISTANCES in constants.py. Default is True.
		verbose (bool): Whether to return intermediate embeddings. Default is False.
		**kwargs: Additional keyword arguments for specific distance metrics. It can
			contain two keys: "args_emb" and "args_dist". The value of "args_emb" is a
			dict of arguments for the calculation of embeddings, and the value of
			"args_dist" is a dict of arguments for the calculation of distance between
			the embeddings. If embeddings are pre-computed and provided as inputs, the
			"args_emb" will be ignored.


	Returns:
		float |  tuple[np.ndarray, TYPE_EMB_ALL, TYPE_EMB_ALL, dict[str, float]]:
		Distance between crystals. If verbose is True, also returns the embeddings and
		the computing time.

	Raises:
		ValueError: If an unsupported distance metric is provided.

	.. _tutorial notebook: https://github.com/WMD-group/xtalmet/blob/main/examples/tutorial.ipynb
	"""
	# conversions from Structure to Crystal
	xtal_1 = _to_crystal(xtal_1)
	xtal_2 = _to_crystal(xtal_2)

	# compute embeddings
	if distance not in DIST_WO_EMB and isinstance(xtal_1, Crystal):
		emb_1 = _compute_embeddings(
			distance, xtal_1, False, **(kwargs.get("args_emb", {}))
		)
	else:
		emb_1 = xtal_1
	if distance not in DIST_WO_EMB and isinstance(xtal_2, Crystal):
		emb_2 = _compute_embeddings(
			distance, xtal_2, False, **(kwargs.get("args_emb", {}))
		)
	else:
		emb_2 = xtal_2

	# compute distance
	if distance == "elmd+amd":
		d_elmd_unnorm = _d_elmd(
			emb_1[0], emb_2[0], **(kwargs.get("args_dist", {}).get("elmd", {}))
		)
		d_amd_unnorm = _d_amd(
			emb_1[1], emb_2[1], **(kwargs.get("args_dist", {}).get("amd", {}))
		)
		d = _blend_elmd_amd_distance(
			d_elmd_unnorm, d_amd_unnorm, kwargs.get("args_dist", {})
		)
	elif distance in _DIST_FUNCS:
		d = _DIST_FUNCS[distance](emb_1, emb_2, **(kwargs.get("args_dist", {})))
	else:
		raise ValueError(f"Unsupported distance metric: {distance}")

	# normalize distance if needed
	if distance in CONTINUOUS_UNNORMALIZED_DISTANCES and normalize:
		d = _normalize(d)

	# return results
	if not verbose:
		return d
	else:
		return d, emb_1, emb_2


TYPE_D_MTX_RETURN = (
	np.ndarray
	| tuple[np.ndarray, list[TYPE_EMB_ALL], dict[str, float]]
	| tuple[
		np.ndarray,
		list[TYPE_EMB_ALL],
		list[TYPE_EMB_ALL],
		dict[str, float],
	]
)  #: Return type for distance_matrix function.


def distance_matrix(
	distance: str,
	xtals_1: list[Structure | Crystal | TYPE_EMB_ALL],
	xtals_2: list[Structure | Crystal | TYPE_EMB_ALL] | None = None,
	normalize: bool = True,
	multiprocessing: bool = False,
	n_processes: int | None = None,
	verbose: bool = False,
	**kwargs,
) -> TYPE_D_MTX_RETURN:
	"""Compute the distance matrix between two sets of crystals.

	If xtals_2 is None, compute the distance matrix within xtals_1.

	Args:
		distance (str): The distance metric to use. Currently supported metrics are
			listed in SUPPORTED_DISTANCES in constants.py. For more detailed information
			about each distance metric, please refer to the `tutorial notebook`_.
		xtals_1 (list[Structure | Crystal | TYPE_EMB_ALL]): A list of pymatgen
			Structures or Crystals or embeddings.
		xtals_2 (list[Structure | Crystal | TYPE_EMB_ALL] | None): A list of
			pymatgen Structures or Crystals or embeddings, or None. Default is None.
		normalize (bool): Whether to normalize the distances d to [0, 1] by using d'
			= d / (1 + d). This argument is only considered when d is a continuous
			distance that is not normalized to [0, 1]. Such distances are listed in
			CONTINUOUS_UNNORMALIZED_DISTANCES in constants.py. Default is True.
		multiprocessing (bool): Whether to use multiprocessing. Default is False.
		n_processes (int | None): Maximum number of processes for multiprocessing. If
			multiprocessing is False, this argument will be ignored. Default is None.
		verbose (bool): Whether to return embeddings and the computing time. Default is
			False.
		**kwargs: Additional keyword arguments for specific distance metrics. It can
			contain two keys: "args_emb" and "args_dist". The value of "args_emb" is a
			dict of arguments for the calculation of embeddings, and the value of
			"args_dist" is a dict of arguments for the calculation of distance matrix
			using the embeddings. If embeddings are pre-computed and provided as inputs,
			the "args_emb" will be ignored.

	Returns:
		TYPE_D_MTX_RETURN: Distance matrix, the embeddings of xtals_1 (and xtals_2 if
		xtals_2 is not None) and the computing time.

	Raises:
		ValueError: If an unsupported distance metric is provided.

	.. _tutorial notebook: https://github.com/WMD-group/xtalmet/blob/main/examples/tutorial.ipynb
	"""
	given_two_sets = xtals_2 is not None
	times = {}

	# conversions from Structure to Crystal
	xtals_1 = _to_crystal_list(xtals_1)
	if given_two_sets:
		xtals_2 = _to_crystal_list(xtals_2)

	# compute embeddings
	if distance not in DIST_WO_EMB and isinstance(xtals_1[0], Crystal):
		emb_1_start = time.time()
		embs_1 = _compute_embeddings(
			distance,
			xtals_1,
			multiprocessing,
			n_processes,
			**(kwargs.get("args_emb", {})),
		)
		emb_1_end = time.time()
		times["emb_1"] = emb_1_end - emb_1_start
	else:
		embs_1 = xtals_1
		times["emb_1"] = 0.0
	if given_two_sets:
		if distance not in DIST_WO_EMB and isinstance(xtals_2[0], Crystal):
			emb_2_start = time.time()
			embs_2 = _compute_embeddings(
				distance,
				xtals_2,
				multiprocessing,
				n_processes,
				**(kwargs.get("args_emb", {})),
			)
			emb_2_end = time.time()
			times["emb_2"] = emb_2_end - emb_2_start
		else:
			embs_2 = xtals_2
			times["emb_2"] = 0.0
	else:
		embs_2 = None

	# compute distances
	d_mtx_start = time.time()
	if distance == "elmd+amd":
		d_mtx_elmd_unnorm = _distance_matrix_d_elmd(
			[emb[0] for emb in embs_1],
			[emb[0] for emb in embs_2] if embs_2 is not None else None,
			multiprocessing,
			n_processes,
			**(kwargs.get("args_dist", {}).get("elmd", {})),
		)
		d_mtx_amd_unnorm = _distance_matrix_d_amd(
			[emb[1] for emb in embs_1],
			[emb[1] for emb in embs_2] if embs_2 is not None else None,
			multiprocessing,
			n_processes,
			**(kwargs.get("args_dist", {}).get("amd", {})),
		)
		d_mtx = _blend_elmd_amd_distance(
			d_mtx_elmd_unnorm, d_mtx_amd_unnorm, kwargs.get("args_dist", {})
		)
	elif distance in _DIST_MTX_FUNCS:
		d_mtx = _DIST_MTX_FUNCS[distance](
			embs_1,
			embs_2,
			multiprocessing,
			n_processes,
			**(kwargs.get("args_dist", {})),
		)
	else:
		raise ValueError(f"Unsupported distance metric: {distance}")
	d_mtx_end = time.time()
	times["d_mtx"] = d_mtx_end - d_mtx_start

	# normalize distances if needed
	if distance in CONTINUOUS_UNNORMALIZED_DISTANCES and normalize:
		d_mtx = _normalize(d_mtx)

	# return results
	if not verbose:
		return d_mtx
	elif given_two_sets:
		return d_mtx, embs_1, embs_2, times
	else:
		return d_mtx, embs_1, times
