# plot_functions.py

import numpy as np
import matplotlib.pyplot as plt
import os
from linear_equilibrium_calculations import calculate_linear_equilibrium  # Updated import

def plot_equilibrium(demand_intercept, slope_demand, supply_intercept, supply_slope):
    """Generates a supply & demand equilibrium plot (Linear Model)."""

    Q_range = np.linspace(0, 20, 100)
    demand_curve = demand_intercept + (slope_demand * Q_range)
    supply_curve = supply_intercept + (supply_slope * Q_range)

    plt.figure(figsize=(6, 4))
    plt.plot(Q_range, demand_curve, label="Demand Curve", color="blue")
    plt.plot(Q_range, supply_curve, label="Supply Curve", color="green")

    # Calculate equilibrium using the new function
    Q_eq, P_eq = calculate_linear_equilibrium(demand_intercept, slope_demand, supply_intercept, supply_slope)

    # Calculate CS & PS
    CS_top = demand_intercept  
    PS_bottom = supply_intercept  

    CS = 0.5 * (CS_top - P_eq) * Q_eq  
    PS = 0.5 * (P_eq - PS_bottom) * Q_eq  

    # Highlight CS (shaded area above equilibrium price)
    plt.fill_between([0, Q_eq], [CS_top, P_eq], P_eq, color="blue", alpha=0.3, label="Consumer Surplus")

    # Highlight PS (shaded area below equilibrium price)
    plt.fill_between([0, Q_eq], P_eq, [PS_bottom, P_eq], color="green", alpha=0.3, label="Producer Surplus")

    # Mark equilibrium point
    plt.scatter(Q_eq, P_eq, color="red", label=f"Equilibrium ({Q_eq}, {P_eq})")

    plt.xlabel("Quantity")
    plt.ylabel("Price")
    plt.legend()
    plt.title("Supply & Demand Equilibrium (Linear Model)")
    plt.grid()

    plot_filename = "equilibrium_plot.png"
    plt.savefig(plot_filename)
    plt.close()

    if not os.path.exists(plot_filename):
        raise FileNotFoundError("Error: Plot file not found after saving!")

    return plot_filename, CS, PS