from importlib import resources
import logging
import gzip
import os.path as op
import os
import re

import numpy as np
import ctypes as C

from kineticsTools.sharedArray import SharedArray

# XXX can this vary?
LINUX_SO_FILE = "tree_predict.cpython-37m-x86_64-linux-gnu.so"

byte = np.dtype('byte')
uint8 = np.dtype('uint8')

# Map for ascii encoded bases to integers 0-3 -- will be used to define a 24-bit lookup code
# for fetching predicted IPDs from the kinetic LUT.

# We start everything at 0, so anything will map to 'A' unless it appears
# in this table
lutCodeMap = np.zeros(256, dtype=uint8)
maps = {'a': 0, 'A': 0, 'c': 1, 'C': 1, 'g': 2, 'G': 2, 't': 3, 'T': 3}
for k in maps:
    lutCodeMap[ord(k)] = maps[k]
lutReverseMap = {0: 'A', 1: 'C', 2: 'G', 3: 'T'}

seqCodeMap = np.ones(256, dtype=uint8) * 4
for k in maps:
    seqCodeMap[ord(k)] = maps[k]
seqMap = {0: 'A', 1: 'C', 2: 'G', 3: 'T', 4: 'N'}
seqMapNp = np.array(['A', 'C', 'G', 'T', 'N'])

seqMapComplement = {0: 'T', 1: 'G', 2: 'C', 3: 'A', 4: 'N'}
seqMapComplementNp = np.array(['T', 'G', 'C', 'A', 'N'])

# Base letters for modification calling
# 'H' : m6A, 'I' : m5C, 'J' : m4C, 'K' : m5C/TET
baseToCode = {'N': 0, 'A': 0, 'C': 1, 'G': 2,
              'T': 3, 'H': 4, 'I': 5, 'J': 6, 'K': 7}
baseToCanonicalCode = {'N': 0, 'A': 0, 'C': 1,
                       'G': 2, 'T': 3, 'H': 0, 'I': 1, 'J': 1, 'K': 1}

codeToBase = dict([(y, x) for (x, y) in baseToCode.items()])


def _getAbsPath(fname):
    try:
        with resources.as_file(resources.files('kineticsTools') / fname) as path:
            return str(path)
    except:
        return os.path.join(os.path.dirname(__file__), path)

class GbmContextModel(object):

    """
    Class for computing ipd predictions on contexts. Evaluate the GBM tree model for a list of contexts
    Contexts may contain arbitrary combinations of modified bases
    """

    def __init__(self, modelH5Group, modelIterations=-1):

        # This will hold the ctypes function pointer
        # It will be lazily initialized
        self.nativeInnerPredict = None
        self.nativeInnerPredictCtx = None

        def ds(name):
            return modelH5Group[name][:]

        self.varNames = [v.decode("ascii") for v in ds("VarNames")]
        self.modFeatureIdx = dict((int(self.varNames[x][1:]), x) for x in range(
            len(self.varNames)) if self.varNames[x][0] == 'M')
        self.canonicalFeatureIdx = dict((int(self.varNames[x][1:]), x) for x in range(
            len(self.varNames)) if self.varNames[x][0] == 'R')

        self.pre = 10
        self.post = 4
        self.ctxSize = self.pre + self.post + 1

        self.splitVar = ds("Variables")
        self.leftNodes = ds("LeftNodes")
        self.rightNodes = ds("RightNodes")
        self.missingNodes = ds("MissingNodes")

        self.splitVar16 = self.splitVar.astype(np.int16)

        self.splitCodes = ds("SplitCodes").astype(np.float32)

        self.cSplits = ds("CSplits")
        self.maxCSplits = self.cSplits.shape[1]

        self.initialValue = ds("InitialValue").astype(np.float32)[0]

        exp = 2 ** np.arange(self.cSplits.shape[1] - 1, -1, -1)
        self.bSplits = ((self.cSplits > 0) * exp).sum(1)

        # total number of trees in model
        self.nTrees = self.splitVar.shape[0]
        self.treeSize = self.splitVar.shape[1]

        offsets = np.floor(np.arange(0, self.leftNodes.size) /
                           self.treeSize) * self.treeSize
        offsets = offsets.astype(np.int32)

        self.leftNodesOffset = self.leftNodes.flatten().astype(np.int32) + offsets
        self.rightNodesOffset = self.rightNodes.flatten().astype(np.int32) + offsets
        self.missingNodesOffset = self.missingNodes.flatten().astype(np.int32) + offsets

        self.splitCodesCtx = self.splitCodes.copy().flatten()

        splitCodesCtxView = self.splitCodesCtx.view()
        splitCodesCtxView.dtype = np.uint32

        # Pack the cSplits as a bit array directly into the splitCode array
        # using an uin32 view of the splitCode array
        flatSplitVar = self.splitVar.flatten()

        powOfTwo = 2 ** np.arange(self.maxCSplits)

        for i in range(self.splitCodesCtx.shape[0]):

            if flatSplitVar[i] != -1:
                # This is a pointer to a cSplit row -- pack the csplit into a unit32, then overwirte
                # this slot of the ctxSplitCodes
                cs = self.cSplits[int(self.splitCodesCtx[i]), :]
                v = (powOfTwo * (cs > 0)).sum()

                splitCodesCtxView[i] = v

        # If the user has requested fewer iterations, update nTrees
        if modelIterations > 0:
            self.nTrees = modelIterations

    def _initNativeTreePredict(self):
        """
        Initialization routine the C tree-predict method
        Needs to be invoked lazily because the native function pointer cannot be pickled
        """

        dll_path = op.dirname(_getAbsPath(LINUX_SO_FILE))
        if os.path.exists(dll_path):
            self._lib = np.ctypeslib.load_library("tree_predict", dll_path)
        else:
            raise ImportError(
                "can't find tree_predict.so at '{}'".format(dll_path))

        lpb = self._lib

        lpb.init_native.argtypes = [C.c_int]

        fp = C.POINTER(C.c_float)
        fpp = C.POINTER(fp)
        ip = C.POINTER(C.c_int)
        sp = C.POINTER(C.c_int16)
        ui64p = C.POINTER(C.c_uint64)

        args = [fp, fpp, C.c_int, ip, ip, ip, fp, ip,
                ip, ip, C.c_float, C.c_int, C.c_int, C.c_int]
        lpb.innerPredict.argtypes = args
        self.nativeInnerPredict = lpb.innerPredict

        # Fast version

        # void innerPredictCtx(
        #    int ctxSize, float radPredF[], uint64_t contextPack[], int cRows,
        #    int16 left[], int16 right[], int16 missing[], float splitCode[], int16 splitVar[],
        #    int varTypes[], float initialValue, int treeSize, int numTrees, int maxCSplitSize)

        args = [C.c_int, fp, ui64p, C.c_int, ip, ip, ip, fp,
                sp, ip, C.c_float, C.c_int, C.c_int, C.c_int]
        lpb.innerPredictCtx.argtypes = args
        self.nativeInnerPredictCtx = lpb.innerPredictCtx

    def getPredictionsSlow(self, ctxStrings, nTrees=None):
        """Compute IPD predictions for arbitrary methylation-containing contexts."""
        # C prototype that we call:
        # void innerPredict(
        #   float[] radPredF,
        #   IntPtr[] dataMatrix,
        #   int cRows, int[] left, int[] right, int[] missing,
        #   float[] splitCode, int[] splitVar, int[] cSplits,
        #   int[] varTypes, float initialValue,
        #   int treeSize, int numTrees, int maxCSplitSize);

        # Make sure native library is initialized
        if self.nativeInnerPredict is None:
            self._initNativeTreePredict()

        def fp(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_float))

        def ip(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_int))

        if nTrees is None:
            nTrees = self.nTrees

        n = len(ctxStrings)

        mCols = [np.zeros(n, dtype=np.float32) for x in range(self.ctxSize)]
        rCols = [np.zeros(n, dtype=np.float32) for x in range(self.ctxSize)]

        for stringIdx in range(len(ctxStrings)):
            s = ctxStrings[stringIdx]

            for i in range(len(s)):
                mCols[i][stringIdx] = baseToCode[s[i]]
                rCols[i][stringIdx] = baseToCanonicalCode[s[i]]

        dataPtrs = (C.POINTER(C.c_float) * (2 * self.ctxSize))()

        varTypes = np.zeros(2 * self.ctxSize, dtype=np.int32)

        for i in range(self.ctxSize):
            dataPtrs[self.modFeatureIdx[i]] = mCols[i].ctypes.data_as(
                C.POINTER(C.c_float))
            dataPtrs[self.canonicalFeatureIdx[i]
                     ] = rCols[i].ctypes.data_as(C.POINTER(C.c_float))

            varTypes[self.modFeatureIdx[i]] = 8
            varTypes[self.canonicalFeatureIdx[i]] = 4

        self.predictions = np.zeros(len(ctxStrings), dtype=np.float32)

        self.nativeInnerPredict(
            fp(self.predictions), dataPtrs,
            n, ip(self.leftNodes), ip(self.rightNodes), ip(self.missingNodes),
            fp(self.splitCodes), ip(self.splitVar), ip(self.cSplits),
            ip(varTypes), self.initialValue, self.treeSize, nTrees, self.maxCSplits)

        return np.exp(self.predictions)

    def getPredictions(self, ctxStrings, nTrees=None):
        """Compute IPD predictions for arbitrary methylation-containing contexts."""
        # C prototype that we call:
        # void innerPredictCtx(
        #   int ctxSize, float[] radPredF,
        #   int[] contextPack,
        #   int cRows, int[] left, int[] right, int[] missing,
        #   float[] splitCode, int[] splitVar,
        #   int[] varTypes, float initialValue,
        #   int treeSize, int numTrees, int maxCSplitSize);

        # Make sure native library is initialized
        if self.nativeInnerPredictCtx is None:
            self._initNativeTreePredict()

        def fp(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_float))

        def ip(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_int))

        def ulp(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_uint64))

        def sp(arr):
            return arr.ctypes.data_as(C.POINTER(C.c_int16))

        n = len(ctxStrings)

        if nTrees is None:
            nTrees = self.nTrees

        packCol = np.zeros(n, dtype=np.uint64)

        for stringIdx in range(len(ctxStrings)):
            s = ctxStrings[stringIdx]  # .strip().strip('\x00')
            code = 0

            for i in range(len(s)):
                modBits = baseToCode[s[i]]

                slotForPosition = self.modFeatureIdx[i]

                code = code | (modBits << (4 * slotForPosition))

            packCol[stringIdx] = code

        # print packed base codes
        # for v in packCol.flatten():
        #    print v
        #    for i in np.arange(12):
        #        print "%d: %o" % (i,  (v.item() >> (5*i)) & 0x1f)

        varTypes = np.zeros(2 * self.ctxSize, dtype=np.int32)

        for i in range(self.ctxSize):
            varTypes[self.modFeatureIdx[i]] = 8
            varTypes[self.canonicalFeatureIdx[i]] = 4

        self.predictions = np.zeros(len(ctxStrings), dtype=np.float32)

        self.nativeInnerPredictCtx(
            self.ctxSize, fp(self.predictions), ulp(packCol),
            n, ip(self.leftNodesOffset), ip(
                self.rightNodesOffset), ip(self.missingNodesOffset),
            fp(self.splitCodesCtx), sp(self.splitVar16),
            ip(varTypes), self.initialValue, self.treeSize, nTrees, self.maxCSplits)

        return np.exp(self.predictions)


class IpdModel:

    """
    Predicts the IPD of an any context, possibly containing multiple modifications.
    We use a 4^12 entry LUT to get the predictions for contexts without modifications,
    then we use the GbmModel to get predictions in the presence of arbitrary mods.
    Note on the coding scheme.  For each contig we store a byte-array that has size = contig.length + 2*self.pad
    The upper 4 bits contain a lookup into seqReverseMap, which can contains N's. This is used for giving
    template snippets that may contains N's if the reference sequence does, or if the snippet
    The lowe 4 bits contain a lookup into lutReverseMap, which
    """

    def __init__(self, fastaRecords, modelFile, modelIterations=-1):
        """
        Load the reference sequences and the ipd lut into shared arrays that can be
        used as numpy arrays in worker processes.
        fastaRecords is a list of FastaRecords, in the alignments file order
        """

        self.pre = 10
        self.post = 4

        self.pad = 30
        self.base4 = 4 ** np.array(range(self.pre + self.post + 1))

        self.refDict = {}
        self.refLengthDict = {}

        for contig in fastaRecords:
            if contig.alignmentID is None:
                # This contig has no mapped reads -- skip it
                continue

            rawSeq = contig.sequence[:]
            refSeq = np.frombuffer(rawSeq.encode("utf-8"), dtype=byte)

            # Store the reference length
            self.refLengthDict[contig.alignmentID] = len(rawSeq)

            # Make a shared array
            sa = SharedArray(dtype='B', shape=len(rawSeq) + self.pad * 2)
            saWrap = sa.getNumpyWrapper()

            # Lut Codes convert Ns to As so that we don't put Ns into the Gbm Model
            # Seq Codes leaves Ns as Ns for getting reference snippets out
            innerLutCodes = lutCodeMap[refSeq]
            innerSeqCodes = seqCodeMap[refSeq]
            innerCodes = np.bitwise_or(
                innerLutCodes, np.left_shift(innerSeqCodes, 4))

            saWrap[self.pad:(len(rawSeq) + self.pad)] = innerCodes

            # Padding codes -- the lut array is padded with 0s the sequence
            # array is padded with N's (4)
            outerCodes = np.left_shift(np.ones(self.pad, dtype=uint8) * 4, 4)
            saWrap[0:self.pad] = outerCodes
            saWrap[(len(rawSeq) + self.pad):(len(rawSeq) + 2 * self.pad)] = outerCodes

            self.refDict[contig.alignmentID] = sa

        # No correction factor for IPDs everything is normalized to 1
        self.meanIpd = 1

        # Find and open the ipd model file
        self.lutPath = modelFile
        if os.path.exists(self.lutPath):
            with gzip.open(self.lutPath, "rb") as npz_in:
                gbmModelData = np.load(npz_in, allow_pickle=True)
                self.gbmModel = GbmContextModel(gbmModelData, modelIterations)

            # We always use the model -- no more LUTS
            self.predictIpdFunc = self.predictIpdFuncModel
            self.predictManyIpdFunc = self.predictManyIpdFuncModel
        else:
            logging.info("Couldn't find model file: %s" % self.lutPath)

    def _loadIpdTable(self, nullModelGroup):
        """
        Read the null kinetic model into a shared numpy array dataset
        """
        nullModelDataset = nullModelGroup["KineticValues"]

        # assert that the dataset is a uint8
        assert(nullModelDataset.dtype == uint8)

        # Construct a 'shared array' (a numpy wrapper around some shared memory
        # Read the LUT into this table
        self.sharedArray = SharedArray('B', nullModelDataset.shape[0])
        lutArray = self.sharedArray.getNumpyWrapper()
        nullModelDataset.read_direct(lutArray)

        # Load the second-level LUT
        self.floatLut = nullModelGroup["Lut"][:]

    def refLength(self, refId):
        return self.refLengthDict[refId]

    def cognateBaseFunc(self, refId):
        """
        Return a function that returns a snippet of the reference sequence around a given position
        """

        # FIXME -- what is the correct strand to return?!
        # FIXME -- what to do about padding when the snippet runs off the end of the reference
        # how do we account for / indicate what is happening
        refArray = self.refDict[refId].getNumpyWrapper()

        def f(tplPos, tplStrand):

            # skip over the padding
            tplPos += self.pad

            # Forward strand
            if tplStrand == 0:
                slc = refArray[tplPos]
                slc = np.right_shift(slc, 4)
                return seqMap[slc]

            # Reverse strand
            else:
                slc = refArray[tplPos]
                slc = np.right_shift(slc, 4)
                return seqMapComplement[slc]

        return f

    def snippetFunc(self, refId, pre, post):
        """
        Return a function that returns a snippet of the reference sequence around a given position
        """

        refArray = self.refDict[refId].getNumpyWrapper()

        def f(tplPos, tplStrand):
            """Closure for returning a reference snippet. The reference is padded with N's for bases falling outside the extents of the reference"""
            # skip over the padding
            tplPos += self.pad

            # Forward strand
            if tplStrand == 0:
                slc = refArray[(tplPos - pre):(tplPos + 1 + post)]
                slc = np.right_shift(slc, 4)
                return "".join(c for c in seqMapNp[slc])

            # Reverse strand
            else:
                slc = refArray[(tplPos + pre):(tplPos - post - 1):-1]
                slc = np.right_shift(slc, 4)
                return "".join(c for c in seqMapComplementNp[slc])

        return f

    def getReferenceWindow(self, refId, tplStrand, start, end):
        """
        Return  a snippet of the reference sequence
        """

        refArray = self.refDict[refId].getNumpyWrapper()

        # adjust position for reference padding
        start += self.pad
        end += self.pad

        # Forward strand
        if tplStrand == 0:
            slc = refArray[start:end]
            slc = np.right_shift(slc, 4)
            return "".join(seqMap[x] for x in slc)

        # Reverse strand
        else:
            slc = refArray[end:start:-1]
            slc = np.right_shift(slc, 4)
            return "".join(seqMapComplement[x] for x in slc)

    def predictIpdFuncLut(self, refId):
        """
        Each (pre+post+1) base context gets mapped to an integer
        by converting each nucleotide to a base-4 number A=0, C=1, etc,
        and treating the 'pre' end of the context of the least significant
        digit.  This code is used to lookup the expected IPD in a
        pre-computed table.  Contexts near the ends of the reference
        are coded by padding the context with 0
        """

        # Materialized the numpy wrapper around the shared data
        refArray = self.refDict[refId].getNumpyWrapper()
        lutArray = self.sharedArray.getNumpyWrapper()
        floatLut = self.floatLut

        def f(tplPos, tplStrand):

            # skip over the padding
            tplPos += self.pad

            # Forward strand
            if tplStrand == 0:
                slc = np.bitwise_and(
                    refArray[(tplPos + self.pre):(tplPos - self.post - 1):-1], 0xf)

            # Reverse strand
            else:
                slc = 3 - \
                    np.bitwise_and(
                        refArray[(tplPos - self.pre):(tplPos + 1 + self.post)], 0xf)

            code = (self.base4 * slc).sum()
            return floatLut[max(1, lutArray[code])]

        return f

    def predictIpdFuncModel(self, refId):
        """
        Each (pre+post+1) base context gets mapped to an integer
        by converting each nucleotide to a base-4 number A=0, C=1, etc,
        and treating the 'pre' end of the context of the least significant
        digit.  This code is used to lookup the expected IPD in a
        pre-computed table.  Contexts near the ends of the reference
        are coded by padding the context with 0
        """

        # Materialized the numpy wrapper around the shared data
        snipFunction = self.snippetFunc(refId, self.post, self.pre)

        def f(tplPos, tplStrand):
            # Get context string
            context = snipFunction(tplPos, tplStrand)

            # Get prediction
            return self.gbmModel.getPredictions([context])[0]

        return f

    def predictManyIpdFuncModel(self, refId):
        """
        Each (pre+post+1) base context gets mapped to an integer
        by converting each nucleotide to a base-4 number A=0, C=1, etc,
        and treating the 'pre' end of the context of the least significant
        digit.  This code is used to lookup the expected IPD in a
        pre-computed table.  Contexts near the ends of the reference
        are coded by padding the context with 0
        """

        # Materialized the numpy wrapper around the shared data
        snipFunction = self.snippetFunc(refId, self.post, self.pre)

        def fMany(sites):
            contexts = [snipFunction(x[0], x[1]) for x in sites]
            return self.gbmModel.getPredictions(contexts)

        return fMany