import pandas as pd
from utils import DATA_DIR, DEFAULT_MECH_FEE
from tqdm import tqdm


def update_roi(row: pd.DataFrame) -> float:
    new_value = row.net_earnings / (
        row.collateral_amount
        + row.trade_fee_amount
        + row.num_mech_calls * DEFAULT_MECH_FEE
    )
    return new_value


def create_unknown_traders_df(trades_df: pd.DataFrame) -> pd.DataFrame:
    """filter trades coming from non-Olas traders that are placing no mech calls"""
    no_mech_calls_mask = (trades_df["staking"] == "non_Olas") & (
        trades_df["num_mech_calls"] == 0
    )
    no_mech_calls_df = trades_df.loc[no_mech_calls_mask]
    trades_df = trades_df.loc[~no_mech_calls_mask]
    return no_mech_calls_df, trades_df


def update_trade_nr_mech_calls(non_agents: bool = False):
    try:
        all_trades_df = pd.read_parquet(DATA_DIR / "all_trades_profitability.parquet")
        tools = pd.read_parquet(DATA_DIR / "tools.parquet")
    except Exception as e:
        print(f"Error reading the profitability and tools parquet files")

    traders = list(all_trades_df.trader_address.unique())
    if non_agents:
        traders = list(
            all_trades_df.loc[
                all_trades_df["staking"] == "non_agent"
            ].trader_address.unique()
        )

    print("before updating")
    print(
        all_trades_df.loc[
            all_trades_df["staking"] == "non_agent"
        ].num_mech_calls.describe()
    )
    for trader in tqdm(traders, desc=f"Updating Traders mech calls", unit="traders"):
        tools_usage = tools[tools["trader_address"] == trader]
        if len(tools_usage) == 0:
            tqdm.write(f"trader with no tools usage found {trader}")
            all_trades_df.loc[
                all_trades_df["trader_address"] == trader, "nr_mech_calls"
            ] = 0
    # update roi
    all_trades_df["roi"] = all_trades_df.apply(lambda x: update_roi(x), axis=1)
    print("after updating")
    print(
        all_trades_df.loc[
            all_trades_df["staking"] == "non_agent"
        ].num_mech_calls.describe()
    )

    # saving
    all_trades_df.to_parquet(DATA_DIR / "all_trades_profitability.parquet", index=False)
    # print("Summarising trades...")
    # summary_df = summary_analyse(all_trades_df)
    # summary_df.to_parquet(DATA_DIR / "summary_profitability.parquet", index=False)


if __name__ == "__main__":
    # update_trade_nr_mech_calls(non_agents=True)
    trades_df = pd.read_parquet(DATA_DIR / "all_trades_profitability.parquet")
    print("before filtering")
    print(trades_df.staking.value_counts())
    unknown_df, trades_df = create_unknown_traders_df(trades_df=trades_df)
    print("after filtering")
    print(trades_df.staking.value_counts())
    print("saving unknown traders")
    unknown_df.to_parquet(DATA_DIR / "unknown_traders.parquet", index=False)
    trades_df.to_parquet(DATA_DIR / "all_trades_profitability.parquet", index=False)