"""rigidez.py: funciones del Método de la Rigidez (MECCOM).

Este módulo se construye a lo largo de las Secciones I a V del curso.

Convenciones:
- Los nudos y los GDL se numeran desde 0 en el código.
- Armaduras 2D: 2 GDL por nudo (ux, uy); el nudo i usa los índices 2*i y 2*i + 1.
- Vigas: 2 GDL por nudo (v, θ). Marcos: 3 GDL por nudo (ux, uy, θ).
- Unidades del SI (m, N, Pa).
"""
import numpy as np
import matplotlib.pyplot as plt


# ============================ Sección I ============================

def geometria_barra(xy1, xy2):
    """Largo y cosenos directores de un elemento que va del nudo 1 al nudo 2.

    Retorna (L, c, s), con c = cos(θ) y s = sin(θ).
    """
    dx, dy = np.subtract(xy2, xy1)
    L = np.hypot(dx, dy)
    return L, dx / L, dy / L


def k_barra(A, E, xy1, xy2):
    """Matriz de rigidez 4x4 de una barra axial en ejes globales: Tᵀ k' T."""
    L, c, s = geometria_barra(xy1, xy2)
    k_local = A * E / L * np.array([[ 1, -1],
                                    [-1,  1]])
    T = np.array([[c, s, 0, 0],
                  [0, 0, c, s]])
    return T.T @ k_local @ T


# ============================ Sección II ============================

def gdl_barra(i, j):
    """Índices de GDL (base 0) de una barra que une los nudos i y j (2 GDL por nudo)."""
    return [2*i, 2*i + 1, 2*j, 2*j + 1]


def ensamblar(K, k_elem, gdl):
    """Suma la matriz de un elemento en la matriz K de la estructura (modifica K)."""
    K[np.ix_(gdl, gdl)] += k_elem


def ensamblar_armadura(nudos, barras, A, E):
    """Matriz K de una armadura 2D.

    nudos: lista de coordenadas (x, y). barras: lista de pares (i, j) de nudos.
    A y E pueden ser un solo valor o una lista con un valor por barra.
    """
    A = np.broadcast_to(A, len(barras))
    E = np.broadcast_to(E, len(barras))
    K = np.zeros((2*len(nudos), 2*len(nudos)))
    for (i, j), A_e, E_e in zip(barras, A, E):
        ensamblar(K, k_barra(A_e, E_e, nudos[i], nudos[j]), gdl_barra(i, j))
    return K


def resolver(K, f, gdl_restringidos, u_R=None):
    """Resuelve K u = f imponiendo los desplazamientos de los GDL restringidos.

    u_R: desplazamientos impuestos en gdl_restringidos (por defecto, ceros).
    Retorna el vector u completo. Si K_LL es (casi) singular, la estructura es
    un mecanismo y se lanza un error en lugar de entregar resultados absurdos.
    """
    gdl_libres = [g for g in range(len(f)) if g not in gdl_restringidos]
    u = np.zeros(len(f))
    if u_R is not None:
        u[gdl_restringidos] = u_R
    K_LL = K[np.ix_(gdl_libres, gdl_libres)]
    K_LR = K[np.ix_(gdl_libres, gdl_restringidos)]
    if np.linalg.cond(K_LL) > 1e12:
        raise np.linalg.LinAlgError("K_LL es singular: revisa los apoyos (¿la estructura es un mecanismo?)")
    u[gdl_libres] = np.linalg.solve(K_LL, f[gdl_libres] - K_LR @ u[gdl_restringidos])
    return u


# ============================ Sección III ============================

def reacciones(K, u, f, gdl_restringidos):
    """Reacciones R = K_RL u_L + K_RR u_R - f_R^ap.

    Es lo mismo que evaluar (K u - f) en los GDL restringidos.
    """
    return (K @ u - f)[gdl_restringidos]


def fuerzas_barras(nudos, barras, A, E, u):
    """Fuerza axial P de cada barra de una armadura 2D (positiva en tracción)."""
    A = np.broadcast_to(A, len(barras))
    E = np.broadcast_to(E, len(barras))
    P = np.zeros(len(barras))
    for b, ((i, j), A_e, E_e) in enumerate(zip(barras, A, E)):
        L, c, s = geometria_barra(nudos[i], nudos[j])
        ux1, uy1, ux2, uy2 = u[gdl_barra(i, j)]
        delta = (ux2 - ux1)*c + (uy2 - uy1)*s   # u'_2 - u'_1
        P[b] = A_e * E_e / L * delta
    return P


def graficar_armadura(nudos, barras, u=None, P=None, escala=1.0, ax=None):
    """Dibuja una armadura 2D, numerando los nudos desde 1.

    Si se entrega u, dibuja también la deformada (desplazamientos multiplicados
    por escala). Si además se entrega P, colorea las barras deformadas:
    rojo = tracción, azul = compresión.
    """
    nudos = np.asarray(nudos, dtype=float)
    if ax is None:
        _, ax = plt.subplots()
    for b, (i, j) in enumerate(barras):
        ax.plot(*nudos[[i, j]].T, "k--" if u is not None else "k-", lw=1)
        if u is not None:
            deformada = nudos + escala * np.reshape(u, (-1, 2))
            color = "k" if P is None else ("tab:red" if P[b] > 0 else "tab:blue")
            ax.plot(*deformada[[i, j]].T, color=color, lw=3)
    for n, (x, y) in enumerate(nudos, start=1):
        ax.annotate(str(n), (x, y), textcoords="offset points", xytext=(5, 5))
    ax.set_aspect("equal")
    ax.set_xlabel("x (m)")
    ax.set_ylabel("y (m)")
    return ax


# ============================ Sección IV ============================

def k_marco(A, E, I, xy1, xy2):
    """Matriz de rigidez 6x6 de un elemento de marco en ejes globales.

    GDL por nudo: (ux, uy, θ). Combina la rigidez axial y la de flexión.
    """
    L, c, s = geometria_barra(xy1, xy2)
    a = A * E / L
    b = E * I / L**3
    k_local = np.array([
        [ a,      0,          0, -a,      0,          0],
        [ 0,   12*b,     6*L*b,   0,  -12*b,     6*L*b],
        [ 0,  6*L*b,  4*L**2*b,   0, -6*L*b,  2*L**2*b],
        [-a,      0,          0,  a,      0,          0],
        [ 0,  -12*b,    -6*L*b,   0,   12*b,    -6*L*b],
        [ 0,  6*L*b,  2*L**2*b,   0, -6*L*b,  4*L**2*b],
    ])
    T = np.zeros((6, 6))
    T[:3, :3] = T[3:, 3:] = [[ c, s, 0],
                             [-s, c, 0],
                             [ 0, 0, 1]]
    return T.T @ k_local @ T


# ============================ Sección V ============================

def k_viga(E, I, L):
    """Matriz de rigidez local de una viga (GDL: v1, θ1, v2, θ2)."""
    return E * I / L**3 * np.array([
        [ 12,    6*L,  -12,    6*L],
        [6*L, 4*L**2, -6*L, 2*L**2],
        [-12,   -6*L,   12,   -6*L],
        [6*L, 2*L**2, -6*L, 4*L**2],
    ])


def f_eq_uniforme(w, L):
    """Cargas nodales equivalentes de una carga uniforme w (N/m, hacia abajo)."""
    return np.array([-w*L/2, -w*L**2/12, -w*L/2, w*L**2/12])


def f_eq_puntual(P, a, L):
    """Cargas nodales equivalentes de una carga puntual P (N, hacia abajo) a una distancia a del nudo 1."""
    b = L - a
    return np.array([-P*b**2*(3*a + b)/L**3, -P*a*b**2/L**2,
                     -P*a**2*(a + 3*b)/L**3,  P*a**2*b/L**2])


def hermite(xi, L):
    """Funciones de forma de Hermite evaluadas en xi = x/L (filas: v1, θ1, v2, θ2)."""
    return np.array([1 - 3*xi**2 + 2*xi**3,
                     L*(xi - 2*xi**2 + xi**3),
                     3*xi**2 - 2*xi**3,
                     L*(-xi**2 + xi**3)])