from boo.dataframe.canonic import is_numeric_column

# columns types


def numeric_columns(df):
    return list(filter(is_numeric_column, df.columns))


def text_columns(df):
    s = set(df.columns) - set(numeric_columns(df))
    return list(s)


def split_columns(df):
    return text_columns(df), numeric_columns(df), df.columns


# change rows


def is_alive(df):
    return (df.sales != 0) | (df.cf != 0) | (df.profit_before_tax != 0)
    # MAYBE: add large deviations from accounting identity


# change values


def change_numeraire(df, unit):
    """Change unit of account (numeraire), eg thousands to billion.
       Assumes *df* units are thousand rubles.
    """
    text_cols, num_cols, all_cols = split_columns(df)
    return df.loc[:, num_cols].divide(unit).round(1).join(df[text_cols])[all_cols]


def to_bln(df):
    return change_numeraire(df, unit=1_000_000)


def to_mln(df):  # from thousands, default
    return change_numeraire(df, unit=1_000)


# export


def large_companies(df):
    _df = (
        df.loc[is_alive(df), :]
        .query("ta > 1_000_000")
        .sort_values("ta", ascending=False)
    )
    return to_bln(_df)


def medium_companies(df):
    _df = df.query("sales >= 1_000").sort_values("sales", ascending=False)
    return to_mln(_df)


def small_companies(df):
    _df = df.query("sales <= 1_000").sort_values("sales", ascending=False)
    return to_mln(_df)


def shorthand(df):
    return df.rename(
        columns={
            "profit_before_tax": "p",
            "profit_before_tax_lag": "p_lag",
            "tp_capital": "cap",
            "tp_capital_lag": "cap_lag",
        }
    )


def no_lags(df):
    return df[[c for c in df.columns if ("_lag" not in c)]]


class Columns:
    MINIMAL = ["region", "ok1", "ok2", "title"] + [
        "ta",
        "of",
        "sales",
        "profit_before_tax",
        "cf",
    ]
    CF = ["cf_oper", "cf_inv", "cf_fin"]


def minimal_columns(df):
    return shorthand(df[Columns.MINIMAL])