import streamlit as st
import bruges as b
from einops import repeat
import numpy as np
from bruges.reflection.reflection import zoeppritz_rpp as zrpp
import pandas as pd
import altair as alt


def vs_from_poisson(vp, poisson):
    '''compute pressure velocity from
    shear velocity and poissons ratio'''
    return np.sqrt((vp**2 - 2*poisson*vp**2)/(2 - 2*poisson))


def gardners(vp):
    '''compute density via gardners relation'''
    return 1000*0.31*np.power(vp, 0.25)/1.33


def cartesian_product(*arrays):
    '''return cartesion product of several arrays'''
    ndim = len(arrays)
    return (np.stack(np.meshgrid(*arrays), axis=-1)
              .reshape(-1, ndim))


def loop_zrpp(vp1,vs1,rho1,vp2,vs2,rho2,theta1):
    '''compute reflectivity for many values of
    input parameters.
    "1" denotes layer one, "2" denotes layer two.
    outside this function, I tend to use 0-indexing though.'''
    refl_loop = np.empty(len(vp1), dtype=complex)
    for i in range(vp1.shape[0]):
        refl_loop[i] = zrpp(vp1=vp1[i], vs1=vs1[i], rho1=rho1[i],
                         vp2=vp2[i], vs2=vs2[i], rho2=rho2[i],
                         theta1=theta1[i])
    return refl_loop

def wb_ava_fig(vwater=1520., rwater = 1025., poissons=0.49):
    '''
    AVO: Amplitude Versus Offset
    AVA: Amplitude Versus Angle
    compute AVA at the WB for a range of values,
    then plot using altair
    '''
    ### Set up some variables
    VP1 = np.arange(1530,1820,50)
    THE = np.arange(0.0, 88., 1.)
    POI = np.array([poissons])
    
    ### Set variables that depend on those variables
    params = cartesian_product(VP1, THE, POI)
    VP1, THE, POI = [a.ravel() for a in np.hsplit(params, 3)]
    VP0 = np.full(VP1.shape, vwater)
    # V-RMS for ~200 m water depth
    VS0 = np.full(VP1.shape, 0.)
    RH0 = np.full(VP1.shape, rwater)
    # 1025 kg/m^3 per Inversion of the physical properties of seafloor surface, South China Sea, Zhou et al 2021
    VP1 = VP1
    VS1 = vs_from_poisson(VP1,POI)
    RH1 = gardners(VP1)
    
    ### Shove into a dict to pass to zrpp
    params = {"vp1": VP0, "vs1": VS0, "rho1": RH0,
          "vp2": VP1, "vs2": VS1, "rho2": RH1,
          "theta1":THE}

    # Compute zoeppritz equation
    r = loop_zrpp(**params)
    
    # Put the results into a DataFrame
    df = pd.DataFrame({"Vp sub-WB": VP1, "Poisson_s ratio": POI, "Vs sub-WB": VS1,
                   "Angle": THE, "Amplitude": np.real(r)})
    
    # Select only points pre-critical angle
    df["Ang_Crit"] = np.degrees(np.arcsin(1500 / df["Vp sub-WB"].values))
    df = df[df["Angle"] < df["Ang_Crit"]]
    
    # Create the altair figure
    highlight = alt.selection(type='single', on='mouseover',
                          fields=["Vp sub-WB:Q"], nearest=True)

    base = alt.Chart(df).encode(
            x="Angle",
            y="Amplitude",
            color="Vp sub-WB:Q",
            tooltip=["Vp sub-WB", "Vs sub-WB", "Angle", "Amplitude"]
    )

    points = base.mark_circle().encode(
        opacity=alt.value(0)
    ).add_selection(
        highlight
    ).properties(
        width=600,
        height=450
    )

    lines = base.mark_line().encode(
        size=alt.condition(~highlight, alt.value(3), alt.value(7))
    )
    
    return points + lines

vwater = st.slider("Select a value for water velocity", min_value=1495, max_value=1545, value=1520)
rwater = st.slider("Select a value for density of water", min_value=1005, max_value=1045, value=1025)
poissons = st.slider("Select a value for Poisson's ratio", min_value=0.4, max_value=0.5, value=0.48)
# st.write("Poisson's ratio is:", p)


st.altair_chart(wb_ava_fig(vwater, rwater, poissons))