"""Frequency-controlled OHLCV features computed per symbol-session from the raw 1-minute export.

Every intraday family is computed identically on bars of length 1m, 5m and 30m, obtained by aggregating the *same* raw
1-minute bars within the regular session (09:30 -> closing print). Daily features come from the same aggregation (D-003).
Zero-volume minutes are absent at the source and are never invented; the share of expected buckets present is a feature.
Bar statistics are labelled as such; nothing here is order flow, spread or informed trading.
"""
from __future__ import annotations

import math
from pathlib import Path

import numpy as np
import polars as pl

from .calendar import _xnys

FREQS = {"1m": 1, "5m": 5, "30m": 30}
INTRADAY_FAMILIES = ["rv", "bpv", "rj", "rsneg_share", "sjv", "rskew", "rkurt", "max_share", "ac1", "drift",
                     "vol_first30", "vol_last30", "vol_gini", "down_vol_share", "present_share", "rv_mktdown_share",
                     "sxy", "syy", "n_bars"]


def session_bounds_table(start, end) -> pl.DataFrame:
    """session_date, open_utc (09:30 ET bar open), close_utc (closing-print bar open), close_minute (minutes after 09:30)."""
    cal = _xnys()
    s = cal.sessions
    s = s[(s >= str(start)) & (s <= str(end))]
    opens = cal.opens.loc[s] if hasattr(cal, "opens") else None
    rows = []
    for d in s:
        o, c = cal.session_open(d), cal.session_close(d)
        rows.append({"session_date": d.date(), "open_utc": o.tz_convert("UTC").to_pydatetime(),
                     "close_utc": c.tz_convert("UTC").to_pydatetime(),
                     "close_minute": int((c - o).total_seconds() // 60)})
    df = pl.DataFrame(rows)
    return df.with_columns([pl.col("open_utc").dt.replace_time_zone("UTC"), pl.col("close_utc").dt.replace_time_zone("UTC")])


def load_regular_session(paths: list[Path], bounds: pl.DataFrame) -> pl.DataFrame:
    """Raw 1m bars restricted to the regular session, with minute-of-session index m (0 = 09:30, close_minute = closing print)."""
    lf = pl.concat([pl.scan_parquet(str(p)) for p in paths]).select(
        ["symbol_as_reported", "timestamp_utc", "session_date", "open", "high", "low", "close", "volume", "quality_flags"])
    df = lf.collect()
    if df.height == 0:
        return df
    # the lake's base+delta layouts (and chunk boundaries) can publish the same bar twice: keep the latest ingestion
    df = pl.concat([pl.scan_parquet(str(p)) for p in paths]).select(
        ["symbol_as_reported", "timestamp_utc", "session_date", "open", "high", "low", "close", "volume", "quality_flags",
         "ingested_at"]).collect().sort("ingested_at").unique(subset=["timestamp_utc"], keep="last").drop("ingested_at")
    df = df.with_columns(pl.col("timestamp_utc").dt.convert_time_zone("UTC"))
    df = df.join(bounds, on="session_date", how="inner")
    df = df.filter((pl.col("timestamp_utc") >= pl.col("open_utc")) & (pl.col("timestamp_utc") <= pl.col("close_utc")))
    df = df.with_columns(((pl.col("timestamp_utc") - pl.col("open_utc")).dt.total_minutes()).cast(pl.Int32).alias("m"))
    df = df.filter(pl.col("quality_flags").list.len() == 0)  # drop rows flagged by FinData (ohlc_inconsistent, price_jump, ...)
    return df.sort(["session_date", "m"])


def aggregate_buckets(reg: pl.DataFrame, freq_min: int) -> pl.DataFrame:
    """Aggregate regular-session 1m bars into buckets of `freq_min` minutes; the closing print joins the last bucket."""
    if freq_min == 1:
        b = reg.with_columns(pl.col("m").alias("b"))
        return b.select(["session_date", "b", "open", "high", "low", "close", "volume", "m", "close_minute"])
    nb = (pl.col("close_minute") / freq_min).ceil().cast(pl.Int32)  # buckets per session
    b = reg.with_columns(pl.min_horizontal((pl.col("m") / freq_min).floor().cast(pl.Int32), nb - 1).alias("b"))
    agg = b.group_by(["session_date", "b"], maintain_order=True).agg([
        pl.col("open").first(), pl.col("high").max(), pl.col("low").min(), pl.col("close").last(), pl.col("volume").sum(),
        pl.col("m").first(), pl.col("close_minute").first(),
    ])
    return agg.sort(["session_date", "b"])


def bucket_returns(bk: pl.DataFrame) -> pl.DataFrame:
    """Within-session log returns: first bucket open->close, then close->close."""
    prev_close = pl.col("close").shift(1).over("session_date")
    r = pl.when(prev_close.is_null()).then((pl.col("close") / pl.col("open")).log()).otherwise((pl.col("close") / prev_close).log())
    return bk.with_columns(r.alias("r"))


def _gini_expr() -> pl.Expr:
    v = pl.col("volume").sort()
    n = pl.len()
    i = pl.int_range(1, n + 1).cast(pl.Float64)
    return ((2.0 * (i * v).sum()) / (n.cast(pl.Float64) * v.sum() + 1e-12) - (n.cast(pl.Float64) + 1.0) / n.cast(pl.Float64))


def session_features(bk: pl.DataFrame, freq_min: int, spy: pl.DataFrame | None) -> pl.DataFrame:
    """One row per session with the intraday families for this frequency (prefix f{freq}_)."""
    df = bucket_returns(bk)
    if spy is not None:
        df = df.join(spy.select(["session_date", "b", pl.col("r").alias("r_spy")]), on=["session_date", "b"], how="left")
    else:
        df = df.with_columns(pl.lit(None, dtype=pl.Float64).alias("r_spy"))
    r = pl.col("r")
    r2 = r.pow(2)
    rv = r2.sum()
    n = pl.len().cast(pl.Float64)
    prev = r.shift(1).over("session_date")
    expected = pl.when(pl.lit(freq_min) == 1).then(pl.col("close_minute").first() + 1).otherwise(
        (pl.col("close_minute").first() / freq_min).ceil()).cast(pl.Float64)
    first30 = pl.col("m") < 30
    last30 = pl.col("m") > (pl.col("close_minute") - 30)
    out = df.group_by("session_date", maintain_order=True).agg([
        n.alias("n_bars"),
        rv.alias("rv"),
        ((math.pi / 2.0) * (r.abs() * prev.abs()).sum()).alias("bpv"),
        (r2.filter(r < 0).sum()).alias("rs_neg"),
        (r2.filter(r > 0).sum()).alias("rs_pos"),
        (n.sqrt() * r.pow(3).sum() / (rv.pow(1.5) + 1e-18)).alias("rskew"),
        (n * r.pow(4).sum() / (rv.pow(2) + 1e-18)).alias("rkurt"),
        (r2.max() / (rv + 1e-18)).alias("max_share"),
        pl.corr(r, prev).alias("ac1"),
        (pl.col("close").last() / pl.col("open").first()).log().alias("drift"),
        (pl.col("volume").filter(first30).sum() / (pl.col("volume").sum() + 1e-9)).alias("vol_first30"),
        (pl.col("volume").filter(last30).sum() / (pl.col("volume").sum() + 1e-9)).alias("vol_last30"),
        _gini_expr().alias("vol_gini"),
        (pl.col("volume").filter(r < 0).sum() / (pl.col("volume").sum() + 1e-9)).alias("down_vol_share"),
        (n / expected).alias("present_share"),
        (r2.filter(pl.col("r_spy") < 0).sum() / (rv + 1e-18)).alias("rv_mktdown_share"),
        (r * pl.col("r_spy")).sum().alias("sxy"),
        pl.col("r_spy").pow(2).sum().alias("syy"),
        pl.col("volume").sum().alias("volume"),
        pl.col("open").first().alias("open"), pl.col("close").last().alias("close"),
        pl.col("high").max().alias("high"), pl.col("low").min().alias("low"),
    ])
    out = out.with_columns([
        (pl.max_horizontal(pl.col("rv") - pl.col("bpv"), 0.0) / (pl.col("rv") + 1e-18)).alias("rj"),
        (pl.col("rs_neg") / (pl.col("rv") + 1e-18)).alias("rsneg_share"),
        ((pl.col("rs_pos") - pl.col("rs_neg")) / (pl.col("rv") + 1e-18)).alias("sjv"),
    ])
    pre = f"f{freq_min}m_"
    keep = {c: pre + c for c in INTRADAY_FAMILIES}
    return out.select(["session_date"] + [pl.col(c).alias(keep[c]) for c in INTRADAY_FAMILIES] +
                      ([pl.col("volume").alias("d_volume"), pl.col("open").alias("d_open"), pl.col("close").alias("d_close"),
                        pl.col("high").alias("d_high"), pl.col("low").alias("d_low")] if freq_min == 1 else []))


def daily_block(day: pl.DataFrame, actions: pl.DataFrame | None, symbol: str) -> pl.DataFrame:
    """Daily OHLCV features from the 1m aggregation + split/dividend-adjusted close-to-close returns."""
    d = day.sort("session_date")
    if actions is not None and actions.height:
        a = actions.filter(pl.col("symbol") == symbol).select(["ex_date", "split_ratio", "dividend"])
        d = d.join(a, left_on="session_date", right_on="ex_date", how="left").with_columns([
            pl.col("split_ratio").fill_null(1.0), pl.col("dividend").fill_null(0.0)])
    else:
        d = d.with_columns([pl.lit(1.0).alias("split_ratio"), pl.lit(0.0).alias("dividend")])
    prev_close_adj = (pl.col("d_close").shift(1) / pl.col("split_ratio"))
    d = d.with_columns([
        ((pl.col("d_close") + pl.col("dividend")) / prev_close_adj - 1.0).alias("d_ret_cc"),
        (pl.col("d_open") / prev_close_adj - 1.0).alias("d_ret_co"),
        (pl.col("d_close") / pl.col("d_open") - 1.0).alias("d_ret_oc"),
        ((pl.col("d_high") / pl.col("d_low")).log().pow(2) / (4.0 * math.log(2.0))).alias("d_parkinson"),
        (pl.col("d_volume") * pl.col("d_close")).alias("d_dollar_vol"),
    ])
    return d


def compute_symbol_features(paths: list[Path], symbol: str, bounds: pl.DataFrame, spy_buckets: dict[str, pl.DataFrame] | None,
                            actions: pl.DataFrame | None) -> pl.DataFrame:
    reg = load_regular_session(paths, bounds)
    if reg.height == 0:
        return pl.DataFrame()
    blocks = []
    for name, fmin in FREQS.items():
        bk = aggregate_buckets(reg, fmin)
        spy = spy_buckets.get(name) if spy_buckets else None
        blocks.append(session_features(bk, fmin, spy))
    out = blocks[0]
    for b in blocks[1:]:
        out = out.join(b, on="session_date", how="left")
    out = daily_block(out, actions, symbol)
    return out.with_columns(pl.lit(symbol).alias("symbol")).sort("session_date")


def compute_spy_buckets(paths: list[Path], bounds: pl.DataFrame) -> dict[str, pl.DataFrame]:
    reg = load_regular_session(paths, bounds)
    out = {}
    for name, fmin in FREQS.items():
        out[name] = bucket_returns(aggregate_buckets(reg, fmin)).select(["session_date", "b", "r"]).unique(
            subset=["session_date", "b"], keep="last").sort(["session_date", "b"])
    return out