from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import pandas as pd
DEFAULT_DIR = Path('datasets/cicddos2019/processed')

def _load_shard(dir: Path, shard: str) -> tuple[dict, pd.DataFrame]:
    p = np.load(dir / f'packets.{shard}.npz')
    f = pd.read_parquet(dir / f'flows.{shard}.parquet')
    assert set(p.files) == {'packet_tokens', 'packet_lengths', 'flow_id'}, p.files
    assert set(f.columns) == {'flow_id', 'label'}, f.columns
    assert len(p['flow_id']) == len(f), f'row count mismatch in {shard}'
    assert np.array_equal(p['flow_id'], f['flow_id'].to_numpy()), f'flow_id mismatch in {shard}'
    return ({'packet_tokens': p['packet_tokens'], 'packet_lengths': p['packet_lengths'], 'flow_id': p['flow_id']}, f)

def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('--dir', type=Path, default=DEFAULT_DIR)
    ap.add_argument('--out-packets', type=Path, default=None)
    ap.add_argument('--out-flows', type=Path, default=None)
    args = ap.parse_args()
    out_p = args.out_packets or args.dir / 'packets.npz'
    out_f = args.out_flows or args.dir / 'flows.parquet'
    print(f'=== merging shards from {args.dir} ===')
    (p1, f1) = _load_shard(args.dir, '01-12')
    (p3, f3) = _load_shard(args.dir, '03-11')
    n1 = len(f1)
    n3 = len(f3)
    N = n1 + n3
    print(f'01-12 rows: {n1:,}  03-11 rows: {n3:,}  total: {N:,}')
    tokens = np.concatenate([p1['packet_tokens'], p3['packet_tokens']], axis=0)
    lengths = np.concatenate([p1['packet_lengths'], p3['packet_lengths']], axis=0)
    flow_id = np.arange(N, dtype=np.uint64)
    print(f'  tokens  shape={tokens.shape}  dtype={tokens.dtype}')
    print(f'  lengths shape={lengths.shape}  dtype={lengths.dtype}')
    flows = pd.concat([f1.drop(columns=['flow_id']), f3.drop(columns=['flow_id'])], ignore_index=True)
    flows.insert(0, 'flow_id', flow_id)
    print(f"  flows rows={len(flows):,}  label unique={flows['label'].nunique()}")
    assert len(tokens) == N == len(flows)
    assert np.array_equal(flow_id, flows['flow_id'].to_numpy())
    print(f'\n=== writing {out_p} ===')
    out_p.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(out_p, packet_tokens=tokens, packet_lengths=lengths, flow_id=flow_id)
    sz = out_p.stat().st_size / 1000000000.0
    print(f'  wrote {sz:.2f} GB')
    print(f'\n=== writing {out_f} ===')
    flows.to_parquet(out_f, compression='snappy', index=False)
    sz = out_f.stat().st_size / 1000000.0
    print(f'  wrote {sz:.2f} MB')
    print(f'\n=== summary ===')
    print(flows['label'].value_counts().to_string())
if __name__ == '__main__':
    main()
