#!/usr/bin/env python

# Merges inputs csv file into one

import argparse
import csv
import logging
import pandas as pd
import gc

from common.configuration import ConfigurationHandler
from common.workspace import WorkspaceHandler

logger = logging.getLogger('pyccflex')
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
logger.addHandler(ch)

if __name__ == '__main__':

    logger.info("\n#### Running: {}".format(__file__))

    # Parse input parameters
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_files', nargs='+', type=str,
                        help="The list of input files to merge", required=True)
    parser.add_argument("--output_file", help="The name of the output file",
                        default=False, type=str)
    parser.add_argument("--locations_config", help="Path to locations configuration file",
                        type=str, required=False, default="./locations.json")
    parser.add_argument("--files_format_config", help="Path to files format configuration file",
                        type=str, required=False, default="./files_format.json")
    parser.add_argument("--add_decision_class", help="Shall the decision class be added to output?",
                        default=False, action='store_true')
    parser.add_argument("--add_contents", help="Shall the content of the line be added to output?",
                        default=False, action='store_true')
    parser.add_argument("--chunk_size", help="Number of lines to process in a batch",
                        type=int, required=False, default=10 ** 5)

    args = vars(parser.parse_args())
    logger.info("Run parameters: {}".format(str(args)))

    locations_file_path = args['locations_config']
    files_format_file_path = args['files_format_config']
    chunk_size = args['chunk_size']

    input_files = args['input_files']
    if len(input_files) < 2:
        logger.error("At least two input files need to be provided")
        exit(1)

    output_file = args['output_file']

    try:
        locations_config = ConfigurationHandler(locations_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(locations_file_path))
        exit(1)

    try:
        files_format_config = ConfigurationHandler(files_format_file_path)
    except Exception as e:
        logger.error("Couldn't load configuration file {}".format(files_format_file_path))
        exit(1)

    csv_separator = files_format_config.get("csv_sep", ",")

    workspace_dir_conf = locations_config.get('workspace_dir', None)
    workspace_dir_path = workspace_dir_conf.get("path", "")
    workspace_dir = WorkspaceHandler(workspace_dir_path)

    add_decision_class = args['add_decision_class']
    add_contents = args['add_contents']

    logger.info(">>> Starting to merge the input csv files...")

    output_file_path = workspace_dir.get_processing_file_path(output_file)

    input_file_path = workspace_dir.get_processing_file_path(input_files[0])
    first_input_file_df_chunks = pd.read_csv(input_file_path, sep=csv_separator,
                                             encoding="utf-8", chunksize=chunk_size, iterator=False)

    input_files_dfs_chunks = []
    for i, input_file in enumerate(input_files):
        if i > 0:
            input_file_path = workspace_dir.get_processing_file_path(input_file)
            input_df_chunks = pd.read_csv(input_file_path, sep=csv_separator,
                                          encoding="utf-8", chunksize=chunk_size, iterator=False)
            input_files_dfs_chunks.append(input_df_chunks)

    with open(output_file_path, 'w', newline='', encoding="utf-8") as f:

        for i, first_input_file_df_chunk in enumerate(first_input_file_df_chunks):
            contents_column = None
            decision_class_name_column = None
            decision_class_value_column = None
            merged_df = None

            columns = list(first_input_file_df_chunk.columns)
            if decision_class_name_column is None and 'class_value' in columns:
                decision_class_name_column = first_input_file_df_chunk['class_name']
                decision_class_value_column = first_input_file_df_chunk['class_value']

            if contents_column is None and 'contents' in columns:
                contents_column = first_input_file_df_chunk['contents']

            columns_to_drop = []
            if 'class_value' in columns:
                columns_to_drop.append('class_value')
                columns_to_drop.append('class_name')
            if 'contents' in columns:
                columns_to_drop.append('contents')

            first_input_file_df_chunk = first_input_file_df_chunk.drop(columns_to_drop, inplace=False, axis=1)
            chunks_to_merge = [first_input_file_df_chunk]

            for input_file_dfs_chunks in input_files_dfs_chunks:
                try:
                    input_df_chunk = input_file_dfs_chunks.get_chunk()
                except StopIteration as e:
                    pass

                logger.info(">>> Merging {} lines of the file {}".format(input_df_chunk.shape[0], input_file))

                columns = list(input_df_chunk.columns)
                if decision_class_name_column is None and 'class_value' in columns:
                    decision_class_name_column = input_df_chunk['class_name']
                    decision_class_value_column = input_df_chunk['class_value']

                if contents_column is None and 'contents' in columns:
                    contents_column = input_df_chunk['contents']

                columns_to_drop = ['id']
                if 'class_value' in columns:
                    columns_to_drop.append('class_value')
                    columns_to_drop.append('class_name')
                if 'contents' in columns:
                    columns_to_drop.append('contents')
                input_df_chunk.drop(columns_to_drop, inplace=True, axis=1)

                chunks_to_merge.append(input_df_chunk)

            merged_df = pd.concat(chunks_to_merge, axis=1)

            for chunk in chunks_to_merge:
                del chunk
            del chunks_to_merge
            gc.collect()

            if add_decision_class and decision_class_name_column is not None:
                merged_df['class_name'] = decision_class_name_column
                merged_df['class_value'] = decision_class_value_column

            if add_contents and contents_column is not None:
                merged_df['contents'] = contents_column

            logger.info(">>> Saving the merged chunks to file {}".format(output_file))
            merged_df.to_csv(f, sep=csv_separator, index=False, encoding="utf-8", header=(i == 0),
                             quoting=csv.QUOTE_NONNUMERIC)
            del merged_df
            gc.collect()
