from __future__ import absolute_import, division, print_function
import numpy as np
import six

import bcolz
from pybedtools import BedTool
from pybedtools import Interval
from pysam import FastaFile
import pyBigWig

from . import backend
from .util import nan_to_zero
from .util import one_hot_encode_sequence

NUM_SEQ_CHARS = 4


class BaseExtractor(object):
    dtype = np.float32

    def __init__(self, datafile, **kwargs):
        self._datafile = datafile

    def __call__(self, intervals, out=None, **kwargs):
        data = self._check_or_create_output_array(intervals, out)
        self._extract(intervals, data, **kwargs)
        return data

    def _check_or_create_output_array(self, intervals, out):
        width = intervals[0].stop - intervals[0].start
        output_shape = self._get_output_shape(len(intervals), width)

        if out is None:
            out = np.zeros(output_shape, dtype=self.dtype)
        else:
            if out.shape != output_shape:
                raise ValueError('out array has incorrect shape: {} '
                                 '(need {})'.format(out.shape, output_shape))
            if out.dtype != self.dtype:
                raise ValueError('out array has incorrect dtype: {} '
                                 '(need {})'.format(out.dtype, self.dtype))
        return out

    def _extract(self, intervals, out, **kwargs):
        'Subclassses should implement this and return the data'
        raise NotImplementedError

    @staticmethod
    def _get_output_shape(num_intervals, width):
        'Subclasses should implement this and return the shape of output'
        raise NotImplementedError


class ArrayExtractor(BaseExtractor):

    def __init__(self, datafile, in_memory=False, **kwargs):
        super(ArrayExtractor, self).__init__(datafile, **kwargs)
        self._data = backend.load_directory(datafile, in_memory=in_memory)
        self.multiprocessing_safe = in_memory

        arr = next(iter(self._data.values()))
        def _mm_extract(self, intervals, out, **kwargs):
            mm_data = self._data
            for index, interval in enumerate(intervals):
                out[index] = mm_data[interval.chrom][interval.start:interval.stop]

        # output shape method
        shape = arr.shape
        if len(shape) == 1:
            def _get_output_shape(num_intervals, width):
                return (num_intervals, width)
        elif len(shape) == 2:
            def _get_output_shape(num_intervals, width):
                return (num_intervals, width, shape[1])
        else:
            raise ValueError('Can only extract from 1D/2D arrays')

        self._mm_extract = _mm_extract.__get__(self)
        self._extract = self._mm_extract
        self._get_output_shape = staticmethod(_get_output_shape).__get__(self)


class FastaExtractor(BaseExtractor):

    def __init__(self, datafile, use_strand=False, **kwargs):
        """Fasta file extractor
        
        NOTE: The extractor is not thread-save.
        If you with to use it with multiprocessing,
        create a new extractor object in each process.
        
        Args:
          datafile (str): path to the bigwig file
          use_strand (bool): if True, the extracted sequence
            is reverse complemented in case interval.strand == "-"
        """
        super(FastaExtractor, self).__init__(datafile, **kwargs)
        self.use_strand = use_strand
        self.fasta = FastaFile(self._datafile)

    def _extract(self, intervals, out, **kwargs):    
        for index, interval in enumerate(intervals):
            seq = self.fasta.fetch(str(interval.chrom), interval.start,
                                       interval.stop)
            one_hot_encode_sequence(seq, out[index, :, :])

            # reverse-complement seq the negative strand
            if self.use_strand and interval.strand == "-":
                out[index, :, :] = out[index, ::-1, ::-1]

        return out

    @staticmethod
    def _get_output_shape(num_intervals, width):
        return (num_intervals, width, NUM_SEQ_CHARS)


class BigwigExtractor(BaseExtractor):

    def __init__(self, datafile, **kwargs):
        """Big-wig file extractor
        
        NOTE: The extractor is not thread-save.
        If you with to use it with multiprocessing,
        create a new extractor object in each process.
        
        Args:
          datafile: path to the bigwig file
        """
        super(BigwigExtractor, self).__init__(datafile, **kwargs)
        self._verbose = kwargs.get('verbose', False)
        self.bw = pyBigWig.open(datafile)
        
    def _extract(self, intervals, out, **kwargs):
        out[:] = self._bigwig_extractor(self.bw, intervals,
                                        **kwargs)

        return out

    @staticmethod
    def _get_output_shape(num_intervals, width):
        return (num_intervals, width)

    @staticmethod
    def _bigwig_extractor(bw, intervals, out=None, **kwargs):
        nan_as_zero = kwargs.get('nan_as_zero', True)
        if out is None:
            width = intervals[0].stop - intervals[0].start
            out = np.zeros((len(intervals), width), dtype=np.float32)
            
        for index, interval in enumerate(intervals):
            out[index] = bw.values(
                interval.chrom, interval.start, interval.stop)
            if nan_as_zero:
                nan_to_zero(out[index])
        return out
    
    def __del__(self):
        return self.close()
    
    def close(self):
        return self.bw.close()