import numpy as np
import pandas as pd
from scipy.integrate import odeint
from scipy.optimize import curve_fit
from scipy.stats import t as t_dist
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import corner

class Bli:
    """Bio-Layer Interferometry (BLI) data class.

    Organizes BLI experiment data and analysis methods. Mirrors the R S4 class
    from the smoke package.

    Attributes:
        traces: DataFrame with first column 'time' and remaining columns as sensor traces.
        ligand: Array of ligand concentrations (high to low).
        t_exp: Array of two values: start times of association and dissociation.
        status: Dict tracking which processing steps have been completed.
        kinetics: Dict with fitted model results (parameters, covariance, residuals, etc.).
        k_on0: Initial on-rate estimate.
        k_off0: Initial off-rate estimate.
    """

    def __init__(self):
        self.traces = pd.DataFrame()
        self._ligand = np.array([])
        self._t_exp = np.array([])
        self.status = {
            "align_load": False,
            "double_blank": False,
            "baseline": False,
            "estimate": False,
            "fit_kinetics": False,
            "ode_kinetics": False,
        }
        self.kinetics = {}
        self.k_on0 = None
        self.k_off0 = None

    @property
    def ligand(self):
        if len(self._ligand) == 0:
            raise ValueError("please input ligand concentration")
        return self._ligand

    @ligand.setter
    def ligand(self, value):
        self._ligand = np.asarray(value, dtype=float)

    @property
    def t_exp(self):
        if len(self._t_exp) == 0:
            raise ValueError("please input start times of association and dissociation")
        return self._t_exp

    @t_exp.setter
    def t_exp(self, value):
        self._t_exp = np.asarray(value, dtype=float)

    def __repr__(self):
        lines = ["Bli Class", "=========="]
        lines.append("status:")
        for k, v in self.status.items():
            lines.append(f"  {k}: {v}")

        lines.append("--------------------")
        lines.append("traces: DataFrame with 1st column of time")
        if self.traces.empty:
            lines.append("  empty")
        else:
            lines.append(str(self.traces.iloc[:5, :6]))
            n_traces = self.traces.shape[1] - 1
            n_time = self.traces.shape[0]
            lines.append(f"  {n_traces} traces, {n_time} time points")

        lines.append("--------------------")
        lines.append("ligand: concentration from high to low")
        if len(self._ligand) == 0:
            lines.append("  empty")
        else:
            lines.append(f"  {self._ligand}")

        lines.append("--------------------")
        lines.append("t_exp: start time of association and dissociation")
        if len(self._t_exp) == 0:
            lines.append("  empty")
        else:
            lines.append(f"  {self._t_exp}")

        lines.append("--------------------")
        lines.append("k_on0: initial On-rate")
        if self.k_on0 is None:
            lines.append("  (to be estimated or input)")
        else:
            lines.append(f"  {self.k_on0}")

        lines.append("--------------------")
        lines.append("k_off0: initial Off-rate")
        if self.k_off0 is None:
            lines.append("  (to be estimated or input)")
        else:
            lines.append(f"  {self.k_off0}")

        lines.append("--------------------")
        lines.append("kinetics: fitting model")
        if self.kinetics:
            lines.append(str(self.get_kinetics()))
        else:
            lines.append("  no fitting model yet")
        lines.append("--------------------")
        return "\n".join(lines)

    def _get_trace_data(self):
        """Return time array and traces DataFrame (without time column)."""
        t = self.traces.iloc[:, 0].values
        trace_data = self.traces.iloc[:, 1:]
        return t, trace_data

    def _make_palette(self, n):
        """Generate a cyan-to-magenta color palette."""
        cmap = LinearSegmentedColormap.from_list("cm", ["cyan", "magenta"], N=n)
        return [cmap(i / max(1, n - 1)) for i in range(n)]

    # ---- Processing methods ----

    def align_load(self, load_start, load_end):
        """Align traces to loading step.

        Args:
            load_start: Start time of loading step.
            load_end: End time of loading step.
        """
        t, traces = self._get_trace_data()
        n_cols = traces.shape[1]
        if n_cols % 2 != 0:
            raise ValueError("the original traces should be paired")

        n_sample = n_cols // 2
        load0_range = (t < load_start) & (t > (load_start - 60))
        load_range = (t < load_end) & (t > (load_end - 60))

        load0 = []
        load = []
        traces_arr = traces.values

        for i in range(n_sample):
            col_sample = i * 2
            col_ref = i * 2 + 1

            # Predict at load_start using linear fit on load0_range
            for col in [col_sample, col_ref]:
                t_fit = t[load0_range]
                y_fit = traces_arr[load0_range, col]
                coeffs = np.polyfit(t_fit, y_fit, 1)
                load0.append(np.polyval(coeffs, load_start))

            # Predict at load_end using linear fit on load_range
            for col in [col_sample, col_ref]:
                t_fit = t[load_range]
                y_fit = traces_arr[load_range, col]
                coeffs = np.polyfit(t_fit, y_fit, 1)
                load.append(np.polyval(coeffs, load_end))

        load0 = np.array(load0)
        load = np.array(load)

        # Subtract load0 offset
        traces_arr = traces_arr - load0[np.newaxis, :]
        load = load - load0

        # Normalize by loading factor
        load = load.reshape(2, n_sample, order="F")
        mean_load = load.mean(axis=1)
        load_factor = (load / mean_load[:, np.newaxis]).flatten(order="F")
        traces_arr = traces_arr / load_factor[np.newaxis, :]

        col_names = [str(i + 1) for i in range(n_cols)]
        self.traces = pd.DataFrame(
            np.column_stack([t, traces_arr]),
            columns=["time"] + col_names,
        )
        self.status["align_load"] = True

    def double_blank(self):
        """Apply double-blank correction to paired traces.

        Traces must be paired (sample, reference) and the last pair serves
        as the blank reference.
        """
        if self.status["double_blank"]:
            raise RuntimeError("already double-blank!")

        t, traces = self._get_trace_data()
        traces_arr = traces.values
        n_cols = traces_arr.shape[1]
        if n_cols % 2 != 0:
            raise ValueError("the double-referenced traces should be paired")

        # Subtract reference from sample for each pair
        bk_cor = traces_arr[:, ::2] - traces_arr[:, 1::2]
        # Subtract blank pair (last column)
        db_cor = bk_cor[:, :-1] - bk_cor[:, -1:]

        if db_cor.shape[1] != len(self._ligand):
            import warnings
            warnings.warn("ligand concentrations mis-match trace number")
            col_names = [str(i) for i in range(db_cor.shape[1])]
        else:
            col_names = [str(round(x, 6)) for x in np.round(self._ligand, 6)]

        self.traces = pd.DataFrame(
            np.column_stack([t, db_cor]),
            columns=["time"] + col_names,
        )
        self.status["double_blank"] = True

    def baseline_correct(self, t_start, t_end):
        """Subtract a linear baseline fitted over [t_start, t_end].

        Args:
            t_start: Start of baseline window.
            t_end: End of baseline window.
        """
        t, traces = self._get_trace_data()
        traces_arr = traces.values
        base_range = (t > t_start) & (t < t_end)
        base_t = t[base_range]

        baselines = []
        for i in range(traces_arr.shape[1]):
            y = traces_arr[base_range, i]
            coeffs = np.polyfit(base_t, y, 1)
            baselines.append(np.polyval(coeffs, t_end))

        baselines = np.array(baselines)
        db_norm = traces_arr - baselines[np.newaxis, :]

        if db_norm.shape[1] != len(self._ligand):
            import warnings
            warnings.warn("ligand concentrations mis-match trace number")
            col_names = [str(i) for i in range(db_norm.shape[1])]
        else:
            col_names = [str(round(x, 6)) for x in np.round(self._ligand, 6)]

        self.traces = pd.DataFrame(
            np.column_stack([t, db_norm]),
            columns=["time"] + col_names,
        )
        self.status["baseline"] = True

    def estimate(self):
        """Estimate initial kOn and kOff from individual trace fits."""
        t, traces = self._get_trace_data()
        traces_arr = traces.values
        t_exp = self.t_exp
        t_ass = t_exp[1] - t_exp[0]
        lig = self.ligand

        if len(lig) != traces_arr.shape[1]:
            raise ValueError("the trace number mis-match ligand concentrations")

        tot_range = t >= t_exp[0]
        t_tot = t[tot_range]
        d_tot = traces_arr[tot_range, :]
        t_rel = t_tot - t_exp[0]

        t_a = t_rel[t_rel < t_ass]
        t_d = t_rel[t_rel > t_ass]
        d_a = d_tot[t_rel < t_ass, :]
        d_d = d_tot[t_rel > t_ass, :]

        k_on = np.zeros(len(lig))
        k_off = np.zeros(len(lig))

        def assoc_model(t_val, k, A):
            return A * (1 - np.exp(-k * t_val))

        def dissoc_model(t_val, k, y_e):
            return y_e + (d0_val - y_e) * np.exp(-k * (t_val - 600))

        for i in range(len(lig)):
            # Association fit
            y = d_a[:, i]
            a_e = y[-1]
            y_p = 1 - y / a_e
            mask = (y_p > 0) & (t_a < (t_ass / 10))
            t_p = t_a[mask]
            y_p_log = np.log(y_p[mask])
            # Linear fit through origin: log(yP) ~ tP
            k0 = -np.sum(y_p_log * t_p) / np.sum(t_p ** 2)

            popt_a, _ = curve_fit(assoc_model, t_a, y, p0=[k0, a_e])
            k_on[i] = popt_a[0]

            # Dissociation fit
            y = d_d[:, i]
            d0_val = assoc_model(t_a[-1], *popt_a)
            d_e = y[-1]
            y_p = (y - d_e) / (d0_val - d_e)
            mask = (y_p > 0) & (t_d < (t_ass * 1.1))
            t_p = t_d[mask] - 600
            y_p_log = np.log(y_p[mask])
            k0 = -np.sum(y_p_log * t_p) / np.sum(t_p ** 2)

            popt_d, _ = curve_fit(dissoc_model, t_d, y, p0=[k0, d_e])
            k_off[i] = popt_d[0]

        self.k_on0 = np.mean(k_on / lig)
        self.k_off0 = np.mean(k_off)
        self.status["estimate"] = True

    def fit_kinetics(self, negative_response=False):
        """Global fit of association and dissociation kinetics.

        Uses a 1:1 binding model with linear drift terms for each ligand
        concentration. Requires k_on0 and k_off0 to be set (via estimate()
        or manually).

        Parameters:
        -----------
        negative_response : bool, optional (default=False)
            If True, assumes the response is negative (decreasing signal),
            and uses maximum value as initial R_max estimate.
            If False, assumes positive response and uses minimum value.
        """
        t, traces = self._get_trace_data()
        traces_arr = traces.values
        t_exp = self.t_exp
        t_ass = t_exp[1] - t_exp[0]
        lig = self.ligand
        k_on0 = self.k_on0
        k_off0 = self.k_off0

        if k_on0 is None or k_off0 is None:
            raise ValueError("no initial kOn0 and kOff0: input or estimate!")
        if len(lig) != traces_arr.shape[1]:
            raise ValueError("the trace number mis-match ligand concentrations")

        tot_range = t >= t_exp[0]
        t_tot = t[tot_range]
        d_tot = traces_arr[tot_range, :]
        t_rel = t_tot - t_exp[0]
        r_max0 = np.min(d_tot) if negative_response else np.max(d_tot)

        n_lig = len(lig)
        n_t = len(t_rel)

        # Build vectorized data (matching R's rep/unlist approach)
        x_vec = np.repeat(lig, n_t)
        t_vec = np.tile(t_rel, n_lig)
        y_vec = d_tot.flatten(order="F")

        # Indicator matrices for each ligand
        d_lig = np.zeros((len(y_vec), n_lig))
        for i in range(n_lig):
            d_lig[i * n_t:(i + 1) * n_t, i] = 1.0

        m_a = (t_vec < t_ass).astype(float)
        m_d = (t_vec > t_ass).astype(float)

        def global_model(_x, r_max, k_on, k_off, *drift_params):
            sa = np.array(drift_params[:n_lig])
            sd = np.array(drift_params[n_lig:2 * n_lig])
            ja = np.array(drift_params[2 * n_lig:3 * n_lig])
            jd = np.array(drift_params[3 * n_lig:4 * n_lig])

            # Association term
            y_ass = r_max / (1 + k_off / k_on / x_vec) * (1 - np.exp(-(k_on * x_vec + k_off) * t_vec))
            # Drift during association for each ligand
            drift_a = np.zeros_like(t_vec)
            for i in range(n_lig):
                drift_a += d_lig[:, i] * (ja[i] + sa[i] * t_vec)

            # Dissociation term
            y_diss = (r_max / (1 + k_off / k_on / x_vec)
                      * (1 - np.exp(-(k_on * x_vec + k_off) * t_ass))
                      * np.exp(-k_off * (t_vec - t_ass)))
            # Drift during dissociation for each ligand
            drift_d = np.zeros_like(t_vec)
            for i in range(n_lig):
                drift_d += d_lig[:, i] * (jd[i] + sd[i] * (t_vec - t_ass))

            return m_a * (y_ass + drift_a) + m_d * (y_diss + drift_d)

        # Initial parameters
        p0 = [r_max0, k_on0, k_off0] + [0.0] * (4 * n_lig)
        param_names = (
            ["rMax", "kOn", "kOff"]
            + [f"sa{i + 1}" for i in range(n_lig)]
            + [f"sd{i + 1}" for i in range(n_lig)]
            + [f"ja{i + 1}" for i in range(n_lig)]
            + [f"jd{i + 1}" for i in range(n_lig)]
        )

        x_dummy = np.arange(len(y_vec))
        popt, pcov = curve_fit(global_model, x_dummy, y_vec, p0=p0, maxfev=10000)

        # Compute fitted values and residuals
        y_fitted = global_model(x_dummy, *popt)
        residuals = y_vec - y_fitted

        # Standard errors from covariance
        n_obs = len(y_vec)
        n_params = len(popt)
        dof = n_obs - n_params
        se = np.sqrt(np.diag(pcov))

        # t-values and p-values
        t_values = popt / se
        p_values = 2 * (1 - t_dist.cdf(np.abs(t_values), dof))

        self.kinetics = {
            "params": dict(zip(param_names, popt)),
            "se": dict(zip(param_names, se)),
            "t_values": dict(zip(param_names, t_values)),
            "p_values": dict(zip(param_names, p_values)),
            "covariance": pcov,
            "residuals": residuals.reshape(d_tot.shape, order="F"),
            "fitted": y_fitted.reshape(d_tot.shape, order="F"),
            "rss": np.sum(residuals ** 2),
            "dof": dof,
            "param_names": param_names,
        }
        self.status["fit_kinetics"] = True
        self.status["ode_kinetics"] = False

    def ode_kinetics(self, negative_response=True):
        """Global fit of association and dissociation kinetics using ODE.

        Uses a 1:1 binding model with linear drift terms for each ligand
        concentration. Requires k_on0 and k_off0 to be set (via estimate()
        or manually).

        Parameters:
        -----------
        negative_response : bool, optional (default=False)
            If True, assumes the response is negative (decreasing signal),
            and uses maximum value as initial R_max estimate.
            If False, assumes positive response and uses minimum value.
        """
        t, traces = self._get_trace_data()
        traces_arr = traces.values
        t_exp = self.t_exp
        t_ass = t_exp[1] - t_exp[0]
        lig = self.ligand
        k_on0 = self.k_on0
        k_off0 = self.k_off0

        if k_on0 is None or k_off0 is None:
            raise ValueError("no initial kOn0 and kOff0: input or estimate!")
        if len(lig) != traces_arr.shape[1]:
            raise ValueError("the trace number mis-match ligand concentrations")

        tot_range = t >= t_exp[0]
        t_tot = t[tot_range]
        d_tot = traces_arr[tot_range, :]
        t_rel = t_tot - t_exp[0]
        r_max0 = np.min(d_tot) if negative_response else np.max(d_tot)

        n_lig = len(lig)
        n_t = len(t_rel)

        # Build vectorized data (matching R's rep/unlist approach)
        x_vec = np.repeat(lig, n_t)
        t_vec = np.tile(t_rel, n_lig)
        y_vec = d_tot.flatten(order="F")

        # Indicator matrices for each ligand
        d_lig = np.zeros((len(y_vec), n_lig))
        for i in range(n_lig):
            d_lig[i * n_t:(i + 1) * n_t, i] = 1.0

        m_a = (t_vec < t_ass).astype(float)
        m_d = (t_vec > t_ass).astype(float)

        def model(R, t, k_on, k_off, R_max, ligand):
            """ODE model for 1:1 binding interaction"""
            dRdt = k_on * ligand * (R_max - R) - k_off * R
            return dRdt

        def global_model(_x, r_max, k_on, k_off, *drift_params):
            """Global model incorporating ODE solution and drift terms"""
            # Extract drift parameters
            sa = np.array(drift_params[:n_lig])
            sd = np.array(drift_params[n_lig:2 * n_lig])
            ja = np.array(drift_params[2 * n_lig:3 * n_lig])
            jd = np.array(drift_params[3 * n_lig:4 * n_lig])

            # Initialize arrays for storing results
            y_ass = np.zeros_like(t_vec)
            y_diss = np.zeros_like(t_vec)

            # Solve ODE for each ligand concentration
            for i in range(n_lig):
                mask = d_lig[:, i] == 1
                t_mask = t_vec[mask]

                # Association phase
                t_ass_mask = t_mask[t_mask < t_ass]
                if len(t_ass_mask) > 0:
                    R_ass = odeint(model, 0, t_ass_mask,
                                 args=(k_on, k_off, r_max, lig[i]))
                    y_ass[mask & (t_vec < t_ass)] = R_ass.flatten()

                # Dissociation phase
                t_diss_mask = t_mask[t_mask > t_ass] - t_ass
                if len(t_diss_mask) > 0:
                    R0_diss = y_ass[mask & (t_vec < t_ass)][-1] if len(t_ass_mask) > 0 else 0
                    R_diss = odeint(model, R0_diss, t_diss_mask,
                                  args=(k_on, k_off, r_max, 0))  # ligand = 0 for dissociation
                    y_diss[mask & (t_vec > t_ass)] = R_diss.flatten()

            # Add drift terms
            drift_a = np.zeros_like(t_vec)
            drift_d = np.zeros_like(t_vec)
            for i in range(n_lig):
                drift_a += d_lig[:, i] * (ja[i] + sa[i] * t_vec)
                drift_d += d_lig[:, i] * (jd[i] + sd[i] * (t_vec - t_ass))

            return m_a * (y_ass + drift_a) + m_d * (y_diss + drift_d)

        # Initial parameters
        p0 = [r_max0, k_on0, k_off0] + [0.0] * (4 * n_lig)
        param_names = (
            ["rMax", "kOn", "kOff"]
            + [f"sa{i + 1}" for i in range(n_lig)]
            + [f"sd{i + 1}" for i in range(n_lig)]
            + [f"ja{i + 1}" for i in range(n_lig)]
            + [f"jd{i + 1}" for i in range(n_lig)]
        )

        x_dummy = np.arange(len(y_vec))
        popt, pcov = curve_fit(global_model, x_dummy, y_vec, p0=p0, maxfev=10000)

        # Compute fitted values and residuals
        y_fitted = global_model(x_dummy, *popt)
        residuals = y_vec - y_fitted

        # Standard errors from covariance
        n_obs = len(y_vec)
        n_params = len(popt)
        dof = n_obs - n_params
        se = np.sqrt(np.diag(pcov))

        # t-values and p-values
        t_values = popt / se
        p_values = 2 * (1 - t_dist.cdf(np.abs(t_values), dof))

        self.kinetics = {
            "params": dict(zip(param_names, popt)),
            "se": dict(zip(param_names, se)),
            "t_values": dict(zip(param_names, t_values)),
            "p_values": dict(zip(param_names, p_values)),
            "covariance": pcov,
            "residuals": residuals.reshape(d_tot.shape, order="F"),
            "fitted": y_fitted.reshape(d_tot.shape, order="F"),
            "rss": np.sum(residuals ** 2),
            "dof": dof,
            "param_names": param_names,
        }
        self.status["ode_kinetics"] = True
        self.status["fit_kinetics"] = False

    def get_kinetics(self):
        """Return a summary DataFrame of kinetics parameters (KD, rMax, kOn, kOff).

        Returns:
            DataFrame with columns Estimate, Std.Error, t value, Pr(>|t|).
        """
        if not self.kinetics:
            raise RuntimeError("no fitting model yet")

        params = self.kinetics["params"]
        se = self.kinetics["se"]
        t_vals = self.kinetics["t_values"]
        p_vals = self.kinetics["p_values"]

        k_on = params["kOn"]
        k_off = params["kOff"]
        kd = k_off / k_on
        kd_se = np.sqrt(
            (se["kOn"] / k_on) ** 2 + (k_off * se["kOff"] / k_on ** 2) ** 2
        )

        rows = {
            "KD": [kd, kd_se, np.nan, np.nan],
            "rMax": [params["rMax"], se["rMax"], t_vals["rMax"], p_vals["rMax"]],
            "kOn": [k_on, se["kOn"], t_vals["kOn"], p_vals["kOn"]],
            "kOff": [k_off, se["kOff"], t_vals["kOff"], p_vals["kOff"]],
        }
        return pd.DataFrame.from_dict(
            rows, orient="index", columns=["Estimate", "Std.Error", "t value", "Pr(>|t|)"]
        )

    # ---- Plotting methods ----

    def plot_traces(self, **kwargs):
        """Plot all traces with a cyan-to-magenta color gradient."""
        if self.traces.empty:
            raise RuntimeError("traces is empty")

        t, traces = self._get_trace_data()
        colors = self._make_palette(traces.shape[1])

        fig, ax = plt.subplots()
        for i in range(traces.shape[1]):
            c = list(colors[i])
            c[3] = 0.3  # alpha
            ax.plot(t, traces.iloc[:, i], color=c, **kwargs)

        if len(self._t_exp) > 0:
            for tv in self._t_exp:
                ax.axvline(tv, linestyle="--", color="gray")

        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Response (nm)")
        plt.tight_layout()
        return fig, ax

    def plot_kinetics(self, **kwargs):
        """Plot traces with fitted kinetics overlay."""
        if not self.kinetics:
            raise RuntimeError("not fitting kinetics yet")

        t, traces = self._get_trace_data()
        t_exp = self.t_exp
        mask = t >= t_exp[0]
        t_plot = t[mask] - t_exp[0]
        traces_plot = traces.values[mask, :]
        y_fitted = self.kinetics["fitted"]

        colors = self._make_palette(traces_plot.shape[1])

        fig, ax = plt.subplots()
        for i in range(traces_plot.shape[1]):
            c = list(colors[i])
            c_solid = list(colors[i])
            c[3] = 0.3
            ax.plot(t_plot, traces_plot[:, i], color=c, **kwargs)
            ax.plot(t_plot, y_fitted[:, i], color=c_solid, **kwargs)

        ax.axvline(t_exp[1] - t_exp[0], linestyle="--", color="gray")
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Response (nm)")
        plt.tight_layout()
        return fig, ax

    def plot_residuals(self, **kwargs):
        """Plot residuals from the kinetics fit."""
        if not self.kinetics:
            raise RuntimeError("not fitting kinetics yet")

        t, traces = self._get_trace_data()
        t_exp = self.t_exp
        mask = t >= t_exp[0]
        t_plot = t[mask] - t_exp[0]
        traces_plot = traces.values[mask, :]
        residuals = self.kinetics["residuals"]

        bound = np.ptp(traces_plot) / 2
        colors = self._make_palette(residuals.shape[1])

        fig, ax = plt.subplots()
        for i in range(residuals.shape[1]):
            c = list(colors[i])
            c[3] = 0.3
            ax.plot(t_plot, residuals[:, i], color=c, **kwargs)

        ax.axvline(t_exp[1] - t_exp[0], linestyle="--", color="gray")
        ax.axhline(0, color="black")
        ax.set_ylim(-bound, bound)
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("Response (nm)")
        plt.tight_layout()
        return fig, ax


    def plot_confidence_contours(self):
        """
        Generate a confidence contour plot (corner plot) for kinetics parameters.
        """
        if not self.kinetics:
            raise RuntimeError("Kinetics have not been fitted yet.")

        # Extract parameters and covariance matrix from kinetics
        pp = np.array([self.kinetics["params"]["kOn"],
                       self.kinetics["params"]["kOff"],
                       self.kinetics["params"]["rMax"]])

        # Extract the full covariance matrix
        full_pcov = self.kinetics["covariance"]

        # Create a mapping of parameter indices
        param_indices = {
            "kOn": 0,
            "kOff": 1,
            "rMax": 2
        }

        # Extract the covariance matrix for k_on, k_off, and rMax
        pcov = full_pcov[np.ix_([param_indices["kOn"], param_indices["kOff"], param_indices["rMax"]],
                                  [param_indices["kOn"], param_indices["kOff"], param_indices["rMax"]])]

        # Degrees of freedom
        dof = self.kinetics["dof"]
        # Residual variance
        resV = np.sum(self.kinetics["residuals"] ** 2) / dof

        # Standard errors
        cv = np.sqrt(np.diag(pcov)) / pp
        diag = np.sqrt(np.diag(np.diag(pcov)))
        gaid = np.linalg.inv(diag)
        pcor = gaid @ pcov @ gaid

        # Generate samples from the multivariate normal distribution
        samples = np.random.multivariate_normal(pp, pcov, size=20000)

        # Create corner plot for parameters
        fig = corner.corner(samples, labels=["k_on", "k_off", "R_max"],
                            show_titles=True,
                            contour_kwargs={'colors': 'blue'},
                            contour_levels=[0.68, 0.95, 0.997],
                            fill_contours=True)

        plt.tight_layout()
        return fig
