import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from PIL import Image


def response_spectra(h, Tmax=10, delta=0.1, dt=0.02):
    m = 1
    S = []
    Ts = np.arange(delta, Tmax+delta, delta)
    t = np.arange(dt, dt*(len(h)+1), dt)
    
    for T in Ts:
        w = 2 * np.pi / T
        k = m * w**2
        d, v, a = newmark_int(t, -m*h, m, k, 0.05)
        tmp = np.array([np.max(np.abs(a[1:] + h[:-1])), np.max(np.abs(v)), np.max(np.abs(d))])
        S.append(tmp)
    return Ts, np.stack(S, 1)


def newmark_int(t, p, m, k, damping):
    gam = 1/2
    beta = 1/4
    wn = np.sqrt(k/m)
    
    dt = t[1] - t[0]
    c = 2 * damping * wn * m
    kgor = k + gam / (beta * dt) * c + m / (beta * (dt**2))
    a = m / (beta * dt) + gam * c / beta
    b = 0.5 * m / beta + dt * (0.5 * gam / beta - 1) * c
    
    dp = diff(p)
    dis = np.zeros(len(t),)
    vel = np.zeros(len(t),)
    acc = np.zeros(len(t),)
    dis[0] = 0
    vel[0] = 0
    acc[0] = 1 / m * p[0]
    
    for i in range(len(t)-2):
        deltaP = dp[i] + a*vel[i] + b*acc[i]
        du_i = deltaP / kgor
        dv_i = gam/(beta*dt)*du_i - gam/beta*vel[i] + dt*(1-0.5*gam/beta)*acc[i]
        da_i = 1/(beta*(dt**2))*du_i - 1/(beta*dt)*vel[i] - 0.5/beta*acc[i]
        dis[i+1] = du_i + dis[i]
        vel[i+1] = dv_i + vel[i]
        acc[i+1] = da_i + acc[i]
    return dis, vel, acc


def diff(h):
    h0 = np.concatenate(([0], h))[:-1]
    return h - h0


def run(name):
    f = np.load(name)
    a = f[1] / np.max(np.abs(f[1]))
    t, rspec = response_spectra(a, Tmax=10, delta=0.05, dt=0.02)
    
    fig = plt.figure(figsize=(8,12), dpi=300, tight_layout=True)
    plt.subplot(4,1,1)
    plt.plot(np.arange(0, len(a)) * 0.02, a, label='earthquake acceleration')
    plt.ylabel('Acc (m/s^2)', fontsize=14)
    plt.yticks(fontsize=12)
    plt.legend(fontsize=12)
    
    plt.subplot(4,1,2)
    plt.plot(t, rspec[0])
    plt.xlabel('Period (s)', fontsize=14)
    plt.ylabel('Sa (m/s^2)', fontsize=14)
    plt.yticks(fontsize=12)
    
    plt.subplot(4,1,3)
    plt.plot(t, rspec[1])
    plt.xlabel('Period (s)', fontsize=14)
    plt.ylabel('Sv (m/s)', fontsize=14)
    plt.yticks(fontsize=12)
    
    plt.subplot(4,1,4)
    plt.plot(t, rspec[2])
    plt.xlabel('Period (s)', fontsize=14)
    plt.ylabel('Sd (m)', fontsize=14)
    plt.yticks(fontsize=12)
    plt.show()
    
    fig.savefig('sample.png', dpi=300)
    return Image.open('sample.png')


dropdown = gr.Dropdown(['test-{}.npy'.format(i+1) for i in range(4)])
iface = gr.Interface(fn=run, inputs=dropdown, outputs='image')
iface.launch()