import numpy as np
import matplotlib.pyplot as plt

def calculate_position_size(account_equity: float, fraction_risked: float, stop_loss_distance: float) -> float:
    """
    Calculates the position size (number of units) using fixed fractional sizing.

    Inputs:
    - account_equity (float): The current total value of the trading account.
    - fraction_risked (float): The fixed percentage of the account equity to risk on one trade (e.g., 0.01 for 1%).
    - stop_loss_distance (float): The maximum potential loss per unit for the trade.

    Outputs:
    - float: The calculated number of units (position size) to trade.

    Formula Used:
    Position Size = (Account Equity * Fraction Risked) / Stop Loss Distance

    Note: The output is typically rounded down to the nearest whole number for practical trading,
    but for simulation purposes, we might keep it as a float to show the theoretical impact.
    """
    if stop_loss_distance <= 0:
        raise ValueError("Stop loss distance must be positive.")
    if not (0 < fraction_risked < 1):
        raise ValueError("Fraction risked must be between 0 and 1.")
    if account_equity <= 0:
        raise ValueError("Account equity must be positive.")

    dollar_risk_per_trade = account_equity * fraction_risked
    position_size = dollar_risk_per_trade / stop_loss_distance
    return position_size

# Scenario 1: Initial account, moderate risk
initial_equity_1 = 10000.0
fraction_risked_1 = 0.01  # 1% risk
stop_loss_distance_1 = 0.50 # $0.50 per unit
position_size_1 = calculate_position_size(initial_equity_1, fraction_risked_1, stop_loss_distance_1)
print(f"Scenario 1 (Equity: ${initial_equity_1:.2f}, Risk: {fraction_risked_1*100:.0f}%, SL: ${stop_loss_distance_1:.2f}): Position Size = {position_size_1:.2f} units")

# Scenario 2: Account grows, same risk settings
grown_equity_2 = 15000.0
fraction_risked_2 = 0.01
stop_loss_distance_2 = 0.50
position_size_2 = calculate_position_size(grown_equity_2, fraction_risked_2, stop_loss_distance_2)
print(f"Scenario 2 (Equity: ${grown_equity_2:.2f}, Risk: {fraction_risked_2*100:.0f}%, SL: ${stop_loss_distance_2:.2f}): Position Size = {position_size_2:.2f} units")

# Scenario 3: Account shrinks, same risk settings
shrunk_equity_3 = 8000.0
fraction_risked_3 = 0.01
stop_loss_distance_3 = 0.50
position_size_3 = calculate_position_size(shrunk_equity_3, fraction_risked_3, stop_loss_distance_3)
print(f"Scenario 3 (Equity: ${shrunk_equity_3:.2f}, Risk: {fraction_risked_3*100:.0f}%, SL: ${stop_loss_distance_3:.2f}): Position Size = {position_size_3:.2f} units")

# Scenario 4: Same equity, larger stop loss distance (less units)
initial_equity_4 = 10000.0
fraction_risked_4 = 0.01
stop_loss_distance_4 = 1.00 # $1.00 per unit
position_size_4 = calculate_position_size(initial_equity_4, fraction_risked_4, stop_loss_distance_4)
print(f"Scenario 4 (Equity: ${initial_equity_4:.2f}, Risk: {fraction_risked_4*100:.0f}%, SL: ${stop_loss_distance_4:.2f}): Position Size = {position_size_4:.2f} units")

def simulate_trading(initial_equity: float, fraction_risked: float, num_trades: int, win_rate: float, avg_r_multiple: float, fixed_units: int = None):
    """
    Simulates a series of trades and tracks account equity using fixed fractional sizing or fixed unit sizing.

    Inputs:
    - initial_equity (float): Starting capital.
    - fraction_risked (float): Percentage of equity risked per trade (for fixed fractional).
    - num_trades (int): Total number of trades to simulate.
    - win_rate (float): Probability of a trade being a win (e.g., 0.5 for 50%).
    - avg_r_multiple (float): Average reward-to-risk ratio (e.g., 2.0 means wins are 2x losses).
    - fixed_units (int, optional): If provided, simulates with a fixed number of units per trade instead of fixed fractional.

    Outputs:
    - list: A list containing the account equity after each trade.
    """
    equity_history = [initial_equity]
    current_equity = initial_equity

    # Assume a constant stop-loss distance and profit per unit for simplicity in simulation
    # For simplicity, let's assume a 'virtual' stop_loss_distance of 1 unit.
    # This means 'dollar_risk_per_unit' is effectively 1.
    # The actual profit/loss per trade will then be 'position_size * 1' for loss,
    # and 'position_size * avg_r_multiple * 1' for win.

    for _ in range(num_trades):
        if current_equity <= 0: # Account busted
            equity_history.append(0.0)
            continue

        if fixed_units is None:
            # Fixed Fractional Sizing
            # We use a conceptual stop_loss_distance=1 for calculation of position size
            # The actual dollar risk per trade is equity * fraction_risked
            # So, position_size = (equity * fraction_risked) / 1
            # loss_per_trade = position_size * 1
            # win_per_trade = position_size * avg_r_multiple * 1
            position_size = calculate_position_size(current_equity, fraction_risked, stop_loss_distance=1.0) # Use 1.0 as a conceptual unit risk
            # Ensure position size is at least 1 for any trade to occur, or 0 if equity is too low
            position_size = max(1, int(position_size)) if position_size >=1 else 0
            if position_size == 0:
                equity_history.append(current_equity)
                continue

            dollar_risk_per_trade = position_size * 1.0 # Conceptual risk per unit is 1.0
            dollar_profit_per_trade = position_size * avg_r_multiple * 1.0
        else:
            # Fixed Unit Sizing
            position_size = fixed_units
            dollar_risk_per_trade = position_size * 1.0 # Assuming $1 loss per unit for simplicity
            dollar_profit_per_trade = position_size * avg_r_multiple * 1.0

        is_win = np.random.rand() < win_rate

        if is_win:
            current_equity += dollar_profit_per_trade
        else:
            current_equity -= dollar_risk_per_trade

        equity_history.append(current_equity)

    return equity_history

initial_capital = 10000.0
risk_fraction = 0.01 # 1% risk per trade
number_of_trades = 500
probability_of_win = 0.5  # 50% win rate
reward_to_risk = 1.5    # 1.5R trades (win $1.5 for every $1 risked)
fixed_units_per_trade = 100 # For comparison with fixed unit sizing

equity_fixed_fractional = simulate_trading(initial_capital, risk_fraction, number_of_trades, probability_of_win, reward_to_risk)
equity_fixed_units = simulate_trading(initial_capital, 0, number_of_trades, probability_of_win, reward_to_risk, fixed_units=fixed_units_per_trade)

print(f"Final Equity (Fixed Fractional): ${equity_fixed_fractional[-1]:.2f}")
print(f"Final Equity (Fixed Units): ${equity_fixed_units[-1]:.2f}")

plt.figure(figsize=(14, 7))
plt.plot(equity_fixed_fractional, label='Fixed Fractional Sizing', color='green', linewidth=2)
plt.plot(equity_fixed_units, label=f'Fixed {fixed_units_per_trade} Units Sizing', color='red', linestyle='--', linewidth=1.5)
plt.title('Portfolio Growth: Fixed Fractional vs. Fixed Unit Sizing')
plt.xlabel('Number of Trades')
plt.ylabel('Account Equity ($)')
plt.grid(True, linestyle='--', alpha=0.7)
plt.axhline(y=initial_capital, color='gray', linestyle=':', label='Initial Capital')
plt.legend()
plt.tight_layout()
plt.show()

risk_fractions = [0.005, 0.01, 0.02, 0.03]
# Using the same initial_capital, num_trades, probability_of_win, reward_to_risk

equity_curves_by_fraction = {}
for fraction in risk_fractions:
    equity_curves_by_fraction[fraction] = simulate_trading(initial_capital, fraction, number_of_trades, probability_of_win, reward_to_risk)

# Print final equities for comparison
for fraction, curve in equity_curves_by_fraction.items():
    print(f"Final Equity (Fraction: {fraction*100:.1f}%): ${curve[-1]:.2f}")

plt.figure(figsize=(14, 7))
colors = ['blue', 'green', 'orange', 'red']

for i, (fraction, curve) in enumerate(equity_curves_by_fraction.items()):
    plt.plot(curve, label=f'Fraction Risked: {fraction*100:.1f}%', color=colors[i], linewidth=1.5)

plt.title('Portfolio Growth with Different Fixed Risk Fractions')
plt.xlabel('Number of Trades')
plt.ylabel('Account Equity ($)')
plt.grid(True, linestyle='--', alpha=0.7)
plt.axhline(y=initial_capital, color='gray', linestyle=':', label='Initial Capital')
plt.legend()
plt.tight_layout()
plt.show()