#!/usr/bin/env python3

import click
import json
from loguru import logger
from pathlib import Path
from snpScore import tableFromVcf, sample_and_group_for_web
from snpScore import tableFromSelectTable
from snpScore import snpTable, async_batch_sh_jobs
from snpScore import snpScoreBox, qtlSeqr, CHR_SIZE
from snpScore import outdir_suffix_from_params


@click.command()
@click.option(
    '-p',
    '--parameters',
    help='snpScore parameters json string.',
    required=True,
    type=click.STRING
)
@click.option(
    '-d',
    '--vcf_dir',
    help=('vcf table directory, can be more than one '
          '[--vcf_dir dir1 --vcf_dir dir2].'),
    type=click.Path(exists=True, file_okay=False),
    multiple=True)
@click.option(
    '--vcf_ann_file',
    help='snp annotation file.',
    required=True,
    type=click.Path()
)
@click.option('--chr_size',
              help='chr size file, default is chr size of bread wheat.',
              type=click.Path(exists=True, dir_okay=False),
              default=CHR_SIZE)
@click.option(
    '-o',
    '--outdir',
    help='results directory.',
    required=True,
    type=click.Path()
)
def main(parameters,
         vcf_dir,
         vcf_ann_file,
         chr_size,
         outdir):
    parameters_obj = json.loads(parameters)
    sample_list, group_list = sample_and_group_for_web(parameters_obj)

    outdir = Path(outdir)
    outdir_suffix = outdir_suffix_from_params(
        parameters_obj)
    outdir = outdir.resolve()
    full_outdir = outdir / outdir_suffix
    results_dir = full_outdir / 'results'

    input_params = dict()
    for param_i in parameters_obj:
        if parameters_obj[param_i]:
            input_params.update({param_i: parameters_obj[param_i]})

    min_depth = input_params.get('min_depth', 5)
    ref_freq = input_params.get('ref_freq', 0.43)

    snp_table_obj = snpTable(out_dir=results_dir,
                             table_dirs=vcf_dir,
                             samples=sample_list,
                             sample_label=group_list,
                             min_depth=min_depth)

    run_jobs = []
    if input_params.get('qtlseqr') or input_params.get('ed'):
        qtlseqr_obj = qtlSeqr(input_table=snp_table_obj.qtlseqr_snp_table,
                              window=input_params.get('qtlseqr_window', 1e7),
                              ref_freq=ref_freq,
                              pop_stru=input_params.get('pop_stru', 'RIL'),
                              min_sample_dp=min_depth,
                              out_dir=results_dir,
                              run_qtlseqr=input_params.get('qtlseqr'),
                              run_ed=input_params.get('ed'))
        run_jobs.append(qtlseqr_obj.qtlseqr_job)

    snp_score_methods = []
    snp_score_methods.append('var')
    if input_params.get('snp_index'):
        snp_score_methods.append('snp_index')
    if input_params.get('est'):
        snp_score_methods.extend(['est_mut_alt', 'est_mut_ref'])

    snpscore_obj = snpScoreBox(alt_freq_df=snp_table_obj.alt_freq_df,
                               grp_list=group_list,
                               method_list=snp_score_methods,
                               outdir=results_dir,
                               chr_size=chr_size,
                               min_depth=min_depth,
                               snp_number_window=input_params.get(
                                   'snp_number_window', 20),
                               snp_number_step=input_params.get(
                                   'snp_number_step', 5),
                               ref_freq=ref_freq,
                               p_ref_freq=input_params.get('p_ref_freq', 0.3),
                               background_ref_freq=input_params.get(
                                   'background_ref_freq', 0.3),
                               mutant_alt_exp=input_params.get(
                                   'mutant_alt_exp'),
                               wild_alt_exp=input_params.get('wild_alt_exp'),
                               vcf_ann_file=vcf_ann_file)
    run_jobs.extend(snpscore_obj.score_jobs)
    logger.info('Launch snpScore jobs...')
    async_batch_sh_jobs(run_jobs, thread=4)
    logger.info('The End.')
    print(full_outdir)


if __name__ == '__main__':
    main()
