import collections


comp = str.maketrans('ATCGNatcgn','TAGCNtagcn')
    
def reverse_comp(st):
    """ Returns the reverse complement of a DNA sequence; non ACGT bases will be ignored. """
    return str(st)[::-1].translate(comp)



class GenomeSource(object):
    def __init__(self, names_to_contigs, aligner_type="bwa"):
        self.names_to_contigs = collections.OrderedDict(names_to_contigs)
        self._bwa = None
        self._ssw = None
        self._blacklist = None

        self.aligner_type = aligner_type

    def get_seq(self, chrom, start, end, strand):
        seq = self.names_to_contigs[chrom][start:end+1]
        if strand == "-":
            seq = reverse_comp(seq)
        return seq

    def keys(self):
        return list(self.names_to_contigs.keys())



class FastaGenomeSource(GenomeSource):
    """ pickle-able wrapper for pyfaidx.Fasta """
    def __init__(self, path):
        self.path = path
        self._fasta = None
        
    def get_seq(self, chrom, start, end, strand):
        seq = self.fasta[chrom][start:end+1]
        if strand == "-":
            seq = reverse_comp(seq)
        return seq

    def keys(self):
        return list(self.fasta.keys())

    @property
    def fasta(self):
        if self._fasta is None:
            import pyfaidx
            self._fasta = pyfaidx.Fasta(self.path, as_raw=True)
        return self._fasta

    def __getstate__(self):
        state = self.__dict__.copy()
        state["_fasta"] = None
        return state
