import matplotlib.pyplot as plt
import numpy as np

import pysra

%matplotlib inline

# Increased figure sizes
plt.rcParams["figure.dpi"] = 120

"""
Explanation: Example 8: RVT SRA with simulated profiles
Use RVT with simulated profiles to compute the surface response spectrum and surface spectral ratio.
End of explanation
"""


m = pysra.motion.SourceTheoryRvtMotion(6.0, 30, "wna")
m.calc_fourier_amps()

"""
Explanation: Create a point source theory RVT motion
End of explanation
"""


profile = pysra.site.Profile(
    [
        pysra.site.Layer(
            pysra.site.DarendeliSoilType(18.0, plas_index=0, ocr=1, stress_mean=100),
            10,
            400,
        ),
        pysra.site.Layer(
            pysra.site.DarendeliSoilType(18.0, plas_index=0, ocr=1, stress_mean=200),
            10,
            450,
        ),
        pysra.site.Layer(
            pysra.site.DarendeliSoilType(18.0, plas_index=0, ocr=1, stress_mean=400),
            30,
            600,
        ),
        pysra.site.Layer(pysra.site.SoilType("Rock", 24.0, None, 0.01), 0, 1200),
    ]
)

"""
Explanation: Create site profile
This is about the simplest profile that we can create. Linear-elastic soil and rock.
End of explanation
"""


calc = pysra.propagation.EquivalentLinearCalculator()

"""
Explanation: Create the site response calculator
End of explanation
"""


var_thickness = pysra.variation.ToroThicknessVariation()
var_velocity = pysra.variation.ToroVelocityVariation.generic_model("USGS C")
var_soiltypes = pysra.variation.SpidVariation(
    -0.5, std_mod_reduc=0.15, std_damping=0.30
)

"""
Explanation: Initialize the variations
End of explanation
"""


freqs = np.logspace(-1, 2, num=500)

outputs = pysra.output.OutputCollection(
    [
        pysra.output.ResponseSpectrumOutput(
            # Frequency
            freqs,
            # Location of the output
            pysra.output.OutputLocation("outcrop", index=0),
            # Damping
            0.05,
        ),
        pysra.output.ResponseSpectrumRatioOutput(
            # Frequency
            freqs,
            # Location in (denominator),
            pysra.output.OutputLocation("outcrop", index=-1),
            # Location out (numerator)
            pysra.output.OutputLocation("outcrop", index=0),
            # Damping
            0.05,
        ),
        pysra.output.InitialVelProfile(),
    ]
)

"""
Explanation: Specify the output
End of explanation
"""


count = 20
for p in pysra.variation.iter_varied_profiles(
    profile,
    count,
    var_thickness=var_thickness,
    var_velocity=var_velocity,
    var_soiltypes=var_soiltypes,
):
    # Here we auto-descretize the profile for wave propagation purposes
    calc(m, p.auto_discretize(), p.location("outcrop", index=-1))
    outputs(calc)

"""
Explanation: Perform the calculation
End of explanation
"""


for o in outputs[:-1]:
    ax = o.plot(style="stats")

ax = outputs[-1].plot(style="stats")

"""
Explanation: Plot the outputs
Create a few plots of the output.
End of explanation
"""


outputs[-1].calc_stats()

outputs[-1].calc_stats(as_dataframe=True)

"""
Explanation: The statistics of the output can be also retrieved and returned as either a dict or pandas.DataFrame.
End of explanation
"""


var_velocity_dd = pysra.variation.DepthDependToroVelVariation.generic_model("USGS C")

outputs.reset()
count = 20
for p in pysra.variation.iter_varied_profiles(
    profile,
    count,
    var_thickness=var_thickness,
    var_velocity=var_velocity_dd,
    var_soiltypes=var_soiltypes,
):
    calc(m, p.auto_discretize(), p.location("outcrop", index=-1))
    outputs(calc)

for o in outputs[:-1]:
    ax = o.plot(style="stats")

ax = outputs[-1].plot(style="stats")

outputs[-1].calc_stats(as_dataframe=True)

"""
Explanation: Repeat using a Generic Depth-Dependent Model
The generic DepthDependToroVelVariation follows the SPID guidance.
End of explanation
"""


var_velocity_dd = pysra.variation.DepthDependToroVelVariation.generic_model(
    "USGS C", depth=[0, 10, 20], ln_std=[0.25, 0.15, 0.10]
)

outputs.reset()
count = 20
for p in pysra.variation.iter_varied_profiles(
    profile,
    count,
    var_thickness=var_thickness,
    var_velocity=var_velocity_dd,
    var_soiltypes=var_soiltypes,
):
    calc(m, p.auto_discretize(), p.location("outcrop", index=-1))
    outputs(calc)

for o in outputs[:-1]:
    ax = o.plot(style="stats")

ax = outputs[-1].plot(style="stats")

outputs[-1].calc_stats(as_dataframe=True)

"""
Explanation: Repeat using a Specific Depth-Dependent Model
End of explanation
"""


