#!/usr/bin/env python3

import sys
import json
import click
import delegator
import pandas as pd
from loguru import logger
from pathlib import Path
from snpScore import sample_and_group_for_web
from snpScore import async_batch_sh_jobs
from snpScore import CHR_SIZE
from snpScore import outdir_suffix_from_params
from snpScore import merge_split_file, score_plot, wrap_param_arg
from snpScore import add_filter_default_params, params_cfg, is_new_cmd, cp_if_not_exist
from snpScore import var_density_stats, var_density_file_suffix, snpFilterBox

BY_CHR_BIN = "snpFilter-bychr"


@click.command()
@click.option(
    "-p",
    "--parameters",
    help="snpScore parameters json string.",
    required=True,
    type=click.STRING,
)
@click.option(
    "--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(
    "--snpeff_cfg",
    help="snpEff config file.",
    type=click.Path(exists=True, dir_okay=False),
)
@click.option("--snpeff_db", help="snpEff database name.", type=click.STRING)
@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()
)
@click.option("-t", "--thread", help="paralle number", default=8, type=click.INT)
@click.option("--plant", help="for general plant", is_flag=True)
def main(
    parameters,
    vcf_dir,
    snpeff_cfg,
    snpeff_db,
    chr_size,
    outdir,
    thread,
    plant,
):

    chr_size_df = pd.read_csv(chr_size, sep="\t", header=None, index_col=0)
    cmd_list = []

    parameters_obj = json.loads(parameters)
    sample_list, group_list = sample_and_group_for_web(parameters_obj)

    input_params = add_filter_default_params(parameters_obj)

    outdir = Path(outdir)
    outdir = outdir.resolve()
    if plant:
        full_outdir = outdir
    else:
        outdir_suffix = outdir_suffix_from_params(parameters_obj)
        full_outdir = outdir / outdir_suffix

    results_dir = full_outdir / "analysis"
    cmd_history_dir = results_dir / "history"

    for chr_i in chr_size_df.index:
        param_wrap_arg = wrap_param_arg(sys.argv[1:])
        arg_line = " ".join(param_wrap_arg)
        cmd_i = f"{BY_CHR_BIN} {arg_line} --chrom {chr_i}"
        cmd_list.append(cmd_i)

    logger.info("Calculating score ...")
    async_batch_sh_jobs(cmd_list, thread=thread)
    logger.info("Ploting score ...")
    snp_score_methods = []
    snp_score_methods.append("var")

    snpFilter_obj = snpFilterBox(
        alt_freq_df=None,
        group=group_list,
        outdir=outdir,
        min_depth=input_params.get("min_depth", 5),
        mutant_freq=input_params.get("mut_freq", 0.4),
        wild_freq=input_params.get("wild_freq", 0.4),
        pat_mutant_freq=input_params.get("p_mut_freq", 0),
        pat_wild_freq=input_params.get("p_wild_freq", 0),
        afd=input_params.get("afd"),
        afd_deviation=input_params.get("afd_deviation", 0.05),
        parent_afd=input_params.get("p_afd", 1),
        parent_afd_deviation=input_params.get("p_afd_deviation", 0.05),
    )

    plot_cmds = []

    # varBScore
    var_filter_dir = outdir / "variantFilter"
    var_filter_file_pattern = snpFilter_obj.alt_filter_freq_file.name
    var_filter_file = merge_split_file(
        results_dir, var_filter_file_pattern, out_dir=var_filter_dir
    )

    var_filter_df = pd.read_csv(var_filter_file)
    snp_density_window = input_params.get("snp_density_window", 1000_000)
    snp_density_step = input_params.get("snp_density_step", None)
    var_filter_density_df = var_density_stats(
        chr_size,
        var_filter_df,
        window=snp_density_window,
        step=snp_density_step,
    )

    if plant:
        var_filter_density_file = var_filter_dir / "var.filter.density.csv"
    else:
        var_filter_density_suffix = var_density_file_suffix(
            snp_density_window, snp_density_step
        )
        var_filter_density_file = (
            var_filter_dir / f"{var_filter_file.stem}{var_filter_density_suffix}"
        )
    var_filter_density_df.to_csv(var_filter_density_file, index=False)
    plot_cmds.append(
        score_plot(
            var_filter_density_file,
            "density-new",
            var_filter_density_file.stem,
            chr_size,
        )
    )
    plot_cmds = list(filter(None, plot_cmds))
    if plot_cmds:
        logger.info("Launch plot jobs...")
        async_batch_sh_jobs(plot_cmds, thread=thread)

    # write params config
    if not plant:
        if is_new_cmd(input_params, cmd_history_dir):
            params_cfg_file = var_filter_dir / "parameters.txt"
            params_cfg(
                params_cfg_file, input_params, cmd_history_dir, cfg_type="varFilter"
            )

    if plant:
        zip_cmd = f"cd {outdir}; zip -r {var_filter_dir.name}.zip {var_filter_dir.name}"
        delegator.run(zip_cmd)
    logger.info("Fin.")
    print(var_filter_density_file.with_suffix(".plot.jpg"))


if __name__ == "__main__":
    main()
