import pandas as pd
from profitability import summary_analyse
from utils import DATA_DIR
from staking import label_trades_by_staking


def clean_old_data_from_parquet_files(cutoff_date: str):
    print("Cleaning oldest data")
    # Convert the string to datetime64[ns, UTC]
    min_date_utc = pd.to_datetime(cutoff_date, format="%Y-%m-%d", utc=True)

    # clean tools.parquet
    try:
        tools = pd.read_parquet(DATA_DIR / "tools.parquet")

        # make sure creator_address is in the columns
        assert "trader_address" in tools.columns, "trader_address column not found"

        # lowercase and strip creator_address
        tools["trader_address"] = tools["trader_address"].str.lower().str.strip()

        tools["request_time"] = pd.to_datetime(tools["request_time"], utc=True)

        print(f"length before filtering {len(tools)}")
        tools = tools.loc[tools["request_time"] > min_date_utc]
        print(f"length after filtering {len(tools)}")
        tools.to_parquet(DATA_DIR / "tools.parquet", index=False)

    except Exception as e:
        print(f"Error cleaning tools file {e}")

    # clean all_trades_profitability.parquet
    try:
        all_trades = pd.read_parquet(DATA_DIR / "all_trades_profitability.parquet")

        all_trades["creation_timestamp"] = pd.to_datetime(
            all_trades["creation_timestamp"], utc=True
        )

        print(f"length before filtering {len(all_trades)}")
        all_trades = all_trades.loc[all_trades["creation_timestamp"] > min_date_utc]
        print(f"length after filtering {len(all_trades)}")
        all_trades.to_parquet(
            DATA_DIR / "all_trades_profitability.parquet", index=False
        )

    except Exception as e:
        print(f"Error cleaning all trades profitability file {e}")

    # generate again summary_profitability.parquet
    try:
        print("Summarising trades...")
        summary_df = summary_analyse(all_trades)

        # add staking labels
        label_trades_by_staking(trades_df=all_trades, update=False)

        # save to parquet
        all_trades.to_parquet(
            DATA_DIR / "all_trades_profitability.parquet", index=False
        )
        summary_df.to_parquet(DATA_DIR / "summary_profitability.parquet", index=False)
    except Exception as e:
        print(f"Error generating summary and saving all trades profitability file {e}")


if __name__ == "__main__":
    clean_old_data_from_parquet_files("2024-09-22")