# AUTOGENERATED! DO NOT EDIT! File to edit: 00_tspecscores.ipynb (unless otherwise specified).

__all__ = ['log_it', 'tsi', 'spm', 'zscore', 'tau', 'ts_func', 'calc_ts']

# Cell
import pandas as pd
import numpy as np


# Cell
def log_it(data: pd.DataFrame) -> pd.DataFrame:
    df = data.copy()
    return np.log(1 + df)


# Cell
def tsi(data: pd.DataFrame) -> pd.DataFrame:
    df = data.copy()
    foo = np.sum(df.values, axis=1)[:, None]
    return pd.DataFrame(
        np.divide(df.values, foo, out=df.values, where=foo != 0), index=df.index, columns=df.columns
    )


# Cell
def spm(data: pd.DataFrame) -> pd.DataFrame:
    df = data.copy()
    foo = np.linalg.norm(df.values, axis=1)[:, None] * df.values
    return pd.DataFrame(
        np.divide(df.values ** 2, foo, out=df.values, where=foo != 0),
        index=df.index,
        columns=df.columns,
    )


# Cell
def zscore(data: pd.DataFrame, transform=True) -> pd.DataFrame:
    df = data.copy()
    std = np.std(df.values, axis=1, ddof=1)[:, None]
    zs = np.divide(df.values - df.values.mean(axis=1)[:, None], std, where=std != 0)
    zs[np.where(std == 0)[0], :] = 0
    if transform:
        max_zs = (df.values.shape[1] - 1) / np.sqrt(df.values.shape[1])
        zs = np.divide((zs + max_zs), (2 * max_zs), where=zs != 0)
    return pd.DataFrame(zs, index=df.index, columns=df.columns)



# Cell
def tau(data):
    df = data.copy()
    foo = np.max(df.values, axis=1)[:, None]
    vector_r = np.divide(df.values, foo, out=df.values, where=foo != 0)
    tau_index = np.sum(1 - vector_r, axis=1) / (df.values.shape[1] - 1)
    tau_index[(df.sum(axis=1) == 0).values] = 0
    return pd.Series(tau_index, index=df.index)


# Cell
ts_func = {"tsi": tsi, "spm": spm, "zscore": zscore, "tau": tau}


# Cell
def calc_ts(
    ts_df: pd.DataFrame, method: str = "spm", log: bool = True, **kwargs
) -> pd.DataFrame:
    if log:
        ts_df = log_it(ts_df)

    return ts_func[method](ts_df, **kwargs)
