import numpy as np
from scipy.stats import norm
from scipy.optimize import minimize
import warnings

# Python translation from Brian Maniscalco's code at http://www.columbia.edu/~bsm2105/type2sdt/
# Matthias Guggenmos, 2021


def fit_meta_d_MLE(nR_S1, nR_S2, s=1, fncdf=None, fninv=None):
    # fit = fit_meta_d_MLE(nR_S1, nR_S2, s, fncdf, fninv)
    #
    # Given data from an experiment where an observer discriminates between two
    # stimulus alternatives on every trial and provides confidence ratings,
    # provides a type 2 SDT analysis of the data.
    #
    # INPUTS

    # * nR_S1, nR_S2
    # these are vectors containing the total number of responses in
    # each response modality, conditional on presentation of S1 and S2.
    #
    # e.g. if nR_S1 = [100 50 20 10 5 1], then when stimulus S1 was
    # presented, the subject had the following response counts:
    # responded S1, rating=3 : 100 times
    # responded S1, rating=2 : 50 times
    # responded S1, rating=1 : 20 times
    # responded S2, rating=1 : 10 times
    # responded S2, rating=2 : 5 times
    # responded S2, rating=3 : 1 time
    #
    # * s
    # this is the ratio of standard deviations for type 1 distributions, i.e.
    #
    # s = sd(S1) / sd(S2)
    #
    # if not specified, s is set to a default value of 1.
    # the function SDT_MLE_fit available at
    # http://www.columbia.edu/~bsm2105/type2sdt/
    # can be used to get an MLE estimate of s using Matlab's optimization
    # toolbox
    #
    # * fncdf
    # a function handle for the CDF of the type 1 distribution.
    # if not specified, fncdf defaults to @normcdf (i.e. CDF for normal
    # distribution)
    #
    # * fninv
    # a function handle for the inverse CDF of the type 1 distribution.
    # if not specified, fninv defaults to @norminv
    #
    # OUTPUT
    #
    # Output is packaged in the struct "fit."
    # In the following, let S1 and S2 represent the distributions of evidence
    # generated by stimulus classes S1 and S2.
    # Then the fields of "fit" are as follows:
    #
    # fit.da        = mean(S2) - mean(S1), in room-mean-square(sd(S1),sd(S2)) units
    # fit.s         = sd(S1) / sd(S2)
    # fit.meta_da   = meta-d' in RMS units
    # fit.M_diff    = meta_da - da
    # fit.M_ratio   = meta_da / da
    # fit.t1ca      = type 1 criterion for meta-d' fit, RMS units
    # fit.t2ca      = type 2 criteria for meta-d' fit, RMS units
    #
    # fit.S1units   = contains same parameters in sd(S1) units.
    #                 these may be of use since the data-fitting is conducted
    #                 using parameters specified in sd(S1) units.
    #
    # fit.logL          = log likelihood of the data fit
    #
    # fit.obs_HR2_rS1  = actual type 2 hit rates for S1 responses
    # fit.est_HR2_rS1  = estimated type 2 hit rates for S1 responses
    # fit.obs_FAR2_rS1 = actual type 2 false alarm rates for S1 responses
    # fit.est_FAR2_rS1 = estimated type 2 false alarm rates for S1 responses
    #
    # fit.obs_HR2_rS2  = actual type 2 hit rates for S2 responses
    # fit.est_HR2_rS2  = estimated type 2 hit rates for S2 responses
    # fit.obs_FAR2_rS2 = actual type 2 false alarm rates for S2 responses
    # fit.est_FAR2_rS2 = estimated type 2 false alarm rates for S2 responses
    #
    # If there are N ratings, then there will be N-1 type 2 hit rates and false
    # alarm rates.

    # 9/7/10 - bm - wrote it

    ## parse inputs

    # nR_S1 = np.array([171275, 174686, 110922, 43117])
    # nR_S2 = np.array([43117, 111053, 174370, 171460])
    # nR_S1 = np.array([17664, 65846, 87765, 88184, 86502, 64187, 46735, 29635, 12072, 1410])
    # nR_S2 = np.array([1421, 11891, 29805, 46620, 64433, 86248, 88122, 88348, 65901, 17211])

    if not (np.mod(len(nR_S1), 2) == 0):
        raise Exception('input arrays must have an even number of elements')
    if len(nR_S1) != len(nR_S2):
        raise Exception('input arrays must have the same number of elements')

    if fncdf is None:
        fncdf = norm.cdf

    if fninv is None:
        fninv = norm.ppf

    nRatings = int(len(nR_S1) / 2)
    nCriteria = 2 * nRatings - 1

    ## set up constraints for MLE estimation

    # parameters
    # meta-d' - 1
    # t2c     - nCriteria-1

    # constrain type 2 criteria values,
    # such that t2c(i) is always <= t2c(i+1)
    # want t2c(i)   <= t2c(i+1)
    # -->  t2c(i+1) >= c(i) + 1e-5 (i.e. very small deviation from equality)
    # -->  t2c(i) - t2c(i+1) <= -1e-5
    A = np.zeros((nCriteria - 2, nCriteria))
    b = np.full(nCriteria - 2, np.nan)
    for i in range(nCriteria - 2):
        A[i, (i + 1):(i + 3)] = [1, -1]
        b[i] = -1e-5

    # lower bounds on parameters
    LB = np.hstack((-10, -20 * np.ones(int((nCriteria - 1) / 2)), np.zeros((int((nCriteria - 1) / 2)))))

    # upper bounds on parameters
    UB = np.hstack((10, np.zeros((int((nCriteria - 1) / 2))), 20 * np.ones(int((nCriteria - 1) / 2))))

    ## select constant criterion type

    constant_criterion = 'meta_d1 * (t1c1 / d1)'  # relative criterion

    ## set up initial guess at parameter values

    ratingHR = np.full(nRatings * 2 - 1, np.nan)
    ratingFAR = np.full(nRatings * 2 - 1, np.nan)
    for c in range(1, nRatings * 2):
        ratingHR[c - 1] = np.sum(nR_S2[c:]) / np.sum(nR_S2)
        ratingFAR[c - 1] = np.sum(nR_S1[c:]) / np.sum(nR_S1)

    t1_index = nRatings - 1
    t2_index = np.setdiff1d(range(2 * nRatings - 1), t1_index)

    d1 = (1 / s) * fninv(ratingHR[t1_index]) - fninv(ratingFAR[t1_index])
    if np.abs(d1) < 1e-4:
        if d1 == 0:
            d1 = np.random.choice([-1, 1]) * 1e-4
        else:
            d1 = np.sign(d1) * 1e-4
        # d1 = 1e-4
        small_d1 = 1
    else:
        small_d1 = 0
    meta_d1 = d1

    c1 = (-1 / (1 + s)) * (fninv(ratingHR) + fninv(ratingFAR))
    t1c1 = c1[t1_index]
    t2c1 = c1[t2_index]

    guess = np.hstack((meta_d1, t2c1 - eval(constant_criterion)))

    ## find the best fit for type 2 hits and FAs

    bounds = [(lb, ub) for lb, ub in zip(LB, UB)]
    constraint = {'type': 'ineq', 'fun': lambda x: b - np.dot(A, x).squeeze()}
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', category=(UserWarning, RuntimeWarning))
        # res = minimize(fit_meta_d_logL, guess, args=(nR_S1, nR_S2, nRatings, d1, t1c1, s, constant_criterion, fncdf, fninv), method='SLSQP', bounds=bounds, constraints=constraint, options=dict(maxiter=1e4, disp=False))  # trust-constr, L-BFGS-B, TNC, SLSQP
        args = [nR_S1, nR_S2, nRatings, d1, t1c1, s, constant_criterion, fncdf, fninv]
        res = minimize(fit_meta_d_logL, guess, args=tuple(args), method='SLSQP', bounds=bounds, constraints=constraint, options=dict(disp=False))  # trust-constr, L-BFGS-B, TNC, SLSQP

    meta_d1 = res.x[0]
    t2c1 = res.x[1:] + eval(constant_criterion)
    logL = -res.fun



    ## data is fit, now to package it...

    ## find observed t2FAR and t2HR

    # I_nR and C_nR are rating trial counts for incorrect and correct trials
    # element i corresponds to # (in)correct w/ rating i
    I_nR_rS2 = nR_S1[nRatings:]
    I_nR_rS1 = nR_S2[nRatings - 1::-1]

    C_nR_rS2 = nR_S2[nRatings:]
    C_nR_rS1 = nR_S1[nRatings - 1::-1]

    obs_FAR2_rS2, obs_HR2_rS2 = np.full(nRatings - 1, np.nan), np.full(nRatings - 1, np.nan)
    obs_FAR2_rS1, obs_HR2_rS1 = np.full(nRatings - 1, np.nan), np.full(nRatings - 1, np.nan)
    for i in range(1, nRatings):
        obs_FAR2_rS2[i - 1] = np.sum(I_nR_rS2[i:]) / np.sum(I_nR_rS2)
        obs_HR2_rS2[i - 1] = np.sum(C_nR_rS2[i:]) / np.sum(C_nR_rS2)

        obs_FAR2_rS1[i - 1] = np.sum(I_nR_rS1[i:]) / np.sum(I_nR_rS1)
        obs_HR2_rS1[i - 1] = np.sum(C_nR_rS1[i:]) / np.sum(C_nR_rS1)

    ## find estimated t2FAR and t2HR

    S1mu = -meta_d1 / 2
    S1sd = 1
    S2mu = meta_d1 / 2
    S2sd = S1sd / s

    mt1c1 = eval(constant_criterion)

    C_area_rS2 = 1 - fncdf(mt1c1, S2mu, S2sd)
    I_area_rS2 = 1 - fncdf(mt1c1, S1mu, S1sd)

    C_area_rS1 = fncdf(mt1c1, S1mu, S1sd)
    I_area_rS1 = fncdf(mt1c1, S2mu, S2sd)

    est_FAR2_rS2, est_HR2_rS2 = np.full(nRatings - 1, np.nan), np.full(nRatings - 1, np.nan)
    est_FAR2_rS1, est_HR2_rS1 = np.full(nRatings - 1, np.nan), np.full(nRatings - 1, np.nan)
    for i in range(0, nRatings - 1):
        t2c1_lower = t2c1[nRatings - i - 2]
        t2c1_upper = t2c1[nRatings - 1 + i]

        I_FAR_area_rS2 = 1 - fncdf(t2c1_upper, S1mu, S1sd)
        C_HR_area_rS2 = 1 - fncdf(t2c1_upper, S2mu, S2sd)

        I_FAR_area_rS1 = fncdf(t2c1_lower, S2mu, S2sd)
        C_HR_area_rS1 = fncdf(t2c1_lower, S1mu, S1sd)

        est_FAR2_rS2[i] = I_FAR_area_rS2 / I_area_rS2
        est_HR2_rS2[i] = C_HR_area_rS2 / C_area_rS2

        est_FAR2_rS1[i] = I_FAR_area_rS1 / I_area_rS1
        est_HR2_rS1[i] = C_HR_area_rS1 / C_area_rS1

    ## package output

    class Fit:
        def __init__(self):
            self.da = np.sqrt(2 / (1 + s ** 2)) * s * d1
            self.d1 = d1
            if np.abs(self.da) < 1e-4:
                if self.da == 0:
                    self.da = np.random.choice([-1, 1]) * 1e-4
                else:
                    self.da = np.sign(self.da) * 1e-4
                self.small_da = 1
            else:
                self.small_da = 0
            self.s = s
            self.meta_da = np.sqrt(2 / (1 + s ** 2)) * s * meta_d1
            if np.abs(self.meta_da) < 1e-4:
                self.meta_da = np.sign(self.meta_da) * 1e-4
                self.small_meta_da = 1
            else:
                self.small_meta_da = 0
            self.small_d1 = small_d1
            self.M_diff = self.meta_da - self.da
            self.M_ratio = self.meta_da / self.da
            self.meta_ca = (np.sqrt(2) * s / np.sqrt(1 + s ** 2)) * t1c1
            t2ca = (np.sqrt(2) * s / np.sqrt(1 + s ** 2)) * t2c1
            self.t2ca_rS1 = t2ca[:nRatings - 1]
            self.t2ca_rS2 = t2ca[nRatings - 1:]
            self.S1units = dict(d1=d1, meta_d1=meta_d1, s=1, meta_c1=t1c1, t2c1_rS1=t2c1[:nRatings - 1], t2c1_rS2=t2c1[nRatings - 1:])
            self.logL = logL
            self.est_HR2_rS1 = est_HR2_rS1
            self.obs_HR2_rS1 = obs_HR2_rS1
            self.est_FAR2_rS1 = est_FAR2_rS1
            self.obs_FAR2_rS1 = obs_FAR2_rS1
            self.est_HR2_rS2 = est_HR2_rS2
            self.obs_HR2_rS2 = obs_HR2_rS2
            self.est_FAR2_rS2 = est_FAR2_rS2
            self.obs_FAR2_rS2 = obs_FAR2_rS2
            self.fit_nfev = res.nfev
            self.fit_nit = res.nit
            self.fit_njev = res.njev
            self.fit_success = int(res.success)
            self.fit_fun = res.fun

    return Fit()


## function to find the likelihood of parameter values, given observed data
def fit_meta_d_logL(parameters, nR_S1, nR_S2, nRatings, d1, t1c1, s, constant_criterion, fncdf, fninv):
    # set up parameters
    meta_d1 = parameters[0]
    t2c1 = np.array(parameters[1:])

    # loads:
    # nR_S1 nR_S2 nRatings d1 t1c1 s constant_criterion fncdf fninv
    # nR_S1, nR_S2, nRatings, d1, t1c1, s, constant_criterion, fncdf, fninv = pickle.load(open('fit_meta_d_MLE.pkl', 'rb'))

    # define mean and SD of S1 and S2 distributions
    S1mu = -meta_d1 / 2
    S1sd = 1
    S2mu = meta_d1 / 2
    S2sd = S1sd / s

    # adjust so that the type 1 criterion is set at 0
    # (this is just to work with optimization toolbox constraints...
    #  to simplify defining the upper and lower bounds of type 2 criteria)
    S1mu = S1mu - eval(constant_criterion)
    S2mu = S2mu - eval(constant_criterion)

    t1c1 = 0

    ### set up MLE analysis

    # get type 2 response counts
    nC_rS1, nC_rS2 = np.full(nRatings, np.nan), np.full(nRatings, np.nan)
    nI_rS1, nI_rS2 = np.full(nRatings, np.nan), np.full(nRatings, np.nan)
    for i in range(nRatings):
        # S1 responses
        nC_rS1[i] = nR_S1[i]
        nI_rS1[i] = nR_S2[i]

        # S2 responses
        nC_rS2[i] = nR_S2[nRatings + i]
        nI_rS2[i] = nR_S1[nRatings + i]

    # get type 2 probabilities
    C_area_rS1 = fncdf(t1c1, S1mu, S1sd)
    I_area_rS1 = fncdf(t1c1, S2mu, S2sd)

    C_area_rS2 = 1 - fncdf(t1c1, S2mu, S2sd)
    I_area_rS2 = 1 - fncdf(t1c1, S1mu, S1sd)

    t2c1x = np.hstack((-np.inf, t2c1[0:nRatings - 1], t1c1, t2c1[nRatings - 1:], np.inf));

    prC_rS1, prC_rS2 = np.full(nRatings, np.nan), np.full(nRatings, np.nan)
    prI_rS1, prI_rS2 = np.full(nRatings, np.nan), np.full(nRatings, np.nan)
    for i in range(nRatings):
        prC_rS1[i] = (fncdf(t2c1x[i + 1], S1mu, S1sd) - fncdf(t2c1x[i], S1mu, S1sd)) / C_area_rS1
        prI_rS1[i] = (fncdf(t2c1x[i + 1], S2mu, S2sd) - fncdf(t2c1x[i], S2mu, S2sd)) / I_area_rS1

        prC_rS2[i] = ((1 - fncdf(t2c1x[nRatings + i], S2mu, S2sd)) - (1 - fncdf(t2c1x[nRatings + i + 1], S2mu, S2sd))) / C_area_rS2
        prI_rS2[i] = ((1 - fncdf(t2c1x[nRatings + i], S1mu, S1sd)) - (1 - fncdf(t2c1x[nRatings + i + 1], S1mu, S1sd))) / I_area_rS2

    # calculate logL
    logL = 0
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', category=RuntimeWarning)
        for i in range(nRatings):
            logL = logL + nC_rS1[i] * np.log(prC_rS1[i]) + nI_rS1[i] * np.log(prI_rS1[i]) + \
                   nC_rS2[i] * np.log(prC_rS2[i]) + nI_rS2[i] * np.log(prI_rS2[i])

    if np.isnan(logL):
        logL = -np.inf

    logL = -logL

    return logL


def type2_SDT_MLE(stimID, response, rating, nRatings, cellpadding=None, equalVariance=1):
    # out = type2_SDT(input)
    #
    # Given data from an experiment where an observer discriminates between two
    # stimulus alternatives on every trial and provides confidence ratings,
    # provides a type 2 SDT analysis of the data.
    #
    # The function estimates the parameters of the unequal variance SDT model,
    # and uses those estimates to find a maximum likelihood estimate of
    # meta-da.
    #
    # INPUTS
    #
    # format of the input may be either:
    #
    # 1) stimID, response, rating, nRatings, (cellpadding), (equalVariance)
    #    where each of the first 3 inputs is a 1xN vector describing the outcome
    #    of N trials. Contents of input should be as follows.
    #
    #    stimID   : 0=S1 stimulus presented, 1=S2 stimulus presented
    #    response : 0=subject responded S1, 1=subject responded S2
    #    rating   : values ranges from 1 to m where 1 is the lowest rating
    #               and m is the highest.
    #
    #               All trials where any of these prescribed ranges of values
    #               are violated are omitted from analysis.
    #
    #    nRatings : the number of ratings available to the subject (e.g. for a
    #               confidence scale of 1-4, nRatings=4).
    #    cellpadding : if any data cells (e.g. high confidence "S2" responses)
    #               are empty, then the value of cellpadding will be added
    #               to every data cell. If not specified, default = 1/(2*nRatings)
    #    equalVariance : if 1, force analysis to use the equal variance SDT
    #               model. If 0, use an estimate of s = sd(S1) / sd(S2) where
    #               s is the slope of the zROC data (estimated using MLE).
    #               If not specified, default = 0.
    #
    # 2) nR_S1, nR_S2, (cellpadding), (equalVariance)
    #    where these are vectors containing the total number of responses in
    #    each response modality, conditional on presentation of S1 and S2.
    #    size of each array is 2*nRatings, where each element corresponds to a
    #    count of responses in each response modality. Response categories are
    #    ordered as follows:
    #    highest conf "S1" ... lowest conf "S1", lowest conf "S2", ... highest conf "S2"
    #
    #    e.g. if nR_S1 = [100 50 20 10 5 1], then when stimulus S1 was
    #    presented, the subject had the following response counts:
    #    responded S1, rating=3 : 100 times
    #    responded S1, rating=2 : 50 times
    #    responded S1, rating=1 : 20 times
    #    responded S2, rating=1 : 10 times
    #    responded S2, rating=2 : 5 times
    #    responded S2, rating=3 : 1 time
    #
    #    cellpadding and equalVariance are defined as above.
    #
    #
    #
    #
    # OUTPUTS
    #
    # out.d_a       : d_a for input data. If s=1, d_a = d'
    # out.meta_d_a  : meta_d_a for input data
    # out.M_ratio   : meta_d_a / d_a; measure of metacognitive efficiency
    # out.M_diff    : meta_d_a - d_a; measure of metacognitive efficiency
    # out.c_a       : criterion c_a for input data. If s=1, c_a = c.
    # out.cprime    : relative criterion used for type 2 estimates. c' = c_a / d_a
    # out.s         : ratio of evidence distribution standard deviations assumed for the analysis.
    # out.type2_fit : output of fit_meta_d_MLE for the type 2 SDT fit.

    # 9/24/10 - bm - fixed program-crashing bug for (nR_S1, nR_S2) input
    # 9/7/10 - bm - wrote it

    ## parse inputs
    if len(stimID) == 0:
        raise ValueError("Empty data")

    if cellpadding is None:
        cellpadding = 1 / (2 * nRatings)

    # filters bad trials
    f = ((stimID == 0) | (stimID == 1)) & ((response == 0) | (response == 1)) & ((rating >= 1) & (rating <= nRatings))
    stimID = stimID[f]
    response = response[f]
    rating = rating[f]

    # convert to trial count format...
    nR_S1, nR_S2 = np.full(nRatings*2, np.nan), np.full(nRatings*2, np.nan)

    # get tallies of "S1" rating responses for S1 and S2 stim
    for i in range(nRatings):
        nR_S1[i] = np.sum((stimID == 0) & (response == 0) & (rating == nRatings - i))
        nR_S2[i] = np.sum((stimID == 1) & (response == 0) & (rating == nRatings - i))

    # get tallies of "S2" rating responses for S1 and S2 stim
    for i in range(nRatings):
        nR_S1[i + nRatings] = np.sum((stimID == 0) & (response == 1) & (rating == i + 1))
        nR_S2[i + nRatings] = np.sum((stimID == 1) & (response == 1) & (rating == i + 1))

    if np.any(nR_S1 == 0) | np.any(nR_S2 == 0):
        nR_S1 = nR_S1 + cellpadding
        nR_S2 = nR_S2 + cellpadding

    ## standard SDT analysis

    if equalVariance:
        s = 1

    ## type 2 SDT analysis

    fit = fit_meta_d_MLE(nR_S1, nR_S2, s)

    return fit


def type2roc(correct, conf, nbins=5):
    # Calculate area under type 2 ROC
    #
    # correct - vector of 1 x ntrials, 0 for error, 1 for correct
    # conf - vector of continuous confidence ratings between 0 and 1
    # nbins - how many bins to use for discretization

    bs = 1 / nbins
    H2, FA2 = np.full(nbins, np.nan), np.full(nbins, np.nan)
    for c in range(nbins):
        if c:
            H2[nbins - c - 1] = np.sum((conf > c*bs) & (conf <= (c+1)*bs) & (correct).astype(bool)) + 0.5
            FA2[nbins - c - 1] = np.sum((conf > c*bs) & (conf <= (c+1)*bs) & ~(correct).astype(bool)) + 0.5
        else:
            H2[nbins - c - 1] = np.sum((conf >= c * bs) & (conf <= (c + 1) * bs) & (correct).astype(bool)) + 0.5
            FA2[nbins - c - 1] = np.sum((conf >= c * bs) & (conf <= (c + 1) * bs) & ~(correct).astype(bool)) + 0.5

    H2 /= np.sum(H2)
    FA2 /= np.sum(FA2)
    cum_H2 = np.hstack((0, np.cumsum(H2)))
    cum_FA2 = np.hstack((0, np.cumsum(FA2)))

    k = np.full(nbins, np.nan)
    for c in range(nbins):
        k[c] = (cum_H2[c+1] - cum_FA2[c])**2 - (cum_H2[c] - cum_FA2[c+1])**2

    auroc2 = 0.5 + 0.25*np.sum(k)

    return auroc2

def type2roc_disc(correct, conf, nRatings=5):
    # Ratings have to start with 0!!!
    # Calculate area under type 2 ROC
    #
    # correct - vector of 1 x ntrials, 0 for error, 1 for correct
    # conf - vector of 1 x ntrials of confidence ratings taking values 0:Nratings-1
    # nRatings - how many confidence levels available

    correct = np.array(correct)
    conf = np.array(conf)

    H2, FA2 = np.full(nRatings, np.nan), np.full(nRatings, np.nan)
    for c in range(nRatings):
        H2[nRatings - c - 1] = np.sum((conf == c) & (correct).astype(bool)) + 0.5
        FA2[nRatings - c - 1] = np.sum((conf == c) & ~(correct).astype(bool)) + 0.5

    H2 /= np.sum(H2)
    FA2 /= np.sum(FA2)
    cum_H2 = np.hstack((0, np.cumsum(H2)))
    cum_FA2 = np.hstack((0, np.cumsum(FA2)))

    k = np.full(nRatings, np.nan)
    for c in range(nRatings):
        k[c] = (cum_H2[c+1] - cum_FA2[c])**2 - (cum_H2[c] - cum_FA2[c+1])**2

    auroc2 = 0.5 + 0.25*np.sum(k)

    return auroc2

def type2roc_disc2(stim, response, conf, nRatings=5):
    return type2roc_disc((np.array(stim) == np.array(response)).astype(int), conf, nRatings=nRatings)


def convert_for_mratio(stimuli, correct, confidence, nbins=3):
    #  see https://github.com/metacoglab/HMeta-d/wiki/HMeta-d-tutorial
    #  and http://www.columbia.edu/~bsm2105/type2sdt/archive/index.html

    bs = 1 / nbins
    nR_S = np.full((2, nbins * 2), np.nan)
    for s in range(2):  # iterate over presented stimuli
        for c in range(2):  # iterate over correctness
            conf = confidence[(stimuli == s) & (correct == c)]
            xrange = (range(nbins, 2*nbins), range(nbins-1, -1, -1))[(1, 0)[c] if s else c]
            for i in range(nbins):  # iterate over confidence bins
                if i == 0:
                    nR_S[s, xrange[i]] = np.sum((conf >= i * bs) & (conf <= (i + 1) * bs))
                else:
                    nR_S[s, xrange[i]] = np.sum((conf > i * bs) & (conf <= (i + 1) * bs))

    return nR_S[0], nR_S[1]


if __name__ == '__main__':
    nsamples = 10000
    stimulus = np.random.randint(0, 2, nsamples)
    choices = np.random.randint(0, 2, nsamples)
    rand_ind = np.random.permutation(range(nsamples))[:int(nsamples/2)]
    choices[rand_ind] = stimulus[rand_ind]
    confidence = 5 * np.ones(nsamples)
    print(type2_SDT_MLE(stimulus, choices, confidence, 5).M_ratio)
    from confidence.mca.type2_SDT_MLE_matlab import type2_SDT_MLE_matlab
    print(type2_SDT_MLE_matlab(stimulus, choices, confidence, 5).M_ratio)
