import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# =====================================================================
# --- 1. INPUT PARAMETERS (Expandable) ---
# =====================================================================
num_elements = 1
num_nodes = 2

A = np.array([16.0])     # Cross-sectional areas [cm^2]
L = np.array([2.0])      # Element lengths [cm]

# Connectivity: [Element-ID, Start_Node, End_Node]
connectivity = np.array([
    [1, 1, 2]
])

# Load
P_final_val = 750

# =====================================================================
# --- 2. MATERIAL MODEL (Non-linear) ---
# =====================================================================
E0 = 21000.0             # Initial Young's modulus [kN/cm^2]
fy = 23.5                # Yield strength [kN/cm^2]
m = 5                    # Hardening parameter
eps_y = fy / E0          # Yield strain (approx. 1.119e-3)

# =====================================================================
# --- 3. STORAGE FOR NEWTON-RAPHSON DATA ---
# =====================================================================
max_iter = 100
max_load_steps = 10
tol = 1e-8

P = np.linspace(0, P_final_val, max_load_steps + 1)

u_it = np.zeros((num_nodes, max_iter, len(P)))
strain_it = np.zeros((num_elements, max_iter, len(P)))
stress_it = np.zeros((num_elements, max_iter, len(P)))
R_it = np.zeros((num_nodes, max_iter, len(P)))
Et_it = np.zeros((num_elements, max_iter, len(P)))

res_list = []

# =====================================================================
# --- 4. MAIN LOOP OVER LOAD STEPS ---
# =====================================================================
for i, P_current in enumerate(P):
    u_current = np.zeros((num_nodes, 1))
    strain_current = np.zeros((num_elements, 1))
    
    # Initial residual: Load P is applied at the last node
    R = np.zeros((num_nodes, 1))
    R[-1, 0] = P_current 
    
    for iter_step in range(max_iter):
        # 4.1 Calculate tangent modulus (Loop over all elements)
        Et = np.zeros(num_elements)
        for s in range(num_elements):
            eps = abs(strain_current[s, -1])
            if eps <= eps_y:
                Et[s] = E0
            else:
                Et[s] = fy / (m * eps) * (eps * E0 / fy)**(1/m)
            Et_it[s, iter_step, i] = Et[s]
        
        # 4.2 Stiffness matrix & Global assembly
        Kt = np.zeros((num_nodes, num_nodes))
        for s in range(num_elements):
            kt_element = Et[s] * A[s] / L[s] * np.array([[1, -1], [-1, 1]])
            start_node = int(connectivity[s, 1] - 1)
            end_node = int(connectivity[s, 2] - 1)
            
            Kt[start_node, start_node] += kt_element[0, 0]
            Kt[start_node, end_node]   += kt_element[0, 1]
            Kt[end_node, start_node]   += kt_element[1, 0]
            Kt[end_node, end_node]     += kt_element[1, 1]
        
        # 4.3 Solve system (Node 1 / Index 0 is fixed)
        Kt_red = Kt[1:, 1:]
        R_red = R[1:]
        
        du_red = np.linalg.solve(Kt_red, R_red).reshape(-1, 1)
        
        # Update full displacement vector
        u_new = np.zeros((num_nodes, 1))
        u_new[1:, :] = u_current[1:, -1:] + du_red
        u_current = np.hstack([u_current, u_new])
        
        # 4.4 Strains & Stresses
        strain_new = np.zeros((num_elements, 1))
        stress_new = np.zeros(num_elements)
        F_int = np.zeros((num_nodes, 1)) # Internal force vector
        
        for s in range(num_elements):
            start_node = int(connectivity[s, 1] - 1)
            end_node = int(connectivity[s, 2] - 1)
            
            # Strain from nodal displacements
            eps_val = (u_new[end_node, 0] - u_new[start_node, 0]) / L[s]
            strain_new[s, 0] = eps_val
            
            # Stress from material law
            if abs(eps_val) <= eps_y:
                sig_val = E0 * eps_val
            else:
                sig_val = np.sign(eps_val) * fy * (abs(eps_val) * E0 / fy)**(1/m)
            stress_new[s] = sig_val
            
            # Accumulate internal forces on the nodes
            axial_force = sig_val * A[s]
            F_int[start_node, 0] -= axial_force
            F_int[end_node, 0]   += axial_force
            
        strain_current = np.hstack([strain_current, strain_new])
        
        # 4.5 Calculate residual (Equilibrium R = F_ext - F_int)
        F_ext = np.zeros((num_nodes, 1))
        F_ext[-1, 0] = P_current
        R = F_ext - F_int
        
        # Store data
        u_it[:, iter_step, i] = u_new.flatten()
        strain_it[:, iter_step, i] = strain_new.flatten()
        stress_it[:, iter_step, i] = stress_new
        R_it[:, iter_step, i] = R.flatten()

        # Dynamically create result row for the dataframe
        row_data = {"Load P [kN]": P_current, "Iter": iter_step + 1}
        for k in range(1, num_nodes):  # u1 is always 0, start from Node 2
            row_data[f"u{k+1} [cm]"] = u_new[k, 0]
        for s in range(num_elements):
            row_data[f"Stress E{s+1}"] = stress_new[s]
            row_data[f"Et_E{s+1}"] = Et[s]
        
        res_norm = np.linalg.norm(R[1:]) # Norm over free nodes only
        row_data["Res_norm"] = res_norm
        
        res_list.append(row_data)

        # Convergence check
        if res_norm < tol:
            break

# =====================================================================
# --- 5. TERMINAL OUTPUT ---
# =====================================================================
df = pd.DataFrame(res_list)
pd.set_option('display.max_rows', None)
pd.set_option('display.width', 1000)
pd.set_option('display.precision', 6)

print(f"\nDETAILED ITERATION LOG ({num_elements} ELEMENTS):")
print(df.to_string(index=False))

# Filter: Show only converged steps
df_converged = df.groupby("Load P [kN]").tail(1)

print(f"\nCONVERGED RESULTS PER LOAD STEP:")
print(df_converged.to_string(index=False))

# =====================================================================
# --- 6. PLOTTING ---
# =====================================================================

# --- Plot 1: Material Law ---
fig1, ax1 = plt.subplots(figsize=(8, 5))

# Calculate plotting data for the material law
strain_mat1 = np.arange(0, eps_y, 1e-6)
strain_mat2 = np.arange(eps_y, 1.6345e-2, 0.001)
stress_mat1 = E0 * strain_mat1
stress_mat2 = fy * (strain_mat2 * E0 / fy)**(1/m)

ax1.plot(strain_mat1, stress_mat1, 'k-', label='Material Law (S235)')
ax1.plot(strain_mat2, stress_mat2, 'k-')

ax1.set_title('Material Law', fontsize=14, pad=15)
ax1.set_xlabel(r'Strain $\epsilon$ [-]', fontsize=12)
ax1.set_ylabel(r'Stress $\sigma$ [kN/cm²]', fontsize=12)
ax1.legend(loc='lower right')
ax1.grid(True, linestyle=':', alpha=0.5)

fig1.tight_layout()

# --- Plot 2: System Response (Load-Deformation) ---
fig2, ax2 = plt.subplots(figsize=(8, 5))

# Extract final displacements for Node 2
u2_final = []
for i in range(len(P)):
    if P[i] == 0:
        u2_final.append(0)
    else:
        # Filter out trailing zeros from unused iterations
        valid_iters = u_it[1, :, i][u_it[1, :, i] != 0]
        u2_final.append(valid_iters[-1] if len(valid_iters) > 0 else 0)

ax2.plot(u2_final, P, 'o--', color='royalblue', alpha=0.8, label='Load-Deformation (Node 2)')

ax2.set_title('Load-Displacement Diagram', fontsize=14, pad=15)
ax2.set_xlabel(r'Displacement $u_2$ [cm]', fontsize=12)
ax2.set_ylabel(r'Load $P$ [kN]', fontsize=12)
ax2.legend(loc='lower right')
ax2.grid(True, linestyle=':', alpha=0.5)

fig2.tight_layout()

plt.show()