import streamlit as st
import numpy as np
import pandas as pd
import altair as alt

def calculate_pile_loads(P, Mx, My, pile_coords):
    """
    Calculates the load on each pile in a group subjected to a vertical load and moments.

    Args:
        P (float): The total vertical load applied to the pile cap (in kN).
        Mx (float): The moment about the x-axis (in kN-m).
        My (float): The moment about the y-axis (in kN-m).
        pile_coords (list of tuples): A list of (x, y) coordinates for each pile.

    Returns:
        list: A list of the calculated loads on each pile.
    """
    n = len(pile_coords)
    if n == 0:
        return []

    x_coords = np.array([coord[0] for coord in pile_coords])
    y_coords = np.array([coord[1] for coord in pile_coords])

    sum_x_sq = np.sum(x_coords**2)
    sum_y_sq = np.sum(y_coords**2)

    loads = []
    for i in range(n):
        x_i = x_coords[i]
        y_i = y_coords[i]

        load_P = P / n
        load_Mx = (Mx * y_i) / sum_y_sq if sum_y_sq != 0 else 0
        load_My = (My * x_i) / sum_x_sq if sum_x_sq != 0 else 0

        total_load = load_P + load_Mx + load_My
        loads.append(total_load)

    return loads

st.set_page_config(layout="wide", page_title="Pile Load Calculator")

st.title("Pile Load Calculator")
st.write("This application calculates the load on each pile in a foundation based on applied loads and moments.")

# --- Sidebar for Inputs ---
st.sidebar.header("Input Parameters")

# Applied Loads
st.sidebar.subheader("Applied Loads (kN, kN-m)")
P = st.sidebar.number_input("Vertical Point Load (P)", value=4500.0, step=100.0)
Mx = st.sidebar.number_input("Moment about X-axis (Mx)", value=680.0, step=50.0)
My = st.sidebar.number_input("Moment about Y-axis (My)", value=400.0, step=50.0)

# Footing Self-Weight
st.sidebar.subheader("Footing Self-Weight (kN)")
footing_self_weight = st.sidebar.number_input("Footing Self-Weight", value=225.0, step=25.0)

# Pile Coordinates
st.sidebar.subheader("Pile Coordinates (meters)")
if 'pile_coords' not in st.session_state:
    st.session_state.pile_coords = [(-3, 3), (0, 3), (3, 3),(-3, 0), (0, 0), (3, 0),(-3, -3),(0, -3),(3, -3)]

def add_pile():
    st.session_state.pile_coords.append((0.0, 0.0))

def remove_pile(index):
    st.session_state.pile_coords.pop(index)

for i, (x, y) in enumerate(st.session_state.pile_coords):
    cols = st.sidebar.columns([2, 2, 1])
    new_x = cols[0].number_input(f"Pile {i+1} X", value=float(x), key=f"x{i}")
    new_y = cols[1].number_input(f"Pile {i+1} Y", value=float(y), key=f"y{i}")
    st.session_state.pile_coords[i] = (new_x, new_y)
    if cols[2].button("X", key=f"del{i}"):
        remove_pile(i)
        st.rerun()

st.sidebar.button("Add Pile", on_click=add_pile)

# --- Main Panel for Results ---
total_vertical_load = P + footing_self_weight
pile_loads = calculate_pile_loads(total_vertical_load, Mx, My, st.session_state.pile_coords)

# Display Results in a Table
st.header("Calculation Results")

results_df = pd.DataFrame({
    'Pile': [f"Pile {i+1}" for i in range(len(st.session_state.pile_coords))],
    'X-coordinate (m)': [f"{c[0]:.2f}" for c in st.session_state.pile_coords],
    'Y-coordinate (m)': [f"{c[1]:.2f}" for c in st.session_state.pile_coords],
    'Calculated Load (kN)': [f"{l:.2f}" for l in pile_loads]
})

st.dataframe(results_df.set_index('Pile'))

# Display Summary of Loads
st.subheader("Load Summary")
st.write(f"**Total Applied Vertical Load (P + Self-Weight):** {total_vertical_load:.2f} kN")
st.write(f"**Maximum Pile Load (Compression):** {max(pile_loads):.2f} kN")
st.write(f"**Minimum Pile Load (Tension/Uplift):** {min(pile_loads):.2f} kN")

# Visualization of Pile Loads
st.header("Pile Load Visualization")

if st.session_state.pile_coords:
    vis_df = pd.DataFrame({
        'x': [c[0] for c in st.session_state.pile_coords],
        'y': [c[1] for c in st.session_state.pile_coords],
        'load': pile_loads,
        'load_text': [f"{l:.1f} kN" for l in pile_loads]
    })

    # Determine the domain for the color scale to center on zero
    max_abs_load = max(abs(vis_df['load'].min()), abs(vis_df['load'].max()))

    # Calculate padding for the chart domain
    x_range = vis_df['x'].max() - vis_df['x'].min()
    y_range = vis_df['y'].max() - vis_df['y'].min()
    x_buffer = x_range * 0.2  # 20% buffer
    y_buffer = y_range * 0.2  # 20% buffer

    base_chart = alt.Chart(vis_df).encode(
        x=alt.X('x:Q', title='X-coordinate (m)',
                scale=alt.Scale(domain=[vis_df['x'].min() - x_buffer, vis_df['x'].max() + x_buffer]),
                axis=alt.Axis(titleFontSize=14, labelFontSize=12)),
        y=alt.Y('y:Q', title='Y-coordinate (m)',
                scale=alt.Scale(domain=[vis_df['y'].min() - y_buffer, vis_df['y'].max() + y_buffer]),
                axis=alt.Axis(titleFontSize=14, labelFontSize=12)),
        tooltip=[
            alt.Tooltip('x:Q', title='X-coordinate', format='.2f'),
            alt.Tooltip('y:Q', title='Y-coordinate', format='.2f'),
            alt.Tooltip('load:Q', title='Load (kN)', format='.2f')
        ]
    )

    # Points with color scale
    points = base_chart.mark_point(size=300, filled=True, stroke='black', strokeWidth=0.5).encode(
        color=alt.Color('load:Q', title='Load (kN)',
                        scale=alt.Scale(scheme='redblue', domain=[-max_abs_load, max_abs_load], reverse=True))
    )

    # Text labels for the loads
    text = base_chart.mark_text(align='center', dy=-15, fontSize=12).encode(
        text='load_text:N'
    )

    chart = (points + text).properties(
        title=alt.TitleParams(
            text='Pile Location and Load Distribution',
            subtitle='Red for Tension/Uplift, Blue for Compression',
            fontSize=20,
            subtitleFontSize=16
        ),
        padding={"left": 20, "top": 20, "right": 20, "bottom": 20} # Add padding
    ).interactive()

    st.altair_chart(chart, use_container_width=True)