#!/usr/bin/env python3

import csv
from pathlib import Path
import argparse


def parse_arguments():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('input_file',
                        type=Path,
                        help='PRODCOM csv file to split')
    parser.add_argument('output_path',
                        type=Path,
                        help='path to save country csv files')
    return parser.parse_args()

def split_bulk_csv(csvfile, output_path):
    output_name = Path(csvfile).stem 
    with open(csvfile) as fin:    
        csvin = csv.DictReader(fin)
        # Category -> open file lookup
        outputs = {}
        for row in csvin:
            cat = row['decl']
            # if necessary open a new file and write the header
            if cat not in outputs:
                fout = open(output_path / '{}_{}.csv'.format(output_name, cat), 'w')
                dw = csv.DictWriter(fout, fieldnames=csvin.fieldnames)
                dw.writeheader()
                outputs[cat] = fout, dw
            # write the row
            outputs[cat][1].writerow(row)
        # close all the files
        for fout, _ in outputs.values():
            fout.close()

if __name__ == "__main__":
    args = parse_arguments()
    split_bulk_csv(args.input_file, args.output_path)
