#!/usr/bin/env python

"""
Combine a set of GVCFs into a MatrixTable
"""

import os
from typing import List, Optional
import logging
import shutil

import click
import pandas as pd

import hail as hl

from cpg_utils import to_path

from hail.experimental.vcf_combiner import vcf_combiner
from hail.experimental.vcf_combiner.vcf_combiner import CombinerConfig

from larcoh.utils import exists, check_duplicates


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


# The target number of rows per partition during each round of merging
TARGET_RECORDS = 25_000

# Default number of partitions in a matrix table generated by combiner
MAX_PARTITIONS = 2586


@click.command()
@click.option(
    '--sample-gvcf-tsv',
    'sample_gvcf_tsv_path',
    required=True,
    help='Sample data TSV path, to get samples names from the "s" column',
)
@click.option(
    '--out-mt',
    'out_mt_path',
    required=True,
    help='path to write the combined MatrixTable',
)
@click.option(
    '--branch-factor',
    'branch_factor',
    required=False,
    type=click.INT,
    default=CombinerConfig.default_branch_factor,
    help='Combiner branch factor (gvcfs will be split into ~branch_factor batches)',
)
@click.option(
    '--batch-size',
    'batch_size',
    required=False,
    type=click.INT,
    default=CombinerConfig.default_phase1_batch_size,
    help='Combiner batch size for phase1 (size of each batch of gvcfs to process '
         'in 1 job)',
)
@click.option(
    '--existing-mt',
    'existing_mt_path',
    help='optional path to an existing MatrixTable. '
    'If provided, will be read and used as a base to get extended with the '
    'samples in the input sample map. Can be read-only, as it will not '
    'be overwritten, instead the result will be written to the new location '
    'provided with --out-mt',
)
@click.option(
    '--tmp-bucket',
    'tmp_bucket',
    required=True,
    help='path to folder for intermediate output. '
    '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(
    '--n-partitions',
    'n_partitions',
    type=click.INT,
    help='Number of partitions for the output matrix table',
)
@click.option(
    '--is-exome',
    'is_exome',
    is_flag=True,
    default=False,
)
def main(
    sample_gvcf_tsv_path: str,
    out_mt_path: str,
    branch_factor: int,
    batch_size: int,
    existing_mt_path: Optional[str],
    tmp_bucket: str,
    local_tmp_dir: str,
    n_partitions: Optional[int],
    is_exome: bool,
):  # pylint: disable=missing-function-docstring
    if n_partitions:
        if n_partitions >= 2586:
            raise click.BadParameter(
                f'n_partitions must be less than {MAX_PARTITIONS} (combiner\'s default '
                f'number of partitions), got {n_partitions}'
            )

    hl.init(default_reference='GRCh38')
    logger.info(f'Combining GVCFs. Reading input paths from {sample_gvcf_tsv_path}')
    with to_path(sample_gvcf_tsv_path).open() as f:
        df = pd.read_table(f)
    sample_names = df.s
    gvcf_paths = df.gvcf
    check_duplicates(sample_names)
    check_duplicates(gvcf_paths)
    print(f'Combining {len(sample_names)} samples: {", ".join(sample_names)}')

    if n_partitions is not None or existing_mt_path is not None:
        # We are going to re-partition later, so writing the
        # matrix table into tmp:
        new_mt_path = os.path.join(tmp_bucket, 'new.mt')
    else:
        new_mt_path = out_mt_path

    if not exists(new_mt_path):
        combine_gvcfs(
            gvcf_paths=gvcf_paths,
            sample_names=sample_names,
            out_mt_path=new_mt_path,
            tmp_bucket=tmp_bucket,
            branch_factor=branch_factor,
            batch_size=batch_size,
            is_exome=is_exome,
        )
    new_mt = hl.read_matrix_table(new_mt_path)
    _log_mt_write(new_mt, new_mt_path)

    new_plus_existing_mt = None
    if existing_mt_path:
        logger.info(f'Combining with the existing matrix table {existing_mt_path}')
        new_plus_existing_mt = _combine_with_the_existing_mt(
            existing_mt_path=existing_mt_path,
            new_mt_path=new_mt_path,
        )
        
    mt = new_plus_existing_mt or new_mt
    if n_partitions is not None:
        mt = mt.repartition(n_partitions)

    # We need to write only if we repartitioned or combined with existing.
    # Otherwise, we are happy with the initial matrix table.
    if existing_mt_path or n_partitions is not None:
        mt.write(out_mt_path, overwrite=True)
        _log_mt_write(mt, out_mt_path)

    shutil.rmtree(local_tmp_dir)
    
    
def _log_mt_write(mt, path):
    logger.info(
        f'Written {mt.cols().count()} new samples to {path}, '
        f'n_partitions={mt.n_partitions()}'
    )


def _combine_with_the_existing_mt(
    existing_mt_path: str,
    # passing as a path because we are going
    # to re-read it with different intervals
    new_mt_path: str,
) -> hl.MatrixTable:
    # Making sure the tables are keyed by locus only, to make the combiner work.
    existing_mt = hl.read_matrix_table(existing_mt_path).key_rows_by('locus')
    intervals = vcf_combiner.calculate_new_intervals(
        existing_mt.rows(),
        n=TARGET_RECORDS,
        reference_genome='GRCh38',
    )
    new_mt = hl.read_matrix_table(new_mt_path, _intervals=intervals).key_rows_by(
        'locus'
    )
    logger.info(
        f'Combining {new_mt_path} ({new_mt.count_cols()} samples) '
        f'with an existing MatrixTable {existing_mt_path} '
        f'({existing_mt.count_cols()} samples), '
        f'split into {existing_mt.n_partitions()} partitions)'
    )
    out_mt = vcf_combiner.combine_gvcfs([existing_mt, new_mt])
    return out_mt


def combine_gvcfs(
    gvcf_paths: List[str],
    sample_names: List[str],
    out_mt_path: str,
    tmp_bucket: str,
    branch_factor: int,
    batch_size: int,
    is_exome: bool = False,
):
    """
    Combine a set of GVCFs in one go
    """
    hl.experimental.run_combiner(
        gvcf_paths,
        sample_names=sample_names,
        # header must be used with sample_names, otherwise sample_names will be ignored
        # first gvcf path is fine as a header because it will be read until 
        # the last line starting with # 
        header=gvcf_paths[0],
        out_file=out_mt_path,
        reference_genome='GRCh38',
        use_genome_default_intervals=not is_exome,
        use_exome_default_intervals=is_exome,
        tmp_path=tmp_bucket,
        overwrite=True,
        key_by_locus_and_alleles=True,
        branch_factor=branch_factor,
        batch_size=batch_size,
    )


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