tool_get_fill_missing_data_v2.py

import os
import polars as pl
import pandas as pd
from datetime import timedelta, datetime

# Base directory for raw data files
datasets_dir = "datasets/raw/full_scope"

# Map each interval folder to timedelta for accurate frequency
INTERVALS = {
    # '1m': timedelta(minutes=1),
    # '3m': timedelta(minutes=3),
    # '5m': timedelta(minutes=5),
    # '15m': timedelta(minutes=15),
    # '30m': timedelta(minutes=30),
    '1h': timedelta(hours=1),
    # '2h': timedelta(hours=2),
    # '4h': timedelta(hours=4),
    # '6h': timedelta(hours=6),
    # '8h': timedelta(hours=8),
    # '12h': timedelta(hours=12),
}


def create_missing_candle(ts, last_row: dict) -> dict:
    """
    Synthesize a frozen candle at timestamp ts.
    All price fields equal last_row['close'], volumes/trades set to zero.
    """
    data = last_row.copy()
    close_price = float(data['close'])
    data.update({
        'timestamp': ts,
        'open': close_price,
        'high': close_price,
        'low': close_price,
        'close': close_price,
        'volume': 0.0,
        'quote_asset_volume': 0.0,
        'number_of_trades': 0,
        'taker_buy_base_asset_volume': 0.0,
        'taker_buy_quote_asset_volume': 0.0
    })
    return data


def forward_fill_file(path: str, interval: str, freq: timedelta):
    """
    Read CSV, detect missing timestamps at interval freq,
    insert synthetic frozen candles, and overwrite in-place.
    """
    # Load CSV
    df = pl.read_csv(path)
    if 'close' not in df.columns:
        raise KeyError(f"Missing 'close' column in {path}")

    # Parse and sort timestamps - ensure consistent datetime precision
    df = df.with_columns([
        pl.col('timestamp')
        .str.strptime(pl.Datetime, "%Y-%m-%d %H:%M:%S", strict=False)
        .cast(pl.Datetime("us"))  # Force microsecond precision
        .alias('timestamp')
    ])
    df = df.sort('timestamp').unique(subset=['timestamp'])

    # Generate full timeline manually since polars date_range has restrictions
    ts_min = df['timestamp'].min()
    ts_max = df['timestamp'].max()

    # Convert polars scalars to python datetime objects
    # Use item() to extract scalar values from polars
    start_dt = ts_min.to_py() if hasattr(ts_min, 'to_py') else ts_min
    end_dt = ts_max.to_py() if hasattr(ts_max, 'to_py') else ts_max

    # If they're already datetime objects, use them directly
    if not isinstance(start_dt, datetime):
        start_dt = pd.to_datetime(start_dt)
    if not isinstance(end_dt, datetime):
        end_dt = pd.to_datetime(end_dt)

    # Convert timedelta to pandas frequency string
    if freq == timedelta(minutes=1):
        pandas_freq = "1min"
    elif freq == timedelta(minutes=3):
        pandas_freq = "3min"
    elif freq == timedelta(minutes=5):
        pandas_freq = "5min"
    elif freq == timedelta(minutes=15):
        pandas_freq = "15min"
    elif freq == timedelta(minutes=30):
        pandas_freq = "30min"
    elif freq == timedelta(hours=1):
        pandas_freq = "1h"
    elif freq == timedelta(hours=2):
        pandas_freq = "2h"
    elif freq == timedelta(hours=4):
        pandas_freq = "4h"
    elif freq == timedelta(hours=6):
        pandas_freq = "6h"
    elif freq == timedelta(hours=8):
        pandas_freq = "8h"
    elif freq == timedelta(hours=12):
        pandas_freq = "12h"
    else:
        raise ValueError(f"Unsupported frequency: {freq}")

    # Use pandas to generate the full range, then convert back to polars
    full_range = pd.date_range(start=start_dt, end=end_dt, freq=pandas_freq)
    full_ts = pl.DataFrame({'timestamp': full_range}).with_columns([
        pl.col('timestamp').cast(pl.Datetime("us"))  # Match precision
    ])

    # Identify missing timestamps
    existing = df.select('timestamp')
    missing_ts = full_ts.join(existing, on='timestamp', how='anti')['timestamp'].to_list()

    if not missing_ts:
        print(f"✅ No missing candles in {os.path.basename(path)}")
    else:
        print(f"⚠️ Filling {len(missing_ts)} missing candles in {os.path.basename(path)}...")
        # Forward-fill approach: iterate chronologically to avoid data leakage
        df_sorted = df.sort('timestamp')
        df_dict = {row['timestamp']: row for row in df_sorted.to_dicts()}
        sorted_ts = sorted(full_ts['timestamp'].to_list())

        inserts = []
        last_known_row = None

        for ts in sorted_ts:
            if ts in df_dict:
                # Real data point - update our last known state
                last_known_row = df_dict[ts]
            elif last_known_row is not None:
                # Missing data point - forward fill from last known
                inserts.append(create_missing_candle(ts, last_known_row))
            # If last_known_row is None, we're before any real data, so skip

        # Append and re-sort
        if inserts:
            df_insert = pl.DataFrame(inserts)
            df = pl.concat([df, df_insert]).sort('timestamp')

    # Finalize: remove any duplicates and write back
    df = df.unique(subset=['timestamp'], keep='first')

    # Format timestamp back to string for CSV output
    df = df.with_columns([
        pl.col('timestamp').dt.strftime("%Y-%m-%d %H:%M:%S")
    ])

    df.write_csv(path)
    print(f"✅ Saved filled file: {os.path.basename(path)}")


def main():
    for interval, freq in INTERVALS.items():
        folder = os.path.join(datasets_dir, interval)
        if not os.path.isdir(folder):
            continue
        print(f"\n--- Processing interval: {interval} ---")
        for fname in sorted(os.listdir(folder)):
            if not fname.endswith('_historical_data.csv'):
                continue
            file_path = os.path.join(folder, fname)
            try:
                forward_fill_file(file_path, interval, freq)
            except Exception as e:
                print(f"❌ Error on {fname}: {e}")


if __name__ == '__main__':
    main()