#!/usr/bin/env python3
"""Filter YIEDL historical data to CrowdCent universe and timeframe.

Reduces ~9.7GB → ~500MB by keeping only:
- Symbols in CrowdCent universe
- Dates within CrowdCent range (2019-07-27 to latest)

Usage:
    uv run docker/filter_yiedl.py
"""
import polars as pl
from pathlib import Path

# Paths
CC_PATH = Path("docker/data/cc_train.parquet")
YIEDL_RAW_PATH = Path("docker/data/yiedl_historical.parquet")
YIEDL_FILTERED_PATH = Path("docker/data/yiedl_train.parquet")


def main():
    # Load CrowdCent to get universe and date range
    print("Loading CrowdCent data...")
    cc = pl.read_parquet(CC_PATH)
    cc_symbols = set(cc["id"].unique().to_list())
    cc_min_date = cc["date"].min().date()
    cc_max_date = cc["date"].max().date()

    print(f"CrowdCent: {len(cc_symbols)} symbols, {cc_min_date} to {cc_max_date}")

    # Filter YIEDL (lazy scan for memory efficiency)
    print(f"\nFiltering YIEDL from {YIEDL_RAW_PATH}...")
    yiedl_filtered = (
        pl.scan_parquet(YIEDL_RAW_PATH)
        .filter(
            pl.col("symbol").is_in(cc_symbols)
            & (pl.col("date") >= cc_min_date)
            & (pl.col("date") <= cc_max_date)
        )
        .collect()
    )

    print(f"Filtered shape: {yiedl_filtered.shape}")
    print(f"Date range: {yiedl_filtered['date'].min()} to {yiedl_filtered['date'].max()}")
    print(f"Symbols: {yiedl_filtered['symbol'].n_unique()}")

    # Write filtered data
    print(f"\nWriting to {YIEDL_FILTERED_PATH}...")
    yiedl_filtered.write_parquet(YIEDL_FILTERED_PATH)

    # Report size
    size_mb = YIEDL_FILTERED_PATH.stat().st_size / 1024 / 1024
    print(f"Done! File size: {size_mb:.1f} MB")


if __name__ == "__main__":
    main()
