#!/usr/bin/env python

"""
Evaluate random forest and VQSR results against the truth data
"""

from os.path import join
from typing import Optional
import logging
from pprint import pformat
import click
import hail as hl

from gnomad.utils.filtering import filter_low_conf_regions, filter_to_clinvar_pathogenic
from gnomad.variant_qc.evaluation import (
    compute_binned_truth_sample_concordance,
    compute_grouped_binned_ht,
    create_truth_sample_ht,
)
from gnomad.variant_qc.pipeline import create_binned_ht, score_bin_agg
from gnomad.resources.grch38 import (
    na12878_giab,
    na12878_giab_hc_intervals,
    syndip,
    syndip_hc_intervals,
)

from joint_calling.utils import get_validation_callback
from joint_calling import utils
from joint_calling import _version

logger = logging.getLogger('random_forest')
logger.setLevel('INFO')


@click.command()
@click.version_option(_version.__version__)
@click.option(
    '--mt',
    'mt_path',
    callback=get_validation_callback(ext='mt'),
    help='Path to the matrix table',
)
@click.option(
    '--rf-annotations-ht',
    'rf_annotations_ht_path',
    callback=get_validation_callback(ext='ht'),
    help='RF annotations table (created by random_forest.py)',
)
@click.option(
    '--info-split-ht',
    'info_split_ht_path',
    required=True,
    callback=get_validation_callback(ext='ht', must_exist=True),
    help='path to info Table with split multiallelics '
    '(generated by generate_info_ht.py --out-split-info-ht)',
)
@click.option(
    '--fam-stats-ht',
    'fam_stats_ht_path',
    callback=get_validation_callback(ext='ht'),
    help='optional path to a Table with trio stats '
    '(generated by generate_variant_qc_annotations.py)',
)
@click.option(
    '--rf-result-ht',
    'rf_result_ht_path',
    callback=get_validation_callback(ext='ht'),
    help='RF result table (created by random_forest.py)',
)
@click.option(
    '--vqsr-filters-split-ht',
    'vqsr_filters_split_ht_path',
    callback=get_validation_callback(ext='ht'),
    help='Use VQSR training sites to train the RF (generated by load_vqsr.py)',
)
@click.option(
    '--out-bin-ht',
    'out_bin_ht_path',
    required=True,
    help='When set, creates file annotated with bin based on rank of VQSR/RF score.',
)
@click.option(
    '--out-aggregated-bin-ht',
    'out_aggregated_bin_ht_path',
    help='When set, creates a file with aggregate counts of variants based on bins.',
)
@click.option(
    '--run-sanity-checks',
    'run_sanity_checks',
    is_flag=True,
    help='When set, runs ranking sanity checks.',
)
@click.option(
    '--n-bins',
    'n_bins',
    help='Number of bins for the binned file (default: 100).',
    default=100,
    type=click.INT,
)
@click.option(
    '--n-partitions',
    'n_partitions',
    type=click.INT,
    help='Desired base number of partitions for output tables',
    default=5000,
)
@click.option(
    '--bucket',
    'work_bucket',
    required=True,
    help='path to write intermediate output and checkpoints. '
    'Can be a Google Storage URL (i.e. start with `gs://`).',
)
@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.',
)
def main(  # pylint: disable=too-many-arguments,too-many-locals
    mt_path: str,
    rf_annotations_ht_path: Optional[str],
    info_split_ht_path: str,
    fam_stats_ht_path: str,
    rf_result_ht_path: Optional[str],
    vqsr_filters_split_ht_path: Optional[str],
    out_bin_ht_path: str,
    out_aggregated_bin_ht_path: str,
    run_sanity_checks: bool,
    n_bins: int,
    n_partitions: int,
    work_bucket: str,
    local_tmp_dir: str,
    overwrite: bool,
):  # pylint: disable=missing-function-docstring
    local_tmp_dir = utils.init_hail('variant_qc_evaluate', local_tmp_dir)

    if overwrite or not utils.file_exists(out_bin_ht_path):
        scores_ht = create_bin_ht(
            rf_annotations_ht=hl.read_table(rf_annotations_ht_path),
            info_split_ht=hl.read_table(info_split_ht_path),
            n_bins=n_bins,
            rf_results_ht=hl.read_table(rf_result_ht_path)
            if rf_result_ht_path
            else None,
            vqsr_filters_split_ht=hl.read_table(vqsr_filters_split_ht_path)
            if vqsr_filters_split_ht_path
            else None,
        )
        scores_ht = scores_ht.checkpoint(out_bin_ht_path, overwrite=True)
    else:
        scores_ht = hl.read_table(out_bin_ht_path)

    if run_sanity_checks:
        logger.info('Running sanity checks...')
        ht = scores_ht
        logger.info(
            ht.aggregate(
                hl.struct(
                    was_biallelic=hl.agg.counter(~ht.was_split),
                    has_biallelic_rank=hl.agg.counter(hl.is_defined(ht.biallelic_bin)),
                    was_singleton=hl.agg.counter(ht.singleton),
                    has_singleton_rank=hl.agg.counter(hl.is_defined(ht.singleton_bin)),
                    was_biallelic_singleton=hl.agg.counter(
                        ht.singleton & ~ht.was_split
                    ),
                    has_biallelic_singleton_rank=hl.agg.counter(
                        hl.is_defined(ht.biallelic_singleton_bin)
                    ),
                )
            )
        )

    if out_aggregated_bin_ht_path:
        if overwrite or not utils.file_exists(out_aggregated_bin_ht_path):
            logger.warning('Use only workers, it typically crashes with preemptibles')
            agg_ht = create_aggregated_bin_ht(
                ht=scores_ht,
                trio_stats_ht=hl.read_table(fam_stats_ht_path)
                if fam_stats_ht_path
                else None,
                work_bucket=work_bucket,
            )
            agg_ht.write(out_aggregated_bin_ht_path, overwrite=True)

    mt = utils.get_mt(mt_path)

    truth_gvcfs = dict(
        syndip=dict(
            s='syndip',
            gvcf='gs://gnomad-public-requester-pays/resources/grch38/syndip/full.38.20180222.vcf.gz',
        ),
        NA12878=dict(
            s='NA12878',
            gvcf='gs://gnomad-public-requester-pays/resources/grch38/na12878/HG001_GRCh38_GIAB_highconf_CG-IllFB-IllGATKHC-Ion-10X-SOLID_CHROM1-X_v.3.3.2_highconf_PGandRTGphasetransfer.vcf.gz',
        ),
    )

    truth_dict = {
        truth_gvcfs['syndip']['s']: {
            's': truth_gvcfs['syndip']['s'],
            'truth_mt': syndip.mt(),
            'hc_intervals': syndip_hc_intervals.ht(),
            'mt': None,
            'ht': None,
        },
        truth_gvcfs['NA12878']['s']: {
            's': truth_gvcfs['NA12878']['s'],
            'truth_mt': na12878_giab.mt(),
            'hc_intervals': na12878_giab_hc_intervals.ht(),
            'mt': None,
            'ht': None,
        },
    }

    truth_snames = [sn for sn in truth_gvcfs if sn in mt.s.collect()]
    if truth_snames:
        truth_dict = {k: v for k, v in truth_dict.items() if k in truth_snames}
        _truth_concordance(
            mt,
            overwrite,
            work_bucket,
            n_partitions,
            scores_ht,
            info_split_ht_path,
            n_bins,
            truth_dict,
        )


def _truth_concordance(
    mt,
    overwrite,
    work_bucket,
    n_partitions,
    scores_ht,
    info_split_ht_path,
    n_bins,
    truth_dict,
):
    logger.info(f'Extracting truth samples from MT...')
    mt = mt.filter_cols(
        hl.literal([v['s'] for k, v in truth_dict.items()]).contains(mt.s)
    )
    mt = hl.experimental.sparse_split_multi(mt, filter_changed_loci=True)

    # Checkpoint to prevent needing to go through the large table a second time
    checkpoint_mt_path = join(work_bucket, 'tmp', 'genomes_split.mt')
    logger.info(f'Saving checkpoint to {checkpoint_mt_path}')
    mt = mt.checkpoint(checkpoint_mt_path, overwrite=overwrite)

    for truth_sample, truth_data in truth_dict.items():
        truth_samples_mt_path = join(work_bucket, 'truth_samples', f'{truth_sample}.mt')
        if not overwrite and utils.file_exists(truth_samples_mt_path):
            truth_data['mt'] = hl.read_matrix_table(truth_samples_mt_path)
        else:
            called_truth_mt = mt.filter_cols(mt.s == truth_dict[truth_sample]['s'])
            # Filter to variants in truth data
            called_truth_mt = called_truth_mt.filter_rows(
                hl.agg.any(called_truth_mt.GT.is_non_ref())
            )
            logger.info(
                f'Saving {truth_sample} called truth sample data to '
                f'{truth_samples_mt_path}'
            )
            called_truth_mt = called_truth_mt.naive_coalesce(n_partitions)
            called_truth_mt.write(truth_samples_mt_path, overwrite=True)
            truth_data['mt'] = called_truth_mt

    for truth_sample, truth_data in truth_dict.items():
        # Merging with truth data. Computes a table for each truth sample comparing
        # the truth sample in the callset vs the truth.
        truth_ht_path = join(work_bucket, 'truth_samples', f'{truth_sample}.ht')
        if not overwrite and utils.file_exists(truth_ht_path):
            truth_data['ht'] = hl.read_table(truth_ht_path)
        else:
            logger.info(
                f'Creating a merged table with callset truth sample and truth data '
                f'for {truth_sample}...'
            )

            # Load truth data
            mt = truth_data['mt']
            truth_hc_intervals = truth_data['hc_intervals']
            truth_mt = truth_data['truth_mt']
            truth_mt = truth_mt.key_cols_by(s=hl.str(truth_data['s']))

            # Remove low quality sites
            info_ht = hl.read_table(info_split_ht_path)
            mt = mt.filter_rows(~info_ht[mt.row_key].AS_lowqual)

            ht = create_truth_sample_ht(mt, truth_mt, truth_hc_intervals)
            ht.write(truth_ht_path, overwrite=True)
            truth_data['ht'] = ht

        # Bin truth sample concordance. Merges concordance results (callset vs.
        # truth) for a given truth sample with bins from specified model
        logger.info(f'Creating binned concordance table for {truth_sample}')
        info_ht = hl.read_table(info_split_ht_path)
        ht = truth_dict[truth_sample]['ht']
        ht = ht.filter(
            ~info_ht[ht.key].AS_lowqual
            & ~hl.is_defined(hl.read_table(utils.TEL_AND_CENT_HT_PATH)[ht.locus])
        )

        logger.info('Filtering out low confidence regions and segdups...')
        lcr = hl.read_table(utils.LCR_INTERVALS_HT_PATH)
        segdup = hl.read_table(utils.SEG_DUP_INTERVALS_HT_PATH)
        ht = ht.filter(hl.is_missing(lcr[ht.locus]))
        ht = ht.filter(hl.is_missing(segdup[ht.locus]))

        logger.info(
            'Loading HT containing RF or VQSR scores annotated with a bin based '
            'on the rank of score...'
        )
        metric_ht = scores_ht
        ht = ht.filter(hl.is_defined(metric_ht[ht.key]))

        ht = ht.annotate(score=metric_ht[ht.key].score)

        ht = compute_binned_truth_sample_concordance(ht, metric_ht, n_bins)
        binned_concordance_ht_path = join(
            work_bucket,
            'binned_concordance',
            f'{truth_sample}_binned_concordance.ht',
        )
        ht.write(binned_concordance_ht_path, overwrite=True)


def create_bin_ht(
    rf_annotations_ht: hl.Table,
    info_split_ht: hl.Table,
    n_bins: int,
    rf_results_ht: Optional[hl.Table] = None,
    vqsr_filters_split_ht: Optional[hl.Table] = None,
) -> hl.Table:
    """
    Creates a table with bin annotations added for a RF or VQSR run and writes it
    to its correct location in annotations.

    :param rf_annotations_ht: table generated by create_rf_annotations.py
    :param info_split_ht: table generated by generate_info_ht.py
    :param n_bins: Number of bins to bin the data into
    :param rf_results_ht: table generated by random_forest.py
    :param vqsr_filters_split_ht: table generated by load_vqsr.py
    :return: Table with bin annotations
    """
    logger.info(f'Annotating model HT with bins using {n_bins} bins')
    if vqsr_filters_split_ht:
        logger.info(f'Using a VQSR model')

        ht = vqsr_filters_split_ht
        ht = ht.annotate(**rf_annotations_ht[ht.key])
        ht = ht.annotate(
            info=info_split_ht[ht.key].info,
            score=ht.info.AS_VQSLOD,
            positive_train_site=ht.info.POSITIVE_TRAIN_SITE,
            negative_train_site=ht.info.NEGATIVE_TRAIN_SITE,
            AS_culprit=ht.info.AS_culprit,
        )

        # Remove all samples with an undefined ac_raw, because ac_raw was
        # calculated on the high quality samples only, and VQSR was run before
        # sample filtering
        ht = ht.filter(hl.is_defined(ht.ac_raw))

    else:
        logger.info(f'Using an RF model')
        ht = rf_results_ht
        ht = ht.annotate(
            info=info_split_ht[ht.key].info,
            positive_train_site=ht.tp,
            negative_train_site=ht.fp,
            score=ht.rf_probability['TP'],
        )

    ht = ht.filter(
        ~info_split_ht[ht.key].AS_lowqual
        & ~hl.is_defined(hl.read_table(utils.TEL_AND_CENT_HT_PATH)[ht.locus])
    )
    ht_non_lcr = filter_low_conf_regions(
        ht,
        filter_lcr=True,
        filter_decoy=False,  # Set if having decoy path
        filter_segdup=True,
    )
    ht = ht.annotate(non_lcr=hl.is_defined(ht_non_lcr[ht.key]))
    bin_ht = create_binned_ht(ht, n_bins, add_substrat={'non_lcr': ht.non_lcr})
    return bin_ht


def create_aggregated_bin_ht(
    ht: hl.Table,
    trio_stats_ht: Optional[hl.Table],
    work_bucket: str,
) -> hl.Table:
    """
    Aggregates variants into bins, grouped by `bin_id` (rank, bi-allelic, etc.),
    contig, and `snv`, `bi_allelic`, and `singleton` status, using previously
    annotated bin information.

    For each bin, aggregates statistics needed for evaluation plots.
    :param ht: table with score bins
    :param trio_stats_ht: optional, HT generated from a FAM file
    :param work_bucket: bucket to write temporary files to
    :return: Table of aggregate statistics by bin
    """

    # Count variants for ranking
    count_expr = {
        x: hl.agg.filter(
            hl.is_defined(ht[x]),
            hl.agg.counter(
                hl.cond(hl.is_snp(ht.alleles[0], ht.alleles[1]), 'snv', 'indel')
            ),
        )
        for x in ht.row
        if x.endswith('bin')
    }
    bin_variant_counts = ht.aggregate(hl.struct(**count_expr))
    logger.info(f'Found the following variant counts:\n {pformat(bin_variant_counts)}')
    ht = ht.annotate_globals(bin_variant_counts=bin_variant_counts)

    # Load ClinVar pathogenic data
    clinvar_pathogenic_ht = filter_to_clinvar_pathogenic(
        hl.read_table(utils.CLINVAR_HT_PATH)
    )
    ht = ht.annotate(clinvar_path=hl.is_defined(clinvar_pathogenic_ht[ht.key]))

    logger.info(f'Creating grouped bin table...')
    checkpoint_path = join(work_bucket, 'tmp', f'grouped_bin.ht')
    grouped_binned_ht = compute_grouped_binned_ht(ht, checkpoint_path=checkpoint_path)

    logger.info(f'Aggregating grouped bin table...')
    # Getting the source Table of the GroupedTable object
    # ("parent" here has nothing to do with the pedigree!)
    parent_ht = grouped_binned_ht._parent  # pylint: disable=protected-access

    agg_ht = grouped_binned_ht.aggregate(
        n_clinvar_path=hl.agg.count_where(parent_ht.clinvar_path),
        **score_bin_agg(grouped_binned_ht, fam_stats_ht=trio_stats_ht),
    )
    return agg_ht


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