# functions to run velocyto and scvelo
import numpy as np
import pandas as pd

# import velocyto as vcy
# import scvelo as scv
from scipy.sparse import csr_matrix
import matplotlib.pyplot as plt
from .moments import *
from anndata import AnnData

def vlm_to_adata(vlm, n_comps=30, basis="umap", trans_mats=None, cells_ixs=None):
    """ Conversion function from the velocyto world to the dynamo world.
    Code original from scSLAM-seq repository

    Parameters
    ----------
        vlm: VelocytoLoom Object
            The VelocytoLoom object that will be converted into adata.
        n_comps: `int` (default: 30)
            The number of pc components that will be stored.
        basis: `str` (default: `umap`)
            The embedding that will be used to store the vlm.ts attribute. Note that velocyto doesn't usually use
            umap as embedding although `umap` as set as default for the convenience of dynamo itself.
        trans_mats: None or dict
            A dict of all relevant transition matrices
        cell_ixs: list of int
            These are the indices of the subsampled cells

    Returns
    -------
        adata: AnnData object
	"""

    from collections import OrderedDict

    # set obs, var
    obs, var = pd.DataFrame(vlm.ca), pd.DataFrame(vlm.ra)
    if "CellID" in obs.keys():
        obs["obs_names"] = obs.pop("CellID")
    if "Gene" in var.keys():
        var["var_names"] = var.pop("Gene")

    if hasattr(vlm, "q"):
        var["gamma_b"] = vlm.q
    if hasattr(vlm, "gammas"):
        var["gamma"] = vlm.gammas
    if hasattr(vlm, "R2"):
        var["gamma_r2"] = vlm.R2

    # rename clusters to louvain
    try:
        ix = np.where(obs.columns == "Clusters")[0][0]
        obs_names = list(obs.columns)
        obs_names[ix] = "louvain"
        obs.columns = obs_names

        # make louvain a categorical field
        obs["louvain"] = pd.Categorical(obs["louvain"])
    except:
        print("Could not find a filed 'Clusters' in vlm.ca.")

    # set layers basics
    layers = OrderedDict(
        unspliced=csr_matrix(vlm.U.T),
        spliced=csr_matrix(vlm.S.T),
        velocity_S=csr_matrix(vlm.velocity.T),
    )

    # set X_spliced / X_unspliced
    if hasattr(vlm, "S_norm"):
        layers["X_spliced"] = csr_matrix(2**vlm.S_norm - 1).T
    if hasattr(vlm, "U_norm"):
        layers["X_unspliced"] = csr_matrix(2**vlm.U_norm - 1).T
    if hasattr(vlm, "S_sz") and not hasattr(vlm, "S_norm"):
        layers["X_spliced"] = csr_matrix(vlm.S_sz).T
    if hasattr(vlm, "U_sz") and hasattr(vlm, "U_norm"):
        layers["X_unspliced"] = csr_matrix(vlm.U_sz).T

    # set M_s / M_u
    if hasattr(vlm, "Sx"):
        layers["M_s"] = csr_matrix(vlm.Sx).T
    if hasattr(vlm, "Ux"):
        layers["M_u"] = csr_matrix(vlm.Ux).T
    if hasattr(vlm, "Sx_sz") and not hasattr(vlm, "Sx"):
        layers["M_s"] = csr_matrix(vlm.Sx_sz).T
    if hasattr(vlm, "Ux_sz") and hasattr(vlm, "Ux"):
        layers["M_u"] = csr_matrix(vlm.Ux_sz).T

    # set obsm
    obsm = {}
    obsm["X"] = vlm.pcs[:, : min(n_comps, vlm.pcs.shape[1])]
    # set basis and velocity on the basis
    if basis is not None:
        obsm["X_" + basis] = vlm.ts
        obsm["velocity_" + basis] = vlm.delta_embedding

    # set transition matrix:
    uns = {}
    if hasattr(vlm, "corrcoef"):
        uns["transition_matrix"] = vlm.corrcoef
    if hasattr(vlm, "colorandum"):
        uns["louvain_colors"] = list(np.unique(vlm.colorandum))

    # add uns annotations
    if trans_mats is not None:
        for key, value in trans_mats.items():
            uns[key] = trans_mats[key]
    if cells_ixs is not None:
        uns["cell_ixs"] = cells_ixs
    if hasattr(vlm, "embedding_knn"):
        from .connectivity import adj_to_knn

        n_neighbors = np.unique((vlm.embedding_knn > 0).sum(1)).min()
        ind_mat, dist_mat = adj_to_knn(
            vlm.emedding_knn, n_neighbors
        )
        uns["neighbors"] = {"indices": ind_mat}
        obsp = {'distances': dist_mat, "connectivities": vlm.emedding_knn}

    uns["dynamics"] = {
        "filter_gene_mode": None,
        "t": None,
        "group": None,
        "X_data": None,
        "X_fit_data": None,
        "asspt_mRNA": "ss",
        "experiment_type": "conventional",
        "normalized": True,
        "model": "deterministic",
        "est_method": "ols",
        "has_splicing": True,
        "has_labeling": False,
        "has_protein": False,
        "use_smoothed": True,
        "NTR_vel": False,
        "log_unnormalized": True,
    }

    # set X
    if hasattr(vlm, "S_norm"):
        X = csr_matrix(vlm.S_norm.T)
    else:
        X = csr_matrix(vlm.S_sz.T) if hasattr(vlm, "S_sz") else csr_matrix(vlm.S.T)

    # create an anndata object with Dynamo characteristics.
    dyn_adata = AnnData(X=X, obs=obs, obsp=obsp, obsm=obsm, var=var, layers=layers, uns=uns)

    return dyn_adata

def converter(data_in, from_type="adata", to_type="vlm", dir="."):
    """
	convert adata to loom object
	- we may save_fig to a temp directory automatically
	- we may write a on-the-fly converter which doesn't involve saving and reading files
	"""
    if from_type == "adata":
        if to_type == "vlm":
            file = dir + "/data.loom"
            data_in.write_loom(file)
            data_out = vcy.VelocytoLoom(file)
    elif from_type == "vlm":
        if to_type == "adata":
            data_out = vlm_to_adata(vlm)

    data_out.ra["Gene"] = data_out.ra["var_names"]  # required by plot_phase_portraits
    colors20 = np.vstack(
        (
            plt.cm.tab20b(np.linspace(0.0, 1, 20))[::2],
            plt.cm.tab20c(np.linspace(0, 1, 20))[1::2],
        )
    )

    def colormap_fun(x: np.ndarray) -> np.ndarray:
        return colors20[np.mod(x, 20)]

    data_out.colorandum = colormap_fun([1] * data_out.S.shape[1])

    return data_out

def run_velocyto(adata):
    """
	1. convert adata to vlm data
	2. set up PCA, UMAP, etc.
	3. estimate the gamma parameter
	"""
    vlm = converter(adata)

    # U_norm: log2(U_sz + pcount)
    # vlm.U_sz: norm_factor * U
    # S_norm: log2(S_sz + pcount)
    # vlm.S_sz norm_factor * S
    # vlm.Ux: smoothed unspliced
    # vlm.Sx: smoothed spliced
    # vlm.Ux_sz: smoothed unspliced -- old code
    # vlm.Sx_sz: smoothed spliced -- old code

    vlm.normalize()  # add U_norm, U_sz, S_norm, S_sz
    vlm.perform_PCA()
    vlm.knn_imputation()  # Ux, Sx, Ux_sz, Sx_sz
    vlm.pcs = adata.X  # pcs: cell x npcs ndarray

    # vlm.Sx = vlm.S_sz
    # vlm.Ux = vlm.U_sz
    # vlm.Sx_sz = vlm.S_sz
    # vlm.Ux_sz = vlm.U_sz

    # gamma fit
    vlm.fit_gammas()  # limit_gamma = False, fit_offset = True,  use_imputed_data = False, use_size_norm = False

    # estimate velocity
    vlm.predict_U()
    vlm.calculate_velocity()

    # predict future state after dt
    vlm.calculate_shift()  # assumption = 'constant_velocity'
    vlm.extrapolate_cell_at_t()  # delta_t = 1.

    return vlm

def run_scvelo(adata):
    """
	1. set up PCA, UMAP, etc.
	2. estimate gamma and all other parameters
	3. return results (adata.var['velocity_gamma'])
	"""
    # scv.pp.filter_and_normalize(adata, min_counts=2, min_counts_u=1, n_top_genes=3000)
    scv.pp.moments(adata)  # , n_pcs = 12, n_neighbors = 15, mode = 'distances'
    scv.tl.velocity(adata)
    scv.tl.velocity_graph(adata)

    # how to fit other parameters, beta, etc.?

    return adata

def mean_var_by_time(X, Time):
    import pandas as pd

    exp_data = pd.DataFrame(X)
    exp_data["Time"] = Time

    mean_val = exp_data.groupby(["Time"]).mean()
    var_val = exp_data.groupby(["Time"]).var()

    return mean_val.values, var_val.values

def run_dynamo(adata, normalize=True, init_num=1, sample_method="lhs"):
    time = adata.obs["Step"].values
    uniqe_time = list(set(time))
    gene_num = adata.X.shape[1]

    # prepare data
    import numpy as np

    x_data = np.zeros((8, len(uniqe_time), gene_num))  # use unique time
    uu, ul, su, sl = (
        adata.layers["uu"].toarray(),
        adata.layers["ul"].toarray(),
        adata.layers["su"].toarray(),
        adata.layers["sl"].toarray(),
    )
    uu = np.log2(uu + 1) if normalize else uu
    ul = np.log2(ul + 1) if normalize else ul
    su = np.log2(su + 1) if normalize else su
    sl = np.log2(sl + 1) if normalize else sl

    x_data[0], x_data[4] = mean_var_by_time(uu, time)
    x_data[1], x_data[5] = mean_var_by_time(ul, time)
    x_data[2], x_data[6] = mean_var_by_time(su, time)
    x_data[3], x_data[7] = mean_var_by_time(sl, time)

    # estimation all parameters
    p0_range = {
        "a": [0, 1],
        "b": [0, 1],
        "la": [0, 1],
        "alpha_a": [10, 1000],
        "alpha_i": [0, 10],
        "sigma": [0, 1],
        "beta": [0, 10],
        "gamma": [0, 10],
    }

    estm = estimation(list(p0_range.values()))
    param_out = pd.DataFrame(
        index=adata.var.index,
        columns=["a", "b", "la", "alpha_a", "alpha_i", "sigma", "beta", "gamma"],
    )
    for i in range(gene_num):
        cur_x_data = x_data[:, :, i].squeeze()
        param_out.iloc[i, :], cost = estm.fit_lsq(
            uniqe_time, cur_x_data, p0=None, n_p0=init_num, sample_method=sample_method
        )

    # estimate only on the spliced and unspliced dataset

    # estimate on the labeled and unlabeled dataset

    # store the fitting result in adata.uns
    adata.uns.update({"dynamo": param_out})

    return adata

def run_dynamo_simple_fit(adata, log=True):
    ncells, gene_num = adata.X.shape

    # estimation all parameters
    param_out = pd.DataFrame(index=adata.var.index, columns=["alpha", "gamma"])

    u, s = adata.layers["unspliced"], adata.layers["spliced"]
    velocity_u, velocity_s = u, s
    for i in range(gene_num):
        cur_u, cur_s = u[:, i], s[:, i]
        gamma = fit_gamma(cur_u.toarray().squeeze(), cur_s.toarray().squeeze())
        alpha = np.mean(cur_s)

        velocity_u[:, i] = cur_u - cur_s * gamma
        velocity_s[:, i] = cur_s / (1 - np.exp(-1)) - cur_u
        param_out.iloc[i, :] = [alpha, gamma]

    adata.layers["velocity_u"] = velocity_u
    adata.layers["velocity_s"] = velocity_s
    adata.uns.update({"dynamo_simple_fit": param_out})

    return adata

def run_dynamo_labelling(adata, log=True, group=False):
    ncells, gene_num = adata.X.shape

    # estimation all parameters
    T = adata.obs["Time"]

    groups = [""] if group == False else np.unique(adata.obs[group])

    param_out = pd.DataFrame(
        index=adata.var.index,
        columns=[i + "_" + j for j in groups for i in ["alpha", "gamma", "u0", "l0"]],
    )
    L, U = adata.layers["L"], adata.layers["U"]
    velocity_u, velocity_s = L, U

    for i in range(gene_num):
        all_parm = []
        for cur_grp in groups.tolist():
            cur_L, cur_U = (
                (L[:, i], U[:, i])
                if cur_grp == ""
                else (
                    L[adata.obs[group] == cur_grp, i],
                    U[adata.obs[group] == cur_grp, i],
                )
            )
            if log:
                cur_U, cur_L = (
                    np.log(cur_U.toarray().squeeze() + 1),
                    np.log(cur_L.toarray().squeeze() + 1),
                )
            else:
                cur_U, cur_L = cur_U.toarray().squeeze(), cur_L.toarray().squeeze()

            gamma, l0 = fit_gamma_labelling(T, cur_L, mode=None)
            alpha, u0 = fit_alpha_labelling(T, cur_U, gamma, mode=None)
            tmp = [alpha, gamma, u0, l0]
            all_parm.extend(tmp)

            velocity_u[:, i] = (cur_L - cur_U * gamma)[:, None]
            velocity_s[:, i] = (cur_U / (1 - np.exp(-1)) - cur_L)[:, None]
            adata.layers[cur_grp + "velocity_u"] = velocity_u
            adata.layers[cur_grp + "velocity_s"] = velocity_s

        param_out.iloc[i, :] = all_parm

    adata.uns.update({"dynamo_labelling": param_out})

    return adata

def compare_res(
    adata,
    velocyto_res,
    svelo_res,
    dynamo_res,
    a_val,
    b_val,
    la_val,
    alpha_a_val,
    alpha_i_val,
    sigma_val,
    beta_val,
    gamma_val,
):
    """
	function to compare results from velocyto and scvelo with our new method
	0. retrieve gamm or gamma with other parameters from velocyto result or scvelo
	1. plot the correlation between parameters estimated with different methods
	2. calculate the correltion between those parameters
	"""
    # self._offset, self._offset2, self._beta, self._gamma, self._r2, self._velocity_genes

    velocyto_gammas = velocyto_res.gammas
    scvelo_gammas = svelo_res.var["velocity_gamma"]

    # scatter plot the true gammas with our result
    plt.subplots(figsize=(15, 5))
    plt.plot()
    plt.subplot(131)
    plt.plot(gamma_val, velocyto_gammas, "o")
    plt.xlabel(r"True $\gamma$")
    plt.ylabel(r"$\gamma$ (velocyto)")
    plt.subplot(132)
    plt.plot(gamma_val, scvelo_gammas, "o")
    plt.xlabel(r"True $\gamma$")
    plt.ylabel(r"$\gamma$ (scvelo)")
    plt.subplot(133)
    plt.plot(gamma_val, dynamo_res.uns["dynamo"]["gamma"], "o")
    plt.xlabel(r"True $\gamma$")
    plt.ylabel(r"$\gamma$ (dynamo)")

    # what if we only have a small number of parameters?
    plt.subplots(figsize=(15, 5))
    plt.plot()
    plt.subplot(131)
    plt.plot(alpha_a_val, svelo_res.var["fit_alpha"], "o")
    plt.xlabel(r"True alpha")
    plt.ylabel(r"$\alpha$ (scvelo)")
    plt.subplot(132)
    plt.plot(beta_val, svelo_res.var["fit_beta"], "o")
    plt.xlabel(r"True $\beta$")
    plt.ylabel(r"$\beta$ (scvelo)")
    plt.subplot(133)
    plt.plot(gamma_val, svelo_res.var["fit_gamma"], "o")
    plt.xlabel(r"True $\gamma$")
    plt.ylabel(r"$\gamma$ (scvelo)")

    #     param_out = pd.DataFrame(index=adata.var.index, columns=['a', 'b', 'la', 'alpha_a', 'alpha_i', 'sigma', 'beta', 'gamma'])
    # what if we only have a small number of parameters?
    plt.subplots(figsize=(15, 15))
    plt.subplot(331)
    plt.plot(a_val, adata.uns["dynamo"]["a"], "o")
    plt.xlabel(r"True $a$")
    plt.ylabel(r"$a$ (dynamo)")
    plt.subplot(332)
    plt.plot(b_val, adata.uns["dynamo"]["b"], "o")
    plt.xlabel(r"True $b$")
    plt.ylabel(r"$b$ (dynamo)")
    plt.subplot(333)
    plt.plot(la_val, adata.uns["dynamo"]["la"], "o")
    plt.xlabel(r"True $l_a$")
    plt.ylabel(r"$l_a$ (dynamo)")
    plt.subplot(334)
    plt.plot(alpha_a_val, adata.uns["dynamo"]["alpha_a"], "o")
    plt.xlabel(r"True $\alpha_a$")
    plt.ylabel(r"$\alpha_a$ (dynamo)")
    plt.subplot(335)
    plt.plot(alpha_i_val, adata.uns["dynamo"]["alpha_i"], "o")
    plt.xlabel(r"True $\alpha_i$")
    plt.ylabel(r"$\alpha_i$ (dynamo)")
    plt.subplot(336)
    plt.plot(sigma_val, adata.uns["dynamo"]["sigma"], "o")
    plt.xlabel(r"True $\sigma$")
    plt.ylabel(r"$\sigma$ (dynamo)")
    plt.subplot(337)
    plt.plot(beta_val, adata.uns["dynamo"]["beta"], "o")
    plt.xlabel(r"True $\beta$")
    plt.ylabel(r"$\beta$ (dynamo)")
    plt.subplot(338)
    plt.plot(gamma_val, adata.uns["dynamo"]["gamma"], "o")
    plt.xlabel(r"True $\gamma$")
    plt.ylabel(r"$\gamma$ (dynamo)")

    velocyto_coef = {"gamma": np.corrcoef(gamma_val, velocyto_gammas)[1, 0]}
    scvelo_coef = {
        "alpha": np.corrcoef(alpha_a_val, svelo_res.var["fit_alpha"])[1, 0],
        "beta": np.corrcoef(beta_val, svelo_res.var["fit_beta"])[1, 0],
        "gamma": np.corrcoef(gamma_val, svelo_res.var["fit_gamma"])[1, 0],
    }

    dynamo_coef = {
        "a": np.corrcoef(a_val, list(dynamo_res.uns["dynamo"]["a"]))[1, 0],
        "b": np.corrcoef(b_val, list(dynamo_res.uns["dynamo"]["b"]))[1, 0],
        "la": np.corrcoef(la_val, list(dynamo_res.uns["dynamo"]["la"]))[1, 0],
        "alpha_a": np.corrcoef(alpha_a_val, list(dynamo_res.uns["dynamo"]["alpha_a"]))[
            1, 0
        ],
        "alpha_i": np.corrcoef(alpha_i_val, list(dynamo_res.uns["dynamo"]["alpha_i"]))[
            1, 0
        ],
        "sigma": np.corrcoef(sigma_val, list(dynamo_res.uns["dynamo"]["sigma"]))[1, 0],
        "beta": np.corrcoef(beta_val, list(dynamo_res.uns["dynamo"]["beta"]))[1, 0],
        "gamma": np.corrcoef(gamma_val, list(dynamo_res.uns["dynamo"]["gamma"]))[1, 0],
    }

    return {
        "velocyto": pd.DataFrame.from_dict(velocyto_coef, orient="index").T,
        "scvelo": pd.DataFrame.from_dict(scvelo_coef, orient="index").T,
        "dynamo": pd.DataFrame.from_dict(dynamo_coef, orient="index").T,
    }


__author__ = ['sibirrer', 'ajshajib', 'dgilman']

import time

import numpy as np
from lenstronomy.Sampling.Samplers.pso import ParticleSwarmOptimizer
from lenstronomy.Util import sampling_util
import emcee
import schwimmbad
from scipy.optimize import minimize

def choose_pool(mpi=False, processes=1, **kwargs):
    """
    Extends the capabilities of the schwimmbad.choose_pool method.

    It handles the `use_dill` parameters in kwargs, that would otherwise raise an error when processes > 1.
    Any thread in the returned multiprocessing pool (e.g. processes > 1) also default

    Docstring from schwimmbad:

    mpi : bool, optional
        Use the MPI processing pool, :class:`~schwimmbad.mpi.MPIPool`. By
        default, ``False``, will use the :class:`~schwimmbad.serial.SerialPool`.
    processes : int, optional
        Use the multiprocessing pool,
        :class:`~schwimmbad.multiprocessing.MultiPool`, with this number of
        processes. By default, ``processes=1``, will use the
        :class:`~schwimmbad.serial.SerialPool`.
    **kwargs
            Any additional kwargs are passed in to the pool class initializer selected by the arguments.
    """
    if processes == 1 or mpi:
        pool = schwimmbad.choose_pool(mpi=mpi, processes=1, **kwargs)
        is_master = pool.is_master()
    else:
        if 'use_dill' in kwargs:
            # schwimmbad MultiPool does not support dill so we remove this option from the kwargs
            _ = kwargs.pop('use_dill')
        pool = schwimmbad.choose_pool(mpi=False, processes=processes, **kwargs)
        # this MultiPool has no is_master() attribute like the SerialPool and MpiPool
        # all threads will then be 'master'.
        is_master = True
    return pool, is_master

class Sampler(object):
    """
    class which executes the different sampling  methods
    Available are: MCMC with emcee and comsoHammer and a Particle Swarm Optimizer.
    This are examples and depending on your problem, you might find other/better solutions.
    Feel free to sample with your convenient sampler!

    """
    def __init__(self, likelihoodModule):
        """

        :param likelihoodModule: instance of LikelihoodModule class
        """
        self.chain = likelihoodModule
        self.lower_limit, self.upper_limit = self.chain.param_limits

    def simplex(self, init_pos, n_iterations, method, print_key='SIMPLEX'):
        """

        :param init_pos: starting point for the optimization
        :param n_iterations: maximum number of iterations
        :param method: the optimization method, default is 'Nelder-Mead'
        :return: the best fit for the lens model using the optimization routine specified by method
        """
        print('Performing the optimization using algorithm:', method)
        time_start = time.time()

        #negativelogL = lambda x: -1 * self.chain.logL(x)

        result = minimize(self.chain.negativelogL, x0=init_pos, method=method,
                          options={'maxiter': n_iterations, 'disp': True})
        logL = self.chain.logL(result['x'])
        kwargs_return = self.chain.param.args2kwargs(result['x'])
        print(-logL * 2 / (max(self.chain.effective_num_data_points(**kwargs_return), 1)),
              'reduced X^2 of best position')
        print(logL, 'logL')
        print(self.chain.effective_num_data_points(**kwargs_return), 'effective number of data points')
        print(kwargs_return.get('kwargs_lens', None), 'lens result')
        print(kwargs_return.get('kwargs_source', None), 'source result')
        print(kwargs_return.get('kwargs_lens_light', None), 'lens light result')
        print(kwargs_return.get('kwargs_ps', None), 'point source result')
        print(kwargs_return.get('kwargs_special', None), 'special param result')
        time_end = time.time()
        print(time_end - time_start, 'time used for ', print_key)
        print('===================')

        return result['x']

    def pso(self, n_particles, n_iterations, lower_start=None, upper_start=None,
            threadCount=1, init_pos=None, mpi=False, print_key='PSO'):
        """
        Return the best fit for the lens model on catalogue basis with
        particle swarm optimizer.

        :param n_particles: number of particles in the sampling process
        :param n_iterations: number of iterations of the swarm
        :param lower_start: numpy array, lower end parameter of the values of the starting particles
        :param upper_start: numpy array, upper end parameter of the values of the starting particles
        :param threadCount: number of threads in the computation (only applied if mpi=False)
        :param init_pos: numpy array, position of the initial best guess model
        :param mpi: bool, if True, makes instance of MPIPool to allow for MPI execution
        :param print_key: string, prints the process name in the progress bar (optional)
        :return: kwargs_result (of best fit), [lnlikelihood of samples, positions of samples, velocity of sampels)
        """
        if lower_start is None or upper_start is None:
            lower_start, upper_start = np.array(self.lower_limit), np.array(self.upper_limit)
            print("PSO initialises its particles with default values")
        else:
            lower_start = np.maximum(lower_start, self.lower_limit)
            upper_start = np.minimum(upper_start, self.upper_limit)

        pool, is_master = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)

        if mpi is True and is_master:
            print('MPI option chosen for PSO.')

        pso = ParticleSwarmOptimizer(self.chain.logL,
                                     lower_start, upper_start, n_particles,
                                     pool=pool)

        if init_pos is None:
            init_pos = (upper_start - lower_start) / 2 + lower_start

        pso.set_global_best(init_pos, [0]*len(init_pos),
                            self.chain.logL(init_pos))

        if is_master:
            print('Computing the %s ...' % print_key)

        time_start = time.time()

        result, [chi2_list, pos_list, vel_list] = pso.optimize(n_iterations)

        if is_master:
            kwargs_return = self.chain.param.args2kwargs(result)
            print(pso.global_best.fitness * 2 / (max(
                self.chain.effective_num_data_points(**kwargs_return), 1)), 'reduced X^2 of best position')
            print(pso.global_best.fitness, 'logL')
            print(self.chain.effective_num_data_points(**kwargs_return), 'effective number of data points')
            print(kwargs_return.get('kwargs_lens', None), 'lens result')
            print(kwargs_return.get('kwargs_source', None), 'source result')
            print(kwargs_return.get('kwargs_lens_light', None), 'lens light result')
            print(kwargs_return.get('kwargs_ps', None), 'point source result')
            print(kwargs_return.get('kwargs_special', None), 'special param result')
            time_end = time.time()
            print(time_end - time_start, 'time used for ', print_key)
            print('===================')
        return result, [chi2_list, pos_list, vel_list]

    def mcmc_emcee(self, n_walkers, n_run, n_burn, mean_start, sigma_start, mpi=False, progress=False, threadCount=1):
        """
        Run MCMC with emcee.
        For details, please have a look at the documentation of the emcee packager.

        :param n_walkers: number of walkers in the emcee process
        :type n_walkers: integer
        :param n_run: number of sampling (after burn-in) of the emcee
        :type n_run: integer
        :param n_burn: number of burn-in iterations (those will not be saved in the output sample)
        :type n_burn: integer
        :param mean_start: mean of the parameter position of the initialising sample
        :type mean_start: numpy array of length the number of parameters
        :param sigma_start: spread of the parameter values (uncorrelated in each dimension) of the initialising sample
        :type sigma_start: numpy array of length the number of parameters
        :param mpi: if True, initializes an MPIPool to allow for MPI execution of the sampler
        :type mpi: bool
        :param progress: if True, prints the progress bar
        :type progress: bool
        :param threadCount: number of threats in multi-processing (not applicable for MPI)
        :type threadCount: integer
        :return: samples, ln likelihood value of samples
        :rtype: numpy 2d array, numpy 1d array
        """
        num_param, _ = self.chain.param.num_param()
        p0 = sampling_util.sample_ball(mean_start, sigma_start, n_walkers)
        time_start = time.time()

        pool, is_master = choose_pool(mpi=mpi, processes=threadCount, use_dill=True)

        sampler = emcee.EnsembleSampler(n_walkers, num_param, self.chain.logL,
                                        pool=pool)

        sampler.run_mcmc(p0, n_burn + n_run, progress=progress)
        flat_samples = sampler.get_chain(discard=n_burn, thin=1, flat=True)
        dist = sampler.get_log_prob(flat=True, discard=n_burn, thin=1)
        if is_master:
            print('Computing the MCMC...')
            print('Number of walkers = ', n_walkers)
            print('Burn-in iterations: ', n_burn)
            print('Sampling iterations:', n_run)
            time_end = time.time()
            print(time_end - time_start, 'time taken for MCMC sampling')
        return flat_samples, dist


# 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 3 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/>.

# This module was adapted from Scikit-Survival python package to
# extract the adjusted TPR and FPR rates for all classification
# thresholds for the censoring adjusted ROC curve. It depends on
# Scikit-Survival. For the full package please check:
# https://github.com/sebp/scikit-survival

import numpy
from scipy.integrate import trapz
from sklearn.utils import check_consistent_length, check_array

from sksurv.nonparametric import CensoringDistributionEstimator, SurvivalFunctionEstimator
from sksurv.util import check_y_survival

__all__ = [
    'brier_score',
    'concordance_index_censored',
    'concordance_index_ipcw',
    'cumulative_dynamic_auc',
    'integrated_brier_score',
]

def _check_estimate(estimate, test_time):
    estimate = check_array(estimate, ensure_2d=False)
    if estimate.ndim != 1:
        raise ValueError(
            'Expected 1D array, got {:d}D array instead:\narray={}.\n'.format(
                estimate.ndim, estimate))
    check_consistent_length(test_time, estimate)
    return estimate

def _check_inputs(event_indicator, event_time, estimate):
    check_consistent_length(event_indicator, event_time, estimate)
    event_indicator = check_array(event_indicator, ensure_2d=False)
    event_time = check_array(event_time, ensure_2d=False)
    estimate = _check_estimate(estimate, event_time)

    if not numpy.issubdtype(event_indicator.dtype, numpy.bool_):
        raise ValueError(
            'only boolean arrays are supported as class labels for survival analysis, got {0}'.format(
                event_indicator.dtype))

    if len(event_time) < 2:
        raise ValueError("Need a minimum of two samples")

    if not event_indicator.any():
        raise ValueError("All samples are censored")

    return event_indicator, event_time, estimate

def _check_times(test_time, times):
    times = check_array(numpy.atleast_1d(times), ensure_2d=False, dtype=test_time.dtype)
    times = numpy.unique(times)

    if times.max() >= test_time.max() or times.min() < test_time.min():
        raise ValueError(
            'all times must be within follow-up time of test data: [{}; {}['.format(
                test_time.min(), test_time.max()))

    return times

def _get_comparable(event_indicator, event_time, order):
    n_samples = len(event_time)
    tied_time = 0
    comparable = {}
    i = 0
    while i < n_samples - 1:
        time_i = event_time[order[i]]
        start = i + 1
        end = start
        while end < n_samples and event_time[order[end]] == time_i:
            end += 1

        # check for tied event times
        event_at_same_time = event_indicator[order[i:end]]
        censored_at_same_time = ~event_at_same_time
        for j in range(i, end):
            if event_indicator[order[j]]:
                mask = numpy.zeros(n_samples, dtype=bool)
                mask[end:] = True
                # an event is comparable to censored samples at same time point
                mask[i:end] = censored_at_same_time
                comparable[j] = mask
                tied_time += censored_at_same_time.sum()
        i = end

    return comparable, tied_time

def _estimate_concordance_index(event_indicator, event_time, estimate, weights, tied_tol=1e-8):
    order = numpy.argsort(event_time)

    comparable, tied_time = _get_comparable(event_indicator, event_time, order)

    concordant = 0
    discordant = 0
    tied_risk = 0
    numerator = 0.0
    denominator = 0.0
    for ind, mask in comparable.items():
        est_i = estimate[order[ind]]
        event_i = event_indicator[order[ind]]
        w_i = weights[order[ind]]

        est = estimate[order[mask]]

        assert event_i, 'got censored sample at index %d, but expected uncensored' % order[ind]

        ties = numpy.absolute(est - est_i) <= tied_tol
        n_ties = ties.sum()
        # an event should have a higher score
        con = est < est_i
        n_con = con[~ties].sum()

        numerator += w_i * n_con + 0.5 * w_i * n_ties
        denominator += w_i * mask.sum()

        tied_risk += n_ties
        concordant += n_con
        discordant += est.size - n_con - n_ties

    cindex = numerator / denominator
    return cindex, concordant, discordant, tied_risk, tied_time

def concordance_index_censored(event_indicator, event_time, estimate, tied_tol=1e-8):
    """Concordance index for right-censored data

    The concordance index is defined as the proportion of all comparable pairs
    in which the predictions and outcomes are concordant.

    Two samples are comparable if (i) both of them experienced an event (at different times),
    or (ii) the one with a shorter observed survival time experienced an event, in which case
    the event-free subject "outlived" the other. A pair is not comparable if they experienced
    events at the same time.

    Concordance intuitively means that two samples were ordered correctly by the model.
    More specifically, two samples are concordant, if the one with a higher estimated
    risk score has a shorter actual survival time.
    When predicted risks are identical for a pair, 0.5 rather than 1 is added to the count
    of concordant pairs.

    See [1]_ for further description.

    Parameters
    ----------
    event_indicator : array-like, shape = (n_samples,)
        Boolean array denotes whether an event occurred

    event_time : array-like, shape = (n_samples,)
        Array containing the time of an event or time of censoring

    estimate : array-like, shape = (n_samples,)
        Estimated risk of experiencing an event

    tied_tol : float, optional, default: 1e-8
        The tolerance value for considering ties.
        If the absolute difference between risk scores is smaller
        or equal than `tied_tol`, risk scores are considered tied.

    Returns
    -------
    cindex : float
        Concordance index

    concordant : int
        Number of concordant pairs

    discordant : int
        Number of discordant pairs

    tied_risk : int
        Number of pairs having tied estimated risks

    tied_time : int
        Number of comparable pairs sharing the same time

    References
    ----------
    .. [1] Harrell, F.E., Califf, R.M., Pryor, D.B., Lee, K.L., Rosati, R.A,
           "Multivariable prognostic models: issues in developing models,
           evaluating assumptions and adequacy, and measuring and reducing errors",
           Statistics in Medicine, 15(4), 361-87, 1996.
    """
    event_indicator, event_time, estimate = _check_inputs(
        event_indicator, event_time, estimate)

    w = numpy.ones_like(estimate)

    return _estimate_concordance_index(event_indicator, event_time, estimate, w, tied_tol)

def concordance_index_ipcw(survival_train, survival_test, estimate, tau=None, tied_tol=1e-8):
    """Concordance index for right-censored data based on inverse probability of censoring weights.

    This is an alternative to the estimator in :func:`concordance_index_censored`
    that does not depend on the distribution of censoring times in the test data.
    Therefore, the estimate is unbiased and consistent for a population concordance
    measure that is free of censoring.

    It is based on inverse probability of censoring weights, thus requires
    access to survival times from the training data to estimate the censoring
    distribution. Note that this requires that survival times `survival_test`
    lie within the range of survival times `survival_train`. This can be
    achieved by specifying the truncation time `tau`.
    The resulting `cindex` tells how well the given prediction model works in
    predicting events that occur in the time range from 0 to `tau`.

    The estimator uses the Kaplan-Meier estimator to estimate the
    censoring survivor function. Therefore, it is restricted to
    situations where the random censoring assumption holds and
    censoring is independent of the features.

    See [1]_ for further description.

    Parameters
    ----------
    survival_train : structured array, shape = (n_train_samples,)
        Survival times for training data to estimate the censoring
        distribution from.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    survival_test : structured array, shape = (n_samples,)
        Survival times of test data.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    estimate : array-like, shape = (n_samples,)
        Estimated risk of experiencing an event of test data.

    tau : float, optional
        Truncation time. The survival function for the underlying
        censoring time distribution :math:`D` needs to be positive
        at `tau`, i.e., `tau` should be chosen such that the
        probability of being censored after time `tau` is non-zero:
        :math:`P(D > \\tau) > 0`. If `None`, no truncation is performed.

    tied_tol : float, optional, default: 1e-8
        The tolerance value for considering ties.
        If the absolute difference between risk scores is smaller
        or equal than `tied_tol`, risk scores are considered tied.

    Returns
    -------
    cindex : float
        Concordance index

    concordant : int
        Number of concordant pairs

    discordant : int
        Number of discordant pairs

    tied_risk : int
        Number of pairs having tied estimated risks

    tied_time : int
        Number of comparable pairs sharing the same time

    References
    ----------
    .. [1] Uno, H., Cai, T., Pencina, M. J., D’Agostino, R. B., & Wei, L. J. (2011).
           "On the C-statistics for evaluating overall adequacy of risk prediction
           procedures with censored survival data".
           Statistics in Medicine, 30(10), 1105–1117.
    """
    test_event, test_time = check_y_survival(survival_test)

    if tau is not None:
        mask = test_time < tau
        survival_test = survival_test[mask]

    estimate = _check_estimate(estimate, test_time)

    cens = CensoringDistributionEstimator()
    cens.fit(survival_train)
    ipcw_test = cens.predict_ipcw(survival_test)
    if tau is None:
        ipcw = ipcw_test
    else:
        ipcw = numpy.empty(estimate.shape[0], dtype=ipcw_test.dtype)
        ipcw[mask] = ipcw_test
        ipcw[~mask] = 0

    w = numpy.square(ipcw)

    return _estimate_concordance_index(test_event, test_time, estimate, w, tied_tol)

def cumulative_dynamic_auc(survival_train, survival_test, estimate, times, tied_tol=1e-8):
    """Estimator of cumulative/dynamic AUC for right-censored time-to-event data.

    The receiver operating characteristic (ROC) curve and the area under the
    ROC curve (AUC) can be extended to survival data by defining
    sensitivity (true positive rate) and specificity (true negative rate)
    as time-dependent measures. *Cumulative cases* are all individuals that
    experienced an event prior to or at time :math:`t` (:math:`t_i \\leq t`),
    whereas *dynamic controls* are those with :math:`t_i > t`.
    The associated cumulative/dynamic AUC quantifies how well a model can
    distinguish subjects who fail by a given time (:math:`t_i \\leq t`) from
    subjects who fail after this time (:math:`t_i > t`).

    Given an estimator of the :math:`i`-th individual's risk score
    :math:`\\hat{f}(\\mathbf{x}_i)`, the cumulative/dynamic AUC at time
    :math:`t` is defined as

    .. math::

        \\widehat{\\mathrm{AUC}}(t) =
        \\frac{\\sum_{i=1}^n \\sum_{j=1}^n I(y_j > t) I(y_i \\leq t) \\omega_i
        I(\\hat{f}(\\mathbf{x}_j) \\leq \\hat{f}(\\mathbf{x}_i))}
        {(\\sum_{i=1}^n I(y_i > t)) (\\sum_{i=1}^n I(y_i \\leq t) \\omega_i)}

    where :math:`\\omega_i` are inverse probability of censoring weights (IPCW).

    To estimate IPCW, access to survival times from the training data is required
    to estimate the censoring distribution. Note that this requires that survival
    times `survival_test` lie within the range of survival times `survival_train`.
    This can be achieved by specifying `times` accordingly, e.g. by setting
    `times[-1]` slightly below the maximum expected follow-up time.
    IPCW are computed using the Kaplan-Meier estimator, which is
    restricted to situations where the random censoring assumption holds and
    censoring is independent of the features.

    The function also provides a single summary measure that refers to the mean
    of the :math:`\\mathrm{AUC}(t)` over the time range :math:`(\\tau_1, \\tau_2)`.

    .. math::

        \\overline{\\mathrm{AUC}}(\\tau_1, \\tau_2) =
        \\frac{1}{\\hat{S}(\\tau_1) - \\hat{S}(\\tau_2)}
        \\int_{\\tau_1}^{\\tau_2} \\widehat{\\mathrm{AUC}}(t)\\,d \\hat{S}(t)

    where :math:`\\hat{S}(t)` is the Kaplan–Meier estimator of the survival function.

    See [1]_, [2]_, [3]_ for further description.

    Parameters
    ----------
    survival_train : structured array, shape = (n_train_samples,)
        Survival times for training data to estimate the censoring
        distribution from.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    survival_test : structured array, shape = (n_samples,)
        Survival times of test data.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    estimate : array-like, shape = (n_samples,)
        Estimated risk of experiencing an event of test data.

    times : array-like, shape = (n_times,)
        The time points for which the area under the
        time-dependent ROC curve is computed. Values must be
        within the range of follow-up times of the test data
        `survival_test`.

    tied_tol : float, optional, default: 1e-8
        The tolerance value for considering ties.
        If the absolute difference between risk scores is smaller
        or equal than `tied_tol`, risk scores are considered tied.

    Returns
    -------
    auc : array, shape = (n_times,)
        The cumulative/dynamic AUC estimates (evaluated at `times`).
    mean_auc : float
        Summary measure referring to the mean cumulative/dynamic AUC
        over the specified time range `(times[0], times[-1])`.

    References
    ----------
    .. [1] H. Uno, T. Cai, L. Tian, and L. J. Wei,
           "Evaluating prediction rules for t-year survivors with censored regression models,"
           Journal of the American Statistical Association, vol. 102, pp. 527–537, 2007.
    .. [2] H. Hung and C. T. Chiang,
           "Estimation methods for time-dependent AUC models with survival data,"
           Canadian Journal of Statistics, vol. 38, no. 1, pp. 8–26, 2010.
    .. [3] J. Lambert and S. Chevret,
           "Summary measure of discrimination in survival models based on cumulative/dynamic time-dependent ROC curves,"
           Statistical Methods in Medical Research, 2014.
    """

    test_event, test_time = check_y_survival(survival_test)

    estimate = _check_estimate(estimate, test_time)

    times = _check_times(test_time, times)

    # sort by risk score (descending)
    o = numpy.argsort(-estimate)
    test_time = test_time[o]
    test_event = test_event[o]
    estimate = estimate[o]
    survival_test = survival_test[o]

    cens = CensoringDistributionEstimator()
    cens.fit(survival_train)
    ipcw = cens.predict_ipcw(survival_test)

    n_samples = test_time.shape[0]
    scores = numpy.empty(times.shape[0], dtype=float)
    rocs = []

    for k, t in enumerate(times):
        is_case = (test_time <= t) & test_event
        is_control = test_time > t
        n_controls = is_control.sum()

        true_pos = []
        false_pos = []
        tp_value = 0.0
        fp_value = 0.0
        est_prev = numpy.infty

        for i in range(n_samples):
            est = estimate[i]
            if numpy.absolute(est - est_prev) > tied_tol:
                true_pos.append(tp_value)
                false_pos.append(fp_value)
                est_prev = est
            if is_case[i]:
                tp_value += ipcw[i]
            elif is_control[i]:
                fp_value += 1
        true_pos.append(tp_value)
        false_pos.append(fp_value)

        sens = numpy.array(true_pos) / ipcw[is_case].sum()
        fpr = numpy.array(false_pos) / n_controls
        scores[k] = trapz(sens, fpr)

        rocs.append((sens, fpr))

    if times.shape[0] == 1:
        mean_auc = scores[0]
    else:
        surv = SurvivalFunctionEstimator()
        surv.fit(survival_test)
        s_times = surv.predict_proba(times)
        # compute integral of AUC over survival function
        d = -numpy.diff(numpy.concatenate(([1.0], s_times)))
        integral = (scores * d).sum()
        mean_auc = integral / (1.0 - s_times[-1])

    return rocs, scores, mean_auc

def brier_score(survival_train, survival_test, estimate, times):
    """Estimate the time-dependent Brier score for right censored data.

    The time-dependent Brier score is the mean squared error at time point :math:`t`:

    .. math::

        \\mathrm{BS}^c(t) = \\frac{1}{n} \\sum_{i=1}^n I(y_i \\leq t \\land \\delta_i = 1)
        \\frac{(0 - \\hat{\\pi}(t | \\mathbf{x}_i))^2}{\\hat{G}(y_i)} + I(y_i > t)
        \\frac{(1 - \\hat{\\pi}(t | \\mathbf{x}_i))^2}{\\hat{G}(t)} ,

    where :math:`\\hat{\\pi}(t | \\mathbf{x})` is the predicted probability of
    remaining event-free up to time point :math:`t` for a feature vector :math:`\\mathbf{x}`,
    and :math:`1/\\hat{G}(t)` is a inverse probability of censoring weight, estimated by
    the Kaplan-Meier estimator.

    See [1]_ for details.

    Parameters
    ----------
    survival_train : structured array, shape = (n_train_samples,)
        Survival times for training data to estimate the censoring
        distribution from.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    survival_test : structured array, shape = (n_samples,)
        Survival times of test data.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    estimate : array-like, shape = (n_samples, n_times)
        Estimated risk of experiencing an event for test data at `times`.
        The i-th column must contain the estimated probability of
        remaining event-free up to the i-th time point.

    times : array-like, shape = (n_times,)
        The time points for which to estimate the Brier score.
        Values must be within the range of follow-up times of
        the test data `survival_test`.

    Returns
    -------
    times : array, shape = (n_times,)
        Unique time points at which the brier scores was estimated.

    brier_scores : array , shape = (n_times,)
        Values of the brier score.

    Examples
    --------
    >>> from sksurv.datasets import load_gbsg2
    >>> from sksurv.linear_model import CoxPHSurvivalAnalysis
    >>> from sksurv.metrics import brier_score
    >>> from sksurv.preprocessing import OneHotEncoder

    Load and prepare data.

    >>> X, y = load_gbsg2()
    >>> X.loc[:, "tgrade"] = X.loc[:, "tgrade"].map(len).astype(int)
    >>> Xt = OneHotEncoder().fit_transform(X)

    Fit a Cox model.

    >>> est = CoxPHSurvivalAnalysis(ties="efron").fit(Xt, y)

    Retrieve individual survival functions and get probability
    of remaining event free up to 5 years (=1825 days).

    >>> survs = est.predict_survival_function(Xt)
    >>> preds = [fn(1825) for fn in survs]

    Compute the Brier score at 5 years.

    >>> times, score = brier_score(y, y, preds, 1825)
    >>> print(score)
    [0.20881843]

    See also
    --------
    integrated_brier_score

    References
    ----------
    .. [1] E. Graf, C. Schmoor, W. Sauerbrei, and M. Schumacher,
           "Assessment and comparison of prognostic classification schemes for survival data,"
           Statistics in Medicine, vol. 18, no. 17-18, pp. 2529–2545, 1999.
    """
    test_event, test_time = check_y_survival(survival_test)
    times = _check_times(test_time, times)

    estimate = check_array(estimate, ensure_2d=False)
    if estimate.ndim == 1 and times.shape[0] == 1:
        estimate = estimate.reshape(-1, 1)

    if estimate.shape[0] != test_time.shape[0]:
        raise ValueError("expected estimate with {} samples, but got {}".format(
            test_time.shape[0], estimate.shape[0]
        ))

    if estimate.shape[1] != times.shape[0]:
        raise ValueError("expected estimate with {} columns, but got {}".format(
            times.shape[0], estimate.shape[1]))

    # fit IPCW estimator
    cens = CensoringDistributionEstimator().fit(survival_train)
    # calculate inverse probability of censoring weight at current time point t.
    prob_cens_t = cens.predict_proba(times)
    prob_cens_t[prob_cens_t == 0] = numpy.inf
    # calculate inverse probability of censoring weights at observed time point
    prob_cens_y = cens.predict_proba(test_time)
    prob_cens_y[prob_cens_y == 0] = numpy.inf

    # Calculating the brier scores at each time point
    brier_scores = numpy.empty(times.shape[0], dtype=float)
    for i, t in enumerate(times):
        est = estimate[:, i]
        is_case = (test_time <= t) & test_event
        is_control = test_time > t

        brier_scores[i] = numpy.mean(numpy.square(est) * is_case.astype(int) / prob_cens_y
                                     + numpy.square(1.0 - est) * is_control.astype(int) / prob_cens_t[i])

    return times, brier_scores

def integrated_brier_score(survival_train, survival_test, estimate, times):
    """The Integrated Brier Score (IBS) provides an overall calculation of
    the model performance at all available times :math:`t_1 \\leq t \\leq t_\\text{max}`.

    The integrated time-dependent Brier score over the interval
    :math:`[t_1; t_\\text{max}]` is defined as

    .. math::

        \\mathrm{IBS} = \\int_{t_1}^{t_\\text{max}} \\mathrm{BS}^c(t) d w(t)

    where the weighting function is :math:`w(t) = t / t_\\text{max}`.
    The integral is estimated via the trapezoidal rule.

    See [1]_ for further details.

    Parameters
    ----------
    survival_train : structured array, shape = (n_train_samples,)
        Survival times for training data to estimate the censoring
        distribution from.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    survival_test : structured array, shape = (n_samples,)
        Survival times of test data.
        A structured array containing the binary event indicator
        as first field, and time of event or time of censoring as
        second field.

    estimate : array-like, shape = (n_samples, n_times)
        Estimated risk of experiencing an event for test data at `times`.
        The i-th column must contain the estimated probability of
        remaining event-free up to the i-th time point.

    times : array-like, shape = (n_times,)
        The time points for which to estimate the Brier score.
        Values must be within the range of follow-up times of
        the test data `survival_test`.

    Returns
    -------
    ibs : float
        The integrated Brier score.

    Examples
    --------
    >>> import numpy
    >>> from sksurv.datasets import load_gbsg2
    >>> from sksurv.linear_model import CoxPHSurvivalAnalysis
    >>> from sksurv.metrics import integrated_brier_score
    >>> from sksurv.preprocessing import OneHotEncoder

    Load and prepare data.

    >>> X, y = load_gbsg2()
    >>> X.loc[:, "tgrade"] = X.loc[:, "tgrade"].map(len).astype(int)
    >>> Xt = OneHotEncoder().fit_transform(X)

    Fit a Cox model.

    >>> est = CoxPHSurvivalAnalysis(ties="efron").fit(Xt, y)

    Retrieve individual survival functions and get probability
    of remaining event free from 1 year to 5 years (=1825 days).

    >>> survs = est.predict_survival_function(Xt)
    >>> times = numpy.arange(365, 1826)
    >>> preds = numpy.asarray([[fn(t) for t in times for fn in survs]])

    Compute the integrated Brier score from 1 to 5 years.

    >>> score = integrated_brier_score(y, y, preds, times)
    >>> print(score)
    0.1815853064627424

    See also
    --------
    brier_score

    References
    ----------
    .. [1] E. Graf, C. Schmoor, W. Sauerbrei, and M. Schumacher,
           "Assessment and comparison of prognostic classification schemes for survival data,"
           Statistics in Medicine, vol. 18, no. 17-18, pp. 2529–2545, 1999.
    """
    # Computing the brier scores
    times, brier_scores = brier_score(survival_train, survival_test, estimate, times)

    if times.shape[0] < 2:
        raise ValueError("At least two time points must be given")

    # Computing the IBS
    ibs_value = trapz(brier_scores, times) / (times[-1] - times[0])

    return ibs_value


from collections import OrderedDict

import numpy as np
import pinocchio
from numpy.linalg import inv
from pinocchio.utils import zero

from .activation import ActivationModelQuad
from .cost import CostDataPinocchio, CostModelPinocchio, CostModelSum
from .state import StatePinocchio
from .utils import a2m

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
class ImpulseModelPinocchio:
    def __init__(self, pinocchioModel, nimpulse):
        assert (hasattr(self, 'ImpulseDataType'))
        self.pinocchio = pinocchioModel
        self.nq, self.nv = pinocchioModel.nq, pinocchioModel.nv
        self.nx = self.nq + self.nv
        self.ndx = 2 * self.nv
        self.nimpulse = nimpulse

    def createData(self, pinocchioData):
        return self.ImpulseDataType(self, pinocchioData)

    def calc(self, data, x):
        assert (False and "This should be defined in the derivative class.")

    def calcDiff(self, data, x, recalc=True):
        assert (False and "This should be defined in the derivative class.")

    def setForces(self, data, forcesArr, forcesVec=None):
        '''
        Convert a numpy array of forces into a stdVector of spatial forces.
        If forcesVec is not none, sum the result in it. Otherwise, reset self.fs
        and put the result there.
        '''
        assert (False and "This should be defined in the derivative class.")
        return self.forces

class ImpulseDataPinocchio:
    def __init__(self, model, pinocchioData):
        nc, nv = model.nimpulse, model.nv
        self.pinocchio = pinocchioData
        self.J = np.zeros([nc, nv])
        self.Jq = np.zeros([nc, nv])
        self.f = np.nan  # not set at construction type
        self.forces = pinocchio.StdVec_Force()
        for i in range(model.pinocchio.njoints):
            self.forces.append(pinocchio.Force.Zero())
        self.Vq = np.zeros([nc, nv])

# --------------------------------------------------------------------------

class ImpulseModel6D(ImpulseModelPinocchio):
    def __init__(self, pinocchioModel, frame):
        self.ImpulseDataType = ImpulseData6D
        ImpulseModelPinocchio.__init__(self, pinocchioModel, nimpulse=6)
        self.frame = frame

    def calc(self, data, x):
        # We suppose forwardKinematics(q,v,a), computeJointJacobian and updateFramePlacement already
        # computed.
        data.J[:, :] = pinocchio.getFrameJacobian(self.pinocchio, data.pinocchio, self.frame,
                                                  pinocchio.ReferenceFrame.LOCAL)

    def calcDiff(self, data, x, recalc=True):
        if recalc:
            self.calc(data, x)
        dv_dq, dv_dv = pinocchio.getJointVelocityDerivatives(self.pinocchio, data.pinocchio, data.joint,
                                                             pinocchio.ReferenceFrame.LOCAL)
        data.Vq[:, :] = data.fXj * dv_dq

    def setForces(self, data, forcesArr, forcesVec=None):
        '''
        Convert a numpy array of forces into a stdVector of spatial forces.
        Side effect: keep the force values in data.
        '''
        # In the dynamic equation, we wrote M*a + J.T*fdyn, while in the ABA it would be
        # M*a + b = tau + J.T faba, so faba = -fdyn (note the minus operator before a2m).
        data.f = forcesArr
        if forcesVec is None:
            forcesVec = data.forces
            data.forces[data.joint] *= 0
        forcesVec[data.joint] += data.jMf * pinocchio.Force(a2m(forcesArr))
        return forcesVec

class ImpulseData6D(ImpulseDataPinocchio):
    def __init__(self, model, pinocchioData):
        ImpulseDataPinocchio.__init__(self, model, pinocchioData)
        frame = model.pinocchio.frames[model.frame]
        self.joint = frame.parent
        self.jMf = frame.placement
        self.fXj = self.jMf.inverse().action

# --------------------------------------------------------------------------

class ImpulseModel3D(ImpulseModelPinocchio):
    def __init__(self, pinocchioModel, frame):
        self.ImpulseDataType = ImpulseData3D
        ImpulseModelPinocchio.__init__(self, pinocchioModel, nimpulse=3)
        self.frame = frame

    def calc(self, data, x):
        # We suppose forwardKinematics(q,v,a), computeJointJacobian and updateFramePlacement already
        # computed.
        data.J[:, :] = pinocchio.getFrameJacobian(self.pinocchio, data.pinocchio, self.frame,
                                                  pinocchio.ReferenceFrame.LOCAL)[:3, :]

    def calcDiff(self, data, x, recalc=True):
        if recalc:
            self.calc(data, x)
        dv_dq, dv_dv = pinocchio.getJointVelocityDerivatives(self.pinocchio, data.pinocchio, data.joint,
                                                             pinocchio.ReferenceFrame.LOCAL)
        data.Vq[:, :] = data.fXj[:3, :] * dv_dq

    def setForces(self, data, forcesArr, forcesVec=None):
        '''
        Convert a numpy array of forces into a stdVector of spatial forces.
        Side effect: keep the force values in data.
        '''
        # In the dynamic equation, we wrote M*a + J.T*fdyn, while in the ABA it would be
        # M*a + b = tau + J.T faba, so faba = -fdyn (note the minus operator before a2m).
        data.f = forcesArr
        if forcesVec is None:
            forcesVec = data.forces
            data.forces[data.joint] *= 0
        forcesVec[data.joint] += data.jMf * pinocchio.Force(a2m(forcesArr), np.zeros((3, 1)))
        return forcesVec

class ImpulseData3D(ImpulseDataPinocchio):
    def __init__(self, model, pinocchioData):
        ImpulseDataPinocchio.__init__(self, model, pinocchioData)
        frame = model.pinocchio.frames[model.frame]
        self.joint = frame.parent
        self.jMf = frame.placement
        self.fXj = self.jMf.inverse().action

# --------------------------------------------------------------------------

class ImpulseModelMultiple(ImpulseModelPinocchio):
    def __init__(self, pinocchioModel, impulses={}):
        self.ImpulseDataType = ImpulseDataMultiple
        ImpulseModelPinocchio.__init__(self, pinocchioModel, nimpulse=0)
        self.impulses = OrderedDict()
        for n, i in impulses.items():
            self.addImpulse(name=n, impulse=i)

    def addImpulse(self, name, impulse):
        self.impulses.update([[name, impulse]])
        self.nimpulse += impulse.nimpulse

    def __getitem__(self, key):
        if isinstance(key, str):
            return self.impulses[key]
        elif isinstance(key, ImpulseModelPinocchio):
            filter = [v for k, v in self.impulses.items() if v.impulse == key]
            assert (len(filter) == 1 and "The given key is not or not unique in the impulse dict. ")
            return filter[0]
        else:
            raise (KeyError("The key should be string or impulsemodel."))

    def calc(self, data, x):
        npast = 0
        for m, d in zip(self.impulses.values(), data.impulses.values()):
            m.calc(d, x)
            data.J[npast:npast + m.nimpulse, :] = d.J
            npast += m.nimpulse

    def calcDiff(self, data, x, recalc=True):
        if recalc:
            self.calc(data, x)
        npast = 0
        for m, d in zip(self.impulses.values(), data.impulses.values()):
            m.calcDiff(d, x, recalc=False)
            data.Vq[npast:npast + m.nimpulse, :] = d.Vq
            npast += m.nimpulse

    def setForces(self, data, fsArr):
        npast = 0
        for i, f in enumerate(data.forces):
            data.forces[i] *= 0
        for m, d in zip(self.impulses.values(), data.impulses.values()):
            m.setForces(d, fsArr[npast:npast + m.nimpulse], data.forces)
            npast += m.nimpulse
        return data.forces

class ImpulseDataMultiple(ImpulseDataPinocchio):
    def __init__(self, model, pinocchioData):
        ImpulseDataPinocchio.__init__(self, model, pinocchioData)
        self.model = model
        self.impulses = OrderedDict([[k, m.createData(pinocchioData)] for k, m in model.impulses.items()])

    def __getitem__(self, key):
        if isinstance(key, str):
            return self.impulses[key]
        elif isinstance(key, ImpulseModelPinocchio):
            filter = [k for k, v in self.model.impulses.items() if v == key]
            assert (len(filter) == 1 and "The given key is not or not unique in the impulse dict. ")
            return self.impulses[filter[0]]
        else:
            raise (KeyError("The key should be string or impulsemodel."))

class CostModelImpactBase(CostModelPinocchio):
    def __init__(self, pinocchioModel, ncost):
        CostModelPinocchio.__init__(self, pinocchioModel, ncost=ncost, nu=0)

    def setImpactData(self, data, vnext):
        data.vnext = vnext

    def setImpactDiffData(self, data, dvnext_dx):
        data.dvnext_dx = dvnext_dx

    def assertImpactDataSet(self, data):
        assert (data.vnext is not None and "vnext should be copied first from impact-data. Call setImpactData first")

    def assertImpactDiffDataSet(self, data):
        assert (data.dvnext_dx is not None and """
        dvnext_dx should be copied first from impact-data. Call setImpactData first""")

class CostDataImpactBase(CostDataPinocchio):
    def __init__(self, model, pinocchioData):
        CostDataPinocchio.__init__(self, model, pinocchioData)
        # These two fields must be informed by ImpactData.
        self.vnext = None
        self.dvnext_dx = None

# --------------------------------------------------------------------------
class CostModelImpactWholeBody(CostModelImpactBase):
    '''
    Penalize the impact on the whole body, i.e. the sum-of-square of ||vnext-v||
    with vnext the velocity after impact and v the velocity before impact.
    '''
    def __init__(self, pinocchioModel, activation=None):
        self.CostDataType = CostDataImpactWholeBody
        CostModelImpactBase.__init__(self, pinocchioModel, ncost=pinocchioModel.nv)
        self.activation = activation if activation is not None else ActivationModelQuad()

    def calc(self, data, x, u=None):
        self.assertImpactDataSet(data)
        nv = self.pinocchio.nv
        data.residuals[:] = data.vnext - x[-nv:]
        data.cost = sum(self.activation.calc(data.activation, data.residuals))
        return data.cost

    def calcDiff(self, data, x, u=None, recalc=True):
        if recalc:
            self.calc(data, x, u)
        self.assertImpactDiffDataSet(data)
        nv = self.pinocchio.nv
        Ax, Axx = self.activation.calcDiff(data.activation, data.residuals)
        data.Rx[:, :] = data.dvnext_dx
        data.Rx[range(nv), range(nv, 2 * nv)] -= 1
        data.Lx[:] = np.dot(data.Rx.T, Ax)
        data.Lxx[:, :] = np.dot(data.Rx.T, Axx * data.Rx)

class CostDataImpactWholeBody(CostDataImpactBase):
    def __init__(self, model, pinocchioData):
        CostDataImpactBase.__init__(self, model, pinocchioData)
        self.activation = model.activation.createData()
        self.Lu = 0
        self.Lxu = 0
        self.Luu = 0
        self.Ru = 0

# --------------------------------------------------------------------------
class CostModelImpactCoM(CostModelImpactBase):
    '''
    Penalize the impact on the com, i.e. the sum-of-square of ||Jcom*(vnext-v)||
    with vnext the velocity after impact and v the velocity before impact.
    '''
    def __init__(self, pinocchioModel, activation=None):
        self.CostDataType = CostDataImpactCoM
        CostModelImpactBase.__init__(self, pinocchioModel, ncost=3)
        self.activation = activation if activation is not None else ActivationModelQuad()

    def calc(self, data, x, u=None):
        self.assertImpactDataSet(data)
        nq, nv = self.pinocchio.nq, self.pinocchio.nv
        pinocchio.centerOfMass(self.pinocchio, data.pinocchio_dv, a2m(x[:nq]), a2m(data.vnext - x[-nv:]))
        data.residuals[:] = data.pinocchio_dv.vcom[0].flat
        data.cost = sum(self.activation.calc(data.activation, data.residuals))
        return data.cost

    def calcDiff(self, data, x, u=None, recalc=True):
        if recalc:
            self.calc(data, x, u)
        self.assertImpactDiffDataSet(data)
        nv = self.pinocchio.nv
        Ax, Axx = self.activation.calcDiff(data.activation, data.residuals)

        # TODO ???
        # r = Jcom(vnext-v)
        # dr/dv = Jcom*(dvnext/dv - I)
        # dr/dq = dJcom_dq*(vnext-v)   + Jcom*dvnext_dq
        #       = dvcom_dq(vq=vnext-v) + Jcom*dvnext_dq
        # Jcom*v = M[:3,:]/mass * v = RNEA(q,0,v)[:3]/mass
        # => dvcom/dq = dRNEA_dq(q,0,v)[:3,:]/mass

        dvc_dq = pinocchio.getCenterOfMassVelocityDerivatives(self.pinocchio, data.pinocchio_dv)
        dvc_dv = pinocchio.jacobianCenterOfMass(self.pinocchio, data.pinocchio_dv)

        # res = vcom(q,vnext-v)
        # dres/dq = dvcom_dq + dvcom_dv*dvnext_dq
        data.Rx[:, :nv] = dvc_dq + np.dot(dvc_dv, data.dvnext_dx[:, :nv])

        # dres/dv = dvcom_dv*(dvnext_dv-I)
        ddv_dv = data.dvnext_dx[:, nv:].copy()
        ddv_dv[range(nv), range(nv)] -= 1
        data.Rx[:, nv:] = np.dot(dvc_dv, ddv_dv)

        data.Lx[:] = np.dot(data.Rx.T, Ax)
        data.Lxx[:, :] = np.dot(data.Rx.T, Axx * data.Rx)

class CostDataImpactCoM(CostDataImpactBase):
    def __init__(self, model, pinocchioData):
        CostDataImpactBase.__init__(self, model, pinocchioData)
        self.activation = model.activation.createData()
        # Those data are ment to be evaluated at v=vnext-v
        self.pinocchio_dv = model.pinocchio.createData()
        self.Lu = 0
        self.Lxu = 0
        self.Luu = 0
        self.Ru = 0

# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------

class ActionModelImpact:
    def __init__(self, pinocchioModel, impulseModel, costModel):
        self.pinocchio = pinocchioModel
        self.State = StatePinocchio(self.pinocchio)
        self.impulse = impulseModel
        self.nq, self.nv = self.pinocchio.nq, self.pinocchio.nv
        self.nx = self.State.nx
        self.ndx = self.State.ndx
        self.nout = self.nx
        self.nu = 0
        self.unone = np.zeros(self.nu)
        self.costs = costModel
        self.impulseWeight = 100.

    @property
    def nimpulse(self):
        return self.impulse.nimpulse

    @property
    def ncost(self):
        return self.costs.ncost

    def createData(self):
        return ActionDataImpact(self)

    def calc(self, data, x, u=None):
        '''
        M(vnext-v) - J^T f = 0
        J vnext = 0

        [MJ^T][vnext] = [Mv]
        [J   ][ -f   ]   [0 ]

        [vnext] = K^-1[Mv], with K = [MJ^T;J0]
        [ -f   ]       [0 ]
        '''
        nq, nv = self.nq, self.nv
        q = a2m(x[:nq])
        v = a2m(x[-nv:])

        pinocchio.computeAllTerms(self.pinocchio, data.pinocchio, q, v)
        pinocchio.updateFramePlacements(self.pinocchio, data.pinocchio)

        self.impulse.calc(data.impulse, x)

        data.K[:nv, :nv] = data.pinocchio.M
        if hasattr(self.pinocchio, 'armature'):
            data.K[range(nv), range(nv)] += self.pinocchio.armature.flat
        data.K[nv:, :nv] = data.impulse.J
        data.K.T[nv:, :nv] = data.impulse.J
        data.Kinv = inv(data.K)
        data.r[:nv] = (data.K[:nv, :nv] * v).flat
        data.r[nv:] = 0

        data.af[:] = np.dot(data.Kinv, data.r)
        data.f[:] *= -1.
        # Convert force array to vector of spatial forces.
        self.impulse.setForces(data.impulse, data.f)

        data.xnext[:nq] = q.flat
        data.xnext[nq:] = data.vnext

        if isinstance(self.costs, CostModelImpactBase):
            self.costs.setImpactData(data.costs, data.vnext)
        if isinstance(self.costs, CostModelSum):
            for cmodel, cdata in zip(self.costs.costs.values(), data.costs.costs.values()):
                if isinstance(cmodel.cost, CostModelImpactBase):
                    cmodel.cost.setImpactData(cdata, data.vnext)

        data.cost = self.costs.calc(data.costs, x, u=None)
        return data.xnext, data.cost

    def calcDiff(self, data, x, u=None, recalc=True):
        '''
        k = [Mv;0]; K = [MJ^T;J0]
        r = [vnext;f] = K^-1 k
        dr/dv = K^-1 [M;0]
        dr/dq = -K^-1 K'K^-1 k + K^-1 k' = -K^-1 (K'r-k')
              = -K^-1 [ M'vnext + J'^T f- M'v ]
                      [ J'vnext               ]
              = -K^-1 [ M'(vnext-v) + J'^T f ]
                      [ J' vnext             ]
        '''
        if recalc:
            xout, cost = self.calc(data, x, u)
        nq, nv = self.nq, self.nv
        q = a2m(x[:nq])
        v = a2m(x[-nv:])
        vnext = a2m(data.vnext)
        fs = data.impulse.forces

        # Derivative M' dv + J'f + b'
        g6bak = self.pinocchio.gravity.copy()
        self.pinocchio.gravity = pinocchio.Motion.Zero()
        pinocchio.computeRNEADerivatives(self.pinocchio, data.pinocchio, q, zero(nv), vnext - v, fs)
        self.pinocchio.gravity = g6bak
        data.did_dq[:, :] = data.pinocchio.dtau_dq

        # Derivative of the impulse constraint
        pinocchio.computeForwardKinematicsDerivatives(self.pinocchio, data.pinocchio, q, vnext, zero(nv))
        # pinocchio.updateFramePlacements(self.pinocchio,data.pinocchio)
        self.impulse.calcDiff(data.impulse, x, recalc=False)
        data.dv_dq = data.impulse.Vq

        data.Fq[:nv, :] = 0
        np.fill_diagonal(data.Fq[:nv, :], 1)  # dq/dq
        data.Fv[:nv, :] = 0  # dq/dv
        data.Fx[nv:, :nv] = -np.dot(data.Kinv[:nv, :], np.vstack([data.did_dq, data.dv_dq]))  # dvnext/dq
        data.Fx[nv:, nv:] = np.dot(data.Kinv[:nv, :nv], data.K[:nv, :nv])  # dvnext/dv

        # data.Rx[:,:] = 0
        # np.fill_diagonal(data.Rv,-1)
        # data.Rx[:,:] += data.Fx[nv:,:]
        # data.Rx *= self.impulseWeight
        # data.Lx [:]   = np.dot(data.Rx.T,data.costResiduals)
        # data.Lxx[:,:] = np.dot(data.Rx.T,data.Rx)

        if isinstance(self.costs, CostModelImpactBase):
            self.costs.setImpactDiffData(data.costs, data.Fx[nv:, :])
        if isinstance(self.costs, CostModelSum):
            for cmodel, cdata in zip(self.costs.costs.values(), data.costs.costs.values()):
                if isinstance(cmodel.cost, CostModelImpactBase):
                    cmodel.cost.setImpactDiffData(cdata, data.Fx[nv:, :])

        self.costs.calcDiff(data.costs, x, u=None, recalc=recalc)

        return data.xnext, data.cost

class ActionDataImpact:
    def __init__(self, model):
        self.pinocchio = model.pinocchio.createData()
        self.impulse = model.impulse.createData(self.pinocchio)
        if model.costs is not None:
            self.costs = model.costs.createData(self.pinocchio)
        self.cost = np.nan
        nx, nu, ndx, nv, nc = model.nx, model.nu, model.State.ndx, model.nv, model.nimpulse
        self.F = np.zeros([ndx, ndx + nu])
        self.Fx = self.F[:, :ndx]
        self.Fu = self.F[:, ndx:]
        self.Fq = self.Fx[:, :nv]
        self.Fv = self.Fx[:, nv:]
        self.Fq[:, :] = 0
        np.fill_diagonal(self.Fq, 1)

        self.costResiduals = np.zeros(nv)
        self.Rx = np.zeros([nv, ndx])
        self.Rq = self.Rx[:, :nv]
        self.Rv = self.Rx[:, nv:]

        self.costResiduals = self.costs.residuals
        self.Lx = self.costs.Lx
        self.Lu = self.costs.Lu
        self.Lxx = self.costs.Lxx
        self.Lxu = self.costs.Lxu
        self.Luu = self.costs.Luu
        self.Rx = self.costs.Rx
        self.Ru = self.costs.Ru

        self.K = np.zeros([nv + nc, nv + nc])  # KKT matrix = [ MJ.T ; J0 ]
        self.r = np.zeros(nv + nc)  # NLE effects =  [ tau-b ; -gamma ]
        self.af = np.zeros(nv + nc)  # acceleration&forces = [ a ; f ]
        self.vnext = self.af[:nv]
        self.f = self.af[nv:]
        self.did_dq = np.zeros([nv, nv])

        self.xnext = np.zeros(nx)