"""Session-aware aggregation of hourly futures to daily bars."""

import os
from datetime import datetime, timedelta
from zoneinfo import ZoneInfo

import numpy as np
import plotly.graph_objects as go
import polars as pl
from plotly.subplots import make_subplots

from data import load_cme_futures
from utils import ML4T_DATA_PATH
from utils.paths import REPO_ROOT, get_chapter_dir
from utils.style import COLORS, ml4t_palette


def _rel(path):
    """Repo-relative display path (keeps absolute machine paths out of outputs)."""
    try:
        return path.relative_to(REPO_ROOT)
    except ValueError:
        return path

# WRITE_TO_DATA=1 materializes the canonical daily parquet under ML4T_DATA_PATH
# for the downstream chapters and case studies; the default (0) writes to a
# chapter-local output dir so a demo run never overwrites production data.
WRITE_TO_DATA = os.environ.get("WRITE_TO_DATA", "0") == "1"

# Output path for the session-aggregated daily data, selected by WRITE_TO_DATA.
OUTPUT_DIR = (
    ML4T_DATA_PATH / "futures" / "market" / "continuous" / "daily"
    if WRITE_TO_DATA
    else get_chapter_dir(2) / "output" / "futures_daily"
)

# Timezone constants
CT = ZoneInfo("America/Chicago")
UTC = ZoneInfo("UTC")

# CME session ends at 4 PM CT
SESSION_END_HOUR_CT = 16  # 4:00 PM


def assign_cme_session_date(ts: datetime) -> datetime:
    """
    Assign CME session date to a UTC timestamp.

    The session date is the date when the session ENDS (4 PM CT).
    A bar at Sunday 11 PM UTC belongs to Monday's session.

    CME closes Friday at 4 PM CT and reopens Sunday 5 PM CT.
    Bars after Friday 4 PM CT still belong to Friday's session —
    they must NOT roll to Saturday.

    Args:
        ts: UTC timestamp

    Returns:
        Session date (as date, no time component)
    """
    # Convert to Central Time
    ts_ct = ts.astimezone(CT)

    # If we're past 4 PM CT, this belongs to tomorrow's session
    if ts_ct.hour >= SESSION_END_HOUR_CT:
        candidate = ts_ct.date() + timedelta(days=1)
        # Friday after 4 PM CT → keep as Friday (no Saturday session)
        # isoweekday: Mon=1, Fri=5, Sat=6
        if candidate.isoweekday() == 6:  # Saturday
            candidate = ts_ct.date()  # Keep as Friday
        session_date = candidate
    else:
        session_date = ts_ct.date()

    return session_date

# Quick test
test_times = [
    datetime(2024, 1, 7, 23, 0, tzinfo=UTC),  # Sunday 11 PM UTC = Sunday 5 PM CT -> Monday
    datetime(2024, 1, 8, 15, 0, tzinfo=UTC),  # Monday 3 PM UTC = Monday 9 AM CT -> Monday
    datetime(2024, 1, 8, 22, 0, tzinfo=UTC),  # Monday 10 PM UTC = Monday 4 PM CT -> Tuesday
    datetime(
        2024, 1, 12, 22, 0, tzinfo=UTC
    ),  # Friday 10 PM UTC = Friday 4 PM CT -> Friday (NOT Saturday)
]

print("Session Assignment Examples:")
for ts in test_times:
    ts_ct = ts.astimezone(CT)
    session = assign_cme_session_date(ts)
    print(f"  {ts} ({ts_ct.strftime('%a %I:%M %p CT')}) -> Session: {session}")

hourly = load_cme_futures(continuous=True, frequency="hourly")
products = sorted(hourly["product"].unique().to_list())

print(f"Loaded {len(hourly):,} hourly bars")
print(f"Products: {hourly['product'].n_unique()}")
print(f"Tenors: {sorted(hourly['tenor'].unique().to_list())}")
print(f"Date range: {hourly['timestamp'].min()} to {hourly['timestamp'].max()}")
print(f"Available products: {', '.join(products)}")

hourly.filter(pl.col("product") == "ES").select(
    "timestamp", "product", "tenor", "open", "high", "low", "close", "volume"
).head(10)

# Vectorized session date assignment using Polars
# Convert to Central Time, then check if hour >= 16 (4 PM)


def add_session_date(df: pl.DataFrame) -> pl.DataFrame:
    """Add session_date column based on CME session boundaries.

    Friday after 4 PM CT stays as Friday — CME has no Saturday session.
    """
    return (
        df.with_columns(pl.col("timestamp").dt.convert_time_zone("America/Chicago").alias("ts_ct"))
        .with_columns(
            pl.col("ts_ct").dt.date().alias("_ct_date"),
            (pl.col("ts_ct").dt.hour() >= SESSION_END_HOUR_CT).alias("_after_close"),
            # isoweekday: Mon=1 ... Fri=5, Sat=6, Sun=7
            (pl.col("ts_ct").dt.weekday() == 5).alias("_is_friday"),
        )
        .with_columns(
            # After 4 PM CT → next day, UNLESS it's Friday (no Saturday session)
            pl.when(pl.col("_after_close") & ~pl.col("_is_friday"))
            .then(pl.col("_ct_date") + pl.duration(days=1))
            .otherwise(pl.col("_ct_date"))
            .alias("session_date")
        )
        .drop("ts_ct", "_ct_date", "_after_close", "_is_friday")
    )

hourly_with_sessions = add_session_date(hourly)

print("Session dates assigned (ES sample):")
hourly_with_sessions.filter(pl.col("product") == "ES").select(
    "timestamp", "session_date", "product", "tenor", "close", "volume"
).head(15)

(
    hourly_with_sessions.filter(
        (pl.col("product") == "ES")
        & (pl.col("tenor") == 0)
        & (pl.col("timestamp").dt.date() == pl.lit("2024-01-08").str.to_date())
    )
    .sort("timestamp")
    .select("timestamp", "session_date", "close", "volume")
)

_boundary_window = (
    hourly_with_sessions.filter(
        (pl.col("product") == "ES")
        & (pl.col("tenor") == 0)
        & (pl.col("timestamp").dt.date() >= pl.lit("2024-01-05").str.to_date())
        & (pl.col("timestamp").dt.date() <= pl.lit("2024-01-09").str.to_date())
    )
    .with_columns(pl.col("timestamp").dt.convert_time_zone("America/Chicago").alias("ts_ct"))
    .sort("timestamp")
)

_sessions = _boundary_window["session_date"].unique().sort().to_list()
_session_colors = dict(zip(_sessions, ml4t_palette(len(_sessions), categorical=True)))

fig = go.Figure()
for sess in _sessions:
    seg = _boundary_window.filter(pl.col("session_date") == sess)
    fig.add_trace(
        go.Scatter(
            x=seg["ts_ct"].to_list(),
            y=seg["close"].to_list(),
            mode="lines+markers",
            line=dict(color=_session_colors[sess], width=1.5),
            marker=dict(size=5),
            name=str(sess),
        )
    )
fig.update_layout(
    title="One color per CME session: Sunday-evening bars belong to Monday",
    xaxis_title="Timestamp (Central Time)",
    yaxis_title="ES front-month close",
    height=420,
    legend_title="Session date",
)
fig.show()

# Sort and detect roll transitions per (product, tenor)
hourly_sorted = hourly_with_sessions.sort(["product", "tenor", "timestamp"])

# Detect instrument_id changes within each (product, tenor) group
hourly_sorted = hourly_sorted.with_columns(
    pl.col("instrument_id").shift(1).over("product", "tenor").alias("_prev_instrument_id"),
    pl.col("close").shift(1).over("product", "tenor").alias("_prev_close"),
)

# Roll points: where instrument_id changes (excluding first row of each group)
rolls = hourly_sorted.filter(
    pl.col("_prev_instrument_id").is_not_null()
    & (pl.col("instrument_id") != pl.col("_prev_instrument_id"))
)

# Ratio = new contract's open / old contract's close (adjacent hourly bars)
roll_ratios = rolls.select(
    "product",
    "tenor",
    "timestamp",
    (pl.col("open") / pl.col("_prev_close")).alias("ratio"),
)

print(f"Roll transitions detected: {len(roll_ratios)}")
print(f"Products with rolls: {roll_ratios['product'].n_unique()}")

# Front month (tenor 0) only — its ~quarterly rolls are what the adjusted
# front-month series below is built from. Deferred tenors roll far more often
# (thin volume flips leadership back and forth), so mixing tenors here would
# overstate the front-month roll count.
es_rolls = roll_ratios.filter((pl.col("product") == "ES") & (pl.col("tenor") == 0)).sort(
    "timestamp"
)
print(f"ES front-month roll ratios ({len(es_rolls)} rolls):")
es_rolls.select("timestamp", "ratio").head(10)

def ratio_adjust(group: pl.DataFrame) -> pl.DataFrame:
    """Apply ratio back-adjustment to a single (product, tenor) group."""
    group = group.sort("timestamp")

    # Get roll ratios for this group
    group_rolls = roll_ratios.filter(
        (pl.col("product") == group["product"][0]) & (pl.col("tenor") == group["tenor"][0])
    ).select("timestamp", "ratio")

    if len(group_rolls) == 0:
        return group.with_columns(pl.lit(1.0).alias("_cumulative_ratio"))

    # Join roll ratios
    group = group.join(group_rolls, on="timestamp", how="left").with_columns(
        pl.col("ratio").fill_null(1.0)
    )

    # Cumulative ratio: product of all FUTURE ratios (reverse cumprod)
    # Bars BEFORE a roll get multiplied; bars ON and AFTER the roll do not
    n = len(group)
    ratios = group["ratio"].to_numpy()
    adj = np.ones(n)
    cumulative = 1.0
    for i in range(n - 1, -1, -1):
        adj[i] = cumulative
        if ratios[i] != 1.0:
            cumulative *= ratios[i]

    return group.with_columns(pl.Series("_cumulative_ratio", adj)).drop("ratio")

# Apply per group
adjusted_groups = []
products_tenors = hourly_sorted.select("product", "tenor").unique().sort("product", "tenor")
n_groups = len(products_tenors)

for i, row in enumerate(products_tenors.iter_rows(named=True)):
    group = hourly_sorted.filter(
        (pl.col("product") == row["product"]) & (pl.col("tenor") == row["tenor"])
    )
    adjusted = ratio_adjust(group)
    adjusted_groups.append(adjusted)
    if (i + 1) % 30 == 0 or i == n_groups - 1:
        print(f"  Adjusted {i + 1}/{n_groups} groups")

hourly_adjusted = pl.concat(adjusted_groups)

# Emit two explicit price series and NO bare OHLC. Ratio adjustment preserves
# within-tenor *returns* but distorts price *levels*: differencing adjusted
# tenor-0/tenor-1 levels reads accumulated roll history, not the curve. Every
# downstream consumer must name which series it wants:
#   - adj_* (returns / momentum / volatility / labels — roll-continuous)
#   - raw_* (carry / term structure / roll yield / notional / costs — contemporaneous)
# cum_ratio is carried through so the two are reconcilable (adj_close == raw_close * cum_ratio).
# Dropping bare open/high/low/close makes the adjusted-vs-raw choice impossible to skip.
hourly_adjusted = hourly_adjusted.with_columns(
    pl.col("open").alias("raw_open"),
    pl.col("high").alias("raw_high"),
    pl.col("low").alias("raw_low"),
    pl.col("close").alias("raw_close"),
)

# Adjusted OHLC (multiply by cumulative ratio, not add).
hourly_adjusted = hourly_adjusted.with_columns(
    (pl.col("open") * pl.col("_cumulative_ratio")).alias("adj_open"),
    (pl.col("high") * pl.col("_cumulative_ratio")).alias("adj_high"),
    (pl.col("low") * pl.col("_cumulative_ratio")).alias("adj_low"),
    (pl.col("close") * pl.col("_cumulative_ratio")).alias("adj_close"),
).drop("open", "high", "low", "close")

print(f"\nRatio adjustment applied to {len(hourly_adjusted):,} hourly bars")

# Show adjustment magnitude for ES front month
es_adj = hourly_adjusted.filter((pl.col("product") == "ES") & (pl.col("tenor") == 0)).sort(
    "timestamp"
)
print(
    f"ES front month cumulative ratio range: "
    f"{es_adj['_cumulative_ratio'].min():.4f} to {es_adj['_cumulative_ratio'].max():.4f}"
)

es_adj_daily = (
    es_adj.group_by("session_date")
    .agg(
        pl.col("raw_close").last(),
        pl.col("adj_close").last(),
        pl.col("_cumulative_ratio").last().alias("cum_ratio"),
    )
    .sort("session_date")
)
_x = es_adj_daily["session_date"].to_list()

fig = make_subplots(
    rows=2,
    cols=1,
    shared_xaxes=True,
    row_heights=[0.68, 0.32],
    vertical_spacing=0.06,
)
fig.add_trace(
    go.Scatter(
        x=_x,
        y=es_adj_daily["raw_close"].to_list(),
        mode="lines",
        line=dict(color=COLORS["copper"], width=1),
        name="Raw (unadjusted)",
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=_x,
        y=es_adj_daily["adj_close"].to_list(),
        mode="lines",
        line=dict(color=COLORS["blue"], width=1),
        name="Ratio-adjusted",
    ),
    row=1,
    col=1,
)
fig.add_trace(
    go.Scatter(
        x=_x,
        y=es_adj_daily["cum_ratio"].to_list(),
        mode="lines",
        line=dict(color=COLORS["slate"], width=1),
        name="Cumulative ratio",
        showlegend=False,
    ),
    row=2,
    col=1,
)
fig.add_hline(y=1.0, line=dict(color=COLORS["neutral"], width=1, dash="dot"), row=2, col=1)
fig.update_layout(
    title="ES front month: ratio back-adjustment removes roll gaps",
    height=560,
    legend_title="Price series",
)
fig.update_yaxes(title_text="Price", row=1, col=1)
fig.update_yaxes(title_text="Cumulative ratio", row=2, col=1)
fig.update_xaxes(title_text="Session date", row=2, col=1)
fig.show()

# Replace hourly_with_sessions with adjusted data for downstream aggregation.
# Keep the cumulative ratio (as cum_ratio) so raw and adjusted stay reconcilable.
hourly_with_sessions = hourly_adjusted.drop("_prev_instrument_id", "_prev_close").rename(
    {"_cumulative_ratio": "cum_ratio"}
)

# Aggregate to daily by session_date, product, tenor
daily = (
    hourly_with_sessions.sort(["product", "tenor", "timestamp"])
    .group_by(["session_date", "product", "tenor"])
    .agg(
        [
            pl.col("adj_open").first(),
            pl.col("adj_high").max(),
            pl.col("adj_low").min(),
            pl.col("adj_close").last(),
            pl.col("raw_open").first(),
            pl.col("raw_high").max(),
            pl.col("raw_low").min(),
            pl.col("raw_close").last(),
            pl.col("cum_ratio").last(),
            pl.col("volume").sum(),
            pl.len().alias("bar_count"),
            pl.col("timestamp").min().alias("session_start"),
            pl.col("timestamp").max().alias("session_end"),
        ]
    )
    .sort(["product", "tenor", "session_date"])
)

print(f"Daily bars: {len(daily):,}")
print(f"Products: {daily['product'].n_unique()}")
print(f"Session date range: {daily['session_date'].min()} to {daily['session_date'].max()}")

es_daily = daily.filter((pl.col("product") == "ES") & (pl.col("tenor") == 0))
print("ES front month daily bars (first 20 sessions):")
es_daily.select(
    "session_date", "adj_open", "adj_high", "adj_low", "adj_close", "volume", "bar_count"
).head(20)

bar_counts = daily.group_by("bar_count").len().sort("bar_count")
typical_sessions = daily.filter(pl.col("bar_count").is_between(20, 24))
print(f"Typical sessions (20-24 bars): {len(typical_sessions):,} / {len(daily):,}")

# Highlight the modal (most common) bucket; a full 23-hour session dominates.
_modal_bars = bar_counts.sort("len", descending=True)["bar_count"][0]
fig = go.Figure(
    go.Bar(
        x=bar_counts["bar_count"].to_list(),
        y=bar_counts["len"].to_list(),
        marker_color=[
            COLORS["amber"] if bc == _modal_bars else COLORS["slate"]
            for bc in bar_counts["bar_count"].to_list()
        ],
    )
)
fig.add_annotation(
    x=_modal_bars,
    y=bar_counts.filter(pl.col("bar_count") == _modal_bars)["len"][0],
    text=f"{_modal_bars}-hour session",
    showarrow=True,
    arrowhead=2,
    yshift=6,
)
fig.update_layout(
    title="Hourly bars per session: the full 23-hour day dominates",
    xaxis_title="Hourly bars in the session",
    yaxis_title="Number of daily bars",
    height=420,
    showlegend=False,
    xaxis=dict(dtick=2),
)
fig.show()

# OHLC invariant check
ohlc_check = daily.with_columns(
    [
        (pl.col("adj_low") <= pl.col("adj_open")).alias("low_le_open"),
        (pl.col("adj_low") <= pl.col("adj_close")).alias("low_le_close"),
        (pl.col("adj_high") >= pl.col("adj_open")).alias("high_ge_open"),
        (pl.col("adj_high") >= pl.col("adj_close")).alias("high_ge_close"),
    ]
)

print("OHLC Invariant Check:")
for col in ["low_le_open", "low_le_close", "high_ge_open", "high_ge_close"]:
    pct = ohlc_check[col].mean() * 100
    status = "[OK]" if pct > 99.9 else "[FAIL]"
    print(f"  {status} {col}: {pct:.2f}%")

# Coverage by product
coverage = (
    daily.group_by("product")
    .agg(
        [
            pl.col("session_date").min().alias("start_date"),
            pl.col("session_date").max().alias("end_date"),
            pl.len().alias("total_bars"),
            pl.col("tenor").n_unique().alias("tenors"),
        ]
    )
    .sort("product")
)

print("Daily data coverage by product:")
coverage

tenor_coverage = (
    daily.group_by("tenor")
    .agg(
        pl.col("product").n_unique().alias("products"),
        pl.len().alias("total_bars"),
    )
    .sort("tenor")
)
print("Coverage by tenor:")
tenor_coverage

# Create output directory
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Save combined daily file
output_path = OUTPUT_DIR / "continuous_daily.parquet"
daily.write_parquet(output_path)
print(f"Saved: {_rel(output_path)}")
print(f"Size: {output_path.stat().st_size / 1e6:.1f} MB")

# Also save per-product files for convenience
per_product_dir = OUTPUT_DIR / "by_product"
per_product_dir.mkdir(exist_ok=True)

for product in products:
    product_df = daily.filter(pl.col("product") == product)
    product_path = per_product_dir / f"{product}.parquet"
    product_df.write_parquet(product_path)

print(f"\nSaved per-product files to: {_rel(per_product_dir)}/")
print(f"Products: {len(products)}")

es_nq_2024 = (
    pl.read_parquet(OUTPUT_DIR / "continuous_daily.parquet")
    .filter(
        pl.col("product").is_in(["ES", "NQ"])
        & (pl.col("tenor") == 0)
        & (pl.col("session_date") >= pl.lit("2024-01-01").str.to_date())
        & (pl.col("session_date") <= pl.lit("2024-12-31").str.to_date())
    )
    .sort("session_date", "product")
)
print(f"ES + NQ front month, 2024: {len(es_nq_2024)} daily bars")
es_nq_2024.head(10)