import numpy as np
import random
import matplotlib.pyplot as plt
from scipy.optimize import fsolve

# Function to compute the slope-zero point from the first code (for bet percentage calculation)
def compute_slope_zero_f(b, a, p, q):
    # Derivative of r with respect to f (using the chain rule)
    def derivative(f, b, a, p, q):
        term1 = p * (1 + f * b) ** (p - 1) * b
        term2 = q * (1 - f * a) ** (q - 1) * (-a)
        return term1 * (1 - f * a) ** q + (1 + f * b) ** p * term2

    # Use fsolve to find the value of f where the derivative is zero
    slope_zero_f = fsolve(derivative, 0.5, args=(b, a, p, q))[0]  # Initial guess: 0.5
    return round(slope_zero_f, 5)

# Function to simulate a coin toss game
def simulate_game(num_people, num_rounds, win_probability, lose_probability, win_percentage, lose_percentage, bet_percentage, starting_money):
    # Initialize everyone's starting money
    money = [starting_money] * num_people  # List to track each person's money

    # Create a matrix to track each person's money after each round
    money_over_time = np.zeros((num_people, num_rounds))

    # Simulate each round of the game for each person
    for round_num in range(num_rounds):
        for i in range(num_people):
            # Calculate how much money to bet for this round
            bet = money[i] * bet_percentage

            # Simulate the coin toss result (True = win, False = lose)
            toss_result = random.random() < win_probability  # Toss result based on win probability

            # Update money based on the result
            if toss_result:
                money[i] += bet * win_percentage  # Win
            else:
                money[i] -= bet * lose_percentage  # Lose

            # Store the money of this person after the current round
            money_over_time[i, round_num] = money[i]

    # Return the matrix of money over time
    return money_over_time

# Function to run the simulation and plot money over rounds
def run_simulation():
    # Get inputs from the user for the first set of code (game parameters)
    win_probability = float(input("Enter the chance of winning (between 0 and 1): "))  # Chance of winning
    win_percentage = float(input("Enter the percentage gained on a win (as a decimal): "))  # Percentage gained on win
    lose_percentage = float(input("Enter the percentage lost on a loss (as a decimal): "))  # Percentage lost on loss
    loss_probability = 1 - win_probability  # Probability of losing is 1 - probability of winning

    # Get inputs from the user for the second set of code (simulation parameters)
    num_people = int(input("Enter the number of players: "))  # Number of people
    num_rounds = int(input("Enter the number of rounds: "))  # Number of rounds each player plays
    starting_money = float(input("Enter the starting amount of money for each player: "))  # Starting amount for each player

    # Calculate the bet percentage from the first set (using the slope-zero point)
    b = win_percentage  # Gain from win
    a = lose_percentage  # Loss from loss
    p = win_probability  # Probability of winning
    q = loss_probability  # Probability of losing
    bet_percentage = compute_slope_zero_f(b, a, p, q)  # Bet percentage from the first set's slope-zero point

    print(f"Calculated bet percentage (fraction of current money to bet): {bet_percentage:.5f}")

    # Simulate the game
    money_over_time = simulate_game(num_people, num_rounds, win_probability, loss_probability, win_percentage, lose_percentage, bet_percentage, starting_money)

    # --- Plot 1: Return Rate and Slope Zero Point ---
    f_values = np.linspace(0, 10000, 1000000)
    r_values = np.round((1 + f_values * b) ** p * (1 - f_values * a) ** q, 5)

    # Create the first plot for return rate r(f)
    plt.figure(figsize=(12, 8))
    plt.plot(f_values, r_values, label=r'$r = ((1 + f \cdot b)^p) \cdot ((1 - f \cdot a)^q)$', color='blue')

    # Add the line where the slope is zero
    plt.axvline(x=bet_percentage, color='green', linestyle='--', label=f'Slope = 0 at f = {bet_percentage:.5f}')

    # Set the x-axis limits (from 0 to 1 as required)
    plt.xlim(0, bet_percentage + (bet_percentage / 10))

    # Dynamically set the y-axis limits based on the range of r_values
    plt.ylim(1, max(r_values) + (max(r_values) / 100))

    plt.title('Return Rate r vs Fraction of Bet f')
    plt.xlabel('Fraction of Bet (f)')
    plt.ylabel('Return Rate (r)')
    plt.grid(True)
    plt.legend()
    plt.show()

    # --- Plot 2: Money Over Time ---
    # Plot money over time for each player
    plt.figure(figsize=(12, 8))

    # Plot the money of the first few players (let's limit it to the first 100 players for clarity)
    num_to_plot = 100  # Number of players to plot for clarity
    for i in range(min(num_to_plot, num_people)):
        plt.plot(money_over_time[i, :], alpha=0.5)  # alpha for transparency to avoid overlap

    # Calculate and plot the average money across all players at each round
    average_money = np.mean(money_over_time, axis=0)
    plt.plot(average_money, label='Average Money', color='red', linewidth=2)

    # Calculate the average money at the end of the simulation
    average_money_final = np.mean(money_over_time[:, -1])  # Average money after all rounds

    # Set the Y-axis to 1.5 times the average final money
    plt.ylim(0, 1.5 * average_money_final)  # Dynamically adjust Y-axis based on average final money

    # Plot a horizontal line at the starting money level
    plt.axhline(starting_money, color='green', linestyle='--', label=f'Starting Money: ${starting_money}')

    plt.title('Money Over Time for Players')
    plt.xlabel('Rounds')
    plt.ylabel('Money ($)')
    plt.grid(True)
    plt.legend(loc='upper right', bbox_to_anchor=(1.1, 1.05), title="Average/Player/Starting Amount")
    plt.show()

    # Calculate and print some statistics
    players_above_avg = np.sum(money_over_time[:, -1] > average_money_final)  # Players with more money than average
    players_below_avg = np.sum(money_over_time[:, -1] < average_money_final)  # Players with less money than average

    # Calculate players above or below the starting amount (dynamic value)
    players_above_start = np.sum(money_over_time[:, -1] > starting_money)  # Players with more than starting money at the end
    players_below_start = np.sum(money_over_time[:, -1] < starting_money)  # Players with less than starting money at the end

    print("\nSimulation results:")
    print(f"Average money per player: ${average_money_final:.2f}")
    print(f"Number of players above the average: {players_above_avg}")
    print(f"Number of players below the average: {players_below_avg}")
    print(f"Number of players above the starting amount of ${starting_money}: {players_above_start}")
    print(f"Number of players below the starting amount of ${starting_money}: {players_below_start}")

    return money_over_time

# Run the simulation
final_money = run_simulation()