import streamlit as st
import numpy as np
import matplotlib.pyplot as plt

def calculate_metrics(a, b, c, d, price_control=None):
    """
    Calculates equilibrium, consumer surplus, producer surplus, and deadweight loss.

    Args:
        a (float): Y-intercept of the supply curve (P = a + bQ).
        b (float): Slope of the supply curve.
        c (float): Y-intercept of the demand curve (P = c - dQ).
        d (float): Slope of the demand curve.
        price_control (float, optional): An enforced price (floor or ceiling). Defaults to None.

    Returns:
        tuple: (Q_eq, P_eq, CS, PS, TS, DWL, Q_traded, price_control)
               Returns (0, 0, 0, 0, 0, 0, 0, price_control) if parameters are invalid.
    """
    # Ensure positive slopes for realistic curves
    b = max(0.01, b)  # Supply slope
    d = max(0.01, d)  # Demand slope

    # Check for valid parameters that lead to an intersection with positive quantity
    if (c - a) <= 0: # Demand intercept must be above supply intercept for Q_eq > 0
        return 0, 0, 0, 0, 0, 0, 0, price_control

    # 1. Equilibrium Calculation
    Q_eq = (c - a) / (d + b)
    P_eq = a + b * Q_eq

    # Ensure valid equilibrium (positive quantity and price)
    if Q_eq < 0.01 or P_eq < 0.01: # Check for near zero or negative values
        return 0, 0, 0, 0, 0, 0, 0, price_control

    # 2. Consumer Surplus (CS)
    # Area of the triangle above equilibrium price and below demand curve
    CS = 0.5 * Q_eq * (c - P_eq)
    CS = max(0, CS) # Ensure non-negative

    # 3. Producer Surplus (PS)
    # Area of the triangle below equilibrium price and above supply curve
    PS = 0.5 * Q_eq * (P_eq - a)
    PS = max(0, PS) # Ensure non-negative

    # 4. Total Surplus
    TS = CS + PS

    # 5. Deadweight Loss (DWL) and Quantity Traded under Price Control
    DWL = 0
    Q_traded = Q_eq  # Initialize quantity traded to equilibrium quantity

    if price_control is not None and price_control > 0: # If a price control is set
        # Calculate quantity demanded and supplied at the controlled price
        Q_demanded_at_pc = (c - price_control) / d
        Q_supplied_at_pc = (price_control - a) / b

        # The actual quantity traded is the minimum of Qd and Qs at the controlled price,
        # ensuring it's non-negative (cannot trade negative quantity).
        Q_traded = min(max(0, Q_demanded_at_pc), max(0, Q_supplied_at_pc))

        # Calculate DWL if trade is restricted from equilibrium
        if Q_traded < Q_eq:
            # Price on the demand curve at the restricted quantity traded
            P_demand_at_Q_traded = c - d * Q_traded
            # Price on the supply curve at the restricted quantity traded
            P_supply_at_Q_traded = a + b * Q_traded

            # DWL is the area of the triangle between demand and supply curves
            # from Q_traded to Q_eq. Height is the vertical distance between D & S at Q_traded.
            DWL = 0.5 * (Q_eq - Q_traded) * (P_demand_at_Q_traded - P_supply_at_Q_traded)
            DWL = max(0, DWL) # Ensure DWL is non-negative

    return Q_eq, P_eq, CS, PS, TS, DWL, Q_traded, price_control


st.set_page_config(layout="wide", page_title="Supply & Demand Model")
st.title("📊 Econ 101: Supply and Demand Model")

st.sidebar.header("⚙️ Adjust Model Parameters")

# Sliders for demand curve: P = c - dQ
st.sidebar.subheader("Demand Curve: P = `c` - `d`Q")
c = st.sidebar.slider("c (Demand Intercept)", 10.0, 200.0, 100.0, 1.0, help="Maximum price consumers are willing to pay for 0 quantity.")
d = st.sidebar.slider("d (Demand Slope)", 0.1, 10.0, 2.0, 0.1, help="How much price decreases for each unit demanded.")

# Sliders for supply curve: P = a + bQ
st.sidebar.subheader("Supply Curve: P = `a` + `b`Q")
a = st.sidebar.slider("a (Supply Intercept)", 0.0, 100.0, 10.0, 1.0, help="Minimum price producers are willing to accept for 0 quantity.")
b = st.sidebar.slider("b (Supply Slope)", 0.1, 10.0, 1.0, 0.1, help="How much price increases for each unit supplied.")

# Price control slider
st.sidebar.subheader("Enforced Price (e.g., Price Floor/Ceiling)")
price_control = st.sidebar.slider("Set Enforced Price (0 for none)", 0.0, 200.0, 0.0, 0.5,
                                  help="Set a price floor or ceiling. If 0, no price control is active.")

# Perform calculations
Q_eq, P_eq, CS, PS, TS, DWL, Q_traded, pc_used = calculate_metrics(a, b, c, d, price_control)

# Display results and plot
col1, col2 = st.columns([1, 2])

with col1:
    st.header("📈 Economic Metrics")
    if Q_eq == 0 and P_eq == 0 and CS == 0: # Indicates invalid parameters from calculate_metrics
        st.error("Invalid parameters: Please adjust 'c' to be greater than 'a' to ensure a valid equilibrium.")
    else:
        st.metric("Equilibrium Quantity (Qe)", f"{Q_eq:.2f}")
        st.metric("Equilibrium Price (Pe)", f"${P_eq:.2f}")
        st.metric("Consumer Surplus (CS)", f"${CS:.2f}")
        st.metric("Producer Surplus (PS)", f"${PS:.2f}")
        st.metric("Total Surplus (TS)", f"${TS:.2f}")

        if price_control > 0: # Only show these metrics if a price control is active
            st.subheader("Metrics with Price Control")
            st.metric("Enforced Price", f"${pc_used:.2f}")
            st.metric("Quantity Traded (Q_traded)", f"{Q_traded:.2f}")
            st.metric("Deadweight Loss (DWL)", f"${DWL:.2f}")
            if DWL > 0:
                st.warning("Deadweight Loss indicates market inefficiency due to the enforced price.")
            else:
                st.info("No Deadweight Loss with current enforced price (or price is ineffective).")
        else:
            st.info("Set 'Enforced Price' to calculate Deadweight Loss.")


with col2:
    st.header("📉 Supply and Demand Graph")
    fig, ax = plt.subplots(figsize=(12, 8))

    # Determine appropriate x and y axis limits for plotting
    max_Q_plot = max(Q_eq * 1.5, (c / d) * 1.1, 20) # Extends beyond equilibrium and demand intercept
    max_P_plot = max(P_eq * 1.5, c * 1.1, 150) # Extends beyond equilibrium and demand intercept

    Q_values = np.linspace(0, max_Q_plot, 400)

    # Demand Curve: P = c - dQ
    P_demand = c - d * Q_values
    # Supply Curve: P = a + bQ
    P_supply = a + b * Q_values

    # Filter out negative prices/quantities for plotting realism
    P_demand[P_demand < 0] = np.nan # Don't plot negative prices
    P_supply[P_supply < 0] = np.nan # Don't plot negative prices
    Q_values[Q_values < 0] = np.nan # Don't plot negative quantities

    ax.plot(Q_values, P_demand, label=f'Demand: P = {c:.1f} - {d:.1f}Q', color='blue', linewidth=2)
    ax.plot(Q_values, P_supply, label=f'Supply: P = {a:.1f} + {b:.1f}Q', color='red', linewidth=2)

    # Plot Equilibrium Point and lines
    if Q_eq > 0 and P_eq > 0:
        ax.plot(Q_eq, P_eq, 'go', markersize=8, label=f'Equilibrium (Q={Q_eq:.2f}, P=${P_eq:.2f})')
        ax.vlines(Q_eq, 0, P_eq, linestyle=':', color='gray', linewidth=1)
        ax.hlines(P_eq, 0, Q_eq, linestyle=':', color='gray', linewidth=1)

        # Consumer Surplus Area
        Q_cs_plot = np.linspace(0, Q_eq, 100)
        P_cs_plot = c - d * Q_cs_plot
        ax.fill_between(Q_cs_plot, P_eq, P_cs_plot, where=P_cs_plot > P_eq, color='skyblue', alpha=0.3, label='Consumer Surplus')

        # Producer Surplus Area
        Q_ps_plot = np.linspace(0, Q_eq, 100)
        P_ps_plot = a + b * Q_ps_plot
        ax.fill_between(Q_ps_plot, P_ps_plot, P_eq, where=P_eq > P_ps_plot, color='lightcoral', alpha=0.3, label='Producer Surplus')

    # Plot Price Control and Deadweight Loss
    if price_control > 0 and Q_traded < Q_eq:
        ax.hlines(pc_used, 0, Q_traded, linestyle='--', color='purple', label=f'Enforced Price (${pc_used:.2f})', linewidth=1.5)
        ax.vlines(Q_traded, 0, pc_used, linestyle='--', color='purple', linewidth=1.5)

        # Identify the points for the DWL triangle
        # Price on demand curve at Q_traded
        P_demand_at_Q_traded_plot = c - d * Q_traded
        # Price on supply curve at Q_traded
        P_supply_at_Q_traded_plot = a + b * Q_traded

        # Shade DWL triangle
        # The points for the DWL triangle are (Q_traded, P_supply_at_Q_traded), (Q_traded, P_demand_at_Q_traded), and (Q_eq, P_eq)
        # Plotting the triangle as a fill_between area between Q_traded and Q_eq
        Q_dwl_fill = np.linspace(Q_traded, Q_eq, 100)
        P_demand_dwl_fill = c - d * Q_dwl_fill
        P_supply_dwl_fill = a + b * Q_dwl_fill

        ax.fill_between(Q_dwl_fill, P_supply_dwl_fill, P_demand_dwl_fill, color='red', alpha=0.5, label='Deadweight Loss')

        # Mark points relevant to price control
        # ax.plot(Q_traded, c - d * Q_traded, 'kx', markersize=8, label=f'Demand at PC ({Q_traded:.2f})')
        # ax.plot(Q_traded, a + b * Q_traded, 'bx', markersize=8, label=f'Supply at PC ({Q_traded:.2f})')

    ax.set_xlabel("Quantity (Q)", fontsize=12)
    ax.set_ylabel("Price (P)", fontsize=12)
    ax.set_title("Supply and Demand Equilibrium", fontsize=14)
    ax.set_ylim(bottom=0, top=max_P_plot)
    ax.set_xlim(left=0, right=max_Q_plot)
    ax.legend(loc='upper right')
    ax.grid(True, linestyle='--', alpha=0.7)
    st.pyplot(fig)