import random
import json

# ---------- PARAMETERS (change these) ----------
NUM_TRIALS = 10000                      # how many times to repeat the experiment per shuffle-count
SHUFFLE_COUNTS = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]  # test these riffle-shuffle counts
MIN_ALTERNATIONS = 10                  # min number of times the hands switch sides during one riffle
MAX_ALTERNATIONS = 20                  # max number of times the hands switch sides during one riffle

MIDDLE_FRACTION = 0.5                  # roughly what fraction of the deck is pulled out as the "heart" (middle) packet
FRACTION_WOBBLE = 0.05                 # how much that fraction is allowed to vary randomly each shuffle

# deck definitions: name -> (ordered "new deck" list, expected rising sequences if truly random)
DECKS = {
    "Standard": {
        "deck": list(range(52)),               # 52 cards, all unique, ace..king order
        "target_random_avg": 26.5,             # (n+1)/2 for n=52
    },
    "Skibbo": {
        "deck": [v for v in range(1, 13) for _ in range(5)],  # 60 cards, values 1-12, 5 copies each
        "target_random_avg": 30.5,              # (n+1)/2 for n=60
    },
}
# one output json per deck, named "<DeckName>_heartcut_shuffle_results.json"
# ------------------------------------------------


def split_pile(pile_size, num_segments):
    """Break pile_size cards into num_segments non-negative chunks (some can be 0),
    using random cut points so any distribution (incl. lumpy or empty) is possible."""
    if num_segments <= 1:
        return [pile_size]
    cuts = sorted(random.uniform(0, pile_size) for _ in range(num_segments - 1))
    cuts = [0] + cuts + [pile_size]
    return [round(cuts[i + 1]) - round(cuts[i]) for i in range(num_segments)]


def heart_cut(deck):
    """Pull the middle ~half of the deck out as one packet ('the heart').
    The leftover top + bottom portions are joined together as the other packet."""
    n = len(deck)
    fraction = MIDDLE_FRACTION + random.uniform(-FRACTION_WOBBLE, FRACTION_WOBBLE)
    fraction = max(0.05, min(0.95, fraction))  # keep it sane

    middle_size = round(n * fraction)
    remaining = n - middle_size
    top_size = remaining // 2 + random.randint(-2, 2)  # imperfect, not exactly centered
    top_size = max(0, min(remaining, top_size))
    bottom_size = remaining - top_size

    top = deck[:top_size]
    middle = deck[top_size:top_size + middle_size]
    bottom = deck[top_size + middle_size:]

    outer = top + bottom  # the two leftover ends, joined into one packet
    heart = middle

    return heart, outer


def clump_riffle_two_piles(pile_a, pile_b, min_alt=MIN_ALTERNATIONS, max_alt=MAX_ALTERNATIONS):
    """Same alternation-based riffle mixing used in the standard method,
    just taking two already-formed piles instead of cutting the deck in half."""
    num_alternations = random.randint(min_alt, max_alt)

    a_segments_count = (num_alternations + 1) // 2
    b_segments_count = num_alternations // 2
    if random.random() < 0.5:
        a_segments_count, b_segments_count = b_segments_count, a_segments_count

    a_sizes = split_pile(len(pile_a), max(a_segments_count, 1))
    b_sizes = split_pile(len(pile_b), max(b_segments_count, 1))

    result = []
    ai = bi = 0
    current = random.randint(0, 1)

    while ai < len(a_sizes) or bi < len(b_sizes):
        if current == 0 and ai < len(a_sizes):
            size = a_sizes[ai]
            result.extend(pile_a[:size])
            pile_a = pile_a[size:]
            ai += 1
        elif current == 1 and bi < len(b_sizes):
            size = b_sizes[bi]
            result.extend(pile_b[:size])
            pile_b = pile_b[size:]
            bi += 1
        current = 1 - current

    result.extend(pile_a)  # leftover due to rounding
    result.extend(pile_b)
    return result


def heart_cut_shuffle(deck):
    heart, outer = heart_cut(deck)
    return clump_riffle_two_piles(heart, outer)


def count_rising_sequences(deck):
    # counts runs of cards still in ascending original order
    count = 1
    for i in range(1, len(deck)):
        if deck[i] < deck[i - 1]:
            count += 1
    return count


def run_experiment_for_deck(deck_name, deck_template, target_random_avg):
    all_results = {}

    for num_shuffles in SHUFFLE_COUNTS:
        rising_seq_counts = []
        for _ in range(NUM_TRIALS):
            deck = list(deck_template)
            for _ in range(num_shuffles):
                deck = heart_cut_shuffle(deck)
            rising_seq_counts.append(count_rising_sequences(deck))

        avg = sum(rising_seq_counts) / len(rising_seq_counts)
        all_results[str(num_shuffles)] = {
            "num_shuffles": num_shuffles,
            "avg_rising_sequences": avg,
            "target_random_avg": target_random_avg,
            "raw_counts": rising_seq_counts,
        }

    settings = {
        "METHOD": "Heart Cut Shuffle",
        "NUM_TRIALS": NUM_TRIALS,
        "MIN_ALTERNATIONS": MIN_ALTERNATIONS,
        "MAX_ALTERNATIONS": MAX_ALTERNATIONS,
        "MIDDLE_FRACTION": MIDDLE_FRACTION,
        "FRACTION_WOBBLE": FRACTION_WOBBLE,
        "DECK": deck_name,
        "DECK_SIZE": len(deck_template),
    }

    output_file = f"{deck_name}_heartcut_shuffle_results.json"
    with open(output_file, "w") as f:
        json.dump({"settings": settings, "results": all_results}, f, indent=2)

    print(f"Done. Results written to {output_file}\n")
    for k, v in all_results.items():
        print(f"{v['num_shuffles']} shuffles -> avg rising sequences: {v['avg_rising_sequences']:.2f} "
              f"(target: {target_random_avg})")
    print()


def run_experiment():
    for deck_name, deck_info in DECKS.items():
        run_experiment_for_deck(deck_name, deck_info["deck"], deck_info["target_random_avg"])


if __name__ == "__main__":
    run_experiment()