"""Calculate heterogeneity and sub-clonal populations for complex input samples.

Use allele frequencies, copy number calls and structural variants to infer
sub-clonal populations within a potentially mixed population. This attempts
to infer these sub-clones to help improve variant calls and interpretation
especially in complex cancer samples.
"""
from __future__ import print_function
import collections
import os

from bcbio import utils
from bcbio.heterogeneity import bubbletree, phylowgs, theta
from bcbio.pipeline import datadict as dd
from bcbio.variation import vcfutils

def _get_calls(data, cnv_only=False):
    """Retrieve calls, organized by name, to use for heterogeneity analysis.
    """
    cnvs_supported = set(["cnvkit", "battenberg"])
    out = {}
    for sv in data.get("sv", []):
        if not cnv_only or sv["variantcaller"] in cnvs_supported:
            out[sv["variantcaller"]] = sv
    return out

def get_variants(data, include_germline=False):
    """Retrieve set of variant calls to use for heterogeneity analysis.
    """
    data = utils.deepish_copy(data)
    supported = ["precalled", "vardict", "vardict-java", "vardict-perl",
                 "freebayes", "octopus", "strelka2"]
    # Right now mutect2 and mutect do not provide heterozygous germline calls
    # to be useful https://github.com/bcbio/bcbio-nextgen/issues/2464
    # supported += ["mutect2", "mutect"]
    if include_germline:
        supported.insert(1, "gatk-haplotype")
    out = []
    # CWL based input
    if isinstance(data.get("variants"), dict) and "samples" in data["variants"]:
        cur_vs = []
        # Unpack single sample list of files
        if (isinstance(data["variants"]["samples"], (list, tuple)) and
              len(data["variants"]["samples"]) == 1 and isinstance(data["variants"]["samples"][0], (list, tuple))):
            data["variants"]["samples"] = data["variants"]["samples"][0]
        for fname in data["variants"]["samples"]:
            variantcaller = utils.splitext_plus(os.path.basename(fname))[0]
            variantcaller = variantcaller.replace(dd.get_sample_name(data) + "-", "")
            for batch in dd.get_batches(data):
                variantcaller = variantcaller.replace(batch + "-", "")
            cur_vs.append({"vrn_file": fname, "variantcaller": variantcaller})
        data["variants"] = cur_vs
    for v in data.get("variants", []):
        if v["variantcaller"] in supported and v.get("vrn_file"):
            out.append((supported.index(v["variantcaller"]), v))
    out.sort()
    return [xs[1] for xs in out]

def _ready_for_het_analysis(items):
    """Check if a sample has input information for heterogeneity analysis.

    We currently require a tumor/normal sample containing both CNV and variant calls.
    """
    paired = vcfutils.get_paired_bams([dd.get_align_bam(d) for d in items], items)
    has_het = any(dd.get_hetcaller(d) for d in items)
    if has_het and paired:
        return get_variants(paired.tumor_data) and _get_calls(paired.tumor_data, cnv_only=True)

def _get_batches(data):
    batches = dd.get_batch(data) or dd.get_sample_name(data)
    if not isinstance(batches, (list, tuple)):
        batches = [batches]
    return batches

def _group_by_batches(items):
    out = collections.OrderedDict()
    for data in (utils.to_single_data(xs) for xs in items):
        for b in _get_batches(data):
            try:
                out[b].append(data)
            except KeyError:
                out[b] = [data]
    return out

def _get_hetcallers(items):
    out = set([])
    for d in items:
        hetcaller = dd.get_hetcaller(d)
        if hetcaller:
            out = out.union(set(hetcaller))
    return sorted(list(out))

def estimate(items, batch, config):
    """Estimate heterogeneity for a pair of tumor/normal samples. Run in parallel.
    """
    hetcallers = {"theta": theta.run,
                  "phylowgs": phylowgs.run,
                  "bubbletree": bubbletree.run}
    paired = vcfutils.get_paired_bams([dd.get_align_bam(d) for d in items], items)
    calls = _get_calls(paired.tumor_data)
    variants = get_variants(paired.tumor_data)
    het_info = []
    for hetcaller in _get_hetcallers(items):
        try:
            hetfn = hetcallers[hetcaller]
        except KeyError:
            hetfn = None
            print("%s not yet implemented" % hetcaller)
        if hetfn:
            hetout = hetfn(variants[0], calls, paired)
            if hetout:
                het_info.append(hetout)
    out = []
    for data in items:
        if batch == _get_batches(data)[0]:
            if dd.get_sample_name(data) == paired.tumor_name:
                if het_info:
                    data["heterogeneity"] = het_info
            out.append([data])
    return out

def run(items, run_parallel):
    """Top level entry point for calculating heterogeneity, handles organization and job distribution.
    """
    to_process = []
    extras = []
    for batch, cur_items in _group_by_batches(items).items():
        if _ready_for_het_analysis(cur_items):
            to_process.append((batch, cur_items))
        else:
            for data in cur_items:
                extras.append([data])
    processed = run_parallel("heterogeneity_estimate", ([xs, b, xs[0]["config"]] for b, xs in to_process))
    return _group_by_sample_and_batch(extras + processed)

def _group_by_sample_and_batch(samples):
    """Group samples split by heterogeneity method back one per sample-batch.

    Groups potentially multiple shared samples (multi batch normals) into
    single items per group.
    """
    out = collections.defaultdict(list)
    for data in [utils.to_single_data(x) for x in samples]:
        out[(dd.get_sample_name(data), dd.get_align_bam(data), tuple(_get_batches(data)))].append(data)
    return [[xs[0]] for xs in out.values()]

