"""Parse train/evaluation CSVs into compact parquet tables (dense int ids).

Outputs in v3/data/:
  rows_train.parquet  : rowid, eid, ii (item idx), clicked
  rows_eval.parquet   : rowid, eid, ii, ord (original row order)
  ev_train.parquet / ev_eval.parquet : eid, u (user idx), ts, tb, dt, h (list[int item idx]), h_ts (list[int])
  items.parquet / users.parquet : id string -> idx
Usage: python prep.py [DATA_DIR]
"""
import polars as pl, sys, time, os
DATA = sys.argv[1] if len(sys.argv) > 1 else '/workspace/huawei/data'
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
os.makedirs(OUT, exist_ok=True)
t0 = time.time()
def log(*a): print(f'[{time.time()-t0:6.1f}s]', *a, flush=True)

def scan(name):
    return pl.scan_csv(f'{DATA}/{name}.csv', schema_overrides={'history_item_ids': pl.Utf8, 'history_timestamps': pl.Utf8,
                       'rowid': pl.Utf8, 'event_id': pl.Utf8, 'user_id': pl.Utf8, 'item_id': pl.Utf8})

raw = {}
for name, lab in (('train', True), ('evaluation', False)):
    lf = scan(name)
    cols = ['rowid', 'event_id', 'item_id'] + (['clicked'] if lab else [])
    rows = lf.select(cols).collect(engine='streaming')
    ev = (lf.select('event_id', 'user_id', 'timestamp', 'time_bucket', 'day_type', 'history_item_ids', 'history_timestamps')
            .unique(subset='event_id', keep='any')
            .with_columns(pl.col('history_item_ids').str.json_decode(pl.List(pl.Utf8)).alias('h_str'),
                          pl.col('history_timestamps').str.json_decode(pl.List(pl.Int64)).alias('h_ts'))
            .drop('history_item_ids', 'history_timestamps').collect(engine='streaming'))
    raw[name] = (rows, ev)
    log(name, rows.shape, ev.shape)

tr_rows, tr_ev = raw['train']; ev_rows, ev_ev = raw['evaluation']
items = pl.concat([tr_rows['item_id'], ev_rows['item_id'], tr_ev['h_str'].explode().drop_nulls(),
                   ev_ev['h_str'].explode().drop_nulls()]).unique().sort()
items = pl.DataFrame({'item_id': items, 'ii': pl.arange(0, len(items), eager=True).cast(pl.Int32)})
users = pl.concat([tr_ev['user_id'], ev_ev['user_id']]).unique().sort()
users = pl.DataFrame({'user_id': users, 'u': pl.arange(0, len(users), eager=True).cast(pl.Int32)})
events = pl.concat([tr_ev['event_id'], ev_ev['event_id']]).unique().sort()
events = pl.DataFrame({'event_id': events, 'eid': pl.arange(0, len(events), eager=True).cast(pl.Int32)})
log('items', len(items), 'users', len(users), 'events', len(events))
items.write_parquet(f'{OUT}/items.parquet'); users.write_parquet(f'{OUT}/users.parquet')

def map_rows(rows, lab):
    r = rows.with_row_index('ord').join(events, on='event_id', how='left').join(items, on='item_id', how='left')
    sel = ['ord', 'rowid', 'eid', 'ii'] + (['clicked'] if lab else [])
    r = r.sort('ord').select(sel)
    if lab: r = r.with_columns(pl.col('clicked').cast(pl.Int8))
    return r

def map_ev(ev):
    e = ev.join(events, on='event_id', how='left').join(users, on='user_id', how='left')
    ex = e.select('eid', 'h_str').with_columns(pl.int_ranges(0, pl.col('h_str').list.len()).alias('k')).explode(['h_str', 'k']).drop_nulls('h_str')
    ex = ex.join(items.rename({'item_id': 'h_str'}), on='h_str', how='left').sort('eid', 'k')
    h = ex.group_by('eid', maintain_order=True).agg(pl.col('ii').alias('h'))
    e = e.join(h, on='eid', how='left').with_columns(pl.col('h').fill_null(pl.lit([], dtype=pl.List(pl.Int32))))
    return e.select('eid', 'u', pl.col('timestamp').cast(pl.Int64).alias('ts'), pl.col('time_bucket').cast(pl.Int8).alias('tb'),
                    pl.col('day_type').cast(pl.Int8).alias('dt'), 'h', 'h_ts').sort('ts', 'eid')

map_rows(tr_rows, True).drop('ord').write_parquet(f'{OUT}/rows_train.parquet')
map_rows(ev_rows, False).write_parquet(f'{OUT}/rows_eval.parquet')
map_ev(tr_ev).write_parquet(f'{OUT}/ev_train.parquet')
map_ev(ev_ev).write_parquet(f'{OUT}/ev_eval.parquet')
log('done')