# -*- coding: utf-8 -*-
"""
Created on Mon Dec 14 19:53:25 2009

Author: josef-pktd

generate arma sample using fft with all the lfilter it looks slow
to get the ma representation first

apply arma filter (in ar representation) to time series to get white noise
but seems slow to be useful for fast estimation for nobs=10000

change/check: instead of using marep, use fft-transform of ar and ma
    separately, use ratio check theory is correct and example works
    DONE : feels much faster than lfilter
    -> use for estimation of ARMA
    -> use pade (scipy.misc) approximation to get starting polynomial
       from autocorrelation (is autocorrelation of AR(p) related to marep?)
       check if pade is fast, not for larger arrays ?
       maybe pade doesn't do the right thing for this, not tried yet
       scipy.pade([ 1.    ,  0.6,  0.25, 0.125, 0.0625, 0.1],2)
       raises LinAlgError: singular matrix
       also doesn't have roots inside unit circle ??
    -> even without initialization, it might be fast for estimation
    -> how do I enforce stationarity and invertibility,
       need helper function

get function drop imag if close to zero from numpy/scipy source, where?

"""

from __future__ import print_function
import numpy as np
import numpy.fft as fft
#import scipy.fftpack as fft
from scipy import signal
#from try_var_convolve import maxabs
from statsmodels.sandbox.archive.linalg_decomp_1 import OneTimeProperty
from statsmodels.tsa.arima_process import ArmaProcess

#trying to convert old experiments to a class

class ArmaFft(ArmaProcess):
    '''fft tools for arma processes

    This class contains several methods that are providing the same or similar
    returns to try out and test different implementations.

    Notes
    -----
    TODO:
    check whether we don't want to fix maxlags, and create new instance if
    maxlag changes. usage for different lengths of timeseries ?
    or fix frequency and length for fft

    check default frequencies w, terminology norw  n_or_w

    some ffts are currently done without padding with zeros

    returns for spectral density methods needs checking, is it always the power
    spectrum hw*hw.conj()

    normalization of the power spectrum, spectral density: not checked yet, for
    example no variance of underlying process is used

    '''

    def __init__(self, ar, ma, n):
        #duplicates now that are subclassing ArmaProcess
        super(ArmaFft, self).__init__(ar, ma)

        self.ar = np.asarray(ar)
        self.ma = np.asarray(ma)
        self.nobs = n
        #could make the polynomials into cached attributes
        self.arpoly = np.polynomial.Polynomial(ar)
        self.mapoly = np.polynomial.Polynomial(ma)
        self.nar = len(ar)  #1d only currently
        self.nma = len(ma)

    def padarr(self, arr, maxlag, atend=True):
        '''pad 1d array with zeros at end to have length maxlag
        function that is a method, no self used

        Parameters
        ----------
        arr : array_like, 1d
            array that will be padded with zeros
        maxlag : int
            length of array after padding
        atend : boolean
            If True (default), then the zeros are added to the end, otherwise
            to the front of the array

        Returns
        -------
        arrp : ndarray
            zero-padded array

        Notes
        -----
        This is mainly written to extend coefficient arrays for the lag-polynomials.
        It returns a copy.

        '''
        if atend:
            return np.r_[arr, np.zeros(maxlag-len(arr))]
        else:
            return np.r_[np.zeros(maxlag-len(arr)), arr]

    def pad(self, maxlag):
        '''construct AR and MA polynomials that are zero-padded to a common length

        Parameters
        ----------
        maxlag : int
            new length of lag-polynomials

        Returns
        -------
        ar : ndarray
            extended AR polynomial coefficients
        ma : ndarray
            extended AR polynomial coefficients

        '''
        arpad = np.r_[self.ar, np.zeros(maxlag-self.nar)]
        mapad = np.r_[self.ma, np.zeros(maxlag-self.nma)]
        return arpad, mapad

    def fftar(self, n=None):
        '''Fourier transform of AR polynomial, zero-padded at end to n

        Parameters
        ----------
        n : int
            length of array after zero-padding

        Returns
        -------
        fftar : ndarray
            fft of zero-padded ar polynomial
        '''
        if n is None:
            n = len(self.ar)
        return fft.fft(self.padarr(self.ar, n))

    def fftma(self, n):
        '''Fourier transform of MA polynomial, zero-padded at end to n

        Parameters
        ----------
        n : int
            length of array after zero-padding

        Returns
        -------
        fftar : ndarray
            fft of zero-padded ar polynomial
        '''
        if n is None:
            n = len(self.ar)
        return fft.fft(self.padarr(self.ma, n))

    #@OneTimeProperty  # not while still debugging things
    def fftarma(self, n=None):
        '''Fourier transform of ARMA polynomial, zero-padded at end to n

        The Fourier transform of the ARMA process is calculated as the ratio
        of the fft of the MA polynomial divided by the fft of the AR polynomial.

        Parameters
        ----------
        n : int
            length of array after zero-padding

        Returns
        -------
        fftarma : ndarray
            fft of zero-padded arma polynomial
        '''
        if n is None:
            n = self.nobs
        return (self.fftma(n) / self.fftar(n))

    def spd(self, npos):
        '''raw spectral density, returns Fourier transform

        n is number of points in positive spectrum, the actual number of points
        is twice as large. different from other spd methods with fft
        '''
        n = npos
        w = fft.fftfreq(2*n) * 2 * np.pi
        hw = self.fftarma(2*n)  #not sure, need to check normalization
        #return (hw*hw.conj()).real[n//2-1:]  * 0.5 / np.pi #doesn't show in plot
        return (hw*hw.conj()).real * 0.5 / np.pi, w

    def spdshift(self, n):
        '''power spectral density using fftshift

        currently returns two-sided according to fft frequencies, use first half
        '''
        #size = s1+s2-1
        mapadded = self.padarr(self.ma, n)
        arpadded = self.padarr(self.ar, n)
        hw = fft.fft(fft.fftshift(mapadded)) / fft.fft(fft.fftshift(arpadded))
        #return np.abs(spd)[n//2-1:]
        w = fft.fftfreq(n) * 2 * np.pi
        wslice = slice(n//2-1, None, None)
        #return (hw*hw.conj()).real[wslice], w[wslice]
        return (hw*hw.conj()).real, w

    def spddirect(self, n):
        '''power spectral density using padding to length n done by fft

        currently returns two-sided according to fft frequencies, use first half
        '''
        #size = s1+s2-1
        #abs looks wrong
        hw = fft.fft(self.ma, n) / fft.fft(self.ar, n)
        w = fft.fftfreq(n) * 2 * np.pi
        wslice = slice(None, n//2, None)
        #return (np.abs(hw)**2)[wslice], w[wslice]
        return (np.abs(hw)**2) * 0.5/np.pi, w

    def _spddirect2(self, n):
        '''this looks bad, maybe with an fftshift
        '''
        #size = s1+s2-1
        hw = (fft.fft(np.r_[self.ma[::-1],self.ma], n)
                / fft.fft(np.r_[self.ar[::-1],self.ar], n))
        return (hw*hw.conj()) #.real[n//2-1:]

    def spdroots(self, w):
        '''spectral density for frequency using polynomial roots

        builds two arrays (number of roots, number of frequencies)
        '''
        return self.spdroots_(self.arroots, self.maroots, w)

    def spdroots_(self, arroots, maroots, w):
        '''spectral density for frequency using polynomial roots

        builds two arrays (number of roots, number of frequencies)

        Parameters
        ----------
        arroots : ndarray
            roots of ar (denominator) lag-polynomial
        maroots : ndarray
            roots of ma (numerator) lag-polynomial
        w : array_like
            frequencies for which spd is calculated

        Notes
        -----
        this should go into a function
        '''
        w = np.atleast_2d(w).T
        cosw = np.cos(w)
        #Greene 5th edt. p626, section 20.2.7.a.
        maroots = 1./maroots
        arroots = 1./arroots
        num = 1 + maroots**2 - 2* maroots * cosw
        den = 1 + arroots**2 - 2* arroots * cosw
        #print 'num.shape, den.shape', num.shape, den.shape
        hw = 0.5 / np.pi * num.prod(-1) / den.prod(-1) #or use expsumlog
        return np.squeeze(hw), w.squeeze()

    def spdpoly(self, w, nma=50):
        '''spectral density from MA polynomial representation for ARMA process

        References
        ----------
        Cochrane, section 8.3.3
        '''
        mpoly = np.polynomial.Polynomial(self.arma2ma(nma))
        hw = mpoly(np.exp(1j * w))
        spd = np.real_if_close(hw * hw.conj() * 0.5/np.pi)
        return spd, w

    def filter(self, x):
        '''
        filter a timeseries with the ARMA filter

        padding with zero is missing, in example I needed the padding to get
        initial conditions identical to direct filter

        Initial filtered observations differ from filter2 and signal.lfilter, but
        at end they are the same.

        See Also
        --------
        tsa.filters.fftconvolve

        '''
        n = x.shape[0]
        if n == self.fftarma:
            fftarma = self.fftarma
        else:
            fftarma = self.fftma(n) / self.fftar(n)
        tmpfft = fftarma * fft.fft(x)
        return fft.ifft(tmpfft)

    def filter2(self, x, pad=0):
        '''filter a time series using fftconvolve3 with ARMA filter

        padding of x currently works only if x is 1d
        in example it produces same observations at beginning as lfilter even
        without padding.

        TODO: this returns 1 additional observation at the end
        '''
        from statsmodels.tsa.filters import fftconvolve3
        if not pad:
            pass
        elif pad == 'auto':
            #just guessing how much padding
            x = self.padarr(x, x.shape[0] + 2*(self.nma+self.nar), atend=False)
        else:
            x = self.padarr(x, x.shape[0] + int(pad), atend=False)

        return fftconvolve3(x, self.ma, self.ar)

    def acf2spdfreq(self, acovf, nfreq=100, w=None):
        '''
        not really a method
        just for comparison, not efficient for large n or long acf

        this is also similarly use in tsa.stattools.periodogram with window
        '''
        if w is None:
            w = np.linspace(0, np.pi, nfreq)[:, None]
        nac = len(acovf)
        hw = 0.5 / np.pi * (acovf[0] +
                            2 * (acovf[1:] * np.cos(w*np.arange(1,nac))).sum(1))
        return hw

    def invpowerspd(self, n):
        '''autocovariance from spectral density

        scaling is correct, but n needs to be large for numerical accuracy
        maybe padding with zero in fft would be faster
        without slicing it returns 2-sided autocovariance with fftshift

        >>> ArmaFft([1, -0.5], [1., 0.4], 40).invpowerspd(2**8)[:10]
        array([ 2.08    ,  1.44    ,  0.72    ,  0.36    ,  0.18    ,  0.09    ,
                0.045   ,  0.0225  ,  0.01125 ,  0.005625])
        >>> ArmaFft([1, -0.5], [1., 0.4], 40).acovf(10)
        array([ 2.08    ,  1.44    ,  0.72    ,  0.36    ,  0.18    ,  0.09    ,
                0.045   ,  0.0225  ,  0.01125 ,  0.005625])
        '''
        hw = self.fftarma(n)
        return np.real_if_close(fft.ifft(hw*hw.conj()), tol=200)[:n]

    def spdmapoly(self, w, twosided=False):
        '''ma only, need division for ar, use LagPolynomial
        '''
        if w is None:
            w = np.linspace(0, np.pi, nfreq)
        return 0.5 / np.pi * self.mapoly(np.exp(w*1j))

    def plot4(self, fig=None, nobs=100, nacf=20, nfreq=100):
        rvs = self.generate_sample(nsample=100, burnin=500)
        acf = self.acf(nacf)[:nacf]  #TODO: check return length
        pacf = self.pacf(nacf)
        w = np.linspace(0, np.pi, nfreq)
        spdr, wr = self.spdroots(w)

        if fig is None:
            import matplotlib.pyplot as plt
            fig = plt.figure()
        ax = fig.add_subplot(2,2,1)
        ax.plot(rvs)
        ax.set_title('Random Sample \nar=%s, ma=%s' % (self.ar, self.ma))

        ax = fig.add_subplot(2,2,2)
        ax.plot(acf)
        ax.set_title('Autocorrelation \nar=%s, ma=%rs' % (self.ar, self.ma))

        ax = fig.add_subplot(2,2,3)
        ax.plot(wr, spdr)
        ax.set_title('Power Spectrum \nar=%s, ma=%s' % (self.ar, self.ma))

        ax = fig.add_subplot(2,2,4)
        ax.plot(pacf)
        ax.set_title('Partial Autocorrelation \nar=%s, ma=%s' % (self.ar, self.ma))

        return fig

def spdar1(ar, w):
    if np.ndim(ar) == 0:
        rho = ar
    else:
        rho = -ar[1]
    return 0.5 / np.pi /(1 + rho*rho - 2 * rho * np.cos(w))

if __name__ == '__main__':
    def maxabs(x,y):
        return np.max(np.abs(x-y))
    nobs = 200  #10000
    ar = [1, 0.0]
    ma = [1, 0.0]
    ar2 = np.zeros(nobs)
    ar2[:2] = [1, -0.9]

    uni = np.zeros(nobs)
    uni[0]=1.
    #arrep = signal.lfilter(ma, ar, ar2)
    #marep = signal.lfilter([1],arrep, uni)
    # same faster:
    arcomb = np.convolve(ar, ar2, mode='same')
    marep = signal.lfilter(ma,arcomb, uni) #[len(ma):]
    print(marep[:10])
    mafr = fft.fft(marep)

    rvs = np.random.normal(size=nobs)
    datafr = fft.fft(rvs)
    y = fft.ifft(mafr*datafr)
    print(np.corrcoef(np.c_[y[2:], y[1:-1], y[:-2]],rowvar=0))

    arrep = signal.lfilter([1],marep, uni)
    print(arrep[:20])  # roundtrip to ar
    arfr = fft.fft(arrep)
    yfr = fft.fft(y)
    x = fft.ifft(arfr*yfr).real  #imag part is e-15
    # the next two are equal, roundtrip works
    print(x[:5])
    print(rvs[:5])
    print(np.corrcoef(np.c_[x[2:], x[1:-1], x[:-2]],rowvar=0))

    # ARMA filter using fft with ratio of fft of ma/ar lag polynomial
    # seems much faster than using lfilter

    #padding, note arcomb is already full length
    arcombp = np.zeros(nobs)
    arcombp[:len(arcomb)] = arcomb
    map_ = np.zeros(nobs)    #rename: map was shadowing builtin
    map_[:len(ma)] = ma
    ar0fr = fft.fft(arcombp)
    ma0fr = fft.fft(map_)
    y2 = fft.ifft(ma0fr/ar0fr*datafr)
    #the next two are (almost) equal in real part, almost zero but different in imag
    print(y2[:10])
    print(y[:10])
    print(maxabs(y, y2))  # from chfdiscrete
    #1.1282071239631782e-014

    ar = [1, -0.4]
    ma = [1, 0.2]

    arma1 = ArmaFft([1, -0.5,0,0,0,00, -0.7, 0.3], [1, 0.8], nobs)

    nfreq = nobs
    w = np.linspace(0, np.pi, nfreq)
    w2 = np.linspace(0, 2*np.pi, nfreq)

    import matplotlib.pyplot as plt
    plt.close('all')

    plt.figure()
    spd1, w1 = arma1.spd(2**10)
    print(spd1.shape)
    _ = plt.plot(spd1)
    plt.title('spd fft complex')

    plt.figure()
    spd2, w2 = arma1.spdshift(2**10)
    print(spd2.shape)
    _ = plt.plot(w2, spd2)
    plt.title('spd fft shift')

    plt.figure()
    spd3, w3 = arma1.spddirect(2**10)
    print(spd3.shape)
    _ = plt.plot(w3, spd3)
    plt.title('spd fft direct')

    plt.figure()
    spd3b = arma1._spddirect2(2**10)
    print(spd3b.shape)
    _ = plt.plot(spd3b)
    plt.title('spd fft direct mirrored')

    plt.figure()
    spdr, wr = arma1.spdroots(w)
    print(spdr.shape)
    plt.plot(w, spdr)
    plt.title('spd from roots')

    plt.figure()
    spdar1_ = spdar1(arma1.ar, w)
    print(spdar1_.shape)
    _ = plt.plot(w, spdar1_)
    plt.title('spd ar1')

    plt.figure()
    wper, spdper = arma1.periodogram(nfreq)
    print(spdper.shape)
    _ = plt.plot(w, spdper)
    plt.title('periodogram')

    startup = 1000
    rvs = arma1.generate_sample(startup+10000)[startup:]
    import matplotlib.mlab as mlb
    plt.figure()
    sdm, wm = mlb.psd(x)
    print('sdm.shape', sdm.shape)
    sdm = sdm.ravel()
    plt.plot(wm, sdm)
    plt.title('matplotlib')

    from nitime.algorithms import LD_AR_est
    #yule_AR_est(s, order, Nfreqs)
    wnt, spdnt = LD_AR_est(rvs, 10, 512)
    plt.figure()
    print('spdnt.shape', spdnt.shape)
    _ = plt.plot(spdnt.ravel())
    print(spdnt[:10])
    plt.title('nitime')

    fig = plt.figure()
    arma1.plot4(fig)

    #plt.show()


import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

try:
    import plotly.graph_objects as go
except ImportError:
    pass

def __calculate_plot_data(stvariogram, **kwargs):
    xx, yy = stvariogram.meshbins
    z = stvariogram.experimental
#    x = xx.flatten()
#    y = yy.flatten()

    # apply the model
    nx = kwargs.get('x_resolution', 100)
    nt = kwargs.get('t_resolution', 100)

    # model spacing
    _xx, _yy = np.mgrid[
        0:np.nanmax(stvariogram.xbins):nx * 1j,
        0:np.nanmax(stvariogram.tbins):nt * 1j
    ]
    model = stvariogram.fitted_model
    lags = np.vstack((_xx.flatten(), _yy.flatten())).T
    # apply the model
    _z = model(lags)

    return xx.T, yy.T, z, _xx, _yy, _z

def matplotlib_plot_3d(stvariogram, kind='scatter', ax=None, elev=30, azim=220, **kwargs):
    # get the data, spanned over a bin meshgrid
    xx, yy, z, _xx, _yy, _z = __calculate_plot_data(stvariogram, **kwargs)
    x = xx.flatten()
    y = yy.flatten()

    # some settings
    c = kwargs.get('color', kwargs.get('c', 'b'))
    cmap = kwargs.get('model_color', kwargs.get('cmap', 'terrain'))
    alpha = kwargs.get('alpha', 0.8)
    depthshade = kwargs.get('depthshade', False)

    # handle the axes
    if ax is not None:
        if not isinstance(ax, Axes3D):
            raise ValueError('The passed ax object is not an instance of mpl_toolkis.mplot3d.Axes3D.')
        fig = ax.get_figure()
    else:
        fig = plt.figure(figsize=kwargs.get('figsize', (10, 10)))
        ax = fig.add_subplot(111, projection='3d')

    # do the plot
    ax.view_init(elev=elev, azim=azim)
    if kind == 'surf':
        ax.plot_trisurf(x, y, z, color=c, alpha=alpha)
    elif kind == 'scatter':
        ax.scatter(x, y, z, c=c, depthshade=depthshade)
    else:
        raise ValueError('%s is not a valid 3D plot' % kind)

    # add the model
    if not kwargs.get('no_model', False):
        ax.plot_trisurf(_xx.flatten(), _yy.flatten(), _z, cmap=cmap, alpha=alpha)

    # labels:
    ax.set_xlabel('space')
    ax.set_ylabel('time')
    ax.set_zlabel('semivariance [%s]' % stvariogram.estimator.__name__)

    # return
    return fig

def plotly_plot_3d(stvariogram, kind='scatter', fig=None, **kwargs):
    # get the data spanned over a bin meshgrid
    xx, yy, z, _xx, _yy, _z = __calculate_plot_data(stvariogram, **kwargs)

    # get some settings
    c = kwargs.get('color', kwargs.get('c', 'black'))
    cmap = kwargs.get('model_color', kwargs.get('colorscale', kwargs.get('cmap', 'Electric')))
    alpha = kwargs.get('opacity', kwargs.get('alpha', 0.6))

    # handle the figue
    if fig is None:
        fig = go.Figure()

    # do the plot
    if kind == 'surf':
        fig.add_trace(
            go.Surface(
                x=xx,
                y=yy,
                z=z.reshape(xx.shape),
                opacity=0.8 * alpha,
                colorscale=[[0, c], [1, c]],
                name='experimental variogram'
            )
        )
    elif kind == 'scatter' or kwargs.get('add_points', False):
        fig.add_trace(
            go.Scatter3d(
                x=xx.flatten(),
                y=yy.flatten(),
                z=z,
                mode='markers',
                opacity=alpha,
                marker=dict(color=c, size=kwargs.get('size', 4)),
                name='experimental variogram'
            )
        )

    # add the model
    if not kwargs.get('no_model', False):
        fig.add_trace(
            go.Surface(
                x=_xx,
                y=_yy,
                z=_z.reshape(_xx.shape),
                opacity=max(1, alpha * 1.2),
                colorscale=cmap,
                name='%s model' % stvariogram.model.__name__
            )
        )

    # set some labels
    fig.update_layout(scene=dict(
        xaxis_title='space',
        yaxis_title='time',
        zaxis_title='semivariance [%s]' % stvariogram.estimator.__name__
    ))

    # return
    return fig


'''

Test the data.py module

'''
import numpy as np
import matplotlib.pyplot as plt

import util
import data as Data
# --- Halotools ---
from halotools.empirical_models import PrebuiltHodModelFactory

from ChangTools.plotting import prettyplot
from ChangTools.plotting import prettycolors

def PlotCovariance(obvs, Mr=21, b_normal=0.25, inference='mcmc'):
    ''' Plot the covariance matrix for a specified obvs
    '''
    # import the covariance matrix
    covar = Data.data_cov(Mr=Mr, b_normal=b_normal, inference=inference)

    if obvs == 'xi':
        obvs_cov = covar[1:16 , 1:16]
        r_bin = Data.xi_binedges()
    elif obvs == 'gmf':
        obvs_cov = covar[17:, 17:]
        binedges = Data.data_gmf_bins()
        r_bin = 0.5 * (binedges[:-1] + binedges[1:])

    n_bin = int(np.sqrt(obvs_cov.size))

    # calculate the reduced covariance for plotting
    red_covar = np.zeros([n_bin, n_bin])
    for ii in range(n_bin):
        for jj in range(n_bin):
            red_covar[ii][jj] = obvs_cov[ii][jj]/np.sqrt(obvs_cov[ii][ii] * obvs_cov[jj][jj])

    prettyplot()
    fig = plt.figure()
    sub = fig.add_subplot(111)
    cont = sub.pcolormesh(r_bin, r_bin, red_covar, cmap=plt.cm.afmhot_r)
    plt.colorbar(cont)

    sub.set_xlim([r_bin[0], r_bin[-1]])
    sub.set_ylim([r_bin[0], r_bin[-1]])
    sub.set_xscale('log')
    sub.set_yscale('log')

    sub.set_xlabel(r'$\mathtt{r}\;[\mathtt{Mpc/h}$]', fontsize=25)
    sub.set_ylabel(r'$\mathtt{r}\;[\mathtt{Mpc/h}$]', fontsize=25)
    fig_file = ''.join([util.fig_dir(),
        obvs.upper(), 'covariance',
        '.Mr', str(Mr),
        '.bnorm', str(round(b_normal,2)),
        '.', inference, '_inf.png'])
    fig.savefig(fig_file, bbox_inches='tight')
    plt.close()
    return None

# ---- Plotting ----
def xi(Mr=20, Nmock=500):
    '''
    Plot xi(r) of the fake observations
    '''
    prettyplot()
    pretty_colors = prettycolors()

    xir, cii = Data.data_xi(Mr=Mr, Nmock=Nmock)
    rbin = Data.data_xi_bins(Mr=Mr)

    fig = plt.figure(1)
    sub = fig.add_subplot(111)
    sub.plot(rbin, rbin*xir, c='k', lw=1)
    sub.errorbar(rbin, rbin*xir, yerr = rbin*cii**0.5 , fmt="ok", ms=1, capsize=2, alpha=1.)

    sub.set_xlim([0.1, 15])
    sub.set_ylim([1, 10])
    sub.set_yscale("log")
    sub.set_xscale("log")

    sub.set_xlabel(r'$\mathtt{r}\; (\mathtt{Mpc})$', fontsize=25)
    sub.set_ylabel(r'$\mathtt{r} \xi_{\rm gg}$', fontsize=25)

    fig_file = ''.join([util.fig_dir(),
        'xi.Mr', str(Mr), '.Nmock', str(Nmock), '.png'])
    fig.savefig(fig_file, bbox_inches='tight')
    plt.close()
    return None

def gmf(Mr=20, Nmock=500):
    '''
    Plot Group Multiplicty Function of fake observations
    '''
    prettyplot()
    pretty_colors = prettycolors()

    # import fake obs GMF
    gmf, sig_gmf = Data.data_gmf(Mr=Mr, Nmock=Nmock)
    # group richness bins
    gmf_bin = Data.data_gmf_bins()

    fig = plt.figure(1)
    sub = fig.add_subplot(111)

    sub.errorbar(
            0.5*(gmf_bin[:-1]+gmf_bin[1:]), gmf, yerr=sig_gmf,
            fmt="ok", capsize=1.0
            )
    sub.set_xlim([1, 60])
    sub.set_yscale('log')
    sub.set_ylabel(r"Group Multiplicity Function (h$^{3}$ Mpc$^{-3}$)", fontsize=20)
    sub.set_xlabel(r"$\mathtt{Group\;\;Richness}$", fontsize=20)
    # save to file
    fig_file = ''.join([util.fig_dir(),
        'gmf.Mr', str(Mr), '.Nmock', str(Nmock), '.png'])
    fig.savefig(fig_file, bbox_inches='tight')
    return None

# ---- tests -----
def xi_binning_tests(Mr=20):
    model = PrebuiltHodModelFactory('zheng07', threshold = -1.0*np.float(Mr))

    rbins = np.concatenate([np.array([0.1]), np.logspace(np.log10(0.5), np.log10(20.), 15)])
    print 'R bins = ', rbins
    for ii in xrange(10):
        model.populate_mock() # population mock realization

        #rbins = np.logspace(-1, np.log10(20.), 16)
        r_bin, xi_r = model.mock.compute_galaxy_clustering(rbins=rbins)
        print xi_r

def test_nbar(Mr=21, b_normal=0.25):
    print Data.data_nbar(Mr=Mr, b_normal=b_normal)

if __name__=='__main__':
    PlotCovariance('gmf', inference='mcmc')

    #test_nbar()
    #xi_cov(Mr=20, Nmock=500)
    #xi_binning_tests(Mr=20)


"""
Copyright (c) 2015 Andreea Georgescu

Created on Wed Mar  4 00:47:37 2015

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

# TODO! This only works for CDMSSi!

from __future__ import division
from __future__ import absolute_import
from __future__ import print_function
from experiment_HaloIndep import *
import interp_uniform as unif
# from interp import interp1d
from scipy import interpolate
from scipy.optimize import brentq, minimize
from basinhopping import *
import matplotlib.pyplot as plt
import os   # for speaking
import parallel_map as par

DEBUG = F
DEBUG_FULL = F
USE_BASINHOPPING = T
ADAPT_KWARGS = F
ALLOW_MOVE = T

class ConstraintsFunction(object):
    """ Class to implement the constraints function that will be passed as an argunent
    to the minimization routines.
    Input:
        args: Arguments needed for calculating the constraints:
            vminStar, logetaStar, vminStar_index
    """
    def __init__(self, *args):
        self.vminStar = args[0]
        self.logetaStar = args[1]
        self.vminStar_index = args[2]
        self.vmin_max = 2000

    def __call__(self, x, close=True):
        """
        Input:
            x: ndarray
        Returns:
            constraints: ndarray
                Constraints vector, where each value must be >= 0 for the
                constraint to be specified. Contains:
             0 -  8: bounds: 3 * (x.size/2) constraints = 9 for x.size/2 = 3
             9 - 12: sorted array: 2 * (x.size/2 - 1) constraints = 4 for x.size/2 = 3
            13 - 15: vminStar_index: x.size/2 constraints = 3 for x.size/2 = 3
            16 - 18: vminStar and logetaStar: x.size/2 constraints = 3 for x.size/2 = 3
        """
        constraints = np.concatenate([x[:x.size/2], self.vmin_max - x[:x.size/2], -x[x.size/2:],
                                      np.diff(x[:x.size/2]), np.diff(-x[x.size/2:]),
                                      (x[:x.size/2] - self.vminStar) * (-x[x.size/2:] + self.logetaStar),
                                      self.vminStar - x[:self.vminStar_index],
                                      x[self.vminStar_index: x.size/2] - self.vminStar,
                                      x[x.size/2: x.size/2 + self.vminStar_index] - self.logetaStar,
                                      self.logetaStar - x[x.size/2 + self.vminStar_index:]])
        if close:
            is_not_close = np.logical_not(np.isclose(constraints, np.zeros_like(constraints), atol=1e-5))
            is_not_close[:3 * (x.size/2)] = True
            constraints = np.where(is_not_close, constraints, np.abs(constraints))
        if np.any(np.isnan(constraints)):
            raise ValueError
        return constraints

class Experiment_EHI(Experiment_HaloIndep):
    """ Class implementing the extended maximum likelihood halo-independent (EHI)
    method to obtain the confidence band for experiments with potential signals and
    unbinned data (arXiv:1507.03902).
    Input:
        expername: string
            The name of the experiment.
        scattering_type: string
            The type of scattering. Can be
                - 'SI' (spin-independent)
                - 'SDAV' (spin-dependent, axial-vector)
                - 'SDPS' (spin-dependent, pseudo-scalar)
        mPhi: float, optional
            The mass of the mediator.
        method: str, optional
            Type of minimization solver to be passed as a parameter to the minimization
                routine. Can be 'SLSQP' or 'COBYLA'.
    """
    def __init__(self, expername, scattering_type, mPhi=mPhiRef, method='SLSQP'):
        super().__init__(expername, scattering_type, mPhi)
        module = import_file(INPUT_DIR + expername + ".py")
        self.ERecoilList = module.ERecoilList
        self.mu_BKG_i = module.mu_BKG_i
        self.NBKG = module.NBKG
        self.method = method

    def _VMinSortedList(self, mx, fp, fn, delta):
        """ Computes the list of vmin corresponsing to measured recoil energies,
        sorted in increasing order. Will be useful as starting guesses.
        """
        self.vmin_sorted_list = np.sort(VMin(self.ERecoilList, self.mT[0], mx, delta))
        return

    def ResponseTables(self, vmin_min, vmin_max, vmin_step, mx, fp, fn, delta,
                       output_file_tail):
        """ Computes response tables
            - self.diff_response_tab is a table of [vmin, DifferentialResponse(Eee_i)]
        pairs for each vmin in the range [vminmin, vminmax], corresponding to measured
        recoil energies Eee_i. It is a 3D matrix where
                axis = 0 has dimension self.ERecoilList.size()
                axis = 1 has dimension vmin_list.size() + 1 (where + 1 is because we
                    prepend zeros for vmin = 0)
                axis = 2 has dimension 2 for the pairs of [vmin, diff_response].
            - self.response_tab is a table of [vmin, Response] pairs for each vmin
        in the range [vminmin, vminmax], corresponding to DifferentialResponse
        integrated over the full energy range. It is a 2D matrix where
                axis = 1 has dimension vmin_list.size() + 1 (where +1 is because we
                    prepend zeros for vmin = 0)
                axis = 2 has dimension 2 for the pairs of [vmin, diff_response].
        Input:
            vmin_min, vmin_max, vmin_step: float
                Vmin range and vmin step size.
            mx, fp, fn, delta: float
            output_file_tail: string
                Tag to be added to the file name since the results for
                self.vmin_sorted_list, self.diff_response_tab and self.response_tab
                are each written to files.
        """
        self._VMinSortedList(mx, fp, fn, delta)
        file = output_file_tail + "_VminSortedList.dat"
        print(file)
        np.savetxt(file, self.vmin_sorted_list)

        if delta == 0:
            branches = [1]
        else:
            branches = [1, -1]
        self.vmin_linspace = np.linspace(vmin_min, vmin_max,
                                         (vmin_max - vmin_min)/vmin_step + 1)
        self.diff_response_tab = np.zeros((self.ERecoilList.size, 1))
        self.response_tab = np.zeros(1)
        self.curly_H_tab = np.zeros((self.ERecoilList.size, 1))
        self.xi_tab = np.zeros(1)
        xi = 0
        vmin_prev = 0
        for vmin in self.vmin_linspace:
            print("vmin =", vmin)
            diff_resp_list = np.zeros((1, len(self.ERecoilList)))
            resp = 0
            curly_H = np.zeros((1, len(self.ERecoilList)))

            for sign in branches:

                (ER, qER, const_factor) = self.ConstFactor(vmin, mx, fp, fn, delta, sign)
                v_delta = min(VminDelta(self.mT, mx, delta))
                diff_resp_list += np.array([self.DifferentialResponse(Eee, qER, const_factor)
                                            for Eee in self.ERecoilList])
                resp += integrate.quad(self.DifferentialResponse, self.Ethreshold, self.Emaximum,
                                       args=(qER, const_factor), epsrel=PRECISSION, epsabs=0)[0]

                curly_H += np.array([[integrate.quad(self.DifferentialResponse_Full, v_delta, vmin,
                                                     args=(Eee, mx, fp, fn, delta, sign),
                                                     epsrel=PRECISSION, epsabs=0)[0]
                                      for Eee in self.ERecoilList]])
            xi += self.Exposure * \
                self.IntegratedResponse(vmin_prev, vmin,
                                        self.Ethreshold, self.Emaximum,
                                        mx, fp, fn, delta)
            vmin_prev = vmin
            self.diff_response_tab = \
                np.append(self.diff_response_tab, diff_resp_list.transpose(), axis=1)
            self.response_tab = np.append(self.response_tab, [resp], axis=0)
            self.curly_H_tab = np.append(self.curly_H_tab, curly_H.transpose(), axis=1)
            # counts/kg/keVee
            self.xi_tab = np.append(self.xi_tab, [xi], axis=0)
            # counts * day
        self.vmin_linspace = np.insert(self.vmin_linspace, 0., 0)
        file = output_file_tail + "_VminLinspace.dat"
        print(file)
        np.savetxt(file, self.vmin_linspace)
        file = output_file_tail + "_DiffRespTable.dat"
        print(file)
        np.savetxt(file, self.diff_response_tab)
        file = output_file_tail + "_RespTable.dat"
        print(file)
        np.savetxt(file, self.response_tab)
        file = output_file_tail + "_CurlyHTable.dat"
        print(file)
        np.savetxt(file, self.curly_H_tab)
        file = output_file_tail + "_XiTable.dat"
        print(file)
        np.savetxt(file, self.xi_tab)
        os.system("say Finished response tables.")
        return

    def PlotTable(self, func, dimension=0, xlim=None, ylim=None,
                  title=None, plot_close=True, plot_show=True, show_zero_axis=False):
        """ Plots response tables.
        Input:
            func: callable
                Function or list of functions of v that should be plotted.
            dimension: int
                0 (if there's only one function) or
                1 (if there are a list of functions).
            xlim, ylim: float
                Axis limits for the plots.
            title: string
                Plot title.
            plot_close, plot_show: bool
                Whether to call plt.close() before and plt.show() after.
            show_zero_axis: bool
                Whether to show a horizontal line at zero.
        """
        if plot_close:
            plt.close()
        if dimension == 0:
            # only one function
            plt.plot(self.vmin_linspace, np.array([func(v)
                     for v in self.vmin_linspace]))
        elif dimension == 1:
            # list of interpolated functions for each energy in self.ERecoilList
            for i in range(self.ERecoilList.size):
                plt.plot(self.vmin_linspace, np.array([func[i](v)
                         for v in self.vmin_linspace]))
        else:
            print("Wrong dimension")
            raise TypeError
        if show_zero_axis:
            plt.plot(self.vmin_linspace, np.zeros(self.vmin_linspace.size))
        if xlim is not None:
            plt.xlim(xlim)
        if ylim is not None:
            plt.ylim(ylim)
        if title is not None:
            plt.title(title)
        if plot_show:
            plt.show()

    def ImportResponseTables(self, output_file_tail, plot=True):
        """ Imports the data for the response tables from files.
        """
        file = output_file_tail + "_VminSortedList.dat"
        with open(file, 'r') as f_handle:
            self.vmin_sorted_list = np.loadtxt(f_handle)
        file = output_file_tail + "_VminLinspace.dat"
        with open(file, 'r') as f_handle:
            self.vmin_linspace = np.loadtxt(f_handle)
        file = output_file_tail + "_DiffRespTable.dat"
        with open(file, 'r') as f_handle:
            self.diff_response_tab = np.loadtxt(f_handle)
        file = output_file_tail + "_RespTable.dat"
        with open(file, 'r') as f_handle:
            self.response_tab = np.loadtxt(f_handle)
        file = output_file_tail + "_CurlyHTable.dat"
        with open(file, 'r') as f_handle:
            self.curly_H_tab = np.loadtxt(f_handle)
        file = output_file_tail + "_XiTable.dat"
        with open(file, 'r') as f_handle:
            self.xi_tab = np.loadtxt(f_handle)
        self.diff_response_interp = np.array([unif.interp1d(self.vmin_linspace, dr)
                                              for dr in self.diff_response_tab])
        self.response_interp = unif.interp1d(self.vmin_linspace, self.response_tab)
        self.curly_H_interp = np.array([unif.interp1d(self.vmin_linspace, h)
                                        for h in self.curly_H_tab])

        if plot:
            self.PlotTable(self.diff_response_interp, dimension=1)
            self.PlotTable(self.response_interp, dimension=0)
            self.PlotTable(self.curly_H_interp, dimension=1, title='Curly H')
        return

    def VminIntegratedResponseTable(self, vmin_list):
        return np.array([[integrate.quad(self.diff_response_interp[i],
                                         vmin_list[a], vmin_list[a + 1],
                                         epsrel=PRECISSION, epsabs=0)[0]
                        for a in range(vmin_list.size - 1)]
                        for i in range(self.ERecoilList.size)])

    def IntegratedResponseTable(self, vmin_list):
        return np.array([integrate.quad(self.response_interp,
                                        vmin_list[a], vmin_list[a + 1],
                                        epsrel=PRECISSION, epsabs=0)[0]
                        for a in range(vmin_list.size - 1)])

    def _MinusLogLikelihood(self, vars_list, vminStar=None, logetaStar=None,
                            vminStar_index=None):
        """ Compute -log(L)
        Input:
            vars_list: ndarray
                List of variables [vmin_1, ..., vmin_No, log(eta_1), ..., log(eta_No)]
            vminStar, logetaStar: float, optional
                Values of fixed vmin^* and log(eta)^*.
        Returns:
            -log(L): float
        """
        if vminStar is None:
            vmin_list_w0 = vars_list[: vars_list.size/2]
            logeta_list = vars_list[vars_list.size/2:]
        else:
            vmin_list_w0 = np.insert(vars_list[: vars_list.size/2],
                                     vminStar_index, vminStar)
            logeta_list = np.insert(vars_list[vars_list.size/2:],
                                    vminStar_index, logetaStar)
        vmin_list_w0 = np.insert(vmin_list_w0, 0, 0)
        vmin_resp_integr = self.VminIntegratedResponseTable(vmin_list_w0)
        resp_integr = self.IntegratedResponseTable(vmin_list_w0)
        mu_i = self.Exposure * np.dot(vmin_resp_integr, 10**logeta_list)
        Nsignal = self.Exposure * np.dot(10**logeta_list, resp_integr)
        if vminStar is None:
            self.gamma_i = (self.mu_BKG_i + mu_i) / self.Exposure
            # counts/kg/keVee/days
        result = self.NBKG + Nsignal - np.log(self.mu_BKG_i + mu_i).sum()
        if np.any(self.mu_BKG_i + mu_i < 0):
            raise ValueError
        return result

    def MinusLogLikelihood(self, vars_list, constr_func=None, vminStar=None,
                           logetaStar=None, vminStar_index=None):
        """ Computes -log(L) and tests whether constraints are satisfied.
        Input:
            vars_list: ndarray
                List of variables [vmin_1, ..., vmin_No, log(eta_1), ..., log(eta_No)].
            constr_func: callable, optional
                Ffunction of vars_list giving an array of values each corresponding to
                a constraint. If the values are > 0 the constraints are satisfied.
            vminStar, logetaStar: float, optional
                Values of fixed vmin^* and log(eta)^*.
            vminStar_index: int, optional
                Index corresponding to the position of vminStar in the array of vmin
                steps.
        Returns:
            -log(L) if all constraints are valid, and the result of an artificial
                function that grows with the invalid constraints if not all constraints
                are valid.
        """
        constraints = constr_func(vars_list)
        constr_not_valid = constraints < 0
        if DEBUG_FULL:
            print("*** vars_list =", repr(vars_list))
        if DEBUG_FULL:
            print("vminStar =", vminStar)
            print("logetaStar =", logetaStar)
            print("constraints =", repr(constraints))
            print("constr_not_valid =", repr(constr_not_valid))
        try:
            return self._MinusLogLikelihood(vars_list, vminStar=vminStar,
                                            logetaStar=logetaStar,
                                            vminStar_index=vminStar_index)
        except:
            if np.any(constr_not_valid):
                constr_list = constraints[constr_not_valid]
                if DEBUG_FULL:
                    print("Constraints not valid!!")
                    print("constr sum =", -constr_list.sum())
                return min(max(-constr_list.sum(), 0.001) * 1e6, 1e6)
            else:
                print("Error!!")
                raise

    def OptimalLikelihood(self, output_file_tail, logeta_guess):
        """ Finds the best-fit piecewise constant eta function corresponding to the
        minimum MinusLogLikelihood, and prints the results to file (value of the minimum
        MinusLogLikelihood and the corresponding values of vmin, logeta steps.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            logeta_guess: float
                Guess for the value of log(eta) in the minimization procedure.
        """
        self.ImportResponseTables(output_file_tail, plot=False)
        vars_guess = np.append(self.vmin_sorted_list,
                               logeta_guess * np.ones(self.vmin_sorted_list.size))
        print("vars_guess =", vars_guess)
        vmin_max = self.vmin_linspace[-1]

        def constr_func(x, vmin_max=vmin_max):
            """ 0 -  8: bounds: 3 * (x.size/2) constraints = 9 for x.size/2 = 3
                9 - 12: sorted array: 2 * (x.size/2 - 1) constraints = 4 for x.size/2 = 3
            """
            constraints = np.concatenate([x[:x.size/2], vmin_max - x[:x.size/2],
                                          -x[x.size/2:],
                                          np.diff(x[:x.size/2]), np.diff(-x[x.size/2:])])
            is_not_close = np.logical_not(
                np.isclose(constraints, np.zeros_like(constraints), atol=1e-5))
            is_not_close[:3 * (x.size/2)] = T
            constr = np.where(is_not_close, constraints, np.abs(constraints))
            if DEBUG:
                print("***constr =", repr(constr))
                print("tf =", repr(constr < 0))
            return constr
        constr = ({'type': 'ineq', 'fun': constr_func})

        np.random.seed(0)
        if USE_BASINHOPPING:
            minimizer_kwargs = {"constraints": constr, "args": (constr_func,)}
            optimum_log_likelihood = basinhopping(self.MinusLogLikelihood, vars_guess,
                                                  minimizer_kwargs=minimizer_kwargs,
                                                  niter=30, stepsize=0.1)
        else:
            optimum_log_likelihood = minimize(self.MinusLogLikelihood, vars_guess,
                                              args=(constr_func,), constraints=constr)

        print(optimum_log_likelihood)
        print("MinusLogLikelihood =", self._MinusLogLikelihood(optimum_log_likelihood.x))
        print("vars_guess =", repr(vars_guess))
        file = output_file_tail + "_GloballyOptimalLikelihood.dat"
        print(file)
        np.savetxt(file, np.append([optimum_log_likelihood.fun],
                                   optimum_log_likelihood.x))
        os.system("say 'Finished finding optimum'")
        return

    def ImportOptimalLikelihood(self, output_file_tail, plot=False):
        """ Import the minumum -log(L) and the locations of the steps in the best-fit
        logeta function.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            plot: bool, optional
                Whether to plot response tables.
        """
        self.ImportResponseTables(output_file_tail, plot=False)
        file = output_file_tail + "_GloballyOptimalLikelihood.dat"
        with open(file, 'r') as f_handle:
            optimal_result = np.loadtxt(f_handle)
        self.optimal_logL = optimal_result[0]
        self.optimal_vmin = optimal_result[1: optimal_result.size/2 + 1]
        self.optimal_logeta = optimal_result[optimal_result.size/2 + 1:]
        print("optimal result =", optimal_result)

        if plot:
            self._MinusLogLikelihood(optimal_result[1:])  # to get self.gamma_i
            self.xi_interp = unif.interp1d(self.vmin_linspace, self.xi_tab)
            self.h_sum_tab = np.sum([self.curly_H_tab[i] / self.gamma_i[i]
                                     for i in range(self.optimal_vmin.size)], axis=0)
            self.q_tab = 2 * (self.xi_tab - self.h_sum_tab)
            self.h_sum_interp = unif.interp1d(self.vmin_linspace, self.h_sum_tab)
            self.q_interp = unif.interp1d(self.vmin_linspace, self.q_tab)

            file = output_file_tail + "_HSumTable.dat"
            print(file)
            np.savetxt(file, self.h_sum_tab)
            file = output_file_tail + "_QTable.dat"
            print(file)
            np.savetxt(file, self.q_tab)

            self.PlotTable(self.xi_interp, dimension=0, plot_show=False)
            self.PlotTable(self.h_sum_interp, dimension=0,
                           xlim=[0, 2000], ylim=[-2e24, 2e24],
                           title='Xi, H_sum', plot_close=False)
            self.PlotTable(self.q_interp, dimension=0,
                           xlim=[0, 2000], ylim=[-2e24, 2e24],
                           title='q', show_zero_axis=True)
        return

    def _PlotStepFunction(self, vmin_list, logeta_list,
                          xlim_percentage=(0., 1.1), ylim_percentage=(1.01, 0.99),
                          mark=None, color=None, linewidth=1,
                          plot_close=True, plot_show=True):
        """ Plots a step-like function, given the location of the steps.
        """
        if plot_close:
            plt.close()
        print(vmin_list)
        print(logeta_list)
        x = np.append(np.insert(vmin_list, 0, 0), vmin_list[-1] + 0.1)
        y = np.append(np.insert(logeta_list, 0, logeta_list[0]), -80)
        if color is not None:
            plt.step(x, y, color=color, linewidth=linewidth)
            if mark is not None:
                plt.plot(x, y, mark, color=color)
        else:
            plt.step(x, y, linewidth=linewidth)
            if mark is not None:
                plt.plot(x, y, mark)
#        plt.xlim([vmin_list[0] * xlim_percentage[0], vmin_list[-1] * xlim_percentage[1]])
        plt.xlim([0, 1000])
        plt.ylim([max(logeta_list[-1] * ylim_percentage[0], -60),
                  max(logeta_list[0] * ylim_percentage[1], -35)])
        if plot_show:
            plt.show()
        return

    def PlotOptimum(self, xlim_percentage=(0., 1.1), ylim_percentage=(1.01, 0.99),
                    color='red', linewidth=1,
                    plot_close=True, plot_show=True):
        """ Plots the best-fit eta(vmin) step function.
        """
        self._PlotStepFunction(self.optimal_vmin, self.optimal_logeta,
                               xlim_percentage=xlim_percentage,
                               ylim_percentage=ylim_percentage,
                               color=color, linewidth=linewidth,
                               plot_close=plot_close, plot_show=plot_show)
        return

    def PlotConstrainedOptimum(self, vminStar, logetaStar, vminStar_index,
                               xlim_percentage=(0., 1.1), ylim_percentage=(1.01, 0.99),
                               plot_close=True, plot_show=True):
        """ Plots the eta(vmin) function given the location of vminStar and logetaStar.
        """
        self._PlotStepFunction(self.optimal_vmin, self.optimal_logeta,
                               plot_close=plot_close, plot_show=False)
        x = np.insert(self.constr_optimal_vmin, vminStar_index, vminStar)
        y = np.insert(self.constr_optimal_logeta, vminStar_index, logetaStar)
        self._PlotStepFunction(x, y,
                               xlim_percentage=xlim_percentage,
                               ylim_percentage=ylim_percentage,
                               plot_close=False, plot_show=False, mark='x', color='k')
        plt.plot(vminStar, logetaStar, '*')
        if plot_show:
            plt.show()
        return

    def _ConstrainedOptimalLikelihood(self, vminStar, logetaStar, vminStar_index):
        """ Finds the constrained minimum MinusLogLikelihood for given vminStar,
        logetaStar and vminStar_index.
        Input:
            vminStar, logetaStar: float
                Location of the constrained step.
            vminStar_index: int
                Index of vminStar in the list of vmin steps of the constrained optimum
                logeta function.
        Returns:
            constr_optimal_logl: float
                The constrained minimum MinusLogLikelihood
        """
        if DEBUG:
            print("~~~~~ vminStar_index =", vminStar_index)
        vmin_guess_left = np.array([self.optimal_vmin[ind]
                                    if self.optimal_vmin[ind] < vminStar
                                    else vminStar * (1 - 0.001*(vminStar_index - ind))
                                    for ind in range(vminStar_index)])
        vmin_guess_right = np.array([self.optimal_vmin[ind]
                                    if self.optimal_vmin[ind] > vminStar
                                    else vminStar * (1 + 0.001*(ind - vminStar_index - 1))
                                    for ind in range(vminStar_index, self.optimal_vmin.size)])
        vmin_guess = np.append(vmin_guess_left, vmin_guess_right)
        logeta_guess = self.optimal_logeta
        logeta_guess_left = np.maximum(logeta_guess[:vminStar_index],
                                       np.ones(vminStar_index)*logetaStar)
        logeta_guess_right = np.minimum(logeta_guess[vminStar_index:],
                                        np.ones(logeta_guess.size - vminStar_index) *
                                        logetaStar)
        logeta_guess = np.append(logeta_guess_left, logeta_guess_right)
        vars_guess = np.append(vmin_guess, logeta_guess)

        constr_func = ConstraintsFunction(vminStar, logetaStar, vminStar_index)
        constr = ({'type': 'ineq', 'fun': constr_func})
        args = (constr_func, vminStar, logetaStar, vminStar_index)

        sol_not_found = True
        attempts = 3
        np.random.seed(1)
        random_variation = 1e-5

        if USE_BASINHOPPING:
            class TakeStep(object):
                def __init__(self, stepsize=0.1):
                    pass
                    self.stepsize = stepsize

                def __call__(self, x):
                    x[:x.size/2] += np.random.uniform(-5. * self.stepsize,
                                                      5. * self.stepsize,
                                                      x[x.size/2:].shape)
                    x[x.size/2:] += np.random.uniform(-self.stepsize,
                                                      self.stepsize, x[x.size/2:].shape)
                    return x
            take_step = TakeStep()

            class AdaptiveKwargs(object):
                def __init__(self, kwargs, random_variation=random_variation):
                    self.kwargs = kwargs
                    self.random_variation = random_variation

                def __call__(self):
                    new_kwargs = {}
                    random_factor_vminStar = \
                        (1 + self.random_variation * np.random.uniform(-1, 1))
                    random_factor_logetaStar = \
                        (1 + self.random_variation * np.random.uniform(-1, 1))
                    constr_func_args = (self.kwargs['args'][1] * random_factor_vminStar,
                                        self.kwargs['args'][2] * random_factor_logetaStar,
                                        self.kwargs['args'][3])
                    constr_func = ConstraintsFunction(*constr_func_args)
                    new_kwargs['args'] = (constr_func,) + constr_func_args
                    new_kwargs['constraints'] = ({'type': 'ineq', 'fun': constr_func})
                    if 'method' in self.kwargs:
                        new_kwargs['method'] = self.kwargs['method']
                    return new_kwargs

            minimizer_kwargs = {"constraints": constr, "args": args, "method": self.method}
            if ADAPT_KWARGS:
                adapt_kwargs = AdaptiveKwargs(minimizer_kwargs, random_variation)
            else:
                adapt_kwargs = None

        while sol_not_found and attempts > 0:
            try:
                if USE_BASINHOPPING:
                    constr_optimum_log_likelihood = \
                        basinhopping(self.MinusLogLikelihood, vars_guess,
                                     minimizer_kwargs=minimizer_kwargs, niter=5,
                                     take_step=take_step, adapt_kwargs=adapt_kwargs,
                                     stepsize=0.2)
                else:
                    constr_optimum_log_likelihood = \
                        minimize(self.MinusLogLikelihood, vars_guess,
                                 args=args, constraints=constr, method=self.method)
                constraints = constr_func(constr_optimum_log_likelihood.x)
                is_not_close = np.logical_not(np.isclose(constraints,
                                                         np.zeros_like(constraints)))
                constr_not_valid = np.logical_and(constraints < 0, is_not_close)
                sol_not_found = np.any(constr_not_valid)
            except ValueError:
                sol_not_found = True
                pass

            attempts -= 1
            args = (constr_func,
                    vminStar * (1 + random_variation * np.random.uniform(-1, 1)),
                    logetaStar * (1 + random_variation * np.random.uniform(-1, 1)),
                    vminStar_index)
            if USE_BASINHOPPING:
                minimizer_kwargs = {"constraints": constr, "args": args}

            if DEBUG and sol_not_found:
                print(attempts, "attempts left! ####################################" +
                      "################################################################")
                print("sol_not_found =", sol_not_found)
        if sol_not_found:
            if DEBUG:
                print("ValueError: sol not found")
            raise ValueError

        if DEBUG:
            print(constr_optimum_log_likelihood)
            print("kwargs =", constr_optimum_log_likelihood.minimizer.kwargs)
            print("args =", constr_optimum_log_likelihood.minimizer.kwargs['args'])
            print("optimum_logL =", self.optimal_logL)
            print("constraints=", repr(constraints))
            print("constr_not_valid =", repr(constr_not_valid))
            print("vars_guess =", repr(vars_guess))
            print("optimum_logL =", self.optimal_logL)
            print("vminStar_index =", vminStar_index)

        return constr_optimum_log_likelihood

    def ConstrainedOptimalLikelihood(self, vminStar, logetaStar, plot=False):
        """ Finds the constrained minimum MinusLogLikelihood for given vminStar,
        logetaStar. Finds the minimum for all vminStar_index, and picks the best one.
        Input:
            vminStar, logetaStar: float
                Location of constrained step.
            plot: bool, optional
                Whether to plot the constrained piecewice-constant logeta function.
        Returns:
            constr_optimal_logl: float
                The constrained minimum MinusLogLikelihood
        """
        vminStar_index = 0
        while vminStar_index < self.optimal_vmin.size and \
                vminStar > self.optimal_vmin[vminStar_index]:
            vminStar_index += 1

        try:
            constr_optimum_log_likelihood = \
                self._ConstrainedOptimalLikelihood(vminStar, logetaStar, vminStar_index)
        except ValueError:
            optim_logL = 10**6
            pass
        else:
            optim_logL = constr_optimum_log_likelihood.fun
            original_optimum = constr_optimum_log_likelihood

        vminStar_index_original = vminStar_index
        index = vminStar_index
        while ALLOW_MOVE and index > 0:
            try:
                index -= 1
                new_optimum = \
                    self._ConstrainedOptimalLikelihood(vminStar, logetaStar, index)
            except ValueError:
                pass
            else:
                if new_optimum.fun < optim_logL:
                    os.system("say Moved left")
                    print("Moved left, index is now", index)
                    print("############################################################" +
                          "############################################################")
                    vminStar_index = index
                    constr_optimum_log_likelihood = new_optimum
                    optim_logL = constr_optimum_log_likelihood.fun
        index = vminStar_index_original
        while ALLOW_MOVE and index < self.optimal_vmin.size:
            try:
                index += 1
                new_optimum = self._ConstrainedOptimalLikelihood(vminStar, logetaStar,
                                                                 index)
            except ValueError:
                pass
            else:
                if new_optimum.fun < optim_logL:
                    os.system("say Moved right")
                    print("Moved right, index is now", index)
                    print("############################################################" +
                          "############################################################")
                    vminStar_index = index
                    constr_optimum_log_likelihood = new_optimum
                    optim_logL = constr_optimum_log_likelihood.fun
        if optim_logL == 10**6:
            raise ValueError

        self.constr_optimal_logl = constr_optimum_log_likelihood.fun
        vars_result = constr_optimum_log_likelihood.x

        self.constr_optimal_vmin = vars_result[: vars_result.size/2]
        self.constr_optimal_logeta = vars_result[vars_result.size/2:]

        if plot:
            print("vminStar =", vminStar)
            print("logetaStar =", logetaStar)
            print("vminStar_index =", vminStar_index)
            try:
                print("original:", original_optimum)
            except:
                print("Original failed.")
                pass
            try:
                print("new:", constr_optimum_log_likelihood)
                print(constr_optimum_log_likelihood.minimizer.kwargs['args'])
            except:
                print("All attepts failed.")
                pass
            try:
                vminStar_rand = constr_optimum_log_likelihood.minimizer.kwargs['args'][1]
                logetaStar_rand = constr_optimum_log_likelihood.minimizer.kwargs['args'][2]
                constr_func = ConstraintsFunction(vminStar_rand, logetaStar_rand,
                                                  vminStar_index)
                constraints = constr_func(constr_optimum_log_likelihood.x)
                is_not_close = np.logical_not(np.isclose(constraints,
                                                         np.zeros_like(constraints)))
                constr_not_valid = np.logical_and(constraints < 0, is_not_close)
                sol_not_found = np.any(constr_not_valid)
                print("random vminStar =", vminStar_rand)
                print("random logetaStar =", logetaStar_rand)
                print("x =", constr_optimum_log_likelihood.x)
                print("constraints =", constraints)
                print("is_not_close =", is_not_close)
                print("constr_not_valid =", constr_not_valid)
                print("sol_not_found =", sol_not_found)
            except:
                print("Error")
                pass
            os.system("say 'Finished plot'")
            self.PlotConstrainedOptimum(vminStar_rand, logetaStar_rand, vminStar_index,
                                        xlim_percentage=(0., 1.1),
                                        ylim_percentage=(1.2, 0.8))
        return self.constr_optimal_logl

    def VminSamplingList(self, output_file_tail, vmin_min, vmin_max, vmin_num_steps,
                         steepness_vmin=1.5, steepness_vmin_center=2.5, plot=False):
        """ Finds a non-linear way to sample the vmin range, such that more points are
        sampled near the location of the steps of the best-fit logeta function, and
        fewer in between. This is done by building a function of vmin that is steeper
        near the steps and flatter elsewhere, and the steeper this function the more
        samplings are done in this region.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            vmin_min, vmin_max: float
                Range in vmin where the sampling should be made.
            vmin_num_steps: int
                Number of samples in vmin (approximate, the final number of steps is
                not exact, due to taking floor() in some places.
            steepness_vmin: float, optional
                Parameter related to the steepness of this function to the left of the
                leftmost step and to the right of the rightmost step.
            steepness_vmin_center: float, optional
                Similar parameter, but for the steepness in between the leftmost step
                and the rightmost step.
            plot: bool, optional
                Whether to plot intermediate results such as the sampling function.
        """
        self.ImportOptimalLikelihood(output_file_tail)
        xmin = vmin_min
        xmax = vmin_max
        # TODO! This +4 is to compensate for a loss of ~4 points (not always 4 though),
        # and it's due to taking floor later on.
        # Find a better way to deal with this.
        x_num_steps = vmin_num_steps  # + 4
        s = steepness_vmin
        sc = steepness_vmin_center

        x_lin = np.linspace(xmin, xmax, 1000)
        x0_list = self.optimal_vmin
        numx0 = x0_list.size

        print("x0 =", x0_list)

        def UnitStep(x): return (np.sign(x) + 1) / 2

        def g1(x, x0, s0, xmin=xmin):
            return np.log10(UnitStep(x - x0) +
                            UnitStep(x0 - x) *
                            (x0 - xmin) / (x + 10**s0 * (-x + x0) - xmin))

        def g2(x, x0, s0, xmax=xmax):
            return np.log10(UnitStep(x0 - x) +
                            UnitStep(x - x0) *
                            (x + 10**s0 * (-x + x0) - xmax) / (x0 - xmax))

        def g(x, x0, s1, s2): return g1(x, x0, s1) + g2(x, x0, s2)

        s_list = np.array([[s, sc]] + [[sc, sc]] * (numx0 - 2) + [[sc, s]])

        def g_total(x, sign=1, x0=x0_list, s_list=s_list):
            return np.array([sign * g(x, x0_list[i], s_list[i, 0], s_list[i, 1])
                             for i in range(x0_list.size)]).prod(axis=0)

        g_lin = g_total(x_lin)

        xT_guess = (x0_list[:-1] + x0_list[1:]) / 2
        bounds = np.array([(x0_list[i], x0_list[i + 1])
                           for i in range(x0_list.size - 1)])
        x_turns_max = np.array([minimize(g_total, np.array(xT_guess[i]),
                                args=(-1,), bounds=[bounds[i]]).x
                                for i in range(0, xT_guess.size, 2)])
        x_turns_min = np.array([minimize(g_total, np.array(xT_guess[i]),
                                         bounds=[bounds[i]]).x
                                for i in range(1, xT_guess.size, 2)])
        x_turns = np.sort(np.append(x_turns_max, x_turns_min))
        x_turns = np.append(np.insert(x_turns, 0, xmin), [xmax])
        y_turns = g_total(x_turns)
        print("x_turns =", x_turns)
        print("y_turns =", y_turns)

        def g_inverse(y, x1, x2):
            return brentq(lambda x: g_total(x) - y, x1, x2)

        def g_inverse_list(y_list, x1, x2):
            return np.array([g_inverse(y, x1, x2) for y in y_list])

        y_diff = np.diff(y_turns)
        y_diff_sum = np.abs(y_diff).sum()
        print("y_diff =", y_diff)
        num_steps = np.array([max(1, np.floor(x_num_steps * np.abs(yd)/y_diff_sum))
                              for yd in y_diff])
        print("num_steps =", num_steps)
        y_list = np.array([np.linspace(y_turns[i], y_turns[i+1], num_steps[i])
                           for i in range(num_steps.size)])
        x_list = np.array([g_inverse_list(y_list[i], x_turns[i], x_turns[i+1])
                           for i in range(y_list.size)])
        x_list = np.concatenate(x_list)
        y_list = np.concatenate(y_list)
        x_list = x_list[np.array([x_list[i] != x_list[i+1]
                        for i in range(x_list.size - 1)] + [True])]
        y_list = y_list[np.array([y_list[i] != y_list[i+1]
                        for i in range(y_list.size - 1)] + [True])]
        self.vmin_sampling_list = x_list

        if plot:
            plt.close()
            plt.plot(x_lin, g_lin)
            plt.plot(x_turns, y_turns, 'o')
            plt.plot(x_list, y_list, '*')
            plt.xlim([xmin, xmax])
            plt.ylim([min(-s * sc**(numx0 - 1), np.min(y_turns)),
                      max(s * sc**(numx0 - 1), np.max(y_turns))])
            plt.show()

        return

    def OptimumStepFunction(self, vmin):
        """ Best-fit logeta as a function of vmin for the optimal log(L).
        Input:
            vmin: float
                Value of vmin for which to evaluate logeta.
        Returns:
            logeta: float
                log(eta(vmin)) for the best-fit piecewise constant function.
        """
        index = 0
        while index < self.optimal_vmin.size and vmin > self.optimal_vmin[index]:
            index += 1
        if index == self.optimal_vmin.size:
            return self.optimal_logeta[-1]*10
        return self.optimal_logeta[index]

    def VminLogetaSamplingTable(self, output_file_tail, logeta_percent_minus,
                                logeta_percent_plus, logeta_num_steps,
                                linear_sampling=True, steepness_logeta=1, plot=False):
        """ Finds a non-linear way to sample both the vmin and logeta range, such that
        more points are sampled near the location of the steps of the best-fit logeta
        function, and fewer in between. This uses the sampling in vmin done by
        VminSamplingList, and computes a non-linear sampling in logeta in a similar way
        (by building a function of logeta that is steeper near the steps and flatter
        elsewhere, and the steeper this function the more samplings are done in this
        region).
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            logeta_percent_minus, logeta_percent_plus: float
                Range in logeta where the sampling should be made, given as percentage
                in the negative and positive direction of the best-fit logeta.
            logeta_num_steps: int
                Number of samples in logeta.
            steepness_logeta: float, optional
                Parameter related to the steepness of this sampling function in logeta.
            plot: bool, optional
                Whether to plot intermediate results such as the sampling function.
        """
        print(self.optimal_vmin)
        print(self.optimal_logeta)

        logeta_num_steps_minus = logeta_num_steps * \
            logeta_percent_minus / (logeta_percent_minus + logeta_percent_plus)
        logeta_num_steps_plus = logeta_num_steps * \
            logeta_percent_plus / (logeta_percent_minus + logeta_percent_plus)

        s = steepness_logeta

        def f(x, xm, i, s0=s):
            return (xm - x) / (10**s0 - 1) * 10**i + (10**s0 * x - xm) / (10**s0 - 1)

        self.vmin_logeta_sampling_table = []
        vmin_last_step = self.optimal_vmin[-1]
        if linear_sampling:
            for vmin in self.vmin_sampling_list:
                logeta_opt = self.OptimumStepFunction(min(vmin, vmin_last_step))
                if vmin < self.optimal_vmin[0]:
                    logeta_min = logeta_opt * (1 + 0.6 * logeta_percent_minus)
                    logeta_max = logeta_opt * (1 - logeta_percent_plus)
                else:
                    if vmin < 600:
                        logeta_min = logeta_opt * (1 + logeta_percent_minus)
                    else:
                        logeta_min = logeta_opt * (1 + 0.6 * logeta_percent_minus)
                    logeta_max = logeta_opt * (1 - 0.5 * logeta_percent_plus)
                logeta_list = [[vmin, logeta]
                               for logeta in np.linspace(logeta_min, logeta_max,
                                                         logeta_num_steps)]
                self.vmin_logeta_sampling_table += [logeta_list]
        else:
            for vmin in self.vmin_sampling_list:
                logeta_opt = self.OptimumStepFunction(min(vmin, vmin_last_step))
                logeta_min = logeta_opt * (1 + logeta_percent_minus)
                logeta_max = logeta_opt * (1 - logeta_percent_plus)
                logeta_list_minus = [[vmin, f(logeta_opt, logeta_min, i)]
                                     for i in np.linspace(s, 0, logeta_num_steps_minus)]
                logeta_list_plus = [[vmin, f(logeta_opt, logeta_max, i)]
                                    for i in np.linspace(s / logeta_num_steps_plus, s,
                                                         logeta_num_steps_plus)]
                self.vmin_logeta_sampling_table += [logeta_list_minus + logeta_list_plus]

        self.vmin_logeta_sampling_table = np.array(self.vmin_logeta_sampling_table)

        if plot:
            self.PlotSamplingTable(plot_close=True)
        return

    def PlotSamplingTable(self, plot_close=False, plot_show=True, plot_optimum=True):
        """ Plots the sampling points in the vmin-logeta plane.
        """
        if plot_close:
            plt.close()
        print("sampling_size =", self.vmin_logeta_sampling_table.shape)
        for tab in self.vmin_logeta_sampling_table:
            plt.plot(tab[:, 0], tab[:, 1], 'o')
        if plot_optimum:
            self.PlotOptimum(xlim_percentage=(0.9, 1.1), ylim_percentage=(1.2, 0.8),
                             plot_close=False, plot_show=plot_show)
        elif plot_show:
            plt.show()
        return

    def GetLikelihoodTable(self, index, output_file_tail, logeta_index_range, extra_tail):
        """ Prints to file lists of the form [logetaStar_ij, logL_ij] needed for
        1D interpolation, where i is the index corresponding to vminStar_i and j is
        the index for each logetaStar. Each file corresponds to a different index i.
            Here only one file is written for a specific vminStar.
        Input:
            index: int
                Index of vminStar.
            output_file_tail: string
                Tag to be added to the file name.
            logeta_index_range: tuple
                A touple (index0, index1) between which logetaStar will be considered.
                If this is None, then the whole list of logetaStar is used.
            extra_tail: string
                Additional tail to be added to filenames.
        """
        print('index =', index)
        print('output_file_tail =', output_file_tail)
        vminStar = self.vmin_logeta_sampling_table[index, 0, 0]
        logetaStar_list = self.vmin_logeta_sampling_table[index, :, 1]
        plot = False
        if logeta_index_range is not None:
            logetaStar_list = \
                logetaStar_list[logeta_index_range[0]: logeta_index_range[1]]
            plot = True
        print("vminStar =", vminStar)
        table = np.empty((0, 2))
        for logetaStar in logetaStar_list:
            try:
                constr_opt = self.ConstrainedOptimalLikelihood(vminStar, logetaStar,
                                                               plot=plot)
            except:
                print("error")
                os.system("say Error")
                pass
            else:
                print("index =", index, "; vminStar =", vminStar,
                      "; logetaStar =", logetaStar, "; constr_opt =", constr_opt)
                table = np.append(table, [[logetaStar, constr_opt]], axis=0)
#                table = np.append(table, [logetaStar])
        print("vminStar =", vminStar, "; table =", table)
        if True:
            temp_file = output_file_tail + "_" + str(index) + \
                "_LogetaStarLogLikelihoodList" + extra_tail + ".dat"
            print(temp_file)
            np.savetxt(temp_file, table)
        return

    def LogLikelihoodList(self, output_file_tail, extra_tail="", processes=None,
                          vmin_index_list=None, logeta_index_range=None):
        """ Loops thorugh the list of all vminStar and calls GetLikelihoodTable,
        which will print the likelihood tables to files.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            extra_tail: string, optional
                Additional tail to be added to filenames.
            processes: int, optional
                Number of processes for parallel programming.
            vmin_index_list: ndarray, optional
                List of indices in vminStar_list for which we calculate the optimal
                likelihood. If not given, the whole list of vminStars is used.
            logeta_index_range: tuple, optional
                Atuple (index0, index1) between which logetaStar will be considered.
                If not given, then the whole list of logetaStar is used.
        """
        if vmin_index_list is None:
            vmin_index_list = range(0, self.vmin_logeta_sampling_table.shape[0])
        else:
            try:
                len(vmin_index_list)
            except TypeError:
                vmin_index_list = range(vmin_index_list,
                                        self.vmin_logeta_sampling_table.shape[0])
        print("vmin_index_list =", vmin_index_list)
        print("logeta_index_range =", logeta_index_range)
        kwargs = ({'index': index,
                   'output_file_tail': output_file_tail,
                   'logeta_index_range': logeta_index_range,
                   'extra_tail': extra_tail}
                  for index in vmin_index_list)
        par.parmap(self.GetLikelihoodTable, kwargs, processes)
        return

    def _logL_interp(vars_list, constraints):
        constr_not_valid = constraints(vars_list)[:-1] < 0
        if np.any(constr_not_valid):
            constr_list = constraints(vars_list)[constr_not_valid]
            return -constr_list.sum() * 10**2
        return logL_interp(vars_list)

    def ConfidenceBand(self, output_file_tail, delta_logL, interpolation_order,
                       extra_tail="", multiplot=True):
        """ Compute the confidence band.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            delta_logL: float
                Target difference between the constrained minimum and the
                unconstrained global minimum of MinusLogLikelihood.
            interpolation_order: int
                interpolation order for the interpolated constrained minimum of
                MinusLogLikelihood as a function of logeta, for a fixed vmin.
            extra_tail: string, optional
                Additional tail to be added to filenames.
            multiplot: bool, optional
                Whether to plot log(L) as a function of logeta for each vmin, and the
                horizontal line corresponding to a given delta_logL.
        """
        print("self.vmin_sampling_list =", self.vmin_sampling_list)
        self.vmin_logeta_band_low = []
        self.vmin_logeta_band_up = []
        vmin_last_step = self.optimal_vmin[-1]
        if multiplot:
            plt.close()
        for index in range(self.vmin_sampling_list.size):
            print("index =", index)
            print("vmin =", self.vmin_sampling_list[index])
            logeta_optim = self.OptimumStepFunction(min(self.vmin_sampling_list[index],
                                                        vmin_last_step))
            file = output_file_tail + "_" + str(index) + \
                "_LogetaStarLogLikelihoodList" + extra_tail + ".dat"
            try:
                with open(file, 'r') as f_handle:
                    table = np.loadtxt(f_handle)
            except:
                continue
            x = table[:, 0]   # this is logeta
            y = table[:, 1]   # this is logL
            logL_interp = interpolate.interp1d(x, y, kind='cubic')

            def _logL_interp(vars_list, constraints):
                constr_not_valid = constraints(vars_list)[:-1] < 0
                if np.any(constr_not_valid):
                    constr_list = constraints(vars_list)[constr_not_valid]
                    return -constr_list.sum() * 1e2
                return logL_interp(vars_list)

            print(self.optimal_logL - delta_logL)
            print(np.array([table[0, 0]]), " ", table[-1, 0])
            print(logeta_optim)

            def constr_func(logeta, logeta_min=np.array([table[0, 0]]),
                            logeta_max=np.array([table[-1, 0]])):
                return np.concatenate([logeta - logeta_min, logeta_max - logeta])

            constr = ({'type': 'ineq', 'fun': constr_func})
            try:
                logeta_minimLogL = minimize(_logL_interp, np.array([logeta_optim]),
                                            args=(constr_func,), constraints=constr).x[0]
            except ValueError:
                print("ValueError at logeta_minimLogL")
                logeta_minimLogL = logeta_optim
                pass
            print("logeta_minimLogL =", logeta_minimLogL)

            print("x =", x)
            print("y =", y)
            if multiplot:
                plt.close()
                plt.plot(x, y, 'o-')
                plt.plot(x, (self.optimal_logL + 1) * np.ones_like(y))
                plt.plot(x, (self.optimal_logL + 2.7) * np.ones_like(y))
                plt.title("index =" + str(index) + ", v_min =" +
                          str(self.vmin_sampling_list[index]) + "km/s")
                plt.xlim(x[0], x[-1])
                plt.ylim(-5, 20)
                plt.show()

            error = F
            try:
                if y[0] > self.optimal_logL + delta_logL and \
                        logeta_minimLogL < self.optimal_logL + delta_logL:
                    sol = brentq(lambda logeta: logL_interp(logeta) - self.optimal_logL -
                                 delta_logL,
                                 table[0, 0], logeta_minimLogL)
                    self.vmin_logeta_band_low += \
                        [[self.vmin_sampling_list[index], sol]]
            except ValueError:
                print("ValueError: Error in calculating vmin_logeta_band_low")
                error = T
            try:
                if y[-1] > self.optimal_logL + delta_logL and \
                        logeta_minimLogL < self.optimal_logL + delta_logL:
                    sol = brentq(lambda logeta: logL_interp(logeta) - self.optimal_logL -
                                 delta_logL,
                                 logeta_minimLogL, table[-1, 0])
                    self.vmin_logeta_band_up += \
                        [[self.vmin_sampling_list[index], sol]]
            except ValueError:
                print("ValueError: Error in calculating vmin_logeta_band_hi")
                error = T

            if error:
                plt.close()
                plt.plot(x, (self.optimal_logL + 1) * np.ones_like(y))
                plt.plot(x, (self.optimal_logL + 2.7) * np.ones_like(y))
                plt.title("index =" + str(index) + "; v_min =" +
                          str(self.vmin_sampling_list[index]) + "km/s")
                plt.xlim(x[0], x[-1])
                plt.ylim([-5, 20])
                plt.plot(x, y, 'o-', color="r")
                plt.plot(logeta_optim, logL_interp(logeta_optim), '*')
                plt.plot(logeta_optim, self.optimal_logL, '*')
                print("ValueError")
                plt.show()
#                raise
                pass

        if multiplot:
            plt.show()

        self.vmin_logeta_band_low = np.array(self.vmin_logeta_band_low)
        self.vmin_logeta_band_up = np.array(self.vmin_logeta_band_up)

        print("lower band: ", self.vmin_logeta_band_low)
        print("upper band: ", self.vmin_logeta_band_up)

        self.PlotConfidenceBand()

        delta_logL = round(delta_logL, 1)
        file = output_file_tail + "_FoxBand_low_deltalogL_" + str(delta_logL) + ".dat"
        print(file)
        np.savetxt(file, self.vmin_logeta_band_low)
        file = output_file_tail + "_FoxBand_up_deltalogL_" + str(delta_logL) + ".dat"
        print(file)
        np.savetxt(file, self.vmin_logeta_band_up)

        return

    def PlotConfidenceBand(self):
        """ Plot the confidence band and the best-fit function.
        """
        plt.close()
        try:
            plt.plot(self.vmin_logeta_band_low[:, 0], self.vmin_logeta_band_low[:, 1], 'o-')
        except IndexError:
            pass
        try:
            plt.plot(self.vmin_logeta_band_up[:, 0], self.vmin_logeta_band_up[:, 1], 'o-')
        except IndexError:
            pass
        self.PlotOptimum(ylim_percentage=(1.2, 0.8), plot_close=F, plot_show=T)

    def ImportConfidenceBand(self, output_file_tail, delta_logL, extra_tail=""):
        """ Import the confidence band from file.
        Input:
            output_file_tail: string
                Tag to be added to the file name.
            delta_logL: float
                Target difference between the constrained minimum and the
                unconstrained global minimum of MinusLogLikelihood.
            extra_tail: string, optional
                Additional tail to be added to filenames.
        """
        delta_logL = round(delta_logL, 1)
        file = output_file_tail + "_FoxBand_low_deltalogL_" + str(delta_logL) + \
            extra_tail + ".dat"
        print(file)
        with open(file, 'r') as f_handle:
            self.vmin_logeta_band_low = np.loadtxt(f_handle)
        file = output_file_tail + "_FoxBand_up_deltalogL_" + str(delta_logL) + \
            extra_tail + ".dat"
        with open(file, 'r') as f_handle:
            self.vmin_logeta_band_up = np.loadtxt(f_handle)
        return


import swantools.io
import swantools.utils
import swantools.plot

import datetime
import matplotlib.pyplot as plt

import numpy as np

def readtable():
    R = swantools.io.SwanIO()
    P = swantools.plot.SwanPlot()

    # Reading TABLE dada with headers:
    df  = R.read_swantable('data/table.txt')

    y = df["Hsig"]
    x = df.index.values
    P.timeseries(x,y,"Significant Wave Heights")

def readspc():
    # Reading spectral data
    R = swantools.io.SwanIO()
    lat,lon,freqs,dirs,times,factors,spectrum = R.read_swanspc('data/spectrum.spc')
    P = swantools.plot.SwanPlot()
    P.spcplot(freqs,dirs,times[15],spectrum[15,:,:]*factors[15])
    # for t, time in enumerate(times):
    #     P.spcplot(freqs,dirs,times[t],spectrum[t,:,:])

def readblock(mode):
    R    = swantools.io.SwanIO()
    P    = swantools.plot.SwanPlot()

    if  mode == "non-stat":
        # Reading a block file - Non stationary example
        lon,lat,times,hs = R.read_swanblock('data/block.mat','Hsig')
        P.blockplot(lon,lat,hs[0,:,:],"Non-stationary Results")
        # for t, time in enumerate(times):
        #     P.blockplot(lon,lat,hs[t,:,:],time.strftime("%Y%m%d %H:%M"))
    elif mode == "stat":
        # Reading a block file - Non stationary example
        lon,lat,times,hs = R.read_swanblock('data/stat_block.mat','Hsig',stat=True)
        P.blockplot(lon,lat,hs,"Stationary Results")

def writescp():

    # Getting some data to play with
    R = swantools.io.SwanIO()
    lat,lon,freqs,dirs,times,factors,spectrum = R.read_swanspc('data/spectrum.spc')

    # Re-writing the data
    R.write_spectrum("spcout.spc",lat,lon,times,freqs,dirs,factors,spectrum)

    # Plot to confirm
    lat,lon,freqs,dirs,times,factors,spectrum = R.read_swanspc('spcout.spc')
    P = swantools.plot.SwanPlot()
    for t, time in enumerate(times):
        P.spcplot(freqs,dirs,times[t],spectrum[t,:,:])

def netcdf_output():
    R = swantools.io.SwanIO()
    W = swantools.io.Converters()

    lon,lat,times,hs = R.read_swanblock('data/block.mat','Hsig')

    W.np2nc("Hsig.nc",lat,lon,times,hs,"Significant Wave Height")

def spectral_output():
    R = swantools.io.SwanIO()
    W = swantools.io.Converters()

    lon,lat,freqs,dirs,times,factors,spectrum = R.read_swanspc('data/spectrum.spc')

    W.spc2nc("spectrum.nc",lat,lon,freqs,dirs,times,factors,spectrum)

if __name__ == "__main__":

    # # Table data
    # import seaborn as sns
    # with sns.axes_style("darkgrid"):
    #     readtable()

    # Spectral data
    readspc()

    # Field data
    readblock("non-stat")

    # Convertung block to netCDF4
    netcdf_output()

    # Converting spctral file to netCDF4
    spectral_output()

    # Wrinting spctral data
    writescp()