#!/usr/bin/python

from __future__ import print_function

import itertools

import numpy as np

from . import parsers
from . import convolution_dp
from . import hyperparameters


def getPosteriors(quantRowsOrig, params):
    quantRows, quantMatrix = parsers.getQuantMatrix(quantRowsOrig)

    pProteinQuantsList, bayesQuantRow = getPosteriorProteinRatios(
        quantMatrix, quantRows, params
    )
    pProteinGroupQuants = getPosteriorProteinGroupRatios(
        pProteinQuantsList, bayesQuantRow, params
    )
    pProteinGroupDiffs, muGroupDiffs = getProteinGroupsDiffPosteriors(
        pProteinGroupQuants, params
    )

    probsBelowFoldChange = getProbBelowFoldChangeDict(pProteinGroupDiffs, params)
    if params["returnPosteriors"]:
        posteriorDists = (pProteinQuantsList, pProteinGroupQuants, pProteinGroupDiffs)
    else:
        posteriorDists = None

    return bayesQuantRow, muGroupDiffs, probsBelowFoldChange, posteriorDists


def getDummyPosteriors(params):
    bayesQuantRow = [1.0 for g in params["groups"] for x in g]
    numGroups = len(params["groups"])
    probsBelowFoldChange, muGroupDiffs = dict(), dict()
    for groupId1, groupId2 in itertools.combinations(range(numGroups), 2):
        (
            probsBelowFoldChange[(groupId1, groupId2)],
            muGroupDiffs[(groupId1, groupId2)],
        ) = (1.0, 0.0)
    posteriorDists = None
    return bayesQuantRow, muGroupDiffs, probsBelowFoldChange, posteriorDists


def getPosteriorProteinRatios(
    quantMatrix, quantRows, params, maxIterations=50, bayesQuantRow=None
):
    numSamples = len(quantMatrix[0])
    bayesQuantRow = np.array([1.0] * numSamples)
    converged = False
    for iteration in range(maxIterations):
        prevBayesQuantRow = np.copy(bayesQuantRow)
        pProteinQuantsList, bayesQuantRow = getPosteriorProteinRatio(
            quantMatrix, quantRows, bayesQuantRow, params
        )

        bayesQuantRow = parsers.geoNormalize(bayesQuantRow)

        diffInIteration = np.log10(prevBayesQuantRow) - np.log10(bayesQuantRow)
        if np.max(diffInIteration * diffInIteration) < 1e-4:
            converged = True
            # print("Converged after iteration", iteration+1, quantRows[0].protein[0])
            break

    if not converged:
        print("Warning: failed to converge for protein", quantRows[0].protein[0])

    return pProteinQuantsList, bayesQuantRow


def getPosteriorProteinRatio(quantMatrix, quantRows, geoAvgQuantRow, params):
    numSamples = len(quantMatrix[0])

    logGeoAvgs = np.log10(parsers.geomAvg(quantMatrix, axis=1))
    featDiffs = np.log10(quantMatrix) - logGeoAvgs[:, np.newaxis]
    pMissingGeomAvg = pMissing(
        logGeoAvgs, params["muDetect"], params["sigmaDetect"]
    )  # Pr(f_grn = NaN | t_grn = 1)

    pQuantIncorrectId = hyperparameters.funcHypsec(
        featDiffs, params["muFeatureDiff"], params["sigmaFeatureDiff"]
    )  # Pr(f_grn = x | t_grn = 1)
    # pQuantIncorrectIdOld = hyperparameters.funcLogitNormal(np.log10(quantMatrix), params["muDetect"], params["sigmaDetect"], params["muXIC"], params["sigmaXIC"])

    xImpsAll = imputeValues(
        quantMatrix, geoAvgQuantRow, params["proteinQuantCandidates"]
    )
    impDiffs = xImpsAll - np.log10(np.array(quantMatrix))[:, :, np.newaxis]
    pDiffs = hyperparameters.funcHypsec(
        impDiffs, params["muFeatureDiff"], params["sigmaFeatureDiff"]
    )  # Pr(f_grn = x | m_grn = 0, t_grn = 0)

    pProteinQuantsList, bayesQuantRow = list(), list()
    for j in range(numSamples):
        pProteinQuant = params["proteinPrior"].copy()  # log likelihood
        for i, row in enumerate(quantMatrix):
            linkPEP = quantRows[i].linkPEP[j]
            identPEP = quantRows[i].identificationPEP[j]
            if identPEP < 1.0:
                pMissings = pMissing(
                    xImpsAll[i, j, :], params["muDetect"], params["sigmaDetect"]
                )  # Pr(f_grn = NaN | m_grn = 1, t_grn = 0)
                if np.isnan(row[j]):
                    likelihood = pMissings * (1.0 - identPEP) * (
                        1.0 - linkPEP
                    ) + pMissingGeomAvg[i] * (identPEP * (1.0 - linkPEP) + linkPEP)
                else:
                    likelihood = (1.0 - pMissings) * pDiffs[i, j, :] * (
                        1.0 - identPEP
                    ) * (1.0 - linkPEP) + (1.0 - pMissingGeomAvg[i]) * (
                        pQuantIncorrectId[i][j] * identPEP * (1.0 - linkPEP) + linkPEP
                    )

                if np.min(likelihood) == 0.0:
                    likelihood += np.nextafter(0, 1)
                pProteinQuant += np.log(likelihood)

        pProteinQuant -= np.max(pProteinQuant)
        pProteinQuant = np.exp(pProteinQuant) / np.sum(np.exp(pProteinQuant))
        pProteinQuantsList.append(pProteinQuant)

        eValue, confRegion = getPosteriorParams(
            params["proteinQuantCandidates"], pProteinQuant
        )
        bayesQuantRow.append(eValue)

    return pProteinQuantsList, bayesQuantRow


def imputeValues(quantMatrix, proteinRatios, testProteinRatios):
    logIonizationEfficiencies = np.log10(quantMatrix) - np.log10(proteinRatios)

    numNonZeros = np.count_nonzero(~np.isnan(logIonizationEfficiencies), axis=1)[
        :, np.newaxis
    ] - ~np.isnan(logIonizationEfficiencies)
    np.nan_to_num(logIonizationEfficiencies, False)
    meanLogIonEff = (
        np.nansum(logIonizationEfficiencies, axis=1)[:, np.newaxis]
        - logIonizationEfficiencies
    ) / numNonZeros

    logImputedVals = (
        np.tile(meanLogIonEff[:, :, np.newaxis], (1, 1, len(testProteinRatios)))
        + testProteinRatios
    )
    return logImputedVals


def pMissing(x, muLogit, sigmaLogit):
    return 1.0 - hyperparameters.logit(x, muLogit, sigmaLogit) + np.nextafter(0, 1)


def getPosteriorProteinGroupRatios(pProteinQuantsList, bayesQuantRow, params):
    numGroups = len(params["groups"])

    pProteinGroupQuants = list()
    for groupId in range(numGroups):
        filteredProteinQuantsList = np.array(
            [
                x
                for j, x in enumerate(pProteinQuantsList)
                if j in params["groups"][groupId]
            ]
        )
        if "shapeInGroupStdevs" in params:
            pMu = getPosteriorProteinGroupMuMarginalized(
                filteredProteinQuantsList, params
            )
        else:
            pMu = getPosteriorProteinGroupMu(
                params["inGroupDiffPrior"], filteredProteinQuantsList, params
            )
        pProteinGroupQuants.append(pMu)

    return pProteinGroupQuants


def getPosteriorProteinGroupMu(pDiffPrior, pProteinQuantsList, params):
    pMus = np.zeros_like(params["proteinQuantCandidates"])
    for pProteinQuants in pProteinQuantsList:
        pMus += np.log(np.convolve(pDiffPrior, pProteinQuants, mode="valid"))

    # pMus = np.nan_to_num(pMus)
    pMus -= np.max(pMus)
    pMus = np.exp(pMus) / np.sum(np.exp(pMus))
    return pMus


def getPosteriorProteinGroupMuMarginalized(pProteinQuantsList, params):
    pMus = np.zeros(
        (len(params["sigmaCandidates"]), len(params["proteinQuantCandidates"]))
    )
    for pProteinQuants in pProteinQuantsList:
        for idx, pDiffPrior in enumerate(params["inGroupDiffPrior"]):
            pMus[idx, :] += np.log(
                np.convolve(pDiffPrior, pProteinQuants, mode="valid")
            )

    pSigmas = hyperparameters.funcGamma(
        params["sigmaCandidates"],
        params["shapeInGroupStdevs"],
        params["scaleInGroupStdevs"],
    )  # prior
    pMus = np.log(np.dot(pSigmas, np.exp(pMus)))

    pMus -= np.max(pMus)
    pMus = np.exp(pMus) / np.sum(np.exp(pMus))

    return pMus


def getProteinGroupsDiffPosteriors(pProteinGroupQuants, params):
    numGroups = len(params["groups"])
    pProteinGroupDiffs, muGroupDiffs = dict(), dict()
    for groupId1, groupId2 in itertools.combinations(range(numGroups), 2):
        pDifference = np.convolve(
            pProteinGroupQuants[groupId1], pProteinGroupQuants[groupId2][::-1]
        )
        pProteinGroupDiffs[(groupId1, groupId2)] = pDifference
        muGroupDiffs[(groupId1, groupId2)], _ = np.log2(
            getPosteriorParams(params["proteinDiffCandidates"], pDifference)
            + np.nextafter(0, 1)
        )
    return pProteinGroupDiffs, muGroupDiffs


def getProbBelowFoldChangeDict(pProteinGroupDiffs, params):
    probsBelowFoldChange = dict()
    numGroups = len(params["groups"])
    for groupId1, groupId2 in itertools.combinations(range(numGroups), 2):
        probsBelowFoldChange[(groupId1, groupId2)] = getPosteriorProteinGroupDiff(
            pProteinGroupDiffs[(groupId1, groupId2)], params
        )
    # probsBelowFoldChange['ANOVA'] = getProbBelowFoldChangeANOVA(pProteinGroupQuants, params)
    return probsBelowFoldChange


def getPosteriorProteinGroupDiff(pDifference, params):
    return pDifference[
        np.abs(np.log2(10 ** params["proteinDiffCandidates"]))
        < params["foldChangeEval"]
    ].sum()


# this is a "pseudo"-ANOVA test which calculates the probability distribution
# for differences of means between multiple groups. With <=4 groups it seemed
# to return reasonable results, but with 10 groups it called many false positives.
def getProbBelowFoldChangeANOVA(pProteinGroupQuants, params):
    if len(pProteinGroupQuants) > 4:
        print(
            "WARNING: this ANOVA-like test might not behave well if >4 treatment groups are present"
        )

    if len(pProteinGroupQuants) >= 2:
        convProbs = convolution_dp.convolveProbs(pProteinGroupQuants)
        bandwidth = np.searchsorted(
            params["proteinQuantCandidates"],
            params["proteinQuantCandidates"][0]
            + np.log10(2 ** params["foldChangeEval"]),
        )
        probBelowFoldChange = 0.0
        for i in range(bandwidth):
            probBelowFoldChange += np.trace(convProbs, offset=i)
    else:
        probBelowFoldChange = 1.0
    return min([1.0, probBelowFoldChange])


def getPosteriorParams(proteinQuantCandidates, pProteinQuants):
    return 10 ** np.sum(proteinQuantCandidates * pProteinQuants), 0.0
    if False:
        eValue, variance = 0.0, 0.0
        for proteinRatio, pq in zip(proteinQuantCandidates, pProteinQuants):
            if pq > 0.001:
                # print(10**proteinRatio, pq)
                eValue += proteinRatio * pq

        for proteinRatio, pq in zip(proteinQuantCandidates, pProteinQuants):
            if pq > 0.001:
                variance += pq * (proteinRatio - eValue) ** 2
        eValueNew = 10**eValue

        return eValueNew, [
            10 ** (eValue - np.sqrt(variance)),
            10 ** (eValue + np.sqrt(variance)),
        ]
