#!/usr/bin/env python
import os.path
import argparse
import strax
import pandas as pd

def parse_args():
    parser = argparse.ArgumentParser(
        description="Rechunker for FileSytemBackend. Interfaces with strax.rechunker."
                    "Please see the documentation of strax.rechunker for more information: "
                    "github.com/AxFoundation/strax/blob/31c114c5f8329e53289d5127fb2125e71c3d6aae/strax/storage/files.py#L371")
    parser.add_argument(
        '--source',
        type=str,
        help="Target directory to rechunk, should be a folder in a "
             "strax.DataDrictory (one datatype)")
    parser.add_argument(
        '--dest', '--destination',
        default=None,
        dest='dest',
        type=str,
        help='Where to store rechunked data. If nothing is specified, replace the source.',
    )
    parser.add_argument(
        '--compressor',
        choices=list(strax.io.COMPRESSORS.keys()),
        help="Recompress using one of these compressors. If nothing specified, "
             "use the same compressor as the source")
    parser.add_argument(
        '--rechunk',
        default=True,
        choices=[True, False],
        type=bool,
        help="rechunk the data")
    parser.add_argument(
        '--target_size_mb', '--target-size-mb',
        dest='target_size_mb',
        type=int,
        default=strax.default_chunk_size_mb,
        help="Target size MB (uncompressed) of the rechunked data")
    parser.add_argument(
        '--write_stats_to', '--write-stats-to',
        dest='write_stats_to',
        type=str,
        default=None,
        help="Write some information to this file (csv format)")
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    source_mb = strax.utils.dir_size_mb(args.source)
    report = strax.rechunker(source_directory=args.source,
                             dest_directory=args.dest,
                             replace=args.dest is None,
                             compressor=args.compressor,
                             target_size_mb=args.target_size_mb,
                             rechunk=args.rechunk,
                             )
    if args.dest is not None:
        recompressed_mb = strax.utils.dir_size_mb(args.dest)
    else:
        recompressed_mb = strax.utils.dir_size_mb(args.source)
    report.update(dict(source_mb=source_mb,
                       dest_mb=recompressed_mb)
                  )
    if args.write_stats_to:
        if os.path.exists(args.write_stats_to):
            df = pd.read_csv(args.write_stats_to)
        else:
            df = pd.DataFrame()
        df_new = pd.concat(
            [df,
             pd.DataFrame({k: [v] for k, v in report.items()})
             ])
        df_new.to_csv(args.write_stats_to, index=False)

    print(f'Re-compressed {args.source}')
    for k, v in report.items():
        print(f'\t{k:16}\t{v}')


if __name__ == '__main__':
    main()
