#!/usr/bin/env python3

import click
import json
from click.decorators import group
import pandas as pd
from pathlib import Path
from snpScore import sample_and_group_for_web
from snpScore import snpTableMP
from snpScore import CHR_SIZE
from snpScore import outdir_suffix_from_params
from snpScore import format_outfile
from snpScore import snpFilterBox, VAR_TO_VCF_COLUMN_MAP, table2annotation_df


@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(
    "--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("--chrom", help="results directory.", required=True, type=click.STRING)
@click.option("-t", "--thread", help="paralle number", default=4, 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, chrom, plant, thread
):
    parameters_obj = json.loads(parameters)
    sample_list, group_list = sample_and_group_for_web(parameters_obj)
    filter_grp = list(set(group_list))

    outdir = Path(outdir)
    if plant:
        full_outdir = outdir
    else:
        outdir_suffix = outdir_suffix_from_params(parameters_obj)
        outdir = outdir.resolve()
        full_outdir = outdir / outdir_suffix
    results_dir = full_outdir / f"analysis/split/{chrom}"

    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)

    snp_table_obj = snpTableMP(
        out_dir=results_dir,
        table_dirs=vcf_dir,
        samples=sample_list,
        sample_label=group_list,
        min_depth=min_depth,
        filter_dp_grp=filter_grp,
        chrom=chrom,
    )

    snpFilter_obj = snpFilterBox(
        alt_freq_df=snp_table_obj.alt_freq_df,
        group=filter_grp,
        outdir=results_dir,
        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),
    )
    snpFilter_obj.alt_filter_freq_df

    # format output
    chr_df = pd.read_csv(chr_size, sep="\t", header=None, names=["chrom", "size"])
    format_res_dir = results_dir / "fmt"
    format_res_dir.mkdir(parents=True, exist_ok=True)

    chromosome_snpeff_db = f"{snpeff_db}.{chrom}"

    snp_annotation_df = table2annotation_df(
        input_table=snpFilter_obj.alt_filter_freq_file,
        column_map=VAR_TO_VCF_COLUMN_MAP,
        snpeff_db=chromosome_snpeff_db,
        snpeff_cfg=snpeff_cfg,
        annotation_dir=results_dir,
        prefix="qtlseqr",
    )

    format_outfile(
        snpFilter_obj.alt_filter_freq_file,
        format_res_dir,
        ann_df=snp_annotation_df,
        chr_list=chr_df.chrom.astype("str"),
    )


if __name__ == "__main__":
    main()
