"""
Generic routines for generating deals based on a set of constraints
"""

from __future__ import annotations

__all__ = ["generate_deal", "generate_deals"]

import warnings
from collections.abc import Iterator
from typing import Optional, Union

from numpy.random import RandomState  # guaranteed to be stable for numpy>=1.16
from tqdm import trange  # type: ignore

from endplay.dealer.constraint import ConstraintInterpreter, Expr
from endplay.types import Card, Deal, Denom, Player, Rank


class DealNotGeneratedError(RuntimeError):
    pass


class DealNotGeneratedWarning(Warning):
    pass


class InconsistentSwappingAlgorithmWarning(Warning):
    pass


# This module could be better implemented as a C extension, with
# a function like the following assigning bits directly to the
# underlying datatype instead of constructing a list of Python
# objects, shuffling it and then slicing the list
#
# void complete_deal(uint16_t** deal)
# {
# 	// Initialize probabilities and flatten dealt cards
# 	uint16_t predeal[4] = { 0, 0, 0, 0 };
# 	int probs[4] = { 13, 26, 39, 52 };
# 	for (int hand = 0; hand < 4; ++hand) {
# 		for (int suit = 0; suit < 4; ++suit) {
# 			int count = popcount(deal[hand][suit]);
# 			predeal[suit] |= deal[hand][suit];
# 			for (int i = hand; i < 4; ++i)
# 				probs[i] -= count;
# 		}
# 	}
#
# 	// Deal remaining cards
# 	for (int rank = 0; rank < 13; ++rank) {
# 		uint16_t rank_bit = 1u << rank;
# 		for (int suit = 0; suit < 4; ++suit) {
# 			// Skip if card already dealt
# 			if (predeal[suit] & rank_bit)
# 				continue;
# 			// Pick a random number, and whichever range in probs
# 			// the number falls into we add the card to the corresponding
# 			// hand. We then decrement all elements in probs starting with
# 			// this hand. This reduces the probablility of a card being
# 			// dealt to the hand we just added a card too, and then adjusts
# 			// the other ranges to account for the fact that there is one
# 			// fewer card to deal.
# 			int k = rand() % probs[3];
# 			if (k < probs[0]) {
# 				deal[0][suit] |= rank_bit;
# 				--probs[0], --probs[1], --probs[2], --probs[3];
# 			}
# 			else if (k < probs[1]) {
# 				deal[1][suit] |= rank_bit;
# 				--probs[1], --probs[2], --probs[3];
# 			}
# 			else if (k < probs[2]) {
# 				deal[2][suit] |= rank_bit;
# 				--probs[2], --probs[3];
# 			}
# 			else {
# 				deal[3][suit] |= rank_bit;
# 				--probs[3];
# 			}
# 		}
# 	}
# }


def generate_deal(
    *constraints: Union[Expr, str],
    predeal: Deal = Deal(),
    swapping: int = 0,
    seed: Optional[int] = None,
    max_attempts: int = 1000000,
    env: dict = {},
) -> Deal:
    """
    Generates a random deal satisfying the constraints, giving 13 cards to each player.
    The constraints should be supplied as functions taking a whole deal and returning
    a boolean, for example `lambda d: hcp(d.north) > 10`, or as expressions compatible
    with the dealer expression syntax (see https://www.bridgebase.com/tools/dealer/Manual/input.html#expr)

    :param constraints: Constraints, as callables or strings
    :param predeal: A :class:`Deal` object which may be partially filled with cards; these will not
            be shuffled, allowing you to specify that players should have particular holdings.
    :param swapping: An integer representing the type of swapping algorithm to use, either
            0 (no swapping), 2 (swapping EW) or 3 (swapping EWS)
    :param seed: The number to seed the random generator with. A `numpy` random generator is
            used which is guaranteed to be stable between Python releases.
    :param max_attempts: Maximum number of shuffles to perform when finding a deal which
            matches the constraints. Set to -1 for infinite
    :param env: A dictionary of the environment used when evaluating constraints
    """
    # We just call generate deals and return the first deal that is yielded. The parameters
    # are the same and so are just forwarded onwards, with the exception that we always set
    # exhaustion_is_error to True (as it really *is* an error to not generate a single hand)
    # and we set show_progress to False as this would show all zeros right until the one hand
    # we want to produce is produced.
    deals = generate_deals(
        *constraints,
        predeal=predeal,
        swapping=swapping,
        show_progress=False,
        produce=1,
        seed=seed,
        max_attempts=max_attempts,
        env=env,
        strict=True,
    )
    return next(deals)


def generate_deals(
    *constraints: Union[Expr, str],
    predeal: Deal = Deal(),
    swapping: int = 0,
    show_progress: bool = False,
    produce: int = 40,
    seed: Optional[int] = None,
    max_attempts: int = 1000000,
    env: dict = {},
    strict: bool = False,
) -> Iterator[Deal]:
    """
    Generates `produce` random deals satisfying the constraints which should
    be given as for :func:`generate_deal`. `produce` and `max_attemps` are upper limits,
    the first to be reached terminates the program

    :param constraints: Constraints, as callables or strings
    :param predeal: A :class:`Deal` object which may be partially filled with cards; these will not
            be shuffled, allowing you to specify that players should have particular holdings.
    :param swapping: An integer representing the type of swapping algorithm to use, either
            0 (no swapping), 2 (swapping EW) or 3 (swapping EWS)
    :param show_progress: If True, a progress bar is displayed with information on how many hands
            have been generated so far
    :param produce: The total number of hands satisfying the constraints to find.
    :param seed: The number to seed the random generator with. A `numpy` random generator is
            used which is guaranteed to be stable between Python releases.
    :param max_attempts: Maximum number of shuffles to perform when finding a deal which
            matches the constraints. Set to -1 for infinite
    :param env: A dictionary of the environment used when evaluating constraints
    :param strict: If True, a `RuntimeError` is raised if `max_attempts` is reached before
            `produce` hands are produced. Otherwise, a warning is generated
    """
    if swapping == 2 and (len(predeal.west) > 0 or len(predeal.east) > 0):
        warnings.warn(
            "2-way swapping is incompatible with E/W predealt, output may be unexpected",
            InconsistentSwappingAlgorithmWarning,
        )
    elif swapping == 3 and (
        len(predeal.west) > 0 or len(predeal.east) > 0 or len(predeal.south) > 0
    ):
        warnings.warn(
            "3-way swapping is incompatible with E/W/S predealt, output may be unexpected",
            InconsistentSwappingAlgorithmWarning,
        )

    rs = RandomState(seed)

    ci = ConstraintInterpreter()
    for name, val in env.items():
        ci.set_env(name, val)
    compiled_constraints = tuple(
        ci.lambdify(c) if not callable(c) else c for c in constraints
    )
    all_cards = set(
        Card(suit=denom, rank=rank) for denom in Denom.suits() for rank in Rank
    )
    cards = list(all_cards.difference(predeal.to_hand()))
    split = [sum([13 - len(hand) for _, hand in predeal][:i]) for i in range(5)]
    generated = 0
    prange = trange(produce, desc="Produced", unit="deals", disable=not show_progress)
    for p in prange:
        if p > 0:
            prange.set_postfix({"success": f"{100*p/generated:.2f}%"})
        produced = False
        while not produced:
            if generated == max_attempts:
                prange.close()
                message = f"Only {p} out of {produce} hands were generated before max_attempts (set to {max_attempts}) was reached"
                if strict:
                    raise DealNotGeneratedError(message)
                else:
                    warnings.warn(message, DealNotGeneratedWarning)
                    return
            generated += 1
            rs.shuffle(cards)  # type: ignore
            deal = predeal.copy()
            for i, player in enumerate(Player):
                deal[player].extend(cards[split[i] : split[i + 1]])
            for perm in _generate_swaps(deal, swapping):
                if all(c(perm) for c in compiled_constraints):
                    yield perm
                    produced = True
                    break
    return


def _generate_swaps(deal: Deal, swapping: int):
    if swapping == 0:
        yield deal
    elif swapping == 2:
        yield deal
        deal.swap(Player.east, Player.west)
        yield deal
    elif swapping == 3:
        for _ in range(3):
            deal.swap(Player.east, Player.south)
            yield deal
            deal.swap(Player.south, Player.west)
            yield deal
    else:
        raise ValueError(f"Invalid swapping value {swapping} used")
