#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Part of the masterfile package: https://github.com/njvack/masterfile
# Copyright (c) 2017 Board of Regents of the University of Wisconsin System
# Written by Nate Vack <njvack@wisc.edu> at the Center for Healthy Minds
# at the University of Wisconsin-Madison.
# Released under MIT licence; see LICENSE at the package root.

""" Extract masterfile-formatted data from a CSV file.

Usage: extract_masterfile_data [options] <masterfile_path> <data> <outfile>

Takes a CSV file with a mix of data destined for a masterfile and other data,
and extracts and formats the data to masterfile format, writing the result
to stdout.

Extraction is done based on column header format: we'll extract the index
column, as well as any columns with enough underscores to make up all the
parts of settings.json's "components" entry.

If there's a "comments" or "description" row in your data, you can skip it
with the --skip option. Additionally, rows with blank index_column will
always be skipped, since they don't make any sense in the masterfile data.

Options:
  --index_column=<col>  Use <col> as the input's index column
  --skip=<rows>         Skip <rows> data columns [default: 0]
  -v, --verbose         Display debugging output
"""

from __future__ import absolute_import, unicode_literals

import sys
import re
import logging

import pandas as pd

import masterfile
from masterfile.masterfile import Masterfile
from masterfile.vendor.docopt import docopt

logging.basicConfig(level=logging.DEBUG, format='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)


def main(argv=None):
    pargs = docopt(__doc__, argv, version=masterfile.__package_version__)
    if pargs['--verbose']:
        logger.setLevel(logging.DEBUG)
    logger.debug(pargs)
    df = pd.read_csv(pargs['<data>'], dtype=str)
    mf = Masterfile.load_path(pargs['<masterfile_path>'])
    formatted = format_dataframe_for_masterfile(
        df, mf, pargs['--index_column'], int(pargs['--skip']))
    output = pargs['<outfile>']
    if output == '-':
        output = sys.stdout
    formatted.to_csv(output, line_terminator='\r\n')


def format_dataframe_for_masterfile(df, mf, input_index_col, skip_rows):
    """
    Return a copy of df, with:
    * input_index_col renamed to mf.index_col
    * index_col as the first column
    * only the columns that match the format (c1_c2_..._cN) included
    * rows with index_col as blank excluded
    """
    index_col = mf.index_column
    logger.debug('Original columns: {}'.format(df.columns))
    logger.debug('Original has {} rows'.format(len(df)))
    col_rx = col_match_regex(mf, input_index_col)
    logger.debug('Column regex: {}'.format(col_rx))
    col_filtered = df.filter(regex=col_rx)
    row_filtered = _filter_rows(col_filtered, input_index_col, skip_rows)
    logger.debug('Now has {} rows'.format(len(row_filtered)))
    renamed = row_filtered.rename(columns={input_index_col: index_col})
    renamed.set_index(index_col, inplace=True)
    logger.debug('New columns: {}'.format(renamed.columns))
    return renamed


def _filter_rows(df, index_col, skip_rows):
    skipped = df[skip_rows:]
    has_data = ~(
        (skipped[index_col].str.strip() == '') |
        (skipped[index_col].isnull())
    )
    return skipped[has_data]


def col_match_regex(mf, input_index_col):
    index_col = input_index_col or mf.index_column
    return '(^{}$)|(^{}$)'.format(
        re.escape(index_col), component_col_regex(len(mf.components)))


def component_col_regex(component_count):
    """
    Match things like 'foo_bar_baz_corge'
    """
    return '_'.join((['[^_\s]+'] * component_count))


if __name__ == '__main__':
    main()

@@ -1361,16 +1361,54 @@
  logging
+%0Afrom contextlib import contextmanager
 %0A%0Aimport
@@ -2066,138 +2066,380 @@
     
-output = pargs%5B'%3Coutfile%3E'%5D%0A    if output == '-':%0A        output = sys.stdout%0A    formatted.to_csv(output, line_terminator='%5Cr%5Cn')
+with file_or_stdout(pargs%5B'%3Coutfile%3E'%5D) as output:%0A        formatted.to_csv(output, line_terminator='%5Cr%5Cn')%0A%0A%0A@contextmanager%0Adef file_or_stdout(filename):%0A    if filename == '-':%0A        logger.info('Writing to stdout')%0A        yield sys.stdout%0A    else:%0A        with open(filename, 'w') as f:%0A            logger.info('Writing to %7B%7D'.format(filename))%0A            yield f
 %0A%0A%0Ad

Add info message to extract_masterfile_data