import gradio as gr
from squid_game_core import (
    parse_tier_map,
    get_expected_value,
    compute_ev_win_lose_two_extremes,
    finite_squid_game_probabilities,
    dp_ev,
    infinite_squid_game_expected_counts_partial,
)
from typing import List, Tuple

def validate_distribution(dist_str: str) -> Tuple[bool, str, List[int]]:
    """Validate the distribution string and convert to list of integers"""
    try:
        dist = [int(x.strip()) for x in dist_str.split(',')]
        if any(x < 0 for x in dist):
            return False, "Distribution cannot contain negative numbers", []
        return True, "", dist
    except ValueError:
        return False, "Distribution must be comma-separated integers (e.g., '0,1,2')", []

def validate_tier_map(tier_str: str) -> Tuple[bool, str]:
    """Validate the tier map string format"""
    try:
        lines = tier_str.strip().splitlines()
        for line in lines:
            if ':' not in line:
                return False, "Each line must contain a colon (e.g., '1-2:1.5')"
            range_part, mult_part = line.split(':')
            float(mult_part.strip())  # Check multiplier is a valid number
            if '-' in range_part:
                low_str, high_str = range_part.split('-')
                int(low_str), int(high_str)
            else:
                int(range_part.strip())
        return True, ""
    except ValueError:
        return False, "Invalid format. Example: '1:1.0\\n2-4:2.0\\n5-6:3.0'"

def solve_game(distribution: str, total_squids: int, tier_map_str: str) -> str:
    """Main function to solve the game and return formatted results"""
    # Validate distribution
    valid_dist, error_msg, dist = validate_distribution(distribution)
    if not valid_dist:
        return error_msg

    # Validate tier map
    valid_tier, error_msg = validate_tier_map(tier_map_str)
    if not valid_tier:
        return error_msg

    # Validate total squids
    try:
        X = int(total_squids)
        if X < 0:
            return "Total squids cannot be negative"
        if X < sum(dist):
            return "Total squids cannot be less than current distribution sum"
    except ValueError:
        return "Total squids must be an integer"

    # Parse tier map and convert to tuple for caching
    try:
        tier_map = parse_tier_map(tier_map_str)
        tier_map_tuple = tuple((a, b, c) for a, b, c in tier_map)

        # Calculate remaining squids to distribute
        remaining = X - sum(dist)

        # Get unforced expected values (full random assignment)
        get_expected_value.cache_clear()
        unforced_ev = get_expected_value(tuple(dist), remaining, tier_map_tuple)

        result = "Unforced Expected Values:\n"
        for i, ev in enumerate(unforced_ev):
            result += f"Player {i+1}: {ev:.3f}\n"

        # Compute each player's forced win/lose EV extremes:
        win_lose_results = compute_ev_win_lose_two_extremes(tuple(dist), remaining, tier_map_tuple)

        result += "\nForced Win/Lose Results:\n"
        for r in win_lose_results:
            result += (f"Player {r['player']+1}: forcedWinEV = {r['forcedWinEV']:.3f}, "
                       f"forcedLoseEV = {r['forcedLoseEV']:.3f}, Diff = {r['difference']:.3f}\n")

        # Add a human-friendly interpretation of the tier map
        result += "\nTier Map Interpretation:\n"
        for low, high, mult in tier_map:
            if low == high:
                result += f"• {low} squid(s): multiplier = {mult:.1f}\n"
            else:
                result += f"• {low}-{high} squids: multiplier = {mult:.1f}\n"

        return result

    except Exception as e:
        return f"Error occurred: {str(e)}"

def solve_finite_game(distribution: str, total_squids: int) -> str:
    """Calculate the probability of each player getting a squid in the finite variant"""
    # Validate distribution
    valid_dist, error_msg, dist = validate_distribution(distribution)
    if not valid_dist:
        return error_msg

    # Validate that no player has more than 1 squid (finite variant rule)
    if any(x > 1 for x in dist):
        return "Error: In the finite variant, players can only have 0 or 1 squid. Please adjust your input."

    # Validate total squids
    try:
        X = int(total_squids)
        if X < 0:
            return "Total squids cannot be negative"
        
        # Check if total squids exceeds number of players
        if X > len(dist):
            return f"Error: In the finite variant, total squids cannot exceed the number of players ({len(dist)})"
            
        # Check if remaining squids + already distributed squids exceeds number of players
        current_sum = sum(dist)
        if current_sum + (X - current_sum) > len(dist):
            return f"Error: Total squids to distribute ({X}) plus already distributed squids ({current_sum}) cannot exceed the number of players ({len(dist)})"
    except ValueError:
        return "Total squids must be an integer"

    try:
        # Clear cache for new calculation
        dp_ev.cache_clear()
        
        # Calculate probabilities using the finite variant
        n = len(dist)  # Number of players
        
        # Calculate remaining squids to distribute
        remaining = X - sum(dist)
        
        # Create bitmask for players who already have squids
        bitmask_u = (1 << n) - 1  # Start with all 1s
        for i, val in enumerate(dist):
            if val == 1:  # If player already has a squid
                bitmask_u ^= (1 << i)  # Set bit to 0
        
        # Calculate probabilities
        probs = dp_ev(n, bitmask_u, remaining)
        
        result = "Finite Squid Game Probabilities:\n"
        result += "(Probability of each player getting a squid)\n\n"
        
        for i, prob in enumerate(probs):
            result += f"Player {i+1}: {prob:.4f}\n"
            
        return result
        
    except Exception as e:
        return f"Error occurred: {str(e)}"

def solve_infinite_game(distribution: str) -> str:
    """Calculate the expected final squid count for each player in the infinite variant"""
    # Validate distribution
    valid_dist, error_msg, dist = validate_distribution(distribution)
    if not valid_dist:
        return error_msg

    try:
        # Calculate expected final counts
        expected_counts = infinite_squid_game_expected_counts_partial(dist)
        
        result = "Infinite Squid Game Expected Final Counts:\n"
        result += "(Expected number of squids each player will have at the end)\n\n"
        
        for i, count in enumerate(expected_counts):
            if count == float('inf'):
                result += f"Player {i+1}: ∞ (infinite)\n"
            else:
                result += f"Player {i+1}: {count:.4f}\n"
        
        # Add explanation based on zero count
        zero_count = sum(1 for x in dist if x == 0)
        if zero_count == 1:
            result += "\nExplanation: Game ends immediately as there is exactly one player with 0 squids."
        elif zero_count == 0:
            result += "\nExplanation: Game never ends (infinite loop) as there are no players with 0 squids."
        else:
            H_z = sum(1.0 / k for k in range(1, zero_count+1))
            increment = (H_z - 1.0)
            result += f"\nExplanation: Each player is expected to receive {increment:.4f} additional squids before the game ends."
            
        return result
        
    except Exception as e:
        return f"Error occurred: {str(e)}"

# Default value for tier map used in interface
DEFAULT_TIER_MAP = """0-0:0
1-2:1
3-4:2
5-6:4
7-7:8
8-8:16
9-9:32
10-100:64"""

with gr.Blocks(title="Squid Game Calculator") as iface:
    gr.Markdown("""
    # Squid Game Expected Value Calculator
    
    Calculate the expected payoff for each player in the Squid Game.
    
    **Classic Variant Rules:**
    1. Players take turns collecting squids randomly.
    2. The game ends when either:
       - Exactly one player has 0 squids, OR
       - There are no squids left to distribute.
       
    **Finite Variant Rules:**
    1. Players take turns collecting squids randomly.
    2. Once a player gets a squid, they can't get another one.
    3. The game ends when all squids are distributed or only one player remains without a squid.
    
    **Infinite Variant Rules:**
    1. Players take turns collecting squids randomly (unlimited supply).
    2. Players accumulate squids over time.
    3. The game ends only when exactly one player has 0 squids.
    """)
    
    with gr.Row():
        with gr.Column():
            distribution_input = gr.Textbox(
                label="Players' Current Squids",
                placeholder="0,0",
                value="0,0",
                info="""Enter each player's current squids, separated by commas.
    Example: '1,0,1,2,0' represents:
      - Player 1 has 1 squid
      - Player 2 has 0 squids
      - Player 3 has 1 squid
      - Player 4 has 2 squids
      - Player 5 has 0 squids"""
            )
            total_squids_input = gr.Number(
                label="Total Squids in Game (Classic & Finite Variants Only)",
                value=9,
                minimum=0,
                step=1,
                precision=0,
                info="The total number of squids to be distributed (must be ≥ sum of current squids for classic variant)"
            )
            tier_map_input = gr.Textbox(
                label="Squid Value Tiers (Classic Variant Only)",
                placeholder=DEFAULT_TIER_MAP,
                value=DEFAULT_TIER_MAP,
                lines=8,
                info="""Define the value tiers for squids.
    Format: range:multiplier (one per line)
    Example:
    0-0:0
    1-2:1
    3-4:2
    5-6:4
    7-7:8
    8-8:16
    9-9:32
    10-100:64"""
            )
        
        with gr.Column():
            results_output = gr.Textbox(label="Results", lines=15)
            
            with gr.Row():
                classic_btn = gr.Button("Calculate Classic Variant", variant="primary")
                finite_btn = gr.Button("Calculate Finite Variant", variant="stop")
                infinite_btn = gr.Button("Calculate Infinite Variant", variant="secondary")
            
            gr.Examples(
                examples=[
                    ["0,0", 9, DEFAULT_TIER_MAP],
                    ["1,0,1", 12, DEFAULT_TIER_MAP],
                    ["2,0,2,0", 14, DEFAULT_TIER_MAP],
                ],
                inputs=[distribution_input, total_squids_input, tier_map_input],
            )
    
    classic_btn.click(
        fn=solve_game,
        inputs=[distribution_input, total_squids_input, tier_map_input],
        outputs=results_output
    )
    
    finite_btn.click(
        fn=solve_finite_game,
        inputs=[distribution_input, total_squids_input],
        outputs=results_output
    )
    
    infinite_btn.click(
        fn=solve_infinite_game,
        inputs=[distribution_input],
        outputs=results_output
    )

if __name__ == "__main__":
    iface.launch()