import numpy as np
import matplotlib.pyplot as plt


def calculate_position_size(
    total_capital: float,
    risk_per_trade_percent: float,
    entry_price: float,
    stop_loss_price: float,
) -> int:
    """
    Calculates the number of units (position size) for a trade based on risk parameters.

    Inputs:
    - total_capital (float): The total capital available in the trading account.
    - risk_per_trade_percent (float): The percentage of total_capital to risk on this trade (e.g., 0.01 for 1%).
    - entry_price (float): The price at which the trade is entered.
    - stop_loss_price (float): The price at which the trade will be exited to limit losses.

    Outputs:
    - int: The calculated number of units for the position, rounded down to the nearest whole number.

    Formula:
    1. Risk_Amount = total_capital * risk_per_trade_percent
    2. Risk_Per_Unit = abs(entry_price - stop_loss_price)
    3. Position_Size = Risk_Amount / Risk_Per_Unit

    Returns 0 if Risk_Per_Unit is zero to prevent division by zero.
    """
    if risk_per_trade_percent <= 0 or entry_price <= 0 or stop_loss_price <= 0:
        raise ValueError(
            "Risk percentage, entry price, and stop loss price must be positive."
        )
    if total_capital <= 0:
        raise ValueError("Total capital must be positive.")

    risk_amount = total_capital * risk_per_trade_percent
    risk_per_unit = abs(entry_price - stop_loss_price)

    if risk_per_unit == 0:
        print("Warning: Stop-loss price is equal to entry price. Cannot calculate position size.")
        return 0

    position_size = risk_amount / risk_per_unit
    return int(np.floor(position_size))  # Round down to trade whole units


def calculate_portfolio_heat(
    total_capital: float,
    positions: dict,
    current_prices: dict,
    stop_loss_prices: dict,
) -> float:
    """
    Calculates the total "heat" (potential percentage loss) of the entire portfolio
    if all current stop-losses are hit.

    Inputs:
    - total_capital (float): The total capital available in the trading account.
    - positions (dict): A dictionary where keys are asset symbols and values are the
                        number of units held for each asset.
    - current_prices (dict): A dictionary where keys are asset symbols and values are
                             their current market prices.
    - stop_loss_prices (dict): A dictionary where keys are asset symbols and values are
                               their stop-loss prices.

    Outputs:
    - float: The total potential loss as a percentage of total_capital.

    Formula:
    1. For each position: Potential_Loss_Per_Position = units * abs(current_price - stop_loss_price)
    2. Total_Potential_Loss = sum(Potential_Loss_Per_Position for all positions)
    3. Portfolio_Heat_Percentage = (Total_Potential_Loss / total_capital) * 100
    """
    if total_capital <= 0:
        raise ValueError("Total capital must be positive.")

    total_potential_loss = 0.0

    for asset, units in positions.items():
        if units == 0:
            continue

        current_price = current_prices.get(asset)
        stop_loss_price = stop_loss_prices.get(asset)

        if current_price is None or stop_loss_price is None:
            print(f"Warning: Price or stop-loss missing for asset {asset}. Skipping.")
            continue

        potential_loss_per_position = units * abs(current_price - stop_loss_price)
        total_potential_loss += potential_loss_per_position

    portfolio_heat_percentage = (total_potential_loss / total_capital) * 100
    return portfolio_heat_percentage


# --- Demo Helper Function ---
def demo_calculation(
    capital,
    risk_pct,
    entry,
    stop_loss,
    asset_name="Asset",
    show_details=True,
):
    """
    Helper function to demonstrate position size calculation and print details.
    """
    try:
        pos_size = calculate_position_size(capital, risk_pct, entry, stop_loss)
        if show_details:
            print(f"--- {asset_name} Details ---")
            print(f"  Total Capital: ${capital:,.2f}")
            print(f"  Risk per Trade: {risk_pct*100:.2f}%")
            print(f"  Risk Amount: ${capital * risk_pct:,.2f}")
            print(f"  Entry Price: ${entry:,.2f}")
            print(f"  Stop-Loss Price: ${stop_loss:,.2f}")
            print(f"  Risk per Unit: ${abs(entry - stop_loss):,.2f}")
            print(f"  Calculated Position Size: {pos_size} units\n")
        return pos_size
    except ValueError as e:
        print(f"Error calculating for {asset_name}: {e}")
        return 0


print("Functions for position sizing and portfolio heat calculation defined.")

# Define mock portfolio parameters
total_capital = 100_000  # $100,000
risk_per_trade_percent = 0.01  # 1% of total capital per trade

print(f"Total Capital: ${total_capital:,.2f}")
print(f"Risk Amount per Trade: ${total_capital * risk_per_trade_percent:,.2f}\n")

# --- Asset 1: Low Volatility Stock ---
asset1_entry = 100.00
asset1_stop_loss = 98.00  # 2% stop-loss (2 units of risk per share)
asset1_pos_size = demo_calculation(
    total_capital,
    risk_per_trade_percent,
    asset1_entry,
    asset1_stop_loss,
    asset_name="Asset A (Low Volatility Stock)",
)

# --- Asset 2: Medium Volatility Stock ---
asset2_entry = 50.00
asset2_stop_loss = 48.00  # 4% stop-loss (2 units of risk per share, but lower price)
asset2_pos_size = demo_calculation(
    total_capital,
    risk_per_trade_percent,
    asset2_entry,
    asset2_stop_loss,
    asset_name="Asset B (Medium Volatility Stock)",
)

# --- Asset 3: High Volatility Stock ---
asset3_entry = 200.00
asset3_stop_loss = 190.00  # 5% stop-loss (10 units of risk per share)
asset3_pos_size = demo_calculation(
    total_capital,
    risk_per_trade_percent,
    asset3_entry,
    asset3_stop_loss,
    asset_name="Asset C (High Volatility Stock)",
)

# --- Asset 4: Cryptocurrency (very high volatility) ---
asset4_entry = 1000.00
asset4_stop_loss = 950.00  # 5% stop-loss (50 units of risk per share)
asset4_pos_size = demo_calculation(
    total_capital,
    risk_per_trade_percent,
    asset4_entry,
    asset4_stop_loss,
    asset_name="Asset D (Cryptocurrency)",
)

print("------------------------------------------")
print("Summary of Position Sizes:")
print(f"Asset A: {asset1_pos_size} units")
print(f"Asset B: {asset2_pos_size} units")
print(f"Asset C: {asset3_pos_size} units")
print(f"Asset D: {asset4_pos_size} units")


# --- Calculate Portfolio Heat ---
print("\n--- Calculating Total Portfolio Heat ---")

portfolio_positions = {
    "Asset A": asset1_pos_size,
    "Asset B": asset2_pos_size,
    "Asset C": asset3_pos_size,
    "Asset D": asset4_pos_size,
}

current_prices = {
    "Asset A": asset1_entry,
    "Asset B": asset2_entry,
    "Asset C": asset3_entry,
    "Asset D": asset4_entry,
}

stop_loss_prices = {
    "Asset A": asset1_stop_loss,
    "Asset B": asset2_stop_loss,
    "Asset C": asset3_stop_loss,
    "Asset D": asset4_stop_loss,
}

portfolio_heat = calculate_portfolio_heat(
    total_capital, portfolio_positions, current_prices, stop_loss_prices
)

print(
    f"Total Portfolio Heat (Potential Loss if all stops hit): {portfolio_heat:.2f}% of capital"
)
print(
    f"Absolute Total Potential Loss: ${total_capital * (portfolio_heat / 100):,.2f}"
)

# Visualization 1: Position Size vs. Risk per Unit

# Generate a range of risk_per_unit values (e.g., from $0.50 to $10.00)
risk_per_unit_values = np.linspace(0.5, 10, 100)

# Assume a fixed entry price for calculation purposes, though only the difference matters
fixed_entry_price = 100.0

# Calculate corresponding position sizes
position_sizes = [
    calculate_position_size(
        total_capital,
        risk_per_trade_percent,
        fixed_entry_price,
        fixed_entry_price - rpu,
    )
    for rpu in risk_per_unit_values
]

plt.figure(figsize=(10, 6))
plt.plot(risk_per_unit_values, position_sizes, linestyle='-')
plt.title('Position Size vs. Risk per Unit (Fixed Risk Amount)')
plt.xlabel('Risk per Unit (Entry Price - Stop Loss Price) ($)')
plt.ylabel('Calculated Position Size (Units)')
plt.grid(True, linestyle='--', alpha=0.7)
plt.xlim(left=0) # Ensure x-axis starts from 0
plt.ylim(bottom=0) # Ensure y-axis starts from 0
plt.tight_layout()
plt.show()

print("Interpretation: As the 'Risk per Unit' (the distance between your entry and stop-loss) increases, your calculated position size decreases. This is because for a fixed dollar amount of risk you are willing to take, you can afford fewer units of an asset if each unit carries more potential loss. This visually confirms the inverse relationship inherent in heat sizing.")


# Visualization 2: Portfolio Risk Contribution (as a pie chart)

# Calculate individual potential losses for the demo portfolio
individual_potential_losses = {}
for asset, units in portfolio_positions.items():
    current_price = current_prices.get(asset)
    stop_loss_price = stop_loss_prices.get(asset)
    if current_price is not None and stop_loss_price is not None:
        individual_potential_losses[asset] = units * abs(current_price - stop_loss_price)

# Filter out assets with zero potential loss if any (e.g., if pos_size was 0)
filtered_losses = {k: v for k, v in individual_potential_losses.items() if v > 0}

if not filtered_losses:
    print("No active positions with potential loss to visualize.")
else:
    labels = filtered_losses.keys()
    sizes = filtered_losses.values()

    plt.figure(figsize=(10, 7))
    plt.pie(sizes, labels=labels, autopct='%1.1f%%', startangle=140, pctdistance=0.85, wedgeprops={'edgecolor': 'black'})
    plt.title('Portfolio Risk Contribution by Asset (If All Stops Hit)')
    plt.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
    plt.tight_layout()
    plt.show()

    print("Interpretation: This pie chart illustrates the proportional 'heat' or potential loss each asset contributes to the total portfolio risk. Even if each individual trade is sized to risk 1% of capital, the sum of these potential losses can accumulate. This visualization helps in quickly identifying which assets are contributing the most to the overall portfolio drawdown potential, assuming all stop-losses are triggered simultaneously.")