import pandas as pd
import numpy as np


class SimpleBacktester:
    def __init__(self, top_n: int = 10, hold_days: int = 1, commission: float = 0.001):
        self.top_n = top_n
        self.hold_days = hold_days
        self.commission = commission

    def run(self, signals: pd.DataFrame, price_df: pd.DataFrame) -> dict:
        if signals.empty or price_df.empty:
            return {"error": "empty data"}

        results = []
        dates = signals["date"].unique()

        for i, date in enumerate(dates):
            day_signals = signals[signals["date"] == date].nlargest(self.top_n, "score")
            if day_signals.empty:
                continue

            next_dates = [d for d in dates if d > date][:self.hold_days]
            if not next_dates:
                continue

            for _, row in day_signals.iterrows():
                code = row["code"]
                entry_price = row.get("close", row.get("price", 0))
                if entry_price <= 0:
                    continue

                stock_prices = price_df[price_df["code"] == code].sort_values("date")
                exit_row = stock_prices[stock_prices["date"] == next_dates[-1]]
                if exit_row.empty:
                    continue

                exit_price = exit_row.iloc[0]["close"]
                ret = (exit_price - entry_price) / entry_price - self.commission * 2
                results.append({
                    "date": date,
                    "code": code,
                    "entry_price": entry_price,
                    "exit_price": exit_price,
                    "return": ret,
                    "hold_days": len(next_dates),
                })

        if not results:
            return {"error": "no trades"}

        df = pd.DataFrame(results)
        return self._calc_metrics(df)

    def run_from_returns(self, daily_returns: pd.Series) -> dict:
        if daily_returns.empty:
            return {"error": "empty data"}

        total_return = (1 + daily_returns).prod() - 1
        n_days = len(daily_returns)
        annual_return = (1 + total_return) ** (252 / max(n_days, 1)) - 1
        annual_vol = daily_returns.std() * np.sqrt(252)
        sharpe = annual_return / annual_vol if annual_vol > 0 else 0

        cumulative = (1 + daily_returns).cumprod()
        running_max = cumulative.cummax()
        drawdown = (cumulative - running_max) / running_max
        max_drawdown = drawdown.min()

        win_rate = (daily_returns > 0).mean()
        avg_win = daily_returns[daily_returns > 0].mean() if (daily_returns > 0).any() else 0
        avg_loss = daily_returns[daily_returns <= 0].mean() if (daily_returns <= 0).any() else 0
        profit_factor = abs(avg_win / avg_loss) if avg_loss != 0 else float("inf")

        return {
            "total_return": total_return,
            "annual_return": annual_return,
            "annual_volatility": annual_vol,
            "sharpe_ratio": sharpe,
            "max_drawdown": max_drawdown,
            "win_rate": win_rate,
            "avg_win": avg_win,
            "avg_loss": avg_loss,
            "profit_factor": profit_factor,
            "trading_days": n_days,
        }

    def _calc_metrics(self, trades_df: pd.DataFrame) -> dict:
        daily_ret = trades_df.groupby("date")["return"].mean()
        metrics = self.run_from_returns(daily_ret)
        metrics["total_trades"] = len(trades_df)
        metrics["avg_trades_per_day"] = len(trades_df) / trades_df["date"].nunique()
        metrics["trades"] = trades_df
        return metrics


def generate_labels(df: pd.DataFrame, forward_days: int = 1, threshold: float = 0.02) -> pd.Series:
    future_ret = df.groupby("code")["close"].pct_change(forward_days).shift(-forward_days)
    labels = (future_ret >= threshold).astype(int)
    return labels