#!/usr/bin/env python3
"""
This file is part of MitoZ.

THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE
OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
PERFORMANCE OF THIS SOFTWARE.

COPYRIGHT © 2019-2022 Guanliang Meng. ALL RIGHTS RESERVED.

"""

import sys
import os
import argparse
import time
import subprocess
from string import Template # Python >=3.6
from pathlib import Path
import glob


def add_arguments(parser):
    parser.add_argument('--fq1', metavar='<file>', required=True,
        help="Fastq1 file")

    parser.add_argument('--fq2', metavar='<file>',
        help="Fastq2 file")

    parser.add_argument('--phred64', action='store_true',
        help="Are the fastq phred64 encoded? [%(default)s]")

    parser.add_argument('--outprefix', metavar='<str>', default='out',
        help="output prefix [%(default)s]")

    parser.add_argument("--fastq_read_length", metavar="<INT>",  type=int, default=150,
        help='''read length of fastq reads, used to split clean fastq files. [%(default)s]''')

    parser.add_argument("--data_size_for_mt_assembly", metavar="<float>", type=float, default=5,
        help='''Data size (GB) used for mitochondrial genome assembly, it is recommmended to be less than 8 GB. Set to 0 if you want to use ALL clean data. [%(default)s]''')

    parser.add_argument('--filter_other_para', metavar='<str>', default='',
        help="other parar. [%(default)s]")

    parser.add_argument('--thread_number', metavar='<int>', type=int, default=4,
        help="thread number [%(default)s]")

    parser.add_argument('--workdir', metavar='<directory>', default='./',
        help="working directory [%(default)s]")

    parser.add_argument('--workdir_done', metavar='<directory>', default='./done',
        help="done directory [%(default)s]")

    parser.add_argument('--workdir_log', metavar='<directory>', default='./log',
        help="log directory [%(default)s]")

    return parser


def filter_rawdata(fq1=None, fq2=None, phred64=False, outprefix='out', other_para='', thread=4, workdir='./', workdir_done='./done/', workdir_log='./log/', logger=None):
    out1 = outprefix + ".clean_R1.fq.gz"
    out2 = outprefix + ".clean_R2.fq.gz"
    report_json = outprefix + ".QC.json"
    report_html = outprefix + ".QC.html"

    filtered_fq1 = os.path.join(workdir, out1)
    filtered_fq2 = os.path.join(workdir, out2)

    if phred64:
        phred64_value = '--phred64'
    else:
        phred64_value = ''

    command_template = Template('''
    # set -vex
    # set -o pipefail
    set -e

    mkdir -p ${workdir_done}
    mkdir -p ${workdir_log}
    mkdir -p ${workdir}
    cd ${workdir}

    task_id="FastpTask"
    done_file="${workdir_done}/$task_id.done"
    log_o_file="${workdir_log}/$task_id.o.log"
    log_e_file="${workdir_log}/$task_id.e.log"

    if [ -f $done_file ]; then
        exit 0
    fi

    if [ -f "${fq1}" ] && [ -f "${fq2}" ];
    then
        fastp ${other_para} \
        -i ${fq1} \
        -I ${fq2} ${phred64} \
        --length_required 50 \
        -o ${out1} \
        -O ${out2} \
        -R "${outprefix}" \
        -j ${report_json} \
        -h ${report_html} \
        -w ${thread} && 1>$log_o_file 2>$log_e_file && touch $done_file   
 
    else
        fastp ${other_para} \
        -i ${fq1} ${phred64} \
        --length_required 50 \
        -o ${out1} \
        -R "${outprefix}" \
        -j ${report_json} \
        -h ${report_html} \
        -w ${thread} && 1>$log_o_file 2>$log_e_file && touch $done_file
    
    fi
    ''')

    command = command_template.safe_substitute(
        workdir=workdir,
        workdir_done=workdir_done,
        workdir_log=workdir_log,
        fq1=fq1,
        fq2=fq2,
        out1=out1,
        out2=out2,
        other_para=other_para,
        outprefix=outprefix,
        report_json=report_json,
        report_html=report_html,
        thread=thread,
        phred64=phred64_value,
        )

    # logger.info("Starting to filter raw data {}".format(command))
    runcmd(command, logger=logger)

    if fq1 and not fq2:
        return os.path.abspath(filtered_fq1), ''
    else:
        return os.path.abspath(filtered_fq1), os.path.abspath(filtered_fq2)


def runcmd(command, logger=None):
    logger.info(command)
    subprocess.run(command, shell=True)

    '''
    try:
        current_time = time.strftime("%Y-%m-%d %H:%M:%S",
                        time.localtime(time.time()))
        print(current_time, "\n", command, "\n", sep="", flush=True)
        subprocess.check_call(command, shell=True)
    except:
        sys.exit("Error occured when running command:\n%s" % command)
    '''

def get_single_fq_file_lines(fastq_read_length=150, data_size_for_mt_assembly=5, fq1=None, fq2=None, logger=None):
    # do not split
    if not data_size_for_mt_assembly:
        logger.info("No spliting clean fastq!")
        return ''

    # pair-end data
    if fq1 and fq2:
        fq_records_per_file = int(data_size_for_mt_assembly * 10e9 / (fastq_read_length * 2))
        
        single_fq_file_lines = fq_records_per_file * 4
        m = '''
        Paired-end data:
        fastq_read_length: {fastq_read_length}
        data_size_for_mt_assembly: {data_size_for_mt_assembly}
        fq_records_per_file: {fq_records_per_file}
        fq_records_per_file * 4: {line}
        '''.format(fastq_read_length=fastq_read_length,
            data_size_for_mt_assembly=data_size_for_mt_assembly, 
            fq_records_per_file=fq_records_per_file,
            line=single_fq_file_lines)

        logger.info(m)

        return single_fq_file_lines

    # SE data
    elif (fq1 or fq2):
        fq_records_per_file = int(data_size_for_mt_assembly * 10e9 / fastq_read_length)
        single_fq_file_lines = fq_records_per_file * 4
        m = '''
        Single-end data:
        fastq_read_length: {fastq_read_length}
        data_size_for_mt_assembly: {data_size_for_mt_assembly}
        fq_records_per_file: {fq_records_per_file}
        fq_records_per_file * 4: {line}
        '''.format(fastq_read_length=fastq_read_length,
            data_size_for_mt_assembly=data_size_for_mt_assembly, 
            fq_records_per_file=fq_records_per_file,
            line=single_fq_file_lines)

        logger.info(m)
        return single_fq_file_lines

    return 30000000


def extract_certain_amount_fq_records(workdir='./', fastq_read_length=150, data_size_for_mt_assembly=5, fq1=None, fq2=None, logger=None):
    single_fq_file_lines = get_single_fq_file_lines(
        fastq_read_length=fastq_read_length,
        data_size_for_mt_assembly=data_size_for_mt_assembly,
        fq1=fq1,
        fq2=fq2,
        logger=logger)
   
    # no need to subsampling
    if not single_fq_file_lines:
        logger.info("extract_certain_amount_fq_records(): since there is no need to subsampling, returns input {fq1} and {fq2} directly.".format(fq1=fq1, fq2=fq2))
        return os.path.abspath(fq1), os.path.abspath(fq2)

    sub_fq1 = "sampling_{}GB.clean_R1.fq.gz".format(data_size_for_mt_assembly)
    sub_fq2 = "sampling_{}GB.clean_R2.fq.gz".format(data_size_for_mt_assembly)

    sub_fq1 = os.path.join(workdir, sub_fq1)
    sub_fq2 = os.path.join(workdir, sub_fq2)

    # already sampled, return existing resulting files directly
    if Path(sub_fq1).is_file() and Path(sub_fq2).is_file():
        logger.info("extract_certain_amount_fq_records(): Found existing {sub_fq1} and {sub_fq2}, there is no need to further subsampling, returns them directly.".format(sub_fq1=sub_fq1, sub_fq2=sub_fq2))
        return sub_fq1, sub_fq2

    # perform sampling
    # pair-end data
    if fq1 and fq2:
        command = "cd {workdir}\n".format(workdir=workdir)
        command += "seqkit head -n {single_fq_file_lines} -o {sub_fq1} {fq1} \n".format(
            single_fq_file_lines=single_fq_file_lines,
            fq1=fq1,
            sub_fq1=sub_fq1)
        command += "seqkit head -n {single_fq_file_lines} -o {sub_fq2} {fq2} \n".format(
            single_fq_file_lines=single_fq_file_lines,
            fq2=fq2,
            sub_fq2=sub_fq2)
        runcmd(command, logger=logger)

        if Path(sub_fq1).is_file() and Path(sub_fq2).is_file():
            logger.info("extract_certain_amount_fq_records(): return subsampling results:\n{sub_fq1}\n{sub_fq2}".format(sub_fq1=sub_fq1, sub_fq2=sub_fq2))
            return sub_fq1, sub_fq2
        else:
            logger.error("extract_certain_amount_fq_records(): Can not find sub_fq1: {sub_fq1} and sub_fq2 {sub_fq2}".format(sub_fq1=sub_fq1, sub_fq2=sub_fq2))

    elif fq1:
        command = "cd {workdir}\n".format(workdir=workdir)
        command += "seqkit head -n {single_fq_file_lines} -o {sub_fq1} {fq1} \n".format(
            single_fq_file_lines=single_fq_file_lines,
            fq1=fq1,
            sub_fq1=sub_fq1)
        runcmd(command, logger=logger)

        if Path(sub_fq1).is_file():
            logger.info("extract_certain_amount_fq_records(): return subsampling results:\n{sub_fq1}".format(sub_fq1=sub_fq1))
            return sub_fq1, ''
        else:
            logger.error("extract_certain_amount_fq_records(): Can not find sub_fq1: {sub_fq1}".format(sub_fq1=sub_fq1))

    else:
        logger.error("extract_certain_amount_fq_records(): Can not find input fq1 and/or fq2")


def find_resulting_fq_files(fq1='', fq2='', logger=None):
    if fq1 and fq2:
        if Path(fq1).is_file() and Path(fq2).is_file():
            m = '''Found fastq files:
            fq1: {fq1}
            fq2: {fq2}'''.format(fq1=fq1, fq2=fq2)
            logger.info(m)
            return fq1, fq2
        else:
            directory = os.path.dirname(fq1)
            f1_base = os.path.basename(fq1)
            f2_base = os.path.basename(fq2)
            fq1_list = glob.glob(directory + "/*"+f1_base)
            fq1_list = sorted(fq1_list)
            fq2_list = glob.glob(directory + "/*"+f2_base)
            fq2_list = sorted(fq2_list)
            m = '''Found fastq files:
fq1_list: \n{fq1_list}
fq2_list: \n{fq2_list}'''.format(fq1_list="\n".join(fq1_list), fq2_list="\n".join(fq2_list))
            logger.info(m)
            return fq1_list, fq2_list
    elif fq1:
        if Path(fq1).is_file():
            m = '''Found fastq file:
            fq1: {fq1}'''.format(fq1=fq1)
            logger.info(m)
            return fq1, ''
        else:
            directory = os.path.dirname(fq1)
            f1_base = os.path.basename(fq1)
            fq1_list = glob.glob(directory + "/*"+f1_base)
            fq1_list = sorted(fq1_list)
            m = '''Found fastq files:
            fq1_list: \n{fq1_list}'''.format(fq1_list="\n".join(fq1_list))
            logger.info(m)
            return fq1_list, ''

    else:
        logger.error("Can not find any fastq files!")


def main(args):
    logger = args.logger
    logger.info("filter.main() get args: {}".format(args))

    if args.fq1 and Path(args.fq1).is_file():
        args.fq1 = os.path.abspath(args.fq1)
    else:
        args.fq1 = ''
    if args.fq2 and Path(args.fq2).is_file():
        args.fq2 = os.path.abspath(args.fq2)
    else:
        args.fq2 = ''

    args.workdir = os.path.abspath(args.workdir)

    logger.info("filter.main() get updated args: {}".format(args))

    single_fq_file_lines = get_single_fq_file_lines(
        fastq_read_length=args.fastq_read_length,
        data_size_for_mt_assembly=args.data_size_for_mt_assembly,
        fq1=args.fq1,
        fq2=args.fq2,
        logger=logger)

    filtered_fq1, filtered_fq2 = filter_rawdata(
        fq1=args.fq1,
        fq2=args.fq2,
        phred64=args.phred64,
        outprefix=args.outprefix,
        other_para=args.filter_other_para,
        thread=args.thread_number,
        workdir=args.workdir,
        workdir_done=args.workdir_done,
        workdir_log=args.workdir_log,
        logger=logger)
    
    sub_fq1, sub_fq2 = extract_certain_amount_fq_records(
        workdir=args.workdir,
        fastq_read_length=args.fastq_read_length,
        data_size_for_mt_assembly=args.data_size_for_mt_assembly,
        fq1=filtered_fq1,
        fq2=filtered_fq2,
        logger=logger)

    logger.info('filter.main() returns:\n{0}\n{1}'.format(sub_fq1, sub_fq2))

    return sub_fq1, sub_fq2

