import sympy as sp
import numpy as np
import pandas as pd
import gradio as gr
import matplotlib.pyplot as plt

def solve_beam(l1_val, l2_val, q1_val, q2_val):
    # Define the variables
    Mx, l1, l2, q1, q2 = sp.symbols('Mx l1 l2 q1 q2')

    # Equation to solve for Mx
    equation = (Mx*l1/3 + q1*l1**3/24 + Mx*l2/3 + q2*l2**3/24)
    Mx_solution = sp.solve(equation, Mx)

    # Define variables for reactions
    VA, VB1, VB2, VC, HA = sp.symbols('VA VB1 VB2 VC HA')

    # System of equations for the first span
    eq1_span1 = VA + VB1 - q1*l1
    eq2_span1 = VB1*l1 - q1*l1**2/2 + Mx_solution[0]

    # System of equations for the second span
    eq1_span2 = VB2 + VC - q2*l2
    eq2_span2 = VB2*l2 - q2*l2**2/2 + Mx_solution[0]

    # Solve for the reactions for the first span
    reactions_span1 = sp.solve((eq1_span1, eq2_span1), (VA, VB1))

    # Solve for the reactions for the second span
    reactions_span2 = sp.solve((eq1_span2, eq2_span2), (VB2, VC))

    # Define variables for positions on the spans
    x1, x2 = sp.symbols('x1 x2')

    # Bending moment and shear for the first span
    M1_expr = reactions_span1[VA] * x1 - q1 * x1**2 / 2
    V1_expr = reactions_span1[VA] - q1 * x1

    # Bending moment and shear for the second span
    M2_expr = Mx_solution[0] - q2 * x2**2 / 2 + reactions_span2[VB2] * x2
    V2_expr = reactions_span2[VB2] - q2 * x2

    # Substitute the provided numerical values into the expressions
    M1_expr = M1_expr.subs({l1: l1_val, q1: q1_val, l2: l2_val, q2: q2_val})
    V1_expr = V1_expr.subs({l1: l1_val, q1: q1_val, l2: l2_val, q2: q2_val})
    M2_expr = M2_expr.subs({l1: l1_val, q1: q1_val, l2: l2_val, q2: q2_val})
    V2_expr = V2_expr.subs({l1: l1_val, q1: q1_val, l2: l2_val, q2: q2_val})
    Mx_value = Mx_solution[0].subs({l1: l1_val, q1: q1_val, l2: l2_val, q2: q2_val})

    # Generate numerical values for x1 and x2
    x1_vals = np.arange(0, l1_val, 0.1)
    x2_vals = np.arange(0, l2_val, 0.1)

    # Evaluate M1 and V1 at each x1
    M1_vals = [float(M1_expr.subs(x1, val)) for val in x1_vals]
    V1_vals = [float(V1_expr.subs(x1, val)) for val in x1_vals]

    # Evaluate M2 and V2 at each x2
    M2_vals = [float(M2_expr.subs(x2, val)) for val in x2_vals]
    V2_vals = [float(V2_expr.subs(x2, val)) for val in x2_vals]

    # Create dataframes for beam1 and beam2
    beam1 = pd.DataFrame({'x': x1_vals, 'M': M1_vals, 'V': V1_vals})
    beam2 = pd.DataFrame({'x': x2_vals + l1_val, 'M': M2_vals, 'V': V2_vals})

    # Concatenate beam1 and beam2 into one dataframe
    beam = pd.concat([beam1, beam2], ignore_index=True)
    
    return beam

def plot_beam(beam_df):
    fig, axs = plt.subplots(2, 1, figsize=(8, 8))

    # Plot bending moment
    axs[0].plot(beam_df['x'], beam_df['M'], label='Bending Moment (M)')
    axs[0].invert_yaxis()
    axs[0].set_title('Bending Moment Diagram')
    axs[0].set_xlabel('Position along the beam (x)')
    axs[0].set_ylabel('Bending Moment (M)')
    axs[0].legend()

    # Plot shear force
    axs[1].plot(beam_df['x'], beam_df['V'], label='Shear Force (V)')
    axs[1].invert_yaxis()
    axs[1].set_title('Shear Force Diagram')
    axs[1].set_xlabel('Position along the beam (x)')
    axs[1].set_ylabel('Shear Force (V)')
    axs[1].legend()

    plt.tight_layout()
    return fig

def main(l1, l2, q1, q2):
    beam_df = solve_beam(l1, l2, q1, q2)
    csv_file = 'beam_analysis.csv'
    beam_df.to_csv(csv_file, index=False)
    fig = plot_beam(beam_df)
    return beam_df, csv_file, fig

inputs = [
    gr.components.Number(label="Length of span 1 (l1)"),
    gr.components.Number(label="Length of span 2 (l2)"),
    gr.components.Number(label="Load on span 1 (q1)"),
    gr.components.Number(label="Load on span 2 (q2)")
]

outputs = [
    gr.components.Dataframe(label="Beam Analysis Data"),
    gr.components.File(label="Download CSV"),
    gr.components.Plot(label="Bending Moment and Shear Force Diagrams")
]

gr.Interface(fn=main, inputs=inputs, outputs=outputs, title="Beam Analysis Tool").launch()