###############################################################################
# __   _            _____    _____
# | \ | |          / ____|  / ____|
# |  \| |  _   _  | |      | (___
# | . ` | | | | | | |       \___ \
# | |\  | | |_| | | |____   ____) |
# |_| \_|  \__,_|  \_____| |_____/
#
# Fast constraint solving in Python  - https://github.com/yangeorget/nucs
#
# Copyright 2024-2026 - Yan Georget
###############################################################################
import sys

from numba import njit  # type: ignore
from numpy.typing import NDArray

from nucs.constants import MAX, MIN
from nucs.heuristics.max_regret_var_heuristic import regret


@njit(cache=True)
def tsp_var_heuristic(decision_variables: NDArray, domains_stk: NDArray, stks_top: NDArray, params: NDArray) -> int:
    """
    :param decision_variables: the decision variables
    :param domains_stk: the stack of domains
    :param stks_top: the index of the top of the stacks as a Numpy array
    :param params: a two-dimensional (first dimension correspond to variables, second to values) cost array
    :return: the variable
    """
    best_score = -sys.maxsize
    best_variable = -1
    top = stks_top[0]
    for variable in decision_variables:
        domain = domains_stk[top, variable]
        if 0 < domain[MAX] - domain[MIN]:
            score = compute_score(domain, variable, params)
            if best_score < score:
                best_variable = variable
                best_score = score
    return best_variable


@njit(cache=True)
def compute_score(domain: NDArray, variable: int, params: NDArray) -> int:
    """
    Minimize [min(12, size(X)), -regret(X)] for lexicographic order.
    """
    size = min(12, domain[MAX] - domain[MIN] + 1)
    return -size * 1024 + regret(domain, params[variable])
