#!/usr/bin/env python

"""
Generate final annotated, soft-filtered Matrix Table
"""

import logging
from typing import Optional

import click
import hail as hl

from joint_calling.utils import get_validation_callback, file_exists
from joint_calling import utils, _version

logger = logging.getLogger('joint-calling')
logger.setLevel('INFO')


@click.command()
@click.version_option(_version.__version__)
@click.option(
    '--mt',
    'mt_path',
    required=True,
    callback=get_validation_callback(ext='mt', must_exist=True),
    help='path to the raw sparse Matrix Table generated by combine_gvcfs.py',
)
@click.option(
    '--out-mt',
    'out_mt_path',
    required=True,
    callback=get_validation_callback(ext='mt', must_exist=False),
    help='path to write the final annotated soft-filtered Matrix Table',
)
@click.option(
    '--out-nonref-mt',
    'out_nonref_mt_path',
    callback=get_validation_callback(ext='mt', must_exist=False),
    help='write a version of output Matrix Table without reference blocks',
)
@click.option(
    '--meta-ht',
    'meta_ht_path',
    required=True,
    help='Table generated by sample_qc.py',
)
@click.option(
    '--freq-ht',
    'freq_ht_path',
    required=True,
    help='Table with frequency annotations',
)
@click.option(
    '--vqsr-ht',
    'vqsr_ht_path',
    required=True,
    help='Table with AS-VQSR annotations (both input and output)',
)
@click.option(
    '--allele-data-ht',
    'allele_data_ht_path',
    required=True,
    help='Table with allele_data annotation',
)
@click.option(
    '--qc-ac-ht',
    'qc_ac_ht_path',
    required=True,
    help='Table with qc_ac_* annotations',
)
@click.option(
    '--local-tmp-dir',
    'local_tmp_dir',
    help='local directory for temporary files and Hail logs (must be local).',
)
@click.option(
    '--overwrite/--reuse',
    'overwrite',
    is_flag=True,
    help='if an intermediate or a final file exists, skip running the code '
    'that generates it.',
)
@click.option(
    '--hail-billing',
    'hail_billing',
    help='Hail billing account ID.',
)
def main(
    mt_path: str,
    out_mt_path: str,
    out_nonref_mt_path: Optional[str],
    meta_ht_path: str,
    vqsr_ht_path: str,
    freq_ht_path: str,
    allele_data_ht_path: str,
    qc_ac_ht_path: str,
    local_tmp_dir: str,
    overwrite: bool,  # pylint: disable=unused-argument
    hail_billing: str,  # pylint: disable=unused-argument
):  # pylint: disable=missing-function-docstring
    utils.init_hail('make_finalised_mt', local_tmp_dir)

    if file_exists(out_mt_path):
        if overwrite:
            logger.info(f'Output {out_mt_path} exists and will be overwritten')
        else:
            logger.error(
                f'Output file {out_mt_path} exists, use --overwrite to overwrite'
            )
            return

    # This will return a table with split multiallelics, and thus GT as a field
    mt = utils.get_mt(
        mt_path, split=True, add_meta=True, meta_ht=hl.read_table(meta_ht_path)
    )

    vqsr_ht = hl.read_table(vqsr_ht_path)
    mt = annotate_vqsr(mt, vqsr_ht)
    # INFO/AS_* fields
    # AS-VQSR output fields (AS_VQSLOD, AS_culprit, NEGATIVE_TRAIN_SITE, POSITIVE_TRAIN_SITE, AS_FilterStatus)
    # filters (corresponds to VCF's FILTER column. Populated from AS_FilterStatus after splitting by allele)
    
    freq_ht = hl.read_table(freq_ht_path)
    mt = mt.annotate_rows(**freq_ht[mt.row_key])
    # InbreedingCoeff
    # freq
    # faf
    # popmax:
    #   AC
    #   AF
    #   AN
    #   homozygote_count
    #   pop
    #   faf95
    # qual_hists
    # raw_qual_hists
    
    allele_data_ht = hl.read_table(allele_data_ht_path)
    mt = mt.annotate_rows(**allele_data_ht[mt.row_key])
    # allele_data:
    #   nonsplit_alleles
    #   has_star
    #   variant_type
    #   n_alt_alleles
    #   allele_type
    #   was_mixed
    
    qc_ac_ht = hl.read_table(qc_ac_ht_path)
    mt = mt.annotate_rows(**qc_ac_ht[mt.row_key])
    # ac_qc_samples_raw
    # ac_qc_samples_unrelated_raw
    # ac_release_samples_raw
    # ac_qc_samples_adj
    # ac_qc_samples_unrelated_adj
    # ac_release_samples_adj

    for ht in [vqsr_ht, freq_ht, allele_data_ht, qc_ac_ht]:
        mt = mt.annotate_globals(**ht.index_globals())

    mt.write(out_mt_path, overwrite=True)

    if out_nonref_mt_path:
        # hl.len(mt.alleles) > 1 - find sites with an ALT allele
        # hl.agg.any(mt.GT.is_non_ref() - check that there is any sample with non-ref
        #    genotype present, in case if all samples with ALT were filtered out
        mt = mt.filter_rows((hl.len(mt.alleles) > 1) & hl.agg.any(mt.GT.is_non_ref()))
        mt.write(out_nonref_mt_path, overwrite=True)


def annotate_vqsr(mt, vqsr_ht):
    """
    Assuming `vqsr_ht` is a site-only table (split multiallelics) from VQSR VCF, 
    and `mt` is a matrix table (split multiallelics), annotates `mt` rows from `vqsr_ht`
    """
    mt = mt.annotate_rows(**vqsr_ht[mt.row_key])
    
    # vqsr_ht has info annotation split by allele; plus new AS-VQSR annotations
    mt = mt.annotate_rows(info=vqsr_ht[mt.row_key].info)

    # populating filters which is outside of info
    mt = mt.annotate_rows(
        filters=mt.filters.union(vqsr_ht[mt.row_key].filters),
    )
    
    mt = mt.annotate_globals(**vqsr_ht.index_globals())
    return mt


if __name__ == '__main__':
    main()  # pylint: disable=E1120
