"""
Source: https://github.com/kerguler/hoppMCMC (01/12/18)

adaptive basin-hopping Markov-chain Monte Carlo for Bayesian optimisation

This is the python (v2.7) implementation of the hoppMCMC algorithm aiming to
identify and sample from the high-probability regions of a posterior
distribution. The algorithm combines three strategies: (i) parallel MCMC,
(ii) adaptive Gibbs sampling and (iii) simulated annealing.

Overall, hoppMCMC resembles the basin-hopping algorithm implemented in the
optimize module of scipy, but it is developed for a wide range of modelling
approaches including stochastic models with or without time-delay.

"""

import os
# import sys
import numpy as np
from struct import pack, unpack
from scipy.stats import ttest_1samp as ttest

MPI_MASTER = 0
try:
    from mpi4py import MPI
    MPI_SIZE = MPI.COMM_WORLD.Get_size()
    MPI_RANK = MPI.COMM_WORLD.Get_rank()

    def Abort(str):
        print("ERROR: " + str)
        MPI.COMM_WORLD.Abort(1)
except:
    MPI_SIZE = 1
    MPI_RANK = 0

    def Abort(str):
        raise errorMCMC(str)

EPS_PULSE_VAR_M
IN = 1e-12
EPS_VARMAT_MIN = 1e-7
EPS_VARMAT_MAX = 1e7


class errorMCMC(Exception):
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


class binfile():

    def __init__(self, fname, mode, rowsize=1):
        self.headsize = 4
        self.bitsize = 8
        self.fname = fname
        self.mode = mode
        self.rowsize = rowsize
        if self.mode == 'r':
  
          try:
                self.f = open(self.fname, "rb")
            except IOError:
                Abort("File not found: " + self.fname)
            self.rowsize = unpack('<i', self.f.read(self.headsize))[0]
        elif self.mode == 'w':
            try:
                self.f = open(self.fname, "w+b")
            except IOError:
                Abort("File not found: " + self.fname)
            tmp = self.f.read(self.headsize)
            if not tmp:
                self.f.write(pack('<i', self.rowsize))
            else:
                self.rowsize = unpack('<i', tmp)[0]
        else:
            Abort("Wrong i/o mode: " + self.mode)
        self.fmt = '<' + 'd' * self.rowsize
        self.size = self.bitsize * self.rowsize
        print("Success: %s opened for %s size %d double rows" % (
            self.fname, "reading" if self.mode == 'r' else
            "writing", self.rowsize))

    def writeRow(self, row):
        self.f.write(pack(self.fmt, *row))
        self.f.flush()

    def readRows(self):
        self.f.seek(0, os.SEEK_END)
        filesize = self.f.tell()
        self.f.seek(self.headsize)
        tmp = self.f.read(filesize - self.headsize)
        ret = np.array(
            unpack('<' + 'd' * int(np.floor((
                filesize - self.headsize) / self.bitsize)), tmp),
            dtype=np.float64).reshape((int(np.floor(((
                filesize - self.headsize) / self.bitsize) / self.rowsize)),
                self.rowsize))
        return ret

    def close(self):
        self.f.close()
        print("Success: %s closed" % (self.fname if self.fname else "file"))


def readFile(filename):
    """
    Reads a binary output file and returns all rows/columns

    Parameters
    ----------

    filename:
         name of the output file

    Returns
    -------

    a np array with all the rows/columns

    """
    a = binfile(filename, "r")
    b = a.readRows()
    a.close()
    return b


def parMin(parmat):
    prm = min(parmat[:, 0])
    prmi = np.where(parmat[:, 0] == prm)[0][0]
    return {'i': prmi, 'f': prm}


def diagdot(mat, vec):
    for n in np.arange(mat.shape[0]):
        mat[n, n] *= vec[n]


def covariance(mat):
    return np.cov(mat, rowvar=False)


def coefVar(mat):
    mean0 = 1.0 / np.mean(mat, 0)
    return mean0 * (np.cov(mat, rowvar=False).T * mean0).T


def sensVar(mat):
    mean0 = np.mean(mat, 0)
    return mean0 * (np.linalg.inv(np.cov(mat, rowvar=False)).T * mean0).T


cov = covariance
rnorm = np.random.multivariate_normal
determinant = np.linalg.det


def join(a, s):
    return s.join(["%.16g" % (x) for x in a])


def logsumexp(x):
    a = np.max(x)
    return a + np.log(np.sum(np.exp(x - a)))


def compareAUCs(parmats, groups, tol=1.0):
    from scipy.stats import gaussian_kde
    ids = np.unique(groups)
    parmats = parmats[:, np.var(parmats, axis=0) != 0]
    parmat_list = [parmats[groups == n, :] for n in ids]
    counts = [np.sum(n == groups) for n in ids]
    try:
        kde = gaussian_kde(parmats[:, 1:].T)
    except:
        print("Warning: Problem encountered in compareAUC!")
        return {ids[g]: 0 for g in range(len(ids))}

    wt_list = np.array([
        [kde.evaluate(pr) for pr in parmat[:, 1:]] for parmat in parmat_list])
    mn = np.min(wt_list)
    wt_list = np.log(wt_list / mn)
    # Importance sampling for Monte Carlo integration:
    favg_exp = np.array([
        [-parmat_list[n][m, 0] / tol - wt_list[n][m] for m in
         range(parmat_list[n].shape[0])] for n in range(len(parmat_list))])
    favg_logsums = np.array([logsumexp(x) for x in favg_exp]) - np.log(counts)
    favg_list = np.exp(favg_logsums - logsumexp(favg_logsums))
    return {ids[g]: favg_list[g] for g in range(len(ids))}


def compareAUC(parmat0, parmat1, T):
    from scipy.stats import gaussian_kde
    parmats = np.vstack((parmat0, parmat1))
    parmat0cp = parmat0[:, np.var(parmats, axis=0) != 0].copy()
    parmat1cp = parmat1[:, np.var(parmats, axis=0) != 0].copy()
    parmats = parmats[:, np.var(parmats, axis=0) != 0]
    try:
        kde = gaussian_kde(parmats[:, 1:].T)
    except:
        print("Warning: Problem encountered in compareAUC!")
        return {'acc': 0, 'favg0': 0, 'favg1': 0}
    wt0 = np.array([kde.evaluate(pr) for pr in parmat0cp[:, 1:]])
    wt1 = np.array([kde.evaluate(pr) for pr in parmat1cp[:, 1:]])
    mn = np.min([wt0, wt1])
    wt0 /= mn
    wt1 /= mn
    # Importance sampling for Monte Carlo integration:
    favg0 = np.mean([
        np.exp(-parmat0cp[m, 0] / T) / wt0[m] for m in
        range(parmat0cp.shape[0])])
    favg1 = np.mean([
        np.exp(-parmat1cp[m, 0] / T) / wt1[m] for m in
        range(parmat1cp.shape[0])])
    acc = (not np.isnan(favg0) and not np.isnan(favg1) and
           favg0 > 0 and favg1 >= 0 and
           (favg1 >= favg0 or
           (np.random.uniform() < (favg1 / favg0))))
    return {'acc': acc, 'favg0': favg0, 'favg1': favg1}


def anneal_exp(y0, y1, steps):
    return y0 * np.exp(-(np.arange(
        steps, dtype=np.float64) / (steps - 1.)) * np.log(np.float64(y0) / y1))


def anneal_linear(y0, y1, steps):
    return np.append(np.arange(
        y0, y1, (y1 - y0) / (steps - 1), dtype=np.float64), y1)


def anneal_sigma(y0, y1, steps):
    return y0 + (y1 - y0) * (1. - 1. / (1. + np.exp(-(
        np.arange(steps) - (0.5 * steps)))))


def anneal_sigmasoft(y0, y1, steps):
    return y0 + (y1 - y0) * (1. - 1. / (1. + np.exp(-12.5 * (
        np.arange(steps) - (0.5 * steps)) / steps)))


def ldet(mat):
    try:
        det = determinant(mat)
    except:
        return -np.Inf
    if det == 0:
        return -np.Inf
    else:
        return np.log(det)


def finalTest(fitFun, param, testnum=10):
    for n in range(testnum):
        f = fitFun(param)
        if not np.isnan(f) and not np.isinf(f):
            return f
    Abort("Incompatible parameter set (%g): %s" % (f, join(param, ",")))


class hoppMCMC:
    def __init__(
        self, fitFun, param, varmat, inferpar=None, gibbs=True, num_hopp=3,
        num_adapt=25, num_chain=12, chain_length=50, rangeT=None,
            model_comp=1000.0, outfilename=''):
        """
        Adaptive Basin-Hopping MCMC Algorithm

        Parameters
        ----------

        fitFun:
              fitFun(x) - objective function which takes a np array as the
              only argument

        param:
              initial parameter vector

        varmat:
              2-dimensional array of initial covariance matrix

        inferpar:
              an array of indexes of parameter dimensions to be inferred
              (all parameters are inferred by default)

        gibbs:
              i