################################################################################
# Diffusion                                                                    #
#                                                                              #
"""Analyse diffusion in a pore."""
################################################################################


import math
import itertools
import scipy as sp
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import poreana.utils as utils


###################
# Diffusion - Bin #
###################
def cui(link, z_dist=0, ax_area=[0.2, 0.8], intent="", is_fit=False, is_plot=True):
    """This function samples and calculates the diffusion coefficient of a
    molecule group in a pore in both axial and radial direction, as described
    in the paper of `Cui <https://doi.org/10.1063/1.1989314>`_.

    The mean square displacement is sampled in function
    :func:`poreana.sample.Sample._diffusion_bin`.

    The axial diffusion is given by the Einstein relation

    .. math::

        \\langle\\left[z(0)-z(t)\\right]^2\\rangle=2D_\\text{axial}t

    with axial diffusion coefficient :math:`D_\\text{axial}`. Thus the
    coefficient corresponds to the slope of the axial msd

    .. math::

        D_\\text{axial}=\\frac12\\frac{\\text{msd}_\\text{axial}[i]-\\text{msd}_\\text{axial}[j]}{t_i-t_j}

    with bin :math:`i>j`. The radial diffusion is given by

    .. math::

        \\langle\\left[r(0)-r(t)\\right]^2\\rangle=R^2\\left[1-\\sum_{n=1}^\\infty\\frac{8}{\\lambda_{1n}^2(\\lambda_{1n}^2-1)}\\exp\\left(-\\frac{\\lambda_{1n}^2}{R^2}D_\\text{radial}t\\right)\\right].

    with radial diffusion coefficient :math:`D_\\text{radial}`, the maximal
    accessible radial position :math:`R` by an atom

    .. math::

        R = \\frac12d-0.2

    with pore diameter :math:`d`, and the zeros :math:`\\lambda_{1n}^2` of
    the derivative

    .. math::

        \\frac{dJ_1}{dx}

    of the first order Bessel function :math:`J_1`. Therefore the coefficient
    has to be fitted to the sampled radial msd. The author observed that
    the first 20 Bessel function zeros were sufficient for the function fit.

    The final unit transformation is done by

    .. math::

        \\frac{\\text{nm}^2}{\\text{ps}}=10^{-2}\\frac{\\text{cm}^2}{\\text{s}}.

    **Note that the function for the radial diffusion is obtained under the
    assumption that the density inside the pore is uniform.**

    Parameters
    ----------
    link : string
        Link to hdf5 data file generated by the sample routine
        :func:`poreana.sample.Sample._diffusion_bin`
    z_dist : float, optional
        Distance from pore centre to calculate the mean
    ax_area : list, optional
        Bin area percentage to calculate the axial diffusion coefficient
    intent : string, optional
        axial, radial or empty for both
    is_fit : bool, optional
        True to plot the fitted function
    is_plot : bool, optional
        True to create plot in this function
    """
    # Load data object
    sample = utils.load(link)

    # Load data
    pore = sample["pore"]
    inp = sample["inp"]
    len_window = int(inp["len_window"])
    len_step = inp["len_step"]
    len_frame = inp["len_frame"]
    bins = int(inp["bin_num"]) if not z_dist else math.floor(z_dist/sample["data"]["width"])
    msd_z = [0 for x in range(len_window)]
    msd_r = [0 for x in range(len_window)]
    norm_z = [0 for x in range(len_window)]
    norm_r = [0 for x in range(len_window)]

    # Sum up all bins
    for i in range(bins):
        for j in range(len_window):
            msd_z[j] += sample["data"]["z_tot"][i][j]
            norm_z[j] += sample["data"]["n_tot"][i][j]

    for i in range(int(inp["bin_num"])):
        for j in range(len_window):
            msd_r[j] += sample["data"]["r_tot"][i][j]
            norm_r[j] += sample["data"]["n_tot"][i][j]

    # Normalize
    msd_z_n = [msd_z[i]/norm_z[i] if norm_z[i] > 0 else 0 for i in range(len_window)]
    msd_r_n = [msd_r[i]/norm_r[i] if norm_r[i] > 0 else 0 for i in range(len_window)]

    # Define time axis and range
    time_ax = [x*len_step*len_frame for x in range(len_window)]
    t_range = (len_window-1)*len_step*len_frame

    # Calculate axial coefficient
    if not intent or intent == "axial":
        dz = (msd_z_n[int(ax_area[1]*len_window)]-msd_z_n[int(ax_area[0]*len_window)])*1e-9**2/((ax_area[1]-ax_area[0])*t_range)/2*1e2**2*1e5  # 10^-9 m^2s^-1

        print("Diffusion axial:  "+"%.3f" % dz+" 10^-9 m^2s^-1")

    # Calculate radial coefficient
    if not intent or intent == "radial":
        def diff_rad(x, a, b, c):
            # Process input
            x = x if isinstance(x, list) or isinstance(x, np.ndarray) else [x]

            # Get bessel function zeros
            jz = sp.special.jnp_zeros(1, math.ceil(b))
            # Calculate sum
            sm = [[8/(z**2*(z**2-1))*math.exp(-(z/c)**2*a*t) for z in jz] for t in x]
            # Final equation
            return [c**2*(1-sum(s)) for s in sm]

        # Fit function
        popt, pcov = sp.optimize.curve_fit(diff_rad, [x*1e12 for x in time_ax], msd_r_n, p0=[1, 20, float(pore["diam"])/2-0.2], bounds=(0, np.inf))

        print("Diffusion radial: "+"%.3f" % (popt[0]*1e3)+" 10^-9 m^2 s^-1; Number of zeros: "+"%2i" % (math.ceil(popt[1]))+"; Radius: "+"%5.2f" % popt[2])

    # Plot
    if is_plot:
        legend = []

    if not intent or intent == "axial":
        sns.lineplot(x=[x*1e12 for x in time_ax], y=msd_z_n)
        if is_plot:
            legend += ["Axial"]
        if is_fit:
            sns.lineplot(x=[x*1e12 for x in time_ax], y=[dz*2*time_ax[x]/1e5/1e-7**2 for x in range(len_window)])
            legend += ["Fitted Axial"]

    if not intent or intent == "radial":
        sns.lineplot(x=[x*1e12 for x in time_ax], y=msd_r_n)
        if is_plot:
            legend += ["Radial"]
        if is_fit:
            sns.lineplot(x=[x*1e12 for x in time_ax], y=diff_rad([x*1e12 for x in time_ax], *popt))
            legend += ["Fitted Radial"]

    if is_plot:
        plt.xlabel("Time (ps)")
        plt.ylabel(r"Mean square displacement (nm$^2$)")
        plt.legend(legend)


def bins(link, ax_area=[0.2, 0.8], is_norm=False):
    """This function calculates the axial (z-axis) diffusion coefficient as a
    function of the radial distance. This is done by sampling the mean square
    displacement for all molecules in a radial sub volume.

    The mean square displacement is sampled in function
    :func:`poreana.sample.Sample._diffusion_bin`.

    For each bin, the msd is summed up, resulting into a msd slope for each
    bin. Thus, the axial diffusion coefficient can be calculated using

    .. math::

        D_\\text{axial}=\\frac12\\frac{\\text{msd}_\\text{axial}[i]-\\text{msd}_\\text{axial}[j]}{t_i-t_j}.

    Note that the msd is evaluated in the area, where the slope is uniform,
    which means that the first twenty and last twenty percent should be
    neglected.

    If ``is_norm`` is set to **True**, the radius will be normalized in respect
    to the effective radius which means, the last radius that has a
    Diffusion greater than zero is taken

    .. math::

        r_\\text{norm}=\\frac{1}{r_\\text{eff}}r.

    Parameters
    ----------
    link : string
        Link to hdf5 data file generated by the sample routine :func:`poreana.sample.Sample._diffusion_bin`
    ax_area : list, optional
        Bin area percentage to calculate the axial diffusion coefficient
    intent : string, optional
        Set to **plot**, for plotting or set to **line** to only return the
        lineplot, leave empty for nothing
    is_norm : bool, optional
        True to normalize x-axis

    Returns
    -------
    diffusion : list
        List of the slope of the non-normalized diffusion coefficient
    """
    # Load data object
    sample = utils.load(link)

    # Load data
    inp = sample["inp"]
    width = sample["data"]["width"]
    msd_z = sample["data"]["z"]
    norm = sample["data"]["n"]
    len_window = int(inp["len_window"])
    len_step = inp["len_step"]
    len_frame = inp["len_frame"]


    # Normalize
    msd_norm = [[msd_z[i][j]/norm[i][j] if norm[i][j] > 0 else 0 for j in range(len_window)] for i in range(int(inp["bin_num"])+1)]

    # Calculate slope
    f_start = int(ax_area[0]*len_window)
    f_end = int(ax_area[1]*len_window)
    time_ax = [x*len_step*len_frame for x in range(len_window)]
    slope = [(msd_norm[i][f_end]-msd_norm[i][f_start])/(time_ax[f_end]-time_ax[f_start]) for i in range(int(inp["bin_num"])+1)]

    # Calculate diffusion coefficient
    diff = [msd*1e-9**2/2*1e2**2*1e5 for msd in slope]  # 10^-9 m^2s^-1

    return {"width": width, "diff": diff, "is_norm": is_norm}


def bins_plot(data, intent="plot", kwargs={}):
    """This function calculates the axial plots the diffusion profile calculated
    by the function :func:`poreana.diffusion.bins`.

    Parameters
    ----------
    data : string
        Data dictionary from function func:`poreana.diffusion.bins`
    intent : string, optional
        Set to **plot**, for plotting or set to **line** to only return the
        lineplot, leave empty for nothing
    kwargs: dict, optional
        Dictionary with plotting parameters
    """
    # Initialize
    width = data["width"]
    diff = data["diff"]
    is_norm = data["is_norm"]

    # Normalize x-axis
    if is_norm:
        for i in range(len(diff)-1, 0, -1):
            if diff[i] > 0:
                x_max = width[i+1]
                break

        bins_norm = [x/x_max for x in width]

    # Plot
    if intent == "plot" or intent == "line":
        x_axis = bins_norm if is_norm else width
        sns.lineplot(x=x_axis[:-1], y=diff, **kwargs)

    if intent == "plot":
        if is_norm:
            plt.xlabel("Normalized distance from pore center")
        else:
            plt.xlabel("Distance from pore center (nm)")
        plt.ylabel(r"Diffusion coefficient ($10^{-9}$ m${^2}$ s$^{-1}$)")


def mean(diff_data, dens_data, ax_area=[0.2, 0.8], is_print=True):
    """This function uses the diffusion coefficient slope obtained from
    function :func:`poreana.diffusion.bins` and the density slope of function
    :func:`poreana.density.bins` to calculate a weighted diffusion
    coefficient inside the pore

    .. math::

        \\langle D_\\text{axial}\\rangle
        =\\frac{\\int\\rho(r)D_\\text{axial}(r)dA(r)}{\\int\\rho(r)dA(r)}.

    In a discrete form, following formula is evaluated

    .. math::

        \\langle D_\\text{axial}\\rangle=\\frac{\\sum_{i=1}^n\\rho(r_i)D_\\text{axial}(r_i)A(r_i)}{\\sum_{i=1}^n\\rho(r_i)A(r_i)}

    with the partial area

    .. math::

        A(r_i)=\\pi(r_i^2-r_{i-1}^2)

    of radial bin :math:`i`. It is assumed that the discretization of the
    density is finer than the diffusion. Therefore, the diffusion values for
    each density bin are interpolated between the nearest available diffusion
    values.

    Parameters
    ----------
    diff_data : dictionary
        Diffusion data dictionary from function :func:`poreana.diffusion.bins`
    dens_data : dictionary
        Density data dictionary from function :func:`poreana.density.bins`
    ax_area : list, optional
        Bin area percentage to calculate the axial diffusion coefficient
    is_print : bool, optional
        True to print mean diffusion

    Returns
    -------
    diff_weight : float
        Density weighted mean axial diffusion in 10^-9 m^2s^-1
    """
    # Load density
    dens_x = dens_data["sample"]["data"]["in_width"]
    dens_y = dens_data["num_dens"]["in"]

    # Load diffusion
    diff_w = diff_data["width"][1]
    diff_x = diff_data["width"]
    diff_y = diff_data["diff"]

    # Integrate density and density weighted diffusion
    temp_diff = []
    dens_int = 0
    diff_int = 0
    for dens_bin, dens_x_val in enumerate(dens_x[:-1]):
        # only consider effective radius
        if dens_y[dens_bin]>0:
            # Find closest bins in diffusion
            diff_bins = [math.floor(dens_x_val/diff_w), math.ceil(dens_x_val/diff_w)]

            # Check boundary
            if diff_bins[1] <= len(diff_x):
                interp_x = [diff_x[bin_id] for bin_id in diff_bins]
                interp_y = [diff_y[bin_id] for bin_id in diff_bins]

                # Linear interpolation for diffusion
                diff_y_val = np.interp(dens_x_val, interp_x, interp_y)
                temp_diff.append(diff_y_val)
            else:
                diff_y_val = diff_y[diff_bins[0]]

            # Integrate
            int_area = (dens_x[dens_bin+1]**2-dens_x[dens_bin]**2)
            dens_int += int_area*dens_y[dens_bin]
            diff_int += int_area*dens_y[dens_bin]*diff_y_val

    # Normalize
    diff_weight = diff_int/dens_int

    # Output
    if is_print:
        print("Mean Diffusion axial: "+"%.3f" % diff_weight+" 10^-9 m^2s^-1")

    return diff_weight


######################################
# Diffusion - MC - Transition Matrix #
######################################
def mc_trans_mat(link, step, kwargs={}, is_norm=False, is_diagonal=False, is_length=True, limit = 0):
    """This function plots the occupation of a normalized transition matrix as a
    heatmap. To normalize the transition matrix the number of the frame are used.
    This means that all entries in the transition matrix are divided by the
    total number of frames which are used to estimate the matrix.

    Parameters
    ----------
    link: string
        Link to the diffusion hdf5, obj or yml data file generated by the sample routine
        :func:`poreana.sample.Sample.init_diffusion_mc` or the MC run
        :func:`poreana.mc.MC.run`
    step: integer
        Step length for which the transition matrix is to be displayed
    kwargs: dict, optional
        Dictionary with plotting parameters
    is_norm: bool,optional
        Normalized the transition with the number for frames
    is_diagonal: bool, optional
        Set the matrix diagonal to zero
    is_length: bool, optional
        x axis is output in box length
    limit: float, optional
        Set the lowest occupancy value for calculating the width of occupancy
    """
    # Load results from the output object file
    data = utils.load(link)

    # Sample obj file is loaded
    if "data" in data:
        trans_mat = data["data"][step]
        inp = data["inp"]
        frame_num = inp["num_frame"]
        frame_length = float(inp["len_frame"])
        bins = inp["bins"]

    # MC obj file is loaded
    else:
        model = data["model"]
        trans_mat = model["data"][step]
        frame_num = model["num_frame"]
        frame_length = float(model["len_frame"])
        bins = model["bins"]

    # Calculation of the width of the diagonal occupation of the transition matrix
    results=[]
    # Loop over the rows of the transition matrix
    for i in range(len(trans_mat[:])-1):
        # Right side of the matrix (row)
        row_right= trans_mat[i][i:]/frame_num
        # Left side of the matrix (row)
        row_left = trans_mat[i][:i]/frame_num
        # If a one element is zero call index
        if np.where(row_right<=limit)[0].size != 0:
            results.append(np.min(np.where(row_right<=limit)))
        if np.where(row_left<=limit)[0].size != 0:
            results.append(i-np.max(np.where(row_left<=limit)))

    # Get average of index
    idx_avg = np.mean(results)
    print("Width of occupancy: " + str(idx_avg))


    # Normalized transition matrix with frame number
    if is_norm:
        trans_mat = trans_mat/frame_num

    # Set diagonal elements of the transition matrix to zero
    if is_diagonal:
        trans_mat[np.diag_indices_from(trans_mat)] = 0



    # Set x axis from bins to box length
    if is_length:
        trans_new={}
        for i in range(len(bins)-1):
            trans_new[str(round(bins[i],2))] = trans_mat[i]
        trans_mat = pd.DataFrame(trans_new)


    # Set title with selected lag time
    plt.title("Lagtime: "+ str(step * frame_length * 10**12) + " ps", fontsize=10)

    # Plot the normalized transition matrix in a heatmap
    sns.heatmap(trans_mat, **kwargs)
    if is_length:
        plt.xlabel("Box length (nm)")


##################
# Diffusion - MC #
##################
def mc_fit(link, len_step=[], section=[], is_std=True, is_print=True, is_plot=True, kwargs_scatter={}, kwargs_line={}):
    """This function uses the diffusion profiles over box length which are
    calculated in the function :func:`poreana.mc.MC.run` to estimate
    the final diffusion coefficient. For that a line is fitted of the averaged
    diffusion profiles :math:`D_{\\mathrm{mean}}({\\Delta t_{\\alpha}})` for the
    different lag times as a function of :math:`1/\\Delta t_{\\alpha}`. The
    function plots all mean diffusion
    :math:`D_\\mathrm{mean}({\\Delta t_{\\alpha}})` over the
    inverse lag times and the linear fit. Additionally, the final diffusion
    coefficient :math:`D` for a lag time
    :math:`\\Delta t_{\\alpha} \\rightarrow \infty` is printed. Also, a table
    of the selected step length and the belonging
    :math:`D_{\\mathrm{mean}}({\\Delta t_{\\alpha}})` can be displayed.

    In order to be able to estimate how the choice of lag times affects the
    result, an additional error estimation is possible. For this all possible
    fourth tuples of the entire calculated lag times are used to estimate a
    mean diffusion coefficient :math:`\\langle D \\rangle` which is the mean
    value of all fourth tuples fitting results. Also, the standard deviation of
    all fitting results is printed. If there is no big difference between the
    diffusion coefficient :math:`D` and the mean diffusion coefficient
    :math:`\\langle D \\rangle`, all :math:`D_\\mathrm{mean}({\\Delta t_{\\alpha}})`
    are on a straight line. The standard deviation also can be used to check
    if the result fluctuates strongly at the choice of other lag times. This
    error estimate is calculated if :math:`\\mathrm{is}\_{\\mathrm{std}}=True`.

    Parameters
    ----------
    link: string
        Link to the diffusion hdf5 data file generated by the MC Alogrithm
        :func:`poreana.mc.MC.run`
    len_step: integer, list, optional
        List of the different step length, if it is [] all calculated lag times
        will be used to fit the diffusion coefficent
    section : list, string, optional
        If :math:`\\mathrm{section} = \\mathrm{"pore"}` the pore section is
        fitted.
        If :math:`\\mathrm{section} = \\mathrm{"reservoir"}` the left reservoir
        area is fitted.
        If :math:`\\mathrm{section} = \\mathrm{[a,b]}` the box area between a
        and b is fitted.
    is_std : bool, optional
        True to calculate a mean diffusion coefficient
        to assess the dependency of the result on the selected lag time
    is_print : bool, optional
        Print diffusion coefficient
    is_plot : bool, optional
        Show the fitting plot
    kwargs_scatter: dict, optional
        Dictionary with plotting parameters for the points
    kwargs_line: dict, optional
        Dictionary with plotting parameters for the fitting line


    Returns
    -------
    diffusion: float
        Diffusion coefficient :math:`D \ \\left( 10^{-9} \\frac{m^2}{s}\\right)` for the
        calculated system
    diffusion_mean: float
        Mean diffusion coefficient :math:`\\langle D \\rangle \ \\left(10^{-9} \\frac{m^2}{s}\\right)`
        for the calculated system
    diff_table : obj
        Table of used lag times with the associated
        :math:`D_{\\mathrm{mean}}({\\Delta t_{\\alpha}}) \ \\left(10^{-9} \\frac{m^2}{s}\\right)`
    res : float
        residual :math:`\\left(10^{-9} \\frac{m^2}{s}\\right)` for fitting the selected lag times
    """
    # Load data
    data = utils.load(link)


    results = data["output"]
    diff_bin = results["diff_profile"]
    model = data["model"]
    bins = model["bins"]
    dt = model["len_frame"]
    diff_unit = model["diffusion unit"]

    # If no specific step length is chosen take the step length from the object file
    if not len_step:
        len_step = model["len_step"]

    # If a pore system is considered
    if "pore" in data:
        pore = data["pore"]
        res = float(pore["res"])
        box = pore["box"]

    # Set vector
    diff_bin_vec = {}

    # Pore
    if isinstance(section, str) and section== "pore":
        # If only the pore area should be considered
        if "pore" in data:
            # Set section
            area = [res, box[2]-res]

        # If only the pore area should be considered
        else:
            print("obj.-file includes results of a simple box")
            return;

    # Reservoir
    elif isinstance(section, str) and section == "reservoir":
        # If only the pore area should be considered
        # Load pore system
        if "pore" in data:
            # Set section
            area = [0,res]

        # If only the pore area should be considered
        else:
            print("obj.-file includes results of a simple box")
            return;

    elif isinstance(section, str) and section not in ["reservoir", "pore"]:
        print("Wrong input for section! Check documentation for available options")
        return;

    elif isinstance(section, list) and len(section)>=3:
        print("Wrong input for section! Check documentation for available options")
        return;

    #Area section
    elif isinstance(section, list):
        area = section

    # If section is not defined -> entire system
    if not section:
        # Set section
        area = [bins[0], bins[-1]]

    # Cut profile
    # Calculated start and end bin index of the pore area
    index_start = np.digitize(area[0], bins)
    index_end = np.digitize(area[1], bins)


    for i in len_step:
        diff_bin_vec[i] = [diff_bin[i][j] for j in range(index_start, index_end)]

    # Calculate mean diffusion coefficient and standard deviation
    # Mean diffusion means the averaged of all possible fourth tuple fitting results
    if is_std:
        len_step_all = len_step

        # Determine all possible combinations
        a = list(itertools.combinations(len_step_all, 2))

        # Allocate the result vector
        res = np.zeros(len(a))

        # Fit to infinity lag time for all combinations
        for i in range(len(a)):

            # Set the current tuple
            rand = list(a[i])

            # Calculate the mean diffusion (m^2/s) over all bins
            D_mean = [np.mean(np.exp([diff_bin_vec[i][j] + diff_unit for j in range(len(diff_bin_vec[i]))])) * 10**3 for i in rand]


            # Calculate the inverse lag time for the linear fit
            lagtime_inverse = [1 / (i * dt * 10**-12) for i in rand]

            # Fit a linear line
            fit = np.poly1d(np.polyfit(lagtime_inverse, D_mean, 1))

            # Set the fitted diffusion coefficient on the results list
            res[i] = fit(0)

        # Determine standard deviation and mean diffusion of fourth touple method
        diffusion_mean = np.mean(res)
        std = res.std()

    # Calculate the mean diffusion (m^2/s) over all bins
    D_mean = [np.mean(np.exp([diff_bin_vec[i][j] + diff_unit for j in range(len(diff_bin_vec[i]))])) * 10**3 for i in len_step]

    # Calculate the inverse lag time (1/s) for the linear fit
    lagtime_inverse = [1 / (len_step[i] * dt * 10**-12) for i in range(len(len_step))]

    # Fit a linear line
    fit = np.poly1d(np.polyfit(lagtime_inverse, D_mean, 1))

    # Calculate residuals
    res = np.polyfit(lagtime_inverse, D_mean, 1, full=True)

    # Print the diffusion coefficient for an entire system
    if is_print:
        if not section:
            print("\nDiffusion axial: "+"%.4e" % (fit(0) * 10 **-9) + " m^2/s\n")

            # Print resudial for fitting
            print("Residual: "+"%.4e" % (float(res[1]) * 10 **-9) + " m^2/s\n")

            # If is_std true print the results of the calculations
            if is_std:
                print("Mean Diffusion axial: "+"%.4e" % (diffusion_mean * 10 **-9)  + " m^2/s\n")

        # Print the diffusion coefficient in the pore area
        if section=="pore":
            print("\nDiffusion axial (Pore): "+"%.4e" % (fit(0) * 10 **-9) + " m^2/s\n")


            # Print resudial for fitting
            print("Residual: "+"%.4e" % (float(res[1]) * 10 **-9) + " m^2/s\n")

            # If is_std true print the results of the calculations
            if is_std:
                print("Mean Diffusion axial (Pore): "+"%.4e" % (diffusion_mean * 10 **-9)  + " m^2/s\n")

        # Print the diffusion coefficient in the reservoir area
        if section=="reservoir":
            print("\nDiffusion axial (Reservoir): "+"%.4e" % (fit(0) * 10 **-9) + " m^2/s\n")

            # Print resudial for fitting
            print("Residual: "+"%.4e" % (float(res[1]) * 10 **-9) + " m^2/s\n")

            # If is_std true print the results of the calculations
            if is_std:
                print("Mean Diffusion axial (Peservoir): "+"%.4e" % (diffusion_mean * 10 **-9) + " m^2/s\n")

        # Print the diffusion coefficient in a selected section
        if (isinstance(section, list)) and len(section)==2:
            print("\nDiffusion axial ([" + "%.2f" % (area[0]) + ", " + "%.2f" % (area[1]) + "]): "+"%.4e" % (fit(0) * 10 **-9) + " m^2/s\n")

            # Print resudial for fitting
            print("Residual: "+"%.4e" % (float(res[1]) * 10 **-9) + " m^2/s\n")

            # If is_std true print the results of the calculations
            if is_std:
                print("Mean Diffusion axial (["+ "%.2f" % (area[0]) + ", " + "%.2f" % (area[1]) +"]): "+"%.4e" % (diffusion_mean * 10 **-9)  + " m^2/s\n")

        # Print std deviation
        if is_std:
            print("Standard deviation: "+"%.4e" % (std * 10 **-9) + " m^2/s\n")

    # Set data frame for the used lag times
    data = [str("%.2f" % D_mean[i] ) for i in range(len(len_step))]
    diff_table = pd.DataFrame(data, index=list(len_step), columns=list(['$D_\mathrm{mean} \ (10^{-9} \mathrm{m^2s^{-1}})$']))
    diff_table = pd.DataFrame(diff_table.rename_axis('Step Length', axis=1))
    styler = diff_table.style.set_caption('Selected step length')
    diff_table = styler.set_properties(**{'text-align': 'right'})
    diff_table = diff_table.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])

    # Set vectors for plotting
    D_mean_vec = [D_mean[i] for i in range(len(lagtime_inverse))]
    lag_time_vec = [1 / (len_step[i] * dt) for i in range(len(len_step))]
    x_vec = np.arange(0, max(lag_time_vec) * 2, (max(lag_time_vec) * 2) / 5)

    # Fit a linear line and calculated diffusion coefficent
    fit = np.poly1d(np.polyfit(lag_time_vec, D_mean_vec, 1))
    diffusion = fit(0)

    if is_plot:
        # Plot the results
        plt.xlim(0, 1.5*max(lag_time_vec))
        plt.ylim(0, 1.5*max(fit(x_vec)))
        sns.scatterplot(x=lag_time_vec, y=D_mean_vec, **kwargs_scatter)
        sns.lineplot(x=x_vec, y=fit(x_vec), **kwargs_line)
        legend = ["$D_{\mathrm{fit}}$", "$D_{\mathrm{mean}}(\\Delta t_{\\alpha})$"]
        plt.legend(legend)
        plt.xlabel(r"Inverse lag time ($10^{12} \ \mathrm{s^{-1}})$")
        plt.ylabel(r"Diff. coeff. ($10^{-9} \ \mathrm{m^2s^{-1}}$)")

    return diffusion, diffusion_mean, diff_table, float(res[1][0])



def mc_profile(link, len_step=[], section=[], infty_profile=True,  is_plot=True, kwargs={}):
    """This function plots the diffusion profile for an infinity
    lag time (:math:`\\Delta t_{\\alpha} \\rightarrow \\infty`) over the box
    fitted with the specified :math:`\\mathrm{len}\_\\mathrm{step}` list.
    Additionally, it is possible to display the diffusion profiles for the
    calculated lag times. Therefore, the
    :math:`\\mathrm{infty}\_\\mathrm{profile}` has to be false.
    The list :math:`\\mathrm{len}\_\\mathrm{step}`
    contains the calculated step length which should be
    displayed.

    Parameters
    ----------
    link : string
        Link to the diffusion hdf5 data file generated by the MC Algorithm
        :func:`poreana.mc.MC.run`
    len_step : integer list, optional
        List of the different step length, if it is [] all diffusion profiles
        depending on the lag time are used to calculate the diffusion profile
        for an infinite lag time or the individual diffusion profiles for these
        lag times are shown
    section : list, string, optional
        If :math:`\\mathrm{section} = \\mathrm{"pore"}` the pore section is
        fitted.
        If :math:`\\mathrm{section} = \\mathrm{"reservoir"}` the left reservoir
        area is fitted.
        If :math:`\\mathrm{section} = \\mathrm{[a,b]}` the box area between a
        and b is fitted.
    infty_profile : bool, optional
        Set to false to display all individual diffusion profiles for
        the selected lag times
    is_plot : bool, optional
        Show diffusion profile
    kwargs: dict, optional
        Dictionary with plotting parameters

    Returns
    -------
    diff_profile_fit : list
        Diffusion profile :math:`\\left(10^{-9} \\frac{m^2}{s}\\right)` for an infinite
        lag time
    diff_profiles : dict
        Diffusion profiles :math:`\\left(10^{-9} \\frac{m^2}{s}\\right)` for every
        calculated lag time
    bins : list
        bins over the box length
    res : float
        residual :math:`\\left(10^{-9} \\frac{m^2}{s}\\right)` for fitting the infinite diffusion profile
    """
    # Load data
    data = utils.load(link)

    # Load results
    results = data["output"]
    diff_bin = results["diff_profile"]

    # Load model inputs
    model = data["model"]
    diff_unit = model["diffusion unit"]
    bins = model["bins"]
    dt = model["len_frame"]

    # If no specific step length is chosen take the step length from the object file
    if not len_step:
        len_step = model["len_step"]

    # If a pore system is considered
    if "pore" in data:
        pore = data["pore"]
        res = pore["res"]
        box = pore["box"]


    # Set dictionaries
    legend = []
    diff_bin_vec = {}
    diff_profile_fit = []
    res_list = []


    # Pore
    if isinstance(section, str) and section== "pore":
        # If only the pore area should be considered
        # Load pore system
        if "pore" in data:
            # Set section
            area = [res, box[2]-res-2*(bins[1]-bins[0])]


        # If only the pore area should be considered
        else:
            print("obj.-file includes results of a simple box")
            return;

    # Reservoir
    elif isinstance(section, str) and section== "reservoir":
        # If only the pore area should be considered
        # Load pore system
        if "pore" in data:
            # Set section
            area = [0,res]

        # If only the pore area should be considered
        else:
            print("obj.-file includes results of a simple box")
            return;

    elif isinstance(section, str) and section not in ["reservoir", "pore"]:
        print("Wrong input for section! Check documentation for available options")
        return;

    elif isinstance(section, list) and len(section)>=3:
        print("Wrong input for section! Check documentation for available options")
        return;

    #Area section
    elif isinstance(section, list):
        area = section

    # If section is not defined -> entire system
    if not section:
        # Set section
        area = [bins[0], bins[-1]]

    # Cut profile
    # Calculated start and end bin index of the pore area
    index_start = np.digitize(area[0], bins)
    index_end = np.digitize(area[1], bins)

    # Save for all lag times the cutted profile
    for i in len_step:
        diff_bin_vec[i] = [diff_bin[i][j] for j in range(index_start, index_end)]

    # Set bin list
    bins = [bins[i] for i in range(index_start, index_end)]


    # Calculate the inverse lag time (1/s) for the linear fit
    lagtime_inverse = [1 / (len_step[i] * dt * 10**-12) for i in range(len(len_step))]

    # Calculate diffusion profiles
    diff_profiles = {}
    for i in len_step:
        diff_profiles[i] = [np.exp(diff_bin_vec[i][j] + diff_unit) * 10 ** 3 for j in range(len(bins))]

    # If infty_profile is false the profiles for the different lag times are plotted

    if not infty_profile:
        # Plot the profiles for the
        if is_plot:
            for i in len_step:
                sns.lineplot(x=bins, y=(diff_profiles[i]), **kwargs)       # Diffusion in m^2/s

        # Plot the diffusion profiles for the different lag times
        legend = ["$\\Delta t_{\\alpha}$ = " + str(len_step[i] * dt) + " ps" for i in range(len(len_step))]

    # If infty_profile is True the diffusion profile for a infinity lag times is shwon
    if len(len_step) >= 2 and infty_profile:
        # Initialize fit vector

        # Calculate the mean diffusion over all bins
        for i in range(len(bins)):
            diff = [np.exp(diff_bin_vec[step][i] + diff_unit) * 10 ** 3 for step in len_step]   # Diffusion in m^2/s

            # fit a linear line
            fit = np.poly1d(np.polyfit(lagtime_inverse, diff, 1))

            # Append diffusion at t-> infty
            diff_profile_fit.append(fit(0))

            # Calculate residuals for fitting
            res = np.polyfit(lagtime_inverse, diff, 1, full=True)

            # Append residual
            res_list.append(float(res[1]))

        # Plot fitted diffusion profile
        if is_plot:
            sns.lineplot(x=bins, y=diff_profile_fit, **kwargs)       # Diffusion in m^2/s

    # Set legend for lag times
    if is_plot:
        if not infty_profile and len(len_step) >= 2:
            legend.append("$\\Delta t_{\\alpha} \\rightarrow \\infty$ ps")


        # Set plot properties
        # Plot axis title for a entire system
        plt.ylabel(r"Diff. coeff. ($10^{-9} \ \mathrm{m^2s^{-1}}$)")
        plt.xlabel(r"Box length (nm)")
        plt.xlim([min(bins),max(bins)])
        if legend:
            plt.legend(legend)

    # Check if profile was fitted
    if not res_list:
        res_list = 0
    else:
        np.mean(res_list)

    return diff_profile_fit, diff_profiles, bins, res_list




###########################
# Diffusion - MC - Radial #
###########################
# def mc_fit_radial(link, len_step=None):
#     """This function uses the diffusion profiles over box length which are
#     calculated in the function :func:`do_mc_cycles` to estimate the final
#     diffusion coefficient. For that a line is fitted of the averaged diffusion
#     profiles :math:`D_{\\mathrm{mean}}({\\Delta_{ij}t_{\\alpha}})` for the
#     different lag times as a function of :math:`1/\\Delta_{ij}t_{\\alpha}`. The
#     function plots all mean diffusion
#     :math:`D_{mean}({\\Delta_{ij}t_{\\alpha}})` over the inverse lag times and
#     the linear fit. Additional the final mean diffusion coefficient :math:`D`
#     for a lagtime :math:`\\Delta_{ij}t_{\\alpha} \\rightarrow \infty` is
#     printed. Also a table of the selected step length and the belonging
#     :math:`D_{\\mathrm{mean}}({\\Delta_{ij}t_{\\alpha}})` can be displayed.
#
#     In order to be able to estimate how the choice of lag times affects the
#     result, an additional error estimation is possible. For this all possible
#     fourth tuples of the entire calculated lag times are used to estimate a
#     mean diffusion coefficient :math:`\\langle D \\rangle` which is the mean
#     value of all fourth tuples fitting results. Also the standard deviation of
#     all fitting results is printed. If there is no big difference between the
#     diffusion coefficient :math:`D` and the mean diffusion ceofficient
#     :math:`\\langle D \\rangle`, all :math:`D_{mean}({\\Delta_{ij}t_{\\alpha}})`
#     are on a straight line. The standard deviation can also be used to check
#     if the result fluctuates strongly at the choice of other lag times. This
#     error estimate is calculated if :math:`\\mathrm{is}_{\\mathrm{std}}=True`.
#
#     Parameters
#     ----------
#     link: string
#         Link to the diffusion data object generated by the sample routine
#         :func:`MC.do_mc_cycles`
#     len_step: integer list, optional
#         List of the different step length, if it is None the lag times which
#         are contains in the obj file are printed
#     is_std : bool, optional
#         True to calculate a mean diffusion coefficient and a Bootstrap diffusion
#         to assess the dependency of the results on the chosen lag time
#
#     Returns
#     -------
#     diffusion: float
#         Radial diffusion coefficient :math:`D \ \\left(\\frac{m^2}{s}\\right)` for the
#         calculated system
#     diffusion_mean: float
#         Radial mean diffusion coefficient :math:`\\langle D \\rangle \ \\left(\\frac{m^2}{s}\\right)` for the
#         calculated system
#     diff_table : obj
#         Table of used lag times with the associated
#         :math:`D_{\\mathrm{mean}}({\\Delta_{ij}t_{\\alpha}}) \ \\left(\\frac{m^2}{s}\\right)`
#     """
#     # Load Results from the output object file
#     results = utils.load(link)
#     diff_bin = results["diff_radial_profile"]
#     df_bin = results["df_profile"]
#     inp = results["model"]
#     bins = inp["bins"]
#     diff_unit = inp["diffusion radial unit"]
#     dt = inp["len_frame"]
#
#     # If no specific step length is chosen take the step length from the object file
#     if not len_step:
#         len_step = inp["len_step"]
#
#     # Calculate the mean diffusion over all bins
#     D_mean = [np.mean(np.exp(diff_bin[i] + diff_unit)) * 10**-
#               6 for i in len_step]   # Diffusion in m^2/s
#
#     # Calculate the inverse lag time (1/s) for the linear fit
#     lagtime_inverse = [1 / (len_step[i] * dt * 10**-12) for i in range(len(len_step))]
#
#     # Fit a linear line
#     fit = np.poly1d(np.polyfit(lagtime_inverse, D_mean, 1))
#
#     # set the x value for the linear fit
#     x = np.arange(0, max(lagtime_inverse) * 2, (max(lagtime_inverse) * 2) / 5)
#
#     # Calculate the diffusion coefficient as a mean value over the bin
#     print("\nMean Radial Diffusion: "+"%.4e" % fit(0) + " m^2/s\n")
#
#     # Integration
#     bin_width = bins[1]-bins[0]
#
#     # Set data frame for the used lag times
#     D_mean_pd = [str("%.2e" % D_mean[i]) for i in range(len(D_mean))]
#     diff_table = pd.DataFrame({'$D_\mathrm{mean}$': D_mean_pd, '$D_\mathrm{avg} (\mathrm{m^2s^{-1}})$': diff_avg_pd}, index=list(len_step))
#     diff_table = pd.DataFrame(diff_table.rename_axis('Step Length', axis=1))
#     styler = diff_table.style.set_caption('Selected step length')
#     diff_table = styler.set_properties(**{'text-align': 'right'})
#     diff_table = diff_table.set_table_styles([dict(selector='th', props=[('text-align', 'left')])])
#
#     # Plot the results
#     plt.xlim(0, 1.5*max(lagtime_inverse))
#     plt.ylim(0, 1.5*max(fit(x)))
#     sns.scatterplot(x=lagtime_inverse, y=D_mean, color='red')
#     sns.lineplot(x=x, y=fit(x))
#     legend = ["linear fit", "$D_{\mathrm{mean}}$"]
#     plt.legend(legend)
#     plt.xlabel(r"Inverse lagtime ($\mathrm{s^{-1}})$")
#     plt.ylabel(r"radial Diff. coeff. ($\mathrm{m^2s^{-1}}$)")
#
#     return fit(0), diff_table
#
#
# def mc_profile_radial(link, len_step=None, avg=False):
#     """This function plots the radial diffusion profile for a infinity
#     lag time (:math:`\\Delta_{ij}t_{\\alpha} \\rightarrow \\infty`) over the box
#     fitted with the specified :math:`\\mathrm{len}\_\\mathrm{step}` list.
#     Additionally, it is possible to display the diffusion profiles for the
#     calculated lag times. Therefore, the
#     :math:`\\mathrm{infty}\_\\mathrm{profile}` has to be false.
#     The list :math:`\\mathrm{len}\_\\mathrm{step}`
#     contains the calculated step length which should be displayed.
#
#     Parameters
#     ----------
#     link : string
#         Link to the diffusion data object generated by the
#         :func:`mc.MC.do_mc_cycles`
#     len_step : integer list, optional
#         List of the different step length, if it is None the lag times which
#         are contains in the obj file are printed
#     infty_profile : bool, optional
#         If its True only the profile for an infity lagtime is plotted
#     """
#     # Load Results from the output object file
#     results = utils.load(link)
#     diff_bin = results["diff_radial_profile"]
#     inp = results["model"]
#     dt = inp["len_frame"]
#     diff_unit = inp["diffusion radial unit"]
#     bins = inp["bins"]
#     bins = [(bins[i]) * 10**-8 for i in range(len(bins))]
#
#     if len_step is None:
#         len_step = inp["len_step"]
#
#     # Plot the diffusion profiles for the different lag times
#     legend = ["lagtime " + str(len_step[i] * dt) + " ps" for i in range(len(len_step))]
#
#     # Plot the diffusion profiles (m^2/s) for the
#     for i in len_step:
#         sns.lineplot(x=bins, y=np.exp(diff_bin[i] + diff_unit) * 10**-6)
#
#     # Averged profile over the selected lag times
#     if len(len_step) != 1 and avg:
#         diff_mean = np.zeros(len(bins[:-1]))
#         diff = []
#         for i in len_step:
#             diff_mean += diff_bin[i]
#         diff = [diff_mean[i]/len(len_step) for i in range(len(diff_mean))]
#
#         # Plot mean profile (m^2/s) over the selected lag times
#         sns.lineplot(x=bins, y=np.exp(diff + diff_unit) * 10**-6)
#         legend.append("avg profile")
#
#     # Set plot properties
#     plt.xlabel(r"Box length (m)")
#     plt.ylabel(r"radial Diffusion ($\mathrm{m^2s^{-1}}$)")
#     plt.legend(legend)
