"""
Add a cropped-spectrum column (channels 700:1018, 318 channels) to the
unknown-data parquet, processed row-group by row-group to bound memory.

Input : Data Parquet/unknown_data_after_noise_filter.parquet
Output: Data Parquet/unknown_data_after_noise_filter_cropped.parquet
"""
import os

import pyarrow.compute as pc
import pyarrow.parquet as pq

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
SRC = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter.parquet')
DST = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter_cropped.parquet')

START, END = 700, 1018  # channels 700..1017 (inclusive) -> 318 channels


def main():
    pf = pq.ParquetFile(SRC)
    n_rgs = pf.metadata.num_row_groups
    print(f'{SRC}\n  rows: {pf.metadata.num_rows:,}  row groups: {n_rgs}')

    writer = None
    for i in range(n_rgs):
        t = pf.read_row_group(i)
        cropped = pc.list_slice(t.column('spectrum'), START, END)
        t = t.append_column('spectrum_cropped', cropped)
        if writer is None:
            writer = pq.ParquetWriter(DST, t.schema)
        writer.write_table(t)
        if (i + 1) % 4 == 0 or i + 1 == n_rgs:
            print(f'  row group {i + 1}/{n_rgs}')
    writer.close()
    print(f'Wrote {DST}')

    # Verify
    out = pq.ParquetFile(DST)
    print(f'Output rows: {out.metadata.num_rows:,}')
    sample = out.read_row_group(0, columns=['spectrum_cropped']).slice(0, 3).to_pydict()['spectrum_cropped']
    print('sample cropped lengths:', [len(s) for s in sample])


if __name__ == '__main__':
    main()