import polars as pl
pl.Config.set_tbl_rows(15)
pl.Config(float_precision=2)

from datetime import datetime
cutoff_date = datetime(2023, 7, 1)

# List of our tables
tables = ['buy_ohlc', 'sell_ohlc', 'buy_sma', 'sell_sma', 'rsi', 'supply', 'demand']
# Create a dictionary where key is tablename and value is a lazyframe of the data for that table
lazy_dfs = {table: pl.scan_parquet(f'Data/{table}.parquet') for table in tables}

# Loop through tables
for table in tables:
    # If the name, split by _ gives us two elements, we know it is a buy_ or sell_
    name_split = table.split('_')
    # If the length is more than 1 and the second part is ohlc, we'll do the ohlc task,
    if len(name_split) > 1 and name_split[1] == 'ohlc':
        # add buy_ or sell_ to the column name so we can merge them 
        # We'll also convert the time field to datetime and truncate it at 1 day since we know 
        # we don't have more granular data than that
        lazy_dfs[table] = (
            lazy_dfs[table].with_columns(
                pl.from_epoch('time', time_unit='s').alias('date').dt.truncate('1d')
            ).rename({
                'open': f'{name_split[0]}_open',
                'high': f'{name_split[0]}_high',
                'low': f'{name_split[0]}_low',
                'close': f'{name_split[0]}_close'    
            }).filter(pl.col('date') > cutoff_date)
            .drop('time'))
    else:
        # If it's not ohlc, we just need to make the field name the same as the table name 
        # (we also do the time thing here)
        lazy_dfs[table] = (
            lazy_dfs[table].with_columns(
                pl.from_epoch('time', time_unit='s').alias('date').dt.truncate('1d'),
                pl.col('value').alias(table)
        ).filter(pl.col('date') > cutoff_date)
        .drop(['time', 'value']))

# Now we'll join our tables into one
combined = (
    lazy_dfs['buy_ohlc']
    .join(lazy_dfs['sell_ohlc'], on=['date', 'id'], how='inner')
    .join(lazy_dfs['sell_sma'], on=['date', 'id'], how='inner')
    .join(lazy_dfs['buy_sma'], on=['date', 'id'], how='inner')
    .join(lazy_dfs['rsi'], on=['date', 'id'], how='inner')
    .join(lazy_dfs['supply'], on=['date', 'id'], how='inner')
    .join(lazy_dfs['demand'], on=['date', 'id'], how='inner')
)

combined.collect().head()

<style> .dataframe > thead > tr, .dataframe > tbody > tr { text-align: right; white-space: pre-wrap; } </style> <small>shape: (5, 15)</small><thead><th>buy_open</th><th>buy_high</th><th>buy_low</th><th>buy_close</th><th>id</th><th>date</th><th>sell_open</th><th>sell_high</th><th>sell_low</th><th>sell_close</th><th>sell_sma</th><th>buy_sma</th><th>rsi</th><th>supply</th><th>demand</th></thead><tbody></tbody><tbody></tbody>

i64	i64	i64	i64	i64	datetime[μs]	i64	i64	i64	i64	f64	f64	f64	i64	i64
183	183	183	183	24	2023-07-02 00:00:00	197	197	196	196	229.29	182.57	40.74	150971	60642
183	183	183	183	24	2023-07-03 00:00:00	196	258	196	258	229.50	182.36	54.15	150792	60556
183	183	183	183	24	2023-07-04 00:00:00	258	258	256	256	229.43	182.14	53.73	150510	60570
183	183	183	183	24	2023-07-05 00:00:00	253	253	252	252	230.64	182.00	52.84	150651	60446
182	183	182	183	24	2023-07-06 00:00:00	183	253	183	253	232.71	181.86	53.05	150674	60364

combined.sink_parquet(f'Data/all_stats_{cutoff_date.strftime('%m-%d-%Y')}_cutoff.parquet', compression='zstd')