from matplotlib import pyplot as plt
from matplotlib import cm
from matplotlib.colors import Normalize, TwoSlopeNorm
import numpy as np

def plotDeformedStructure2d(coordinates,connectivity,disp,scaleFactor,theTitle):
    # Create a figure and axis
    fig, ax = plt.subplots()

    # Plot undeformed structure
    for i in range(connectivity.shape[0]):
        x=[coordinates[connectivity[i,0],0],coordinates[connectivity[i,1],0]]
        y=[coordinates[connectivity[i,0],1],coordinates[connectivity[i,1],1]]
        ax.plot(x,y,'k--o')
    
    # Plot deformed structure
    for i in range(connectivity.shape[0]):
        x=[coordinates[connectivity[i,0],0]+disp[2*connectivity[i,0]]*scaleFactor,coordinates[connectivity[i,1],0]+disp[2*connectivity[i,1]]*scaleFactor]
        y=[coordinates[connectivity[i,0],1]+disp[2*connectivity[i,0]+1]*scaleFactor,coordinates[connectivity[i,1],1]+disp[2*connectivity[i,1]+1]*scaleFactor]
        ax.plot(x,y,'k-o')

    # Add node numbers
    for i in range(coordinates.shape[0]):
        ax.text(coordinates[i,0], coordinates[i,1], str(i+1), fontsize=12, ha='right')

    # Plot properties
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_aspect('equal')  # Ensure the x and y scales are the same

    # Save the figure 
    #plt.savefig(f'{theTitle}.pdf', format='pdf')
    
    # Show the figure
    plt.show()
    
    return

def plotInternalForces2dTruss(coordinates,connectivity, E, areas, disp, theTitle):
    # Create a figure and axis
    fig, ax = plt.subplots()

    # Plot undeformed structure
    for i in range(connectivity.shape[0]):
        x=[coordinates[connectivity[i,0],0],coordinates[connectivity[i,1],0]]
        y=[coordinates[connectivity[i,0],1],coordinates[connectivity[i,1],1]]
        ax.plot(x,y,'k--o')

    # Compute axial load
    N=np.zeros(connectivity.shape[0])
    for i in range(connectivity.shape[0]):
        l_undeformed=np.linalg.norm(coordinates[connectivity[i,1]]-coordinates[connectivity[i,0]])
        theta_undeformed=np.arctan2(coordinates[connectivity[i,1],1]-coordinates[connectivity[i,0],1],coordinates[connectivity[i,1],0]-coordinates[connectivity[i,0],0])    

        # Global reference coordinates
        u_X1=disp[2*connectivity[i,0]]
        u_Y1=disp[2*connectivity[i,0]+1]
        u_X2=disp[2*connectivity[i,1]]
        u_Y2=disp[2*connectivity[i,1]+1]
        # Local reference coordinates
        u_x1=u_X1*np.cos(theta_undeformed)+u_Y1*np.sin(theta_undeformed)
        u_x2=u_X2*np.cos(theta_undeformed)+u_Y2*np.sin(theta_undeformed)

        deltaL=u_x2-u_x1
        N[i]=E*areas[i]/l_undeformed*deltaL/1000
    
    # Determine the colormap and normalization
    max_abs_N = np.max(np.abs(N))
    norm = TwoSlopeNorm(vmin=-max_abs_N, vcenter=0, vmax=max_abs_N)
    cmap = cm.get_cmap('coolwarm')

    # Plot internal forces with heat map effect
    for i in range(connectivity.shape[0]):
        x = [coordinates[connectivity[i,0],0], coordinates[connectivity[i,1],0]]
        y = [coordinates[connectivity[i,0],1], coordinates[connectivity[i,1],1]]
        color = cmap(norm(N[i]))
        ax.plot(x, y, color=color, linewidth=2)

        # Add the value of N on each bar
        mid_x = (x[0] + x[1]) / 2
        mid_y = (y[0] + y[1]) / 2
        ax.text(mid_x, mid_y, f'{N[i]:.2f}', color='black', fontsize=12, ha='center', va='center')
    
    # Plot properties
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_aspect('equal')  # Ensure the x and y scales are the same

    # Add a single colorbar
    sm = cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    fig.colorbar(sm, ax=ax, label='Axial Force (kN)')

    # Remove the box (frame) around the plot
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['left'].set_visible(False)
    ax.spines['bottom'].set_visible(False)


    # Save the figure 
    #plt.savefig(f'{theTitle}.pdf', format='pdf')
    #plt.savefig(f'{theTitle}.png', format='PNG')
    
    # Show the figure
    plt.show()
    
    return