
from collections import namedtuple, defaultdict
import logging

import numpy as np

from pbreports.statistics import mean_int

log = logging.getLogger(__name__)


class Constants(object):
    LABEL_NONE = "Not Barcoded"
    BIOSAMPLE_NONE = "No Name"


class BarcodeGroup(object):
    """
    Utility class for storing per-barcode metrics from multiple reads
    """
    __slots__ = ["label", "n_bases", "readlengths", "bqs", "srl_max", "idx"]

    def __init__(self,
                 label,
                 n_bases=0,
                 readlengths=(),
                 bqs=(),
                 srl_max=(),
                 idx=None):
        self.label = label
        self.n_bases = n_bases
        self.readlengths = list(readlengths)
        self.bqs = list(bqs)
        self.srl_max = list(srl_max)
        self.idx = idx

    def add_read(self, read_info):
        assert read_info.label == self.label
        self.n_bases += read_info.nbases
        self.readlengths.append(read_info.readlength)
        self.bqs.extend(list(read_info.bq))
        self.srl_max.append(read_info.srl_max)

    @property
    def n_subreads(self):
        return len(self.bqs)

    @property
    def n_reads(self):
        return len(self.readlengths)

    def mean_read_length(self):
        if self.n_reads == 0:
            return 0
        return int(sum(self.readlengths) / self.n_reads)

    def mean_bcqual(self):
        if self.n_subreads == 0:
            return 0
        return int(sum(self.bqs) / self.n_subreads)

    def mean_longest_subread_length(self):
        if self.n_reads == 0:
            return 0
        return mean_int(sum(self.srl_max), self.n_reads)


class ReadInfo(namedtuple("ReadInfo", ["label", "nbases", "readlength", "srl_max", "bq", "bc_idx", "qId", "holeNumber"])):
    """
    Container class for storing information about a single barcoded ZMW read
    and associated subreads.
    """

    @property
    def idx(self):
        if self.is_barcoded():
            return "{f}--{r}".format(f=self.bc_idx[0], r=self.bc_idx[1])
        else:
            return "None"

    def is_barcoded(self):
        return self.bc_idx != (-1, -1)

    @property
    def n_subreads(self):
        return len(self.bq)


def get_biosample_dict(barcoded_subreads):
    biosamples = {}
    for bio_sample in barcoded_subreads.metadata.bioSamples:
        for barcode in bio_sample.DNABarcodes:
            if biosamples.get(barcode.name, bio_sample.name) != bio_sample.name:
                # FIXME this is probably user error, but not guaranteed
                log.warn("DNABarcode %s refers to multiple BioSamples: %s, %s",
                         barcode.name, biosamples[barcode.name], bio_sample.name)
            biosamples[barcode.name] = bio_sample.name
    return biosamples


def _to_zmw_read_info(label, bc_id, qStart, qEnd, bcQual, qId, holeNumber):
    subread_lengths = qEnd - qStart
    readlength = qEnd.max() - qStart.min()
    nbases = np.sum(subread_lengths)
    srl_max = subread_lengths.max()
    return ReadInfo(label, nbases, readlength, srl_max, bcQual, bc_id, qId, holeNumber)


def _iter_zmws_by_barcode(pb_index, bc_key, bc_label, isoseq_mode=False):
    (bc_fw, bc_rev) = bc_key
    bc_sel = (pb_index.bcForward == bc_fw) & (pb_index.bcReverse == bc_rev)
    if isoseq_mode:
        bc_sel |= (pb_index.bcForward == bc_rev) & (
            pb_index.bcReverse == bc_fw)
    qIds = list(set(pb_index.qId[bc_sel]))
    for movie_id in qIds:
        movie_sel = (bc_sel) & (pb_index.qId == movie_id)
        holeNumber = pb_index.holeNumber[movie_sel]
        qStart = pb_index.qStart[movie_sel]
        qEnd = pb_index.qEnd[movie_sel]
        bcQual = pb_index.bcQual[movie_sel]
        zmw_starts = np.where(np.concatenate(
            [[True], holeNumber[1:] != holeNumber[:-1]]))[0]
        zmw_ends = np.concatenate([zmw_starts[1:], [len(qEnd)]])
        for zmw_start, zmw_end in zip(zmw_starts, zmw_ends):
            yield _to_zmw_read_info(
                label=bc_label,
                bc_id=bc_key,
                qStart=qStart[zmw_start:zmw_end],
                qEnd=qEnd[zmw_start:zmw_end],
                bcQual=bcQual[zmw_start:zmw_end],
                qId=movie_id,
                holeNumber=holeNumber[zmw_start])


def _get_unique_barcode_ids(pb_index, isoseq_mode=False):
    """
    Get a list of sorted, unique fw/rev barcode indices from an index object.
    """
    bc_sel = (pb_index.bcForward != -1) & (pb_index.bcReverse != -1)
    bcFw = pb_index.bcForward[bc_sel]
    bcRev = pb_index.bcReverse[bc_sel]
    bc_ids = sorted(list(set(zip(bcFw, bcRev))))
    if isoseq_mode:
        bc_ids = sorted(list(set([tuple(sorted(bc)) for bc in bc_ids])))
    return bc_ids


def get_reads_by_barcode(pb_index, barcodes, isoseq_mode=False,
                         label_none=Constants.LABEL_NONE):
    """
    Given a SubreadSet or ConsensusReadSet and BarcodeSet as input, return an
    iterable of ReadInfo objects
    """
    barcoded_reads = []
    bcs = [bc for bc in barcodes]
    for bc_id in _get_unique_barcode_ids(pb_index, isoseq_mode):
        barcode_fw, barcode_rev = bc_id
        if barcode_fw == -1:
            barcode_label = label_none
        else:
            barcode_label = "{f}--{r}".format(f=bcs[barcode_fw].id,
                                              r=bcs[barcode_rev].id)
        for read_info in _iter_zmws_by_barcode(pb_index,
                                               bc_id,
                                               barcode_label,
                                               isoseq_mode):
            barcoded_reads.append(read_info)
    return barcoded_reads


def _get_movie_names(ds):
    movies = {}
    for bam in ds.resourceReaders():
        movies.update({rg.ID: rg.MovieName for rg in bam.readGroupTable})
    return movies


def _iter_unbarcoded_reads_info(pb_index, movie_names, barcoded_zmws,
                                label_none=Constants.LABEL_NONE):
    movie2qid = {movie_names[qId]: qId for qId in set(pb_index.qId)}
    for movie_name in sorted(movie2qid.keys()):
        qId = movie2qid[movie_name]
        movie_sel = pb_index.qId == qId
        movie_zmws = pb_index.holeNumber[movie_sel]
        qStart = pb_index.qStart[movie_sel]
        qEnd = pb_index.qEnd[movie_sel]
        zmw_starts = np.where(np.concatenate(
            [[True], movie_zmws[1:] != movie_zmws[:-1]]))[0]
        zmw_ends = np.concatenate([zmw_starts[1:], [len(qEnd)]])
        for i_zmw, i_end in zip(zmw_starts, zmw_ends):
            zmw = movie_zmws[i_zmw]
            if not (qId, zmw) in barcoded_zmws:
                yield _to_zmw_read_info(
                    label=label_none,
                    bc_id=(-1, -1),
                    qStart=qStart[i_zmw:i_end],
                    qEnd=qEnd[i_zmw:i_end],
                    bcQual=np.full(i_end - i_zmw, 0.0),
                    qId=qId,
                    holeNumber=zmw)


def get_unbarcoded_reads_info(dataset_in, barcoded_zmws,
                              label_none=Constants.LABEL_NONE):
    # not necessarily subreads
    nonbarcoded_records_by_zmw = defaultdict(list)
    movies = _get_movie_names(dataset_in)
    return _iter_unbarcoded_reads_info(dataset_in.index,
                                       movies,
                                       barcoded_zmws,
                                       label_none=label_none)
