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

"""----------------------------------------------------------------------------
  Script Name: functions.py
  Description: Functions for functional tests
  Input files: Expected and observed folders
  Created By:  Pierre Martin
  Date:        2021-12-16
-------------------------------------------------------------------------------
"""

# Metadata
__author__ = 'Pierre Martin \
- MIAT - PF Genotoul'
__copyright__ = 'Copyright (C) 2021 INRAe'
__license__ = 'GNU General Public License'
__version__ = '1'
__email__ = 'support.bioinfo.genotoul@inra.fr'
__status__ = 'dev'

# Functions of metagwgs functional_tests (main.py)
try:
    import argparse
    import sys
    import re
    import os
    import os.path as path
    from pathlib import Path
    import json
    import subprocess

except ImportError as error:
    print(error)
    exit(1)

# Arguments parsing (3 required / 2 optional)
##  Usage example:
# 
# python functional_tests.py \
#   -step 01_clean_qc \
#   -exp_dir ../../functional_tests/expected \
#   -obs_dir ../../functional_tests/observed/results \
#   --script ../../functional_tests/test.sh
#   --verbose

def parse_arguments():

    parser = argparse.ArgumentParser()

    parser.add_argument('-step', type = str, \
    help = '(at least one required) step(s) of metagwgs you wish to perform a test on (multiple steps can be tested at once)')
    parser.add_argument('-exp_dir', type = str, \
    help = '(required) expected logs dir containing logs from a healthy metagwgs workflow')
    parser.add_argument('-obs_dir', type = str, \
    help = '(required) observed logs dir containing logs from the metagwgs workflow you wish to test')
    parser.add_argument('--script', type = str, \
    help = '(optional) script file containing metagwgs Nextflow launch command')
    parser.add_argument('--verbose', action = "store_true", \
    help = '(optional) print test results in stdout (default: false)')

    if len(sys.argv) == 1:
        parser.print_usage(sys.stderr)
        sys.exit(1)

    return parser.parse_args()

# Launch nextflow from script if given by user
def launch_nextflow(script):

    script_file = path.abspath(script)
    script_dir = path.dirname(script_file)

    print('Launching test run with the provided file:\n\n{}\n'.format(script_file))

    process = subprocess.Popen('sh' + ' {}'.format(script_file), cwd = script_dir, shell = True).wait()

    print('Test run completed')

# Find files in expected directory to test
def files_load(exp_dir, step):

    files_list = []
    files_dir = path.join(path.abspath(exp_dir), step)
    for file in Path(files_dir).rglob('*.[a-z]*'):
        files_list.append(path.relpath(file, start = files_dir))

    return sorted(files_list)

def gz_file(directory, func) :
        files_dir = path.join(path.abspath(directory), "01_clean_qc/01_3_taxonomic_affiliation_reads")

        for file in Path(files_dir).rglob('*_kaiju.out*'):
            os.system( "{} {}".format(func,file))

# Do file comparisons for given step (and write output in stdout if --verbose)
def check_files(exp_dir, obs_dir, step, methods, verbose):

    if ("01_clean_qc" in step):
        for directory in [exp_dir,obs_dir]:
            gz_file(directory,"gunzip")


    # Check existence of expected and observed directories
    if not path.exists(exp_dir) or not path.exists(obs_dir):
        sys.exit('{a} or {b} folder(s) do not exist, please check --exp_dir and --obs_dir parameters'.format(a = exp_dir, b = obs_dir))
    # Load files list to test from exp_dir folder
    files_list = files_load(exp_dir, step)

    # Initiate log for untested files (removes already existing logs)
    not_tested_log = 'ft_{}.not_tested'.format(step)
    os.remove(not_tested_log) if path.exists(not_tested_log) else None

    # Create a new log for this functional test
    tested_log = 'ft_{}.log'.format(step)
    log = open(tested_log, 'w+')

    # Paths to expected and observed parent directories
    expected_prefix = path.join(path.abspath(exp_dir), step)
    observed_prefix = path.join(path.abspath(obs_dir), step)
    log.write('Expected directory: {a}\nvs\nObserved directory: {b}\n'.format(a = expected_prefix, b = observed_prefix))

    # Passed and failed tests count initialization
    global true_cnt, false_cnt
    max_cnt = len(files_list)
    true_cnt = 0
    false_cnt = 0
    nt_cnt = 0

    if verbose: print('\nLaunching {}...\n'.format(step))

    # Test each file
    for file_path in files_list:

        # Metadata on file to find them and know which test to perform
        expected_path = path.join(expected_prefix, file_path)
        observed_path = path.join(observed_prefix, file_path)

        if verbose:
            print("exp:\t",expected_path)
            print("obs:\t",observed_path)

        file_name = path.basename(file_path)
        file_extension = path.splitext(file_name)[1]

        # Find which test to perform on given file (exceptions being "taxo_diff" and "cut_diff")
        method = ''
        for test in methods:
            if type(methods[test]) != list and re.search(methods[test], file_name):
                method = test
                break

            elif type(methods[test]) == list and file_extension in methods[test]:
                method = test
                break

        if method == '':
            sys.exit('Method {} doesn\'t exist for {} in {}'.format(test, file_name, expected_path))


        # Non existing files
        if not path.exists(observed_path):
            nt = open(not_tested_log, 'a')
            nt.write(observed_path + '\n')
            nt.close()
            nt_cnt += 1
            if max_cnt > 0:
                max_cnt -= 1
            file_out = '''
------------------------------------------------------------------------------

File:           {a}
Not tested (list in {b})
'''.format(a = file_path, b = not_tested_log)
            log.write(file_out)
            if verbose: print(file_out)


        # Existing files
        if path.exists(expected_path) and path.exists(observed_path):

            test, out = test_file(expected_path, observed_path, method)

            # Test failed:
            if test == False:
                file_out = '''
------------------------------------------------------------------------------

File:           {a}
Test method:    {b}

'''.format(a = file_path, b = method)
                log.write(file_out)
                if verbose: print(file_out)
                log.write(out)
                if verbose: print(out)

            # Test passed:
            elif test == True:
                file_out = '''
------------------------------------------------------------------------------

File:           {a}
Test method:    {b}
'''.format(a = file_path, b = method)

                log.write(file_out)
                if verbose: print(file_out)
                log.write(out)
                if verbose: print(out)

            continue

    if max_cnt != 0:
        if true_cnt != 0:
            true_perc = round((float(true_cnt) / float(max_cnt) * 100), 2)
        else:
            true_perc = float(0)
        if false_cnt != 0:
            false_perc = round(100 - (float(true_cnt) / float(max_cnt) * 100), 2)
        else:
            false_perc = float(0)
    else:
        true_perc = float(0)
        false_perc = float(0)

    out = '''
=========================================
-----------------------------------------

Testing the {a} step of metagWGS:

Total:      {b}
Passed:     {c} ({d}%)
Missed:     {e} ({f}%)
Not tested: {g}

Find more details in {h}

-----------------------------------------
=========================================
'''.format(a = step, b = max_cnt, c = true_cnt, d = true_perc, e = false_cnt, f = false_perc, g = nt_cnt, h = tested_log)

    log.write(out)
    log.close()

    if ("01_clean_qc" in step):
        for directory in [exp_dir,obs_dir]:
            gz_file(directory,"gzip")

    return out

# Perform a test for each method on given files
def test_file(exp_path, obs_path, method):

    global true_cnt, false_cnt

    if re.search('diff', method):

        if method == 'cut_diff':
            command = 'diff <(tail -n+6 {}) <(tail -n+6 {})'.format(exp_path, obs_path)

        elif method == 'sort_diff':
            command = 'diff <(sort {}) <(sort {})'.format(exp_path, obs_path)

        elif method == 'no_header_diff':
            command = 'diff <(grep -v "^#" {}) <(grep -v "^#" {})'.format(exp_path, obs_path)

        elif method == 'zdiff':
            command = 'zdiff {} {}'.format(exp_path, obs_path)

        process = subprocess.Popen(command, stdout = subprocess.PIPE, shell = True, executable = '/bin/bash')
        diff_out, error = process.communicate()
        
        if not error:
            diff_result = diff_out.decode()
            if diff_result != '':
                test = False
                out = 'Test result:    Failed\nDifferences:\n{}\n'.format(diff_result)
                false_cnt += 1

            elif diff_result == '':
                test = True
                out = 'Test result:    Passed\n'
                true_cnt += 1
        else:
            test = False
            out = 'Test result:    Failed\nSubprocess error:    {}\nCommand:    {}'.format(error, command)
            false_cnt += 1

        return test, out

    elif method == 'not_empty':
        test = path.getsize(obs_path) > 0

        if test:
            test = True
            out = 'Test result:    Passed\n'
            true_cnt += 1
        else:
            test = False
            out = 'Test result:    Failed\n'
            false_cnt += 1

        return test, out
