''' From a directory containing a submission and a directory containing reference files, validate a submission

    Usage: 
        python3 OpenASR_validate_submission.py version
        
        python3 OpenASR_validate_submission.py validate 
        -ref test_suite_validation/reference/ 
        -s test_suite_validation/submissions/PASS_submission/
        -l pashto

'''
import os
import argparse
import pandas as pd
import subprocess
import logging
import csv
import __init__

logger = logging.getLogger('VALIDATION')

def validate_submission_file(subm_file,ref_file,language):
    '''
        Checks that a submission document references the same documents as 
        a reference, that the submission contains either 5 or 6 columns (if CF provided),
        that the 2nd column of the submission document contains either '1' for inLine or '2' for outLine
    '''

    # Run the regular SCTK validator to check the syntax
    if language==None:
        pipe = subprocess.Popen(" ".join(["./ctmValidator.pl", "-i", subm_file]), stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True, universal_newlines=True, shell=True)
    else:
        pipe = subprocess.Popen(" ".join(["./ctmValidator.pl", "-i", subm_file, "-l", language]), stdout=subprocess.PIPE, stderr=subprocess.PIPE, close_fds=True, universal_newlines=True, shell=True)
    
    result,error = pipe.communicate()

    if result[:5] == "ERROR":
        logger.error(result)
        return False

    # Check that the file is a .ctm file (whitespace separated or tab separated hypfile)
    submission_df = None
    try:
        submission_df = pd.read_csv(subm_file, header=None, dtype=str, delim_whitespace=True)
    except Exception as e:
        logger.error('{} is not a valid ctm file'.format(subm_file))
        return False    

    # # 1st check: check that the submission has either 5 or 6 columns (optional confidence factor) ##issue handled by SCTK
    # if not submission_df.shape[1] in [5,6]:
    #     logger.error('{} is not a valid tsv file: invalid columns'.format(subm_file))
    #     return False

    # 1st check: check that the 2nd column of the submission contains only 1 for inLine or 2 for outLine
    if "outLine" in subm_file:
        if not set(submission_df[1].unique()) <= set(['2']):
            logger.error('Invalid file: {}'.format(subm_file))
            logger.error('Second column in file {} should only contain 2 for outLine'.format(subm_file))
            return False
    else:
        if not set(submission_df[1].unique()) <= set(['1']):
            logger.error('Invalid file: {}'.format(subm_file))
            logger.error('Second column in file {} should only contain 1 for inLine'.format(subm_file))
            return False

    # Load ref content to be validated against hyp docID
    rows = []
    csv_header = ['Filename', 'Channel', 'Speaker','BeginTime','EndTime','Transcript']
    with open(ref_file, 'r') as ref:
        for row in csv.DictReader(ref, delimiter='\t', fieldnames=csv_header[:-1], restkey=csv_header[-1], skipinitialspace=True):
            try:
                rows.append([row['Filename'], row['Channel'], row['Speaker'], row['BeginTime'], row['EndTime'], ' '.join(row['Transcript'])])
            except Exception as e:
                rows.append([row['Filename'], row['Channel'], row['Speaker'], row['BeginTime'], row['EndTime'], ' '])

    reference_df = pd.DataFrame(rows)

    subm_files = submission_df[0]
    ref_files = reference_df[0]

    # 2nd check: check that the documents are identical
    if set(ref_files) != set(subm_files):
        logger.error('Invalid file: {}'.format(subm_file))

        # Check if we had not enough docs in the  submission files
        invalid_files = set(ref_files) - set(subm_files)
        if invalid_files:
            logger.error('The following document(s) were not found: {}'.format(invalid_files))
 
        # Check whether we had too many docs in the submission files
        invalid_files = set(subm_files) - set(ref_files)
        if invalid_files:
            logger.error('The following extra document(s) have been found: {}'.format(invalid_files))

        return False

    # 3rd check: ensure that the 5th column doesn't start with a number
    if submission_df[4].str.match(r'^[0-9]').any():
        logger.error('Invalid file: {}'.format(subm_file))
        logger.error('5th column in file {} should not contain a number'.format(subm_file))
        return False

    logger.debug('Submission file {} is valid'.format(subm_file))

    return True

def get_reference_submission_files_pairs(ref_dir, subm_dir, validate=True):

    logger.info('Gathering {} system outputs from {} and reference from {}...'.format(
                    'and validating' if validate else '',
                    subm_dir, ref_dir)
                )

    # Build the reference files list
    ref_files = [ files for _, _, files in os.walk(ref_dir) ]

    if not ref_files:
        logger.error('No such directory {}'.format(ref_dir))
        exit(1)

    ref_files = ref_files[0]
    ref_filebase = [ os.path.splitext(file)[0] for file in ref_files ]

    # Build the submission files list
    subm_files = [ files for _, _, files in os.walk(subm_dir) ]

    if not subm_files:
        logger.error('No such directory {}'.format(subm_dir))
        exit(1)

    subm_files = subm_files[0]
    subm_filebase = [ os.path.splitext(file)[0] for file in subm_files ]

    # Check that both list are identical
    if set(ref_filebase) != set(subm_filebase) and validate:
        logger.error('Invalid submission')

        # If the directory is empty 
        if len(subm_filebase) == 0:
            logger.error('No file found')
            exit(1)

        # Check if we had not enough submission files
        invalid_filebase = set(ref_filebase) - set(subm_filebase)
        invalid_files = [ '{}.ctm'.format(file) for file in invalid_filebase ]
        if invalid_filebase:
            logger.error('The following file(s) were not found: {}'.format(invalid_files))
 
        # Check whether we had too many submission files
        invalid_filebase = set(subm_filebase) - set(ref_filebase)
        invalid_files = [ '{}.ctm'.format(file) for file in invalid_filebase ]
        if invalid_filebase:
            logger.error('The following extra file(s) have been found: {}'.format(invalid_files))

        exit(1)

    # Here, the directory is supposed to be valid, let's build the dict
    # contain the submission file for each reference file
    results = {}
    for ref_file in ref_filebase:

        # Get the file path of each query.tsv file
        ref_file_path = os.path.join(ref_dir, '{}.stm'.format(ref_file))
        subm_file_path = os.path.join(subm_dir, '{}.ctm'.format(ref_file))

        # Update a dictionary with the scoring results for each query file
        results.update({ ref_file_path: subm_file_path })

    return results

def validate_submission_dir(args):

    # Load the submission and reference files
    ref_subm_pairs = get_reference_submission_files_pairs(args.reference_dir, args.submission_dir, validate=True)
    # Here, the directory is supposed to be valid, let's execute the validation for each file
    results = {}
    for ref_file_path, subm_file_path in ref_subm_pairs.items():

        # Update a dictionary with the result of the validation
        results.update({ os.path.basename(ref_file_path): validate_submission_file(subm_file_path,ref_file_path,args.language)})


    # Gather the invalid results
    invalid_files = [file for file in results if not results[file] ]

    if invalid_files:
        logger.error('Validation failed for the following file(s): {}'.format(invalid_files))
        exit(1)
    else:
        logger.info('Validation success')


def main():

    logging.basicConfig(level=logging.INFO)
    
    #validate commands
    parser = argparse.ArgumentParser(description='Validate a submission against a directory containing reference files')
    parser.add_argument('-s','--submission_dir', type=str, required=True, help='directory containing a submission')
    parser.add_argument('-ref','--reference_dir', type=str, required=True, help='reference directory')
    parser.add_argument('-l','--language', type=str, help='indicates the language of the submission files. If not given, the default language is English.')
    parser.add_argument('-v','--version', action='version', version='{version}'.format(version=__init__.__version__))
    parser.set_defaults(func=validate_submission_dir)

    args = parser.parse_args()

    if hasattr(args, 'func') and args.func:
        args.func(args)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()
