main.py

import numpy as np
import math
import itertools as it
import json
import os
import random
from tqdm import tqdm as ProgressDisplay
from scipy.stats import entropy
import multiprocessing
from functools import partial

# Constants
WRONG = np.uint8(0)  # ⚫
MISPLACED = np.uint8(1)  # ⚪
EXACT = np.uint8(2)  # 🟢

class NumberWordle:
    def __init__(self, n=4, positions=4, max_attempts=7):
        """
        Initialize the NumberWordle game
        n: The range of numbers (1 to n)
        positions: Number of positions in the code (default 4)
        max_attempts: Maximum number of attempts allowed (default 7)
        """
        self.n = n
        self.positions = positions
        self.max_attempts = max_attempts
        self.all_possibilities = self._generate_all_possibilities()
        self.answer = None
        self.reset_game()

    def _generate_all_possibilities(self):
        """Generate all possible combinations of n numbers in positions slots"""
        return list(it.permutations(range(1, self.n + 1), self.positions))

    def reset_game(self, answer=None):
        """Reset the game with a new answer"""
        self.guesses = []
        self.patterns = []
        self.possibilities = list(self.all_possibilities)
        if answer is None:
            self.answer = random.choice(self.possibilities)
        else:
            self.answer = answer

    def get_pattern(self, guess, answer=None):
        """
        Calculate the pattern between a guess and the answer
        Returns: (exact_count, misplaced_count, wrong_count)
        """
        if answer is None:
            answer = self.answer

        exact_count = 0
        misplaced_count = 0

        # Count exact matches
        for g, a in zip(guess, answer):
            if g == a:
                exact_count += 1

        # Count misplaced matches
        guess_remaining = [g for i, g in enumerate(guess) if guess[i] != answer[i]]
        answer_remaining = [a for i, a in enumerate(answer) if guess[i] != answer[i]]

        for g in guess_remaining:
            if g in answer_remaining:
                misplaced_count += 1
                answer_remaining.remove(g)

        wrong_count = self.positions - exact_count - misplaced_count
        return (exact_count, misplaced_count, wrong_count)

    def pattern_to_string(self, pattern):
        """Convert a pattern to a human-readable string"""
        exact, misplaced, wrong = pattern
        return f"🟢 {exact} | ⚪ {misplaced} | ⚫ {wrong}"

    def make_guess(self, guess):
        """Make a guess and update the game state"""
        pattern = self.get_pattern(guess)
        self.guesses.append(guess)
        self.patterns.append(pattern)

        # Update possibilities
        self.possibilities = [
            p for p in self.possibilities
            if self.get_pattern(guess, p) == pattern
        ]

        return pattern

    def is_won(self):
        """Check if the game is won"""
        return self.patterns and self.patterns[-1][0] == self.positions

    def is_lost(self):
        """Check if the game is lost"""
        return len(self.guesses) >= self.max_attempts and not self.is_won()


# Entropy calculation functions
def safe_log2(x):
    """Safely calculate log2, returning 0 for x=0"""
    return math.log2(x) if x > 0 else 0

def get_weights(possibilities):
    """Get uniform weights for all possibilities"""
    n = len(possibilities)
    if n == 0:
        return np.array([])
    return np.ones(n) / n


def get_pattern_distributions(allowed_guesses, possible_answers, positions=4):
    """
    Calculate the distribution of patterns for each possible guess
    Returns a matrix of shape (len(allowed_guesses), (positions+1)*(positions+1))
    """
    n_guesses = len(allowed_guesses)
    n_answers = len(possible_answers)

    # We have 3 possible states (EXACT, MISPLACED, WRONG) for each position
    # For 4 positions, we need to represent (exact, misplaced, wrong) where sum = positions
    # This gives us (positions+1)*(positions+1) possible patterns
    n_patterns = (positions + 1) * (positions + 1)  # Upper bound

    distributions = np.zeros((n_guesses, n_patterns))

    # Create a temporary game instance for pattern calculation
    game = NumberWordle(positions=positions)

    # For each guess and possible answer, calculate the pattern and update distribution
    for i, guess in enumerate(allowed_guesses):
        for j, answer in enumerate(possible_answers):
            pattern = game.get_pattern(guess, answer)
            # Convert pattern to a unique index
            pattern_idx = pattern[0] * (positions + 1) + pattern[1]
            distributions[i, pattern_idx] += 1 / n_answers

    return distributions

def entropy_of_distributions(distributions):
    """Calculate the entropy of each distribution"""
    return entropy(distributions, base=2, axis=1)


def get_entropies(allowed_guesses, possible_answers):
    """Calculate the entropy for each possible guess"""
    if len(possible_answers) == 0:
        return np.zeros(len(allowed_guesses))

    # Get the number of positions from the first possible answer
    positions = len(possible_answers[0]) if possible_answers else 4

    distributions = get_pattern_distributions(allowed_guesses, possible_answers, positions)
    return entropy_of_distributions(distributions)

def get_expected_scores(allowed_guesses, possible_answers):
    """
    Calculate the expected score for each guess
    Lower is better
    """
    n = len(possible_answers)
    if n == 0:
        return np.zeros(len(allowed_guesses))

    # Calculate probability that each guess is correct
    probs = np.array([
        1 / n if guess in possible_answers else 0
        for guess in allowed_guesses
    ])

    # Calculate entropy for each guess
    entropies = get_entropies(allowed_guesses, possible_answers)

    # Expected score = P(correct) * 1 + P(incorrect) * (1 + expected_future_score)
    # We estimate expected_future_score based on entropy reduction
    current_entropy = math.log2(n) if n > 0 else 0
    expected_future_scores = np.zeros(len(allowed_guesses))

    for i, ent in enumerate(entropies):
        entropy_reduction = current_entropy - ent
        # Estimate future score based on entropy reduction
        # More entropy reduction -> lower future score
        if entropy_reduction > 0:
            expected_future_scores[i] = 1 + math.log2(n) / entropy_reduction
        else:
            expected_future_scores[i] = 1 + math.log2(n)

    return probs + (1 - probs) * expected_future_scores

def optimal_guess(allowed_guesses, possible_answers):
    """Find the optimal guess that minimizes expected score"""
    if len(possible_answers) == 1:
        return possible_answers[0]

    expected_scores = get_expected_scores(allowed_guesses, possible_answers)
    return allowed_guesses[np.argmin(expected_scores)]

def get_possible_answers_after_guess(guess, pattern, possible_answers, game):
    """Get the list of possible answers after a guess and pattern"""
    return [
        answer for answer in possible_answers
        if game.get_pattern(guess, answer) == pattern
    ]

def simulate_game(game, first_guess=None, verbose=True):
    """Simulate a full game with optimal play"""
    game.reset_game()
    all_guesses = list(game.all_possibilities)

    if first_guess is None:
        # Choose the first guess that maximizes entropy
        entropies = get_entropies(all_guesses, game.possibilities)
        first_guess = all_guesses[np.argmax(entropies)]

    guess = first_guess
    attempts = 0

    while attempts < game.max_attempts:
        attempts += 1

        # Make the guess
        pattern = game.make_guess(guess)

        if verbose:
            print(f"Guess {attempts}: {guess} -> {game.pattern_to_string(pattern)}")
            print(f"Remaining possibilities: {len(game.possibilities)}")

        # Check if won
        if game.is_won():
            if verbose:
                print(f"Won in {attempts} attempts!")
            return attempts

        # Choose next guess based on updated possible answers
        if len(game.possibilities) > 0:
            if len(game.possibilities) == 1:
                # If only one possibility remains, guess it
                guess = game.possibilities[0]
            else:
                # Use entropy to select the best guess
                entropies = get_entropies(all_guesses, game.possibilities)
                guess = all_guesses[np.argmax(entropies)]

    if verbose:
        print("Lost! Maximum attempts reached.")
        print(f"The answer was: {game.answer}")

    return None




def analyze_next_guesses(game, top_n=5):
    """Analyze the top N next guesses based on entropy"""
    if len(game.possibilities) == 0:
        return []

    # If only one possibility remains, return it as the only recommendation
    if len(game.possibilities) == 1:
        return [{
            'guess': game.possibilities[0],
            'entropy': 0.0,
            'probability': 1.0,
            'expected_remaining': 1.0
        }]

    all_guesses = list(game.all_possibilities)
    entropies = get_entropies(all_guesses, game.possibilities)

    # Create a list of (index, entropy, is_possible) tuples
    guess_data = [(i, entropies[i], all_guesses[i] in game.possibilities)
                  for i in range(len(all_guesses))]

    # Sort by entropy (descending), then by whether it's a possible answer (True first)
    guess_data.sort(key=lambda x: (-x[1], not x[2]))

    results = []
    for i in range(min(top_n, len(all_guesses))):
        idx, entropy_value, is_possible = guess_data[i]
        guess = all_guesses[idx]

        # Calculate probability of being correct
        prob_correct = 1 / len(game.possibilities) if is_possible else 0

        results.append({
            'guess': guess,
            'entropy': entropy_value,
            'probability': prob_correct,
            'expected_remaining': 2 ** (math.log2(len(game.possibilities)) - entropy_value)
        })

    return results



def play_interactive_game(n=4, positions=4):
    """Play an interactive game where the user makes guesses"""
    game = NumberWordle(n=n, positions=positions)
    print(f"Welcome to NumberWordle! I'm thinking of {positions} different numbers from 1 to {n}.")
    print(f"You have {game.max_attempts} attempts to guess it.")
    print("After each guess, I'll tell you:")
    print("🟢 - How many numbers are correct and in the right position")
    print("⚪ - How many numbers are correct but in the wrong position")
    print("⚫ - How many numbers are incorrect")

    attempts = 0
    while attempts < game.max_attempts:
        # Show analysis of best next guesses
        if attempts > 0:
            print("\nTop recommended guesses:")
            recommendations = analyze_next_guesses(game, top_n=5)
            for i, rec in enumerate(recommendations):
                guess_str = ', '.join(map(str, rec['guess']))
                print(
                    f"{i + 1}. {guess_str} - Entropy: {rec['entropy']:.2f}, Expected remaining: {rec['expected_remaining']:.1f}")

        # Get user guess
        valid_guess = False
        while not valid_guess:
            try:
                guess_input = input(f"\nGuess {attempts + 1}/{game.max_attempts}: ")

                # Handle both comma-separated and non-comma-separated inputs
                if ',' in guess_input:
                    guess = tuple(map(int, guess_input.split(',')))
                else:
                    # For inputs like "1234", convert to (1,2,3,4)
                    guess = tuple(int(digit) for digit in guess_input if digit.isdigit())

                if len(guess) != positions:
                    print(f"Please enter exactly {positions} numbers separated by commas.")
                    continue

                if any(num < 1 or num > n for num in guess):
                    print(f"All numbers must be between 1 and {n}.")
                    continue

                if len(set(guess)) != positions:
                    print("All numbers must be different.")
                    continue

                valid_guess = True
            except ValueError:
                print("Invalid input. Please enter numbers separated by commas (e.g., 1,2,3,4)")

        # Make the guess
        attempts += 1
        pattern = game.make_guess(guess)

        print(f"Result: {game.pattern_to_string(pattern)}")
        print(f"Remaining possibilities: {len(game.possibilities)}")

        # Check if won
        if game.is_won():
            print(f"Congratulations! You won in {attempts} attempts!")
            return True

    print("Sorry, you've used all your attempts.")
    print(f"The answer was: {', '.join(map(str, game.answer))}")
    return False


def simulate_single_game(n, positions, max_attempts, _):
    """Simulate a single game for parallel processing"""
    game = NumberWordle(n=n, positions=positions, max_attempts=max_attempts)
    game.reset_game()
    score = simulate_game(game, verbose=False)
    return score


def evaluate_strategy(n=4, positions=4, num_games=100, num_processes=None):
    """Evaluate the performance of the optimal strategy using multiprocessing"""
    # Use all available CPUs by default
    if num_processes is None:
        num_processes = multiprocessing.cpu_count()

    # Create a partial function with fixed parameters
    sim_func = partial(simulate_single_game, n, positions, 7)  # 7 is default max_attempts

    # Create a pool of workers
    with multiprocessing.Pool(processes=num_processes) as pool:
        # Map the simulation function to the range of games
        scores = list(ProgressDisplay(
            pool.imap(sim_func, range(num_games)),
            total=num_games,
            desc=f"Simulating games using {num_processes} processes"
        ))

    # Filter out None values (losses)
    valid_scores = [s for s in scores if s is not None]
    win_rate = len(valid_scores) / num_games

    if valid_scores:
        avg_score = sum(valid_scores) / len(valid_scores)
        # Count occurrences of each score
        score_distribution = [(i, scores.count(i)) for i in range(1, 8)]  # 7 is default max_attempts
    else:
        avg_score = 0
        score_distribution = []

    print(f"Results for n={n}, positions={positions}:")
    print(f"Win rate: {win_rate:.2%}")
    print(f"Average score: {avg_score:.2f}")
    print(f"Score distribution: {score_distribution}")

    return {
        'win_rate': win_rate,
        'average_score': avg_score,
        'score_distribution': score_distribution
    }


def helper_mode(n=4, positions=4):
    """Help the user solve an external NumberWordle game"""
    game = NumberWordle(n=n, positions=positions)
    # We don't use the game's answer since the user has their own
    print(f"NumberWordle Helper Mode - I'll help you solve a {positions}-number game (1-{n})")
    print("After each guess, enter the result as three numbers:")
    print("- Number of 🟢 (correct position)")
    print("- Number of ⚪ (correct number, wrong position)")
    print("- Number of ⚫ (incorrect numbers)")

    # Initialize with all possibilities
    game.possibilities = list(game.all_possibilities)

    attempts = 0
    while attempts < game.max_attempts:
        # Show analysis of best next guesses
        if attempts > 0:
            print(f"\nRemaining possibilities: {len(game.possibilities)}")
            if len(game.possibilities) <= 10:
                print("Possible answers:")
                for i, possibility in enumerate(game.possibilities):
                    print(f"  {i+1}. {', '.join(map(str, possibility))}")

        print("\nTop recommended guesses:")
        recommendations = analyze_next_guesses(game, top_n=5)
        for i, rec in enumerate(recommendations):
            guess_str = ', '.join(map(str, rec['guess']))
            print(f"{i+1}. {guess_str} - Entropy: {rec['entropy']:.2f}, Expected remaining: {rec['expected_remaining']:.1f}")

        # Get user guess
        valid_guess = False
        while not valid_guess:
            try:
                guess_input = input(f"\nEnter your guess {attempts+1}/{game.max_attempts}: ")

                # Handle both comma-separated and non-comma-separated inputs
                if ',' in guess_input:
                    guess = tuple(map(int, guess_input.split(',')))
                else:
                    # For inputs like "1234", convert to (1,2,3,4)
                    guess = tuple(int(digit) for digit in guess_input if digit.isdigit())

                if len(guess) != positions:
                    print(f"Please enter exactly {positions} numbers.")
                    continue

                if any(num < 1 or num > n for num in guess):
                    print(f"All numbers must be between 1 and {n}.")
                    continue

                if len(set(guess)) != positions:
                    print("All numbers must be different.")
                    continue

                valid_guess = True
            except ValueError:
                print("Invalid input. Please enter numbers (e.g., 1,2,3,4)")

        # Get result from user
        valid_result = False
        while not valid_result:
            try:
                result_input = input("Enter result (🟢,⚪,⚫): ")

                # Parse result
                if ',' in result_input:
                    pattern = tuple(map(int, result_input.split(',')))
                else:
                    # For inputs like "210", convert to (2,1,0)
                    pattern = tuple(int(digit) for digit in result_input if digit.isdigit())

                if len(pattern) != 3:
                    print("Please enter exactly 3 numbers for the result.")
                    continue

                if sum(pattern) != positions:
                    print(f"The sum of the result should be {positions}.")
                    continue

                valid_result = True
            except ValueError:
                print("Invalid input. Please enter three numbers (e.g., 2,1,1)")

        # Update game state
        attempts += 1
        game.guesses.append(guess)
        game.patterns.append(pattern)

        # Update possibilities
        game.possibilities = [
            p for p in game.possibilities
            if game.get_pattern(guess, p) == pattern
        ]

        # Check if won
        if pattern[0] == positions:
            print(f"Congratulations! You won in {attempts} attempts!")
            return True

        # Check if no possibilities remain
        if len(game.possibilities) == 0:
            print("Error: No possibilities match your inputs. Check for mistakes.")
            return False

    print("Sorry, you've used all your attempts.")
    if len(game.possibilities) == 1:
        print(f"The answer was probably: {', '.join(map(str, game.possibilities[0]))}")
    return False

# Main function to run the game
if __name__ == "__main__":
    while True:
        print("\nNumberWordle - A Wordle-like game with numbers")
        print("Options:")
        print("1. Play interactive game")
        print("2. Simulate a game with optimal play")
        print("3. Evaluate strategy performance")
        print("4. Helper mode (for external games)")

        choice = input("Enter your choice (1-4): ")

        if choice == "1" or choice == "4":
            difficulty = input("Choose difficulty (easy=4, medium=5, hard=6, expert=7): ")
            n = int(difficulty) if difficulty.isdigit() and 4 <= int(difficulty) <= 7 else 4

            # Keep playing the same mode until user quits
            play_again = True
            while play_again:
                if choice == "1":
                    play_interactive_game(n=n)
                else:  # choice == "4"
                    helper_mode(n=n)

                play_again_input = input("\nPlay again with same settings? (y/n): ").lower()
                play_again = play_again_input.startswith('y')

        elif choice == "2":
            difficulty = input("Choose difficulty (easy=4, medium=5, hard=6, expert=7): ")
            n = int(difficulty) if difficulty.isdigit() and 4 <= int(difficulty) <= 7 else 4

            # Keep playing the same mode until user quits
            play_again = True
            while play_again:
                game = NumberWordle(n=n)
                simulate_game(game)

                play_again_input = input("\nSimulate another game with same settings? (y/n): ").lower()
                play_again = play_again_input.startswith('y')

        elif choice == "3":
            difficulty = input("Choose difficulty (easy=4, medium=5, hard=6, expert=7): ")
            n = int(difficulty) if difficulty.isdigit() and 4 <= int(difficulty) <= 7 else 4
            num_games = int(input("Number of games to simulate: ") or "100")
            evaluate_strategy(n=n, num_games=num_games)

        else:
            print("Invalid choice")

        continue_program = input("\nReturn to main menu? (y/n): ").lower()
        if not continue_program.startswith('y'):
            break