import numpy as np
import pandas as pd


def sto06(Hs0, Tp, beta):
    """
    :param Hs0: List or float of offshore significant wave height values
    :param Tp: List or float of peak wave period
    :param beta: List of float of beach slope
    :return: Float or list of R2, setup, S_total, S_inc and S_ig values
    """

    df = pd.DataFrame({"Hs0": Hs0, "Tp": Tp, "beta": beta}, index=[x for x in range(0, np.size(Hs0))])

    df["Lp"] = 9.8 * df["Tp"] ** 2 / 2 / np.pi

    # General equation
    df["S_ig"] = pd.to_numeric(0.06 * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
    df["S_inc"] = pd.to_numeric(0.75 * df["beta"] * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
    df["setup"] = pd.to_numeric(0.35 * df["beta"] * np.sqrt(df["Hs0"] * df["Lp"]), errors="coerce")
    df["S_total"] = np.sqrt(df["S_inc"] ** 2 + df["S_ig"] ** 2)
    df["R2"] = 1.1 * (df["setup"] + df["S_total"] / 2)

    # Dissipative conditions
    dissipative = df["beta"] / (df["Hs0"] / df["Lp"]) ** (0.5) <= 0.3

    df.loc[dissipative, "setup"] = 0.016 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 16
    df.loc[dissipative, "S_total"] = 0.046 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 17
    df.loc[dissipative, "R2"] = 0.043 * (df["Hs0"] * df["Lp"]) ** (0.5)  # eqn 18

    return (
        float_or_list(df["R2"].tolist()),
        float_or_list(df["setup"].tolist()),
        float_or_list(df["S_total"].tolist()),
        float_or_list(df["S_inc"].tolist()),
        float_or_list(df["S_ig"].tolist()),
    )


def float_or_list(a):
    """
    If only one value in the array, return the float, else return a list
    :param a:
    :return:
    """
    if len(a) == 1:
        return a[0]
    else:
        return list(a)


if __name__ == "__main__":
    pass