import pandas as pd


def rolling_beta(
    stock: pd.Series,
    market: pd.Series,
    window: int,
) -> pd.Series:

    cov = stock.rolling(window).cov(market)
    var = market.rolling(window).var()

    return cov / (var + 1e-9)


def add_correlation_features(
    df: pd.DataFrame,
) -> pd.DataFrame:

    results = []

    for symbol, grp in df.groupby(
        "symbol",
        sort=False,
    ):

        grp = grp.sort_values("timestamp").copy()

        stock_ret = grp["ret_1d"]

        # ─────────────────────────────────────────
        # NIFTY
        # ─────────────────────────────────────────

        if "nifty_ret_1d" in grp.columns:

            mkt = grp["nifty_ret_1d"]

            grp["corr_nifty_20d"] = (
                stock_ret.rolling(20).corr(mkt)
            )

            grp["corr_nifty_60d"] = (
                stock_ret.rolling(60).corr(mkt)
            )

            grp["beta_nifty_60d"] = rolling_beta(
                stock_ret,
                mkt,
                60,
            )

            grp["is_high_beta"] = (
                grp["beta_nifty_60d"] > 1.2
            ).astype("int8")

            grp["is_low_beta"] = (
                grp["beta_nifty_60d"] < 0.8
            ).astype("int8")

            grp["corr_breakdown"] = (
                grp["corr_nifty_20d"]
                <
                grp["corr_nifty_60d"] - 0.3
            ).astype("int8")

        # ─────────────────────────────────────────
        # Relative strength
        # ─────────────────────────────────────────

        if "nifty_ret_5d" in grp.columns:

            grp["rel_strength_5d"] = (
                grp["ret_5d"]
                - grp["nifty_ret_5d"]
            )

        if "nifty_ret_20d" in grp.columns:

            grp["rel_strength_20d"] = (
                grp["ret_20d"]
                - grp["nifty_ret_20d"]
            )

        # ─────────────────────────────────────────
        # Macro
        # ─────────────────────────────────────────

        for asset in [
            "gold",
            "brent",
            "usd_inr",
        ]:

            col = f"{asset}_ret_1d"

            if col in grp.columns:

                ar = grp[col]

                corr_col = f"corr_{asset}_60d"

                grp[corr_col] = (
                    stock_ret
                    .rolling(60)
                    .corr(ar)
                )

                grp[
                    f"corr_{asset}_rising"
                ] = (
                    grp[corr_col]
                    >
                    grp[corr_col].shift(20)
                ).astype("int8")

        # ─────────────────────────────────────────
        # US market
        # ─────────────────────────────────────────

        if "us_ret_1d" in grp.columns:

            grp["corr_us_60d"] = (
                stock_ret
                .rolling(60)
                .corr(grp["us_ret_1d"])
            )

        results.append(grp)

    if not results:
        return df

    return (
        pd.concat(results)
        .sort_values(["timestamp", "symbol"])
    )