# -*- coding: utf-8 -*-
"""
Created on Thu Dec 15 15:28:28 2016

@author: Alberto Camjayi
"""

# Parámetros del problema
g = 9.8                 # aceleración de la gravedad
m = 1                  # masa de los péndulos
l_largos = 1           # largo de los primeros péndulos
l_cortos = 0.5         # idem de los del final de la cadena
k = 5                  # constante de los resortes
Gama = 0.1             # constante de amortiguamiento
N_largos = 30          # número de péndulos largos
N_cortos = 200         # número de péndulos cortos
F0_sobre_m = 4         # fuerza sobre la primera masa
Omega = 3              # frecuencia angular de forzado
# ----------------------------------------------------------------
# Importamos algunas herramientas necesarias
from scipy import *
from scipy.linalg import eigh
from scipy.sparse import spdiags, hstack, vstack, bmat, dia_matrix
from pylab import figure, grid, plot, axis, xlabel, ylabel, title, legend, show

def get_M(N1,l1,N2,l2):             # Aquí construímos la matriz
    Nm = N1+N2
    diag0 = zeros(Nm); diag1 = zeros(Nm)
    diag0[:N1] = g/l1+2*k/m          # Elementos de la diagonal
    diag0[N1:] = g/l2+2*k/m          # Elementos de la diagonal
    diag1[1:] = -k/m                 # Elementos extradiagonales

    Mi = spdiags([diag0,diag1],[0,1],Nm,Nm)
    return Mi

def get_sol_estacionaria(w0,Fm,Omega):
    Aes = zeros(w0.size)
    Bes = zeros(w0.size)

    Aes =         Omega*Gama*Fm/((w0**2-Omega**2)**2 + Omega**2*Gama**2)
    Bes = (w0**2 - Omega**2)*Fm/((w0**2-Omega**2)**2 + Omega**2*Gama**2)

    return Aes, Bes


M = get_M(N_largos,l_largos,N_cortos,l_cortos).toarray()
w0_cuadrado, Dmatrix = eigh(M,lower=False)

F = zeros(N_largos+N_cortos); F[0] = F0_sobre_m

F_en_modos = dot(transpose(Dmatrix),F)


for i in range(7):
    Omega = Omega + i*3

    A_en_modos, B_en_modos = get_sol_estacionaria(w0_cuadrado,F_en_modos,Omega)