import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import scipy.io as sio
import scipy.stats as ss

import elfi
import elfi.examples.bdm
import elfi.examples.ma2

%matplotlib inline

"""
Explanation: This tutorial is generated from a Jupyter notebook that can be found here. 
Using non-Python operations
If your simulator or other operations are implemented in a programming language other than Python, you can still use ELFI. This notebook briefly demonstrates how to do this in three common scenarios:

External executable (written e.g. in C++ or a shell script)
R function
MATLAB function

Let's begin by importing some libraries that we will be using:
End of explanation
"""


# Make an external command. {0} {1} are positional arguments and {seed} a keyword argument `seed`.
command = 'echo {0} {1} {seed}'
echo_sim = elfi.tools.external_operation(command)

# Test that `echo_sim` can now be called as a regular python function
echo_sim(3, 1, seed=123)

"""
Explanation: External executables
ELFI supports using external simulators and other operations that can be called from the command-line. ELFI provides some tools to easily incorporate such operations to ELFI models. This functionality is introduced in this tutorial.
We demonstrate here how to wrap executables as ELFI nodes. We will first use elfi.tools.external_operation tool to wrap executables as a Python callables (function). Let's first investigate how it works with a simple shell echo command:
End of explanation
"""


# Vectorize it with elfi tools
echo_sim_vec = elfi.tools.vectorize(echo_sim)

# Make a simple model
m = elfi.ElfiModel(name='echo')
elfi.Prior('uniform', .005, 2, model=m, name='alpha')
elfi.Simulator(echo_sim_vec, m['alpha'], 0, name='echo')

# Test to generate 3 simulations from it
m['echo'].generate(3)

"""
Explanation: The placeholders for arguments in the command string are just Python's format strings.
Currently echo_sim only accepts scalar arguments. In order to work in ELFI, echo_sim needs to be vectorized so that we can pass to it a vector of arguments. ELFI provides a handy tool for this as well:
End of explanation
"""


# Fixed model parameters
delta = 0
tau = 0.198
N = 20

# The zeros are to make the observed population vector have length N
y_obs = np.array([6, 3, 2, 2, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype='int16')

"""
Explanation: So above, the first column draws from our uniform prior for $\alpha$, the second column has constant zeros, and the last one lists the seeds provided to the command by ELFI.
Complex external operations $-$ case BDM
To provide a more realistic example of external operations, we will consider the Birth-Death-Mutation (BDM) model used in Lintusaari at al 2016 [1].
Birth-Death-Mutation process
We will consider here the Birth-Death-Mutation process simulator introduced in Tanaka et al 2006 [2] for the spread of Tuberculosis. The simulator outputs a count vector where each of its elements represents a "mutation" of the disease and the count describes how many are currently infected by that mutation. There are three rates and the population size:

$\alpha$ - (birth rate) the rate at which any infectious host transmits the disease.
$\delta$ - (death rate) the rate at which any existing infectious hosts either recovers or dies.
$\tau$ - (mutation rate) the rate at which any infectious host develops a new unseen mutation of the disease within themselves.
$N$ - (population size) the size of the simulated infectious population

It is assumed that the susceptible population is infinite, the hosts carry only one mutation of the disease and transmit that mutation onward. A more accurate description of the model can be found from the original paper or e.g. Lintusaari at al 2016 [1].
<img src="resources/bdm.png" alt="BDM model illustration from Lintusaari et al. 2016" style="width: 400px;"/>
This simulator cannot be implemented effectively with vectorized operations so we have implemented it with C++ that handles loops efficiently. We will now reproduce Figure 6(a) in Lintusaari at al 2016 [2] with ELFI. Let's start by defining some constants:
End of explanation
"""


m = elfi.ElfiModel(name='bdm')
elfi.Prior('uniform', .005, 2, model=m, name='alpha')

# Get the BDM source directory
sources_path = elfi.examples.bdm.get_sources_path()

# Compile (unix-like systems)
!make -C $sources_path

# Move the executable in to the working directory
!mv $sources_path/bdm .

# Test the executable (assuming we have the executable `bdm` in the working directory)
sim = elfi.tools.external_operation('./bdm {0} {1} {2} {3} --seed {seed} --mode 1')
sim(1, delta, tau, N, seed=123)

"""
Explanation: Let's build the beginning of a new model for the birth rate $\alpha$ as the only unknown
End of explanation
"""


# Assuming we have the executable `bdm` in the working directory
command = './bdm {filename} --seed {seed} --mode 1 > {output_filename}'


# Function to prepare the inputs for the simulator. We will create filenames and write an input file.
def prepare_inputs(*inputs, **kwinputs):
    alpha, delta, tau, N = inputs
    meta = kwinputs['meta']

    # Organize the parameters to an array. The broadcasting works nicely with constant arguments here.
    param_array = np.row_stack(np.broadcast(alpha, delta, tau, N))
    
    # Prepare a unique filename for parallel settings
    filename = '{model_name}_{batch_index}_{submission_index}.txt'.format(**meta)
    np.savetxt(filename, param_array, fmt='%.4f %.4f %.4f %d')

    # Add the filenames to kwinputs
    kwinputs['filename'] = filename
    kwinputs['output_filename'] = filename[:-4] + '_out.txt'
    
    # Return new inputs that the command will receive
    return inputs, kwinputs


# Function to process the result of the simulation
def process_result(completed_process, *inputs, **kwinputs):
    output_filename = kwinputs['output_filename']
    
    # Read the simulations from the file.
    simulations = np.loadtxt(output_filename, dtype='int16')
    
    # Clean up the files after reading the data in
    os.remove(kwinputs['filename'])
    os.remove(output_filename)
    
    # This will be passed to ELFI as the result of the command
    return simulations


# Create the python function (do not read stdout since we will work through files)
bdm = elfi.tools.external_operation(command, 
                                    prepare_inputs=prepare_inputs, 
                                    process_result=process_result, 
                                    stdout=False)

"""
Explanation: The BDM simulator is actually already internally vectorized if you provide it an input file with parameters on the rows. This is more efficient than looping in Python (elfi.tools.vectorize), because one simulation takes very little time and we wish to generate tens of thousands of simulations. We will also here redirect the output to a file and then read the file into a numpy array. 
This is just one possibility among the many to implement this. The most efficient would be to write a native Python module with C++ but it's beyond the scope of this article. So let's work through files which is a fairly common situation especially with existing software.
End of explanation
"""


# Create the simulator
bdm_node = elfi.Simulator(bdm, m['alpha'], delta, tau, N, observed=y_obs, name='sim')

# Ask ELFI to provide the meta dict
bdm_node.uses_meta = True

# Draw the model
elfi.draw(m)

# Test it
data = bdm_node.generate(3)
print(data)

"""
Explanation: Now let's replace the echo simulator with this. To create unique but informative filenames, we ask ELFI to provide the operation some meta information. That will be available under the meta keyword (see the prepare_inputs function above):
End of explanation
"""


def T1(clusters):
    clusters = np.atleast_2d(clusters)
    return np.sum(clusters > 0, 1)/np.sum(clusters, 1)

def T2(clusters, n=20):
    clusters = np.atleast_2d(clusters)
    return 1 - np.sum((clusters/n)**2, axis=1)

# Add the different distances to the model
elfi.Summary(T1, bdm_node, name='T1')
elfi.Distance('minkowski', m['T1'], p=1, name='d_T1')

elfi.Summary(T2, bdm_node, name='T2')
elfi.Distance('minkowski', m['T2'], p=1, name='d_T2')

elfi.Distance('minkowski', m['sim'], p=1, name='d_sim')

elfi.draw(m)

# Save parameter and simulation results in memory to speed up the later inference
pool = elfi.OutputPool(['alpha', 'sim'])
# Fix a seed
seed = 20170511

rej = elfi.Rejection(m, 'd_T1', batch_size=10000, pool=pool, seed=seed)
%time T1_res = rej.sample(5000, n_sim=int(1e5))

rej = elfi.Rejection(m, 'd_T2', batch_size=10000, pool=pool, seed=seed)
%time T2_res = rej.sample(5000, n_sim=int(1e5))

rej = elfi.Rejection(m, 'd_sim', batch_size=10000, pool=pool, seed=seed)
%time sim_res = rej.sample(5000, n_sim=int(1e5))

# Load a precomputed posterior based on an analytic solution (see Lintusaari et al 2016)
matdata = sio.loadmat('./resources/bdm.mat')
x = matdata['likgrid'].reshape(-1)
posterior_at_x = matdata['post'].reshape(-1)

# Plot the reference
plt.figure()
plt.plot(x, posterior_at_x, c='k')

# Plot the different curves
for res, d_node, c in ([sim_res, 'd_sim', 'b'], [T1_res, 'd_T1', 'g'], [T2_res, 'd_T2', 'r']):
    alphas = res.outputs['alpha']
    dists = res.outputs[d_node]
    # Use gaussian kde to make the curves look nice. Note that this tends to benefit the algorithm 1 
    # a lot as it ususally has only a very few accepted samples with 100000 simulations
    kde = ss.gaussian_kde(alphas[dists<=0])
    plt.plot(x, kde(x), c=c)
    
plt.legend(['reference', 'algorithm 1', 'algorithm 2, T1\n(eps=0)', 'algorithm 2, T2\n(eps=0)'])
plt.xlim([-.2, 1.2]);
print('Results after 100000 simulations. Compare to figure 6(a) in Lintusaari et al. 2016.')

"""
Explanation: Completing the BDM model
We are now ready to finish up the BDM model. To reproduce Figure 6(a) in Lintusaari at al 2016 [2], let's add different summaries and discrepancies to the model and run the inference for each of them:
End of explanation
"""


import rpy2.robjects as robj
from rpy2.robjects import numpy2ri as np2ri

# Converts numpy arrays automatically
np2ri.activate()

"""
Explanation: Interfacing with R
It is possible to run R scripts in command line for example with Rscript. However, in Python it may be more convenient to use rpy2, which allows convenient access to the functionality of R from within Python. You can install it with pip install rpy2.
Here we demonstrate how to calculate the summary statistics used in the ELFI tutorial (autocovariances) using R's acf function for the MA2 model.
End of explanation
"""


robj.r('''
    # create a function `f`
    f <- function(x, lag=1) {
        ac = acf(x, plot=FALSE, type="covariance", lag.max=lag, demean=FALSE)
        ac[['acf']][lag+1]
    }
    ''')

f = robj.globalenv['f']

def autocovR(x, lag=1):
    x = np.atleast_2d(x)
    apply = robj.r['apply']
    ans = apply(x, 1, f, lag=lag)
    return np.atleast_1d(ans)

# Test it
autocovR(np.array([[1,2,3,4], [4,5,6,7]]), 1)

"""
Explanation: Let's create a Python function that wraps the R commands (please see the documentation of rpy2 for details):
End of explanation
"""


ma2 = elfi.examples.ma2.get_model(seed_obs=4)
elfi.draw(ma2)

"""
Explanation: Load a ready made MA2 model:
End of explanation
"""


# Replace with R autocov
S1 = elfi.Summary(autocovR, ma2['MA2'], 1)
S2 = elfi.Summary(autocovR, ma2['MA2'], 2)
ma2['S1'].become(S1)
ma2['S2'].become(S2)

# Run the inference
rej = elfi.Rejection(ma2, 'd', batch_size=1000, seed=seed)
rej.sample(100)

"""
Explanation: Replace the summaries S1 and S2 with our R autocovariance function.
End of explanation
"""


import matlab.engine

"""
Explanation: Interfacing with MATLAB
There are a number of options for running MATLAB (or Octave) scripts from within Python. Here, evaluating the distance is demonstrated with a MATLAB function using the official MATLAB Python cd API. (Tested with MATLAB 2016b.)
End of explanation
"""


eng = matlab.engine.start_matlab()  # takes a while...

"""
Explanation: A MATLAB session needs to be started (and stopped) separately:
End of explanation
"""


def euclidean_M(x, y):
    # MATLAB array initialized with Python's list
    ddM = matlab.double((x-y).tolist())
    
    # euclidean distance
    dM = eng.sqrt(eng.sum(eng.power(ddM, 2.0), 2))
    
    # Convert back to numpy array
    d = np.atleast_1d(dM).reshape(-1)
    return d

# Test it
euclidean_M(np.array([[1,2,3], [6,7,8], [2,2,3]]), np.array([2,2,2]))

"""
Explanation: Similarly as with R, we have to write a piece of code to interface between MATLAB and Python:
End of explanation
"""


ma2M = elfi.examples.ma2.get_model(seed_obs=4)
elfi.draw(ma2M)

"""
Explanation: Load a ready made MA2 model:
End of explanation
"""


# Replace with Matlab distance implementation
d = elfi.Distance(euclidean_M, ma2M['S1'], ma2M['S2'])
ma2M['d'].become(d)

# Run the inference
rej = elfi.Rejection(ma2M, 'd', batch_size=1000, seed=seed)
rej.sample(100)

"""
Explanation: Replace the summaries S1 and S2 with our R autocovariance function.
End of explanation
"""


eng.quit()

"""
Explanation: Finally, don't forget to quit the MATLAB session:
End of explanation
"""


