import numpy as np
import matplotlib.pyplot as plt

# Function to calculate results for simply supported beams or continuous beams
def analyze_beam(length, load, beam_type):
    if beam_type == "simply_supported":
        return analyze_simply_supported(length, load)
    elif beam_type == "continuous":
        return analyze_continuous(length, load)

# Simply supported beam analysis
def analyze_simply_supported(length, load):
    # Simplified calculation example for a uniformly loaded simply supported beam
    shear_force = load * length / 2  # Max shear force (at supports)
    bending_moment = load * length**2 / 8  # Max bending moment (at mid-span)
    deflection = (5 * load * length**4) / (384 * 200000 * 100)  # Example deflection formula for steel (in inches)
    
    return {
        "max_shear": shear_force,
        "max_bending_moment": bending_moment,
        "deflection": deflection
    }

# Continuous beam analysis (simplified example)
def analyze_continuous(length, load):
    # Continuous beam analysis (simplified)
    shear_force = load * length / 3  # Approximation for continuous beams
    bending_moment = load * length**2 / 12  # Approximation for continuous beams
    deflection = (5 * load * length**4) / (384 * 200000 * 100)  # Same deflection calculation as simply supported

    return {
        "max_shear": shear_force,
        "max_bending_moment": bending_moment,
        "deflection": deflection
    }

# Function to plot beam diagrams (SFD and BMD)
def plot_beam_diagrams(beam_type, length, load):
    x = np.linspace(0, length, 100)
    shear = np.zeros_like(x)
    moment = np.zeros_like(x)

    if beam_type == "simply_supported":
        # Shear Force and Bending Moment Diagrams for simply supported beam
        shear[:] = load / 2  # Simplified constant shear force
        moment[:] = (load * x * (length - x)) / (2 * length)  # Bending moment equation for simply supported beam
    elif beam_type == "continuous":
        # Shear Force and Bending Moment Diagrams for continuous beam
        shear[:] = load / 3  # Simplified constant shear force for continuous beam
        moment[:] = (load * x * (length - x)) / (3 * length)  # Approximation for continuous beam

    # Plot the diagrams
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
    ax1.plot(x, shear, label="Shear Force")
    ax1.set_title("Shear Force Diagram (SFD)")
    ax1.set_xlabel("Beam Length (ft)")
    ax1.set_ylabel("Shear Force (lb)")
    ax1.grid(True)

    ax2.plot(x, moment, label="Bending Moment", color="orange")
    ax2.set_title("Bending Moment Diagram (BMD)")
    ax2.set_xlabel("Beam Length (ft)")
    ax2.set_ylabel("Bending Moment (lb-ft)")
    ax2.grid(True)

    plt.tight_layout()
    plt.show()