import pandas as pd
from profitability import DATA_DIR, DEFAULT_MECH_FEE, summary_analyse
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 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)