import pandas as pd

try:
    from utils import DATA_DIR, TMP_DIR
except ImportError:
    from scripts.utils import DATA_DIR, TMP_DIR

from datetime import datetime, timezone
from tqdm import tqdm


def transform_to_datetime(x):
    return datetime.fromtimestamp(int(x), tz=timezone.utc)


def get_weekly_total_mech_calls(trader_data: pd.DataFrame) -> int:
    """Function to compute the total weekly number of mech calls for all markets
    that the trader bet upon"""
    try:
        all_mech_calls_df = pd.read_parquet(DATA_DIR / "weekly_mech_calls.parquet")
    except Exception:
        print("Error reading the weekly_mech_calls file")

    trading_weeks = trader_data.month_year_week.unique()
    trader_address = trader_data.trader_address.unique()[0]
    if len(trading_weeks) > 1:
        raise ValueError("The trader data should contain only one week information")
    trading_week = trading_weeks[0]
    try:
        return all_mech_calls_df.loc[
            (all_mech_calls_df["trader_address"] == trader_address)
            & (all_mech_calls_df["month_year_week"] == trading_week),
            "total_mech_calls",
        ].iloc[0]
    except Exception as e:
        print(
            f"Error getting the number of mech calls for the trader {trader_address} and week {trading_week}"
        )
        return 280  # average number 40 mech calls in 7 days


def compute_weekly_total_mech_calls(
    trader: str, week: str, weekly_trades: pd.DataFrame, weekly_tools: pd.DataFrame
) -> dict:
    weekly_total_mech_calls_dict = {}
    weekly_total_mech_calls_dict["trader_address"] = trader
    weekly_total_mech_calls_dict["month_year_week"] = week
    weekly_total_mech_calls_dict["total_trades"] = len(weekly_trades)
    weekly_total_mech_calls_dict["total_mech_calls"] = len(weekly_tools)
    return weekly_total_mech_calls_dict


def compute_total_mech_calls():
    """Function to compute the total number of mech calls for all traders and all markets
    at a weekly level"""
    try:
        print("Reading tools file")
        tools = pd.read_parquet(TMP_DIR / "tools.parquet")
        tools["request_time"] = pd.to_datetime(tools["request_time"])
        tools["request_date"] = tools["request_time"].dt.date
        tools = tools.sort_values(by="request_time", ascending=True)
        tools["month_year_week"] = (
            tools["request_time"].dt.to_period("W").dt.strftime("%b-%d-%Y")
        )

    except Exception as e:
        print(f"Error updating the invalid trades parquet {e}")

    print("Reading trades weekly info file")
    fpmmTrades = pd.read_parquet(DATA_DIR / "fpmmTrades.parquet")
    try:
        fpmmTrades["creationTimestamp"] = fpmmTrades["creationTimestamp"].apply(
            lambda x: transform_to_datetime(x)
        )
    except Exception as e:
        print(f"Transformation not needed")

    fpmmTrades["creation_timestamp"] = pd.to_datetime(fpmmTrades["creationTimestamp"])
    fpmmTrades["creation_date"] = fpmmTrades["creation_timestamp"].dt.date
    fpmmTrades = fpmmTrades.sort_values(by="creation_timestamp", ascending=True)
    fpmmTrades["month_year_week"] = (
        fpmmTrades["creation_timestamp"].dt.to_period("W").dt.strftime("%b-%d-%Y")
    )

    nr_traders = len(fpmmTrades["trader_address"].unique())
    all_mech_calls = []
    for trader in tqdm(
        fpmmTrades["trader_address"].unique(),
        total=nr_traders,
        desc="creating weekly mech calls dataframe",
    ):
        # compute the mech calls estimations for each trader
        all_trades = fpmmTrades[fpmmTrades["trader_address"] == trader]
        all_tools = tools[tools["trader_address"] == trader]
        weeks = fpmmTrades.month_year_week.unique()

        for week in weeks:
            weekly_trades = all_trades.loc[all_trades["month_year_week"] == week]
            weekly_tools = all_tools.loc[all_tools["month_year_week"] == week]

            weekly_mech_calls_dict = compute_weekly_total_mech_calls(
                trader, week, weekly_trades, weekly_tools
            )
            all_mech_calls.append(weekly_mech_calls_dict)

    all_mech_calls_df: pd.DataFrame = pd.DataFrame.from_dict(
        all_mech_calls, orient="columns"
    )
    print("Saving weekly_mech_calls.parquet file")
    print(all_mech_calls_df.total_mech_calls.describe())

    all_mech_calls_df.to_parquet(DATA_DIR / "weekly_mech_calls.parquet", index=False)


if __name__ == "__main__":
    compute_total_mech_calls()