from atlasvar.probes.models import Mutation
from atlasvar.probes import AlleleGenerator
from atlasvar.probes import make_variant_probe
from atlasvar.annotation.genes import GeneAminoAcidChangeToDNAVariants


import logging
logging.basicConfig()
logger = logging.getLogger(__name__)
flatten = lambda l: [item for sublist in l for item in sublist]


class CBGVariantSearch(object):

    def __init__(self, cbg, reference):
        self.cbg = cbg
        self.reference = reference
        self.al = AlleleGenerator(reference_filepath=self.reference,
                                  kmer=self.cbg.kmer_size)

    def search_for_alleles(self, ref_seqs, alt_seqs):
        results = {"ref": [], "alt": []}
        for ref in ref_seqs:
            logger.info(ref)
            res = self.cbg.search(ref, score=False)
            logger.info(str(res))
            results["ref"].extend(res.keys())
        for alt in alt_seqs:
            logger.info(alt)
            res = self.cbg.search(alt, score=False)
            logger.info(str(res))
            results["alt"].extend(res.keys())
        return results

    def make_variant_probe_set(self, var):
        return make_variant_probe(
            self.al, var, self.cbg.kmer_size, DB=None)

    def create_variant_probe_set(self, var_name):
        var = Mutation(var_name=var_name, reference=self.reference).variant
        variant_panel = self.make_variant_probe_set(var)
        return variant_panel

    def search_for_variant(self, ref_base, pos, alt_base="X", alphabet="DNA"):
        if not alphabet in ["DNA", "PROT"]:
            raise ValueError("alphabet must be either DNA or PROT")
        var_name = "".join([ref_base, str(pos), alt_base])
        variant_probe_set = self.create_variant_probe_set(var_name=var_name)
        return {var_name: self.genotype_alleles(variant_probe_set.refs, variant_probe_set.alts)}

    def genotype_alleles(self, refs, alts):
        ref_alt_samples = self.search_for_alleles(refs, alts)
        results = {}
        for sample_id in set(flatten(ref_alt_samples.values())):
            if sample_id in ref_alt_samples["ref"] and sample_id in ref_alt_samples["alt"]:
                results[sample_id] = {"genotype": "0/1"}
            elif sample_id in ref_alt_samples["ref"]:
                results[sample_id] = {"genotype": "0/0"}
            elif sample_id in ref_alt_samples["alt"]:
                results[sample_id] = {"genotype": "1/1"}
        return results


class CBGAminoAcidMutationSearch(CBGVariantSearch):

    def __init__(self, cbg, reference, genbank):
        super(CBGAminoAcidMutationSearch, self).__init__(cbg, reference)
        self.genbank = genbank
        self.aa2dna = GeneAminoAcidChangeToDNAVariants(
            self.reference,
            self.genbank)

    def search_for_amino_acid_variant(self, gene, ref, pos, alt):
        mut_name = "".join([ref, str(pos), alt])
        gene_mut_name = "_".join([gene, mut_name])
        results = {gene_mut_name: {}}
        _results = results[gene_mut_name]

        for var_name in self.aa2dna.get_variant_names(gene, mut_name, True):
            mut = Mutation(reference=self.reference,
                           var_name=var_name,
                           gene=gene,
                           mut=mut_name)
            variant_probe_set = self.create_variant_probe_set(var_name)
            variant_calls = self.genotype_alleles(
                variant_probe_set.refs, variant_probe_set.alts)
            for sample, genotype in variant_calls.items():
                _results[sample] = {"genotype": genotype[
                    "genotype"], "aa_mut": mut.mut, "variant": mut.variant.var_name, "gene": gene}
        print(results)
        return results
        #     if not mut.mut in mutations:
        #         mutations[mut.mut] = [mut]
        #     else:
        #         mutations[mut.mut].append(mut)

        # for amino_acid_mutation, mut_object_list in mutations.items():
        #     refs = []
        #     alts = []
        #     for mut_object in mut_object_list:
        #         variant_probe_set = self.make_variant_probe_set(
        #             mut_object.variant)
        #        	refs.extend(variant_probe_set.refs)
        #        	alts.extend(variant_probe_set.alts)
        #     self.genotype_alleles(refs,alts)

        # print(mut.mut)
        # print(var_name)
        # variant_probe_set = self.create_variant_probe_set(
        #     var_name=var_name)
        # print(self.genotype_alleles(
        #     variant_probe_set.refs, variant_probe_set.alts))

@@ -729,37 +729,8 @@
 qs:%0A
-            logger.info(ref)%0A
     
@@ -781,42 +781,8 @@
 se)%0A
-            logger.info(str(res))%0A
     
@@ -856,37 +856,8 @@
 qs:%0A
-            logger.info(alt)%0A
     
@@ -908,42 +908,8 @@
 se)%0A
-            logger.info(str(res))%0A
     

Remove debuggin logging