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,
    get_expected_value_finite, 
    compute_ev_win_lose_finite_tuple,
)
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", []
        if not dist:
            return False, "Distribution cannot be empty.", []
        return True, "", dist
    except ValueError:
        return False, "Distribution must be comma-separated integers (e.g., '0,1,2')", []

def validate_folded_players(folded_str: str, num_players: int) -> Tuple[bool, str, List[int]]:
    """Validate the folded players string."""
    try:
        folded = [int(x.strip()) for x in folded_str.split(',')]
        if len(folded) != num_players:
            return False, f"Folded players list must have {num_players} entries, but it has {len(folded)}.", []
        if any(x not in [0, 1] for x in folded):
            return False, "Folded players list can only contain 0s and 1s.", []
        return True, "", folded
    except ValueError:
        return False, "Folded players must be comma-separated integers (e.g., '0,0,1,0')", []

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, folded_players_str: 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 folded players
    valid_folded, error_msg, folded = validate_folded_players(folded_players_str, len(dist))
    if not valid_folded:
        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)
        dist_tuple = tuple(dist)
        folded_tuple = tuple(folded)

        # 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(dist_tuple, remaining, tier_map_tuple)

        result = "Unforced Expected Values:\n"
        for i, ev in enumerate(unforced_ev):
            player_status = "(Folded)" if folded[i] == 1 else ""
            result += f"Player {i+1}: {ev:.3f} {player_status}\n"

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

        result += "\nForced Win/Lose Results (for non-folded players):\n"
        for i,r in enumerate(win_lose_results):
            if folded_tuple[r['player']] == 0:
                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, folded_players_str: str, total_squids: int) -> str:
    """Calculate the Expected Value (EV) for each player in the finite variant using tuple-based state."""
    # Validate distribution
    valid_dist, error_msg, dist = validate_distribution(distribution)
    if not valid_dist:
        return error_msg

    # Validate folded players
    valid_folded, error_msg, folded = validate_folded_players(folded_players_str, len(dist))
    if not valid_folded:
        return error_msg

    # Validate finite variant rules
    if any(x > 1 for x in dist):
        return "Error: In the finite variant, players can only have 0 or 1 squid."

    # Validate total squids
    try:
        X = int(total_squids)
        current_squids = sum(dist)
        if X < 0:
            return "Total squids cannot be negative."
        if X > len(dist):
            return f"Error: Total squids ({X}) cannot exceed the number of players ({len(dist)})."
        if X < current_squids:
            return f"Total squids ({X}) cannot be less than current distribution sum ({current_squids})."
    except ValueError:
        return "Total squids must be an integer."

    try:
        get_expected_value_finite.cache_clear()
        
        remaining = X - sum(dist)
        dist_tuple = tuple(dist)
        folded_tuple = tuple(folded)
        
        from squid_game_core import _is_terminal_finite, _compute_final_payout_finite_tuple
        
        if _is_terminal_finite(dist_tuple, remaining):
            final_payoffs = _compute_final_payout_finite_tuple(dist_tuple)
            result = "Finite Squid Game Final Payouts:\n(Game has ended or no squids left)\n\n"
            for i, payoff in enumerate(final_payoffs):
                player_status = "(Folded)" if folded[i] == 1 else ""
                result += f"Player {i+1}: {payoff:.4f} {player_status}\n"
            return result

        base_ev, win_lose_results = compute_ev_win_lose_finite_tuple(dist_tuple, remaining, folded_tuple)
        
        result = "Finite Squid Game Expected Value (EV):\n"
        result += "(Payout Rule: Losers pay for all winners' squids)\n\n"
        
        result += "Baseline EV (fully random):\n"
        for i, ev in enumerate(base_ev):
            player_status = "(Folded)" if folded[i] == 1 else ""
            result += f"Player {i+1}: {ev:.4f} {player_status}\n"

        if win_lose_results:
            result += "\nForced Win/Lose EV (for players without a squid who haven't folded):\n"
            for i,r in enumerate(win_lose_results):
                if folded_tuple[r['player']] == 0:
                    result += (f"Player {r['player']+1}: forcedWinEV = {r['forcedWinEV']:.4f}, "
                           f"forcedLoseEV = {r['forcedLoseEV']:.4f}, Diff = {r['difference']:.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:1
5-6:1
7-7:1
8-8:1
9-9:1
10-100:1"""

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 5 players."""
            )
            folded_players_input = gr.Textbox(
                label="Folded Players (0 = Playing, 1 = Folded)",
                placeholder="0,0",
                value="0,0",
                info="""Enter 0 for playing, 1 for folded. Must match the number of players.
    Example: '0,0,1,0,0' means Player 3 has folded and won't receive any more 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)"""
            )
        
        with gr.Column():
            results_output = gr.Textbox(label="Results", lines=20)
            
            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", "0,0", 9, DEFAULT_TIER_MAP],
                    ["1,0,1", "0,0,0", 12, DEFAULT_TIER_MAP],
                    ["2,0,2,0", "0,1,0,0", 14, DEFAULT_TIER_MAP],
                ],
                inputs=[distribution_input, folded_players_input, total_squids_input, tier_map_input],
            )
    
    classic_btn.click(
        fn=solve_game,
        inputs=[distribution_input, folded_players_input, total_squids_input, tier_map_input],
        outputs=results_output
    )
    
    finite_btn.click(
        fn=solve_finite_game,
        inputs=[distribution_input, folded_players_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()