#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------

__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"

"""Plotting utilities for the WIP-P3M package."""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.colors import TwoSlopeNorm
from matplotlib.gridspec import GridSpec
from mpl_toolkits.axes_grid1 import make_axes_locatable
import cmocean.cm as cm
import plotly.graph_objects as go

from wip3m.plot_params import *  # type: ignore

from wip3m.logger import getCustomLogger, INDENT, UNINDENT

logger = getCustomLogger(__name__)

# Configure global plotting settings
setup_plotting()

fs = GLOBAL_FS_SMALL
fs_titles = GLOBAL_FS_LARGE
cols = COLOUR_LIST
cmap = cm.thermal  # type: ignore

def matplotlib_to_plotly(cmap, n=255):
    """Convert a matplotlib colormap to a Plotly colorscale."""
    colorscale = []
    for i in range(n):
        norm = i / (n - 1)
        r, g, b, _ = cmap(norm)
        colorscale.append([norm, f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'])
    return colorscale

thermal_plotly = matplotlib_to_plotly(cmap)

def plotly_3d(
    field: np.ndarray,
    size: int = 128,
    L: float = None,  # type: ignore
    colormap: str | list = "RdYlBu",
    limits: str = "max",
) -> go.Figure:
    """
    Create an interactive 3D plot of volume slices using Plotly.

    Parameters
    ----------
    field : array-like
        3D data field to visualise.
    size : int, optional
        Size of the field along one dimension. Default is 128.
    L : float, optional
        Physical size of the field in Mpc/h. Used for axis labels only.
    colormap : str, optional
        Colour map for visualisation. Default is 'RdYlBu'.
    limits : str, optional
        Colour scale limits ('max', 'truncate', or 'default'). Default
        is 'max'.

    Returns
    -------
    go.Figure
        Plotly figure object.
    """
    volume = field.T
    rows, cols = volume[0].shape

    # Define colour scale limits
    if limits == "max":
        maxcol = np.max(np.abs(volume))
        mincol = -maxcol
    elif limits == "truncate":
        maxcol = min(np.max(-volume), np.max(volume))
        mincol = -maxcol
    else:
        maxcol = np.max(volume)
        mincol = np.min(volume)
    midcol = np.mean(volume)

    # Generate frames for the animation
    nb_frames = size
    frames = [
        go.Frame(
            data=go.Surface(
                z=(size - k) * np.ones((rows, cols)),
                surfacecolor=np.flipud(volume[cols - 1 - k]),
                cmin=mincol,
                cmid=midcol,
                cmax=maxcol,
            ),
            name=str(k),  # Frames must be named for proper animation
        )
        for k in range(nb_frames)
    ]

    # Initial plot configuration
    fig = go.Figure(
        frames=frames,
        data=go.Surface(
            z=size * np.ones((rows, cols)),
            surfacecolor=np.flipud(volume[cols // 2]),
            colorscale=colormap,
            cmin=mincol,
            cmid=midcol,
            cmax=maxcol,
            colorbar=dict(thickness=20, ticklen=4),
        ),
    )

    def frame_args(duration):
        """Helper function to set animation frame arguments."""
        return {
            "frame": {"duration": duration},
            "mode": "immediate",
            "fromcurrent": True,
            "transition": {"duration": duration, "easing": "linear"},
        }

    # Add animation slider
    sliders = [
        {
            "pad": {"b": 10, "t": 60},
            "len": 0.9,
            "x": 0.1,
            "y": 0,
            "steps": [
                {
                    "args": [[f.name], frame_args(0)],
                    "label": str(k),
                    "method": "animate",
                }
                for k, f in enumerate(fig.frames)
            ],
        }
    ]

    # Configure layout with or without physical size
    layout_config = dict(
        title="Slices in density field",
        width=600,
        height=600,
        scene=dict(
            zaxis=dict(range=[0, size - 1], autorange=False),
            xaxis_title="x [Mpc/h]",
            yaxis_title="y [Mpc/h]",
            zaxis_title="z [Mpc/h]",
            aspectratio=dict(x=1, y=1, z=1),
        ),
        updatemenus=[
            {
                "buttons": [
                    {"args": [None, frame_args(50)], "label": "&#9654;", "method": "animate"},
                    {"args": [[None], frame_args(0)], "label": "&#9724;", "method": "animate"},
                ],
                "direction": "left",
                "pad": {"r": 10, "t": 70},
                "type": "buttons",
                "x": 0.1,
                "y": 0,
            }
        ],
        sliders=sliders,
    )
    if L is not None:
        layout_config["scene"]["xaxis"] = dict(
            ticktext=[0, L / 2, L],
            tickvals=[0, size / 2, size],
            title="x [Mpc/h]",
        )
        layout_config["scene"]["yaxis"] = dict(
            ticktext=[0, L / 2, L],
            tickvals=[0, size / 2, size],
            title="y [Mpc/h]",
        )
        layout_config["scene"]["zaxis"]["ticktext"] = [0, L / 2, L]
        layout_config["scene"]["zaxis"]["tickvals"] = [0, size / 2, size]

    fig.update_layout(**layout_config)
    return fig

def clear_large_plot(fig):
    """
    Clear a figure to free up memory.

    Parameters
    ----------
    fig : matplotlib.figure.Figure
        The figure to clear.
    """
    from IPython.display import clear_output

    del fig
    clear_output()

def load_force_diagnostic(filename):
    """
    Load force diagnostic data.

    Parameters
    ----------
    filename : str
        Path to the CSV file written by `run_force_subtraction_test`
        from p3m.c.

    Returns
    -------
    r : ndarray
        Bin-centred distances.
    fmag : ndarray
        Total force magnitudes |f_after - f_before|.
    data : recarray
        Full structured array with all columns.
    """
    data = np.genfromtxt(filename, delimiter=",", names=True)
    # Append the magnitude of force difference to the data
    fmag = np.sqrt(data["fx"]**2 + data["fy"]**2 + data["fz"]**2)
    data = np.lib.recfunctions.append_fields(data, "fmag", fmag, usemask=False, asrecarray=True)

    return data["distance"], data["fmag"], data

def plot_force_distance(r, fmag, f_max=1e-1, a=None, title=None, save_path=None):
    """
    Plot total force magnitude vs distance and a theoretical 1/r² profile.

    Parameters
    ----------
    r : array_like
        Distance values (bin centers or actual distances).
    fmag : array_like
        Corresponding force magnitudes.
    f_max : float, optional
        Maximum force magnitude for including points in the plot (default: 1e-1).
    a : float or None, optional
        If provided, the theoretical prefactor for the 1/r² profile.
    title : str, optional
        Plot title.
    save_path : str or None, optional
        If provided, path to save the figure as a PDF.
    """
    from scipy.optimize import curve_fit

    # Select points to include in the plot
    mask = fmag < f_max
    rs = r[mask]
    fmag_select = fmag[mask]

    fig, ax = plt.subplots()
    ax.scatter(rs, fmag_select, s=10, alpha=0.8, label="Particle/hole pairs", color="tab:blue")

    # If a is provided, plot the theoretical curve
    if a is not None:
        def inverse_square(r, a):
            return a / r**2
        r = np.linspace(np.min(rs), np.max(rs), 300)
        f_th = inverse_square(r, a)
        ax.plot(r, f_th, "r-", label=r"$\propto 1/r^2$")

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel(r"Distance $r$ [Mpc/$h$]")
    ax.set_ylabel(r"Force magnitude [code units]")
    if title is not None:
        ax.set_title(title)
    ax.legend()
    ax.grid(True, which="both", ls=":", lw=0.5)

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=300)
        print(f"Figure saved to: {save_path}")
    plt.show()

def plot_force_distance_comparison(rr, ff, ll, L, Np, Npm, ss=None, a=None, title=None, save_path=None):
    """
    Plot force magnitude vs distance.

    Parameters
    ----------
    rr : list of array_like
        List of distance arrays.
    ff : list of array_like
        List of force magnitude arrays corresponding to each distance array.
    ll : list of str
        List of labels for each dataset.
    ss : list of str or None, optional
        List of symbols for each dataset (default: None).
    L : float
        Physical size of the field in Mpc/h.
    Np : int
        Number of particles per dimension.
    Npm : int
        Number of PM cells per dimension.
    a : float or None, optional
        If provided, the theoretical prefactor for the 1/r² profile.
    title : str, optional
        Plot title.
    save_path : str or None, optional
        If provided, path to save the figure as a PDF.
    """
    fig, ax = plt.subplots()

    colours = plt.rcParams["axes.prop_cycle"].by_key()["color"]
    handles1 = []

    if ss is None:
        ss = ["o"] * len(rr)
    for i, (r, f, label, symbol) in enumerate(zip(rr, ff, ll, ss)):
        scatter = ax.scatter(r, f, alpha=0.8, s=20, label=label,
                             color=colours[i % len(colours)], marker=symbol)
        handles1.append(scatter)

    loc1="lower right"

    # Theoretical curve
    theory_line = None
    if a is not None and len(rr) > 0:
        loc1="upper right"
        def inverse_square(r, a):
            return a / r**2
        r_min = min(np.min(r) for r in rr)
        r_max = max(np.max(r) for r in rr)
        r_plot = np.linspace(r_min, r_max, 300)
        f_th = inverse_square(r_plot, a)
        theory_line, = ax.plot(r_plot, f_th, "k-", label=r"theory $\propto 1/r^2$")
        handles1.append(theory_line)

    # Characteristic vertical reference scales
    nyquist = 2 * L / Npm
    epsilon = 0.03 * L / Np
    xs = 1.25 * L / Npm
    xr = 4.5 * xs
    line1 = ax.axvline(x=nyquist, color="black", linestyle="-", lw=1, label="Nyquist (PM grid)")
    line2 = ax.axvline(x=2*epsilon, color="gray", linestyle="--", lw=2, label=r"Particle length $2\epsilon$")
    line3 = ax.axvline(x=xs, color="gray", linestyle="-.", lw=2, label=r"Split scale $x_s$")
    line4 = ax.axvline(x=xr, color="gray", linestyle=":", lw=2, label=r"Short-range reach $x_r$")
    print(f"Nyquist: {nyquist:.2f} Mpc/h")
    print(f"Particle length: {2*epsilon:.2f} Mpc/h")
    print(f"Split scale: {xs:.2f} Mpc/h")
    print(f"Short-range reach: {xr:.2f} Mpc/h")

    # Set log-log axes and labels
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel(r"Distance $r$ [Mpc/$h$]")
    ax.set_ylabel(r"Force magnitude [code units]")
    if title is not None:
        ax.set_title(title)

    # Legend for data
    legend1 = ax.legend(handles=handles1, loc=loc1, frameon=True, fontsize=GLOBAL_FS_TINY)
    ax.add_artist(legend1)

    # Legend for vertical lines
    handles2 = [line1, line2, line3, line4]
    fig.subplots_adjust(bottom=0.3)
    legend2 = fig.legend(handles=handles2, loc='lower center', ncol=2, frameon=False)

    ax.grid(True, which="both", ls=":", lw=0.5)

    if save_path:
        plt.savefig(save_path)
        print(f"Figure saved to: {save_path}")
    plt.show()

def plot_timestepping_diagnostics(log_path, aiDrift=None, TimeStepDistribution=None, nsteps=None, ymin=2e-3, save_path=None, exact=True, show=True):
    """
    Plot and compare timestep limiters from log file, as given by the
    `compute_time_step` function in p3m.c.

    Parameters
    ----------
    log_path : str
        Path to the timestep log file.
    aiDrift : array_like or None
        Reference array of scale factors.
    TimeStepDistribution : int or None
        Baseline timestep distribution mode (0, 1, or 2).
    nsteps : int or None
        Total number of baseline steps (used for labelling).
    save_path : str or None
        If provided, path to save the figure.
    show : bool, optional
        Whether to display the plot (default: True).
    """
    logger.info(f"Plotting timestep limiters from {log_path}...")
    INDENT()

    a, da, da_max, da_H_approx, da_H_exact, da_dyn_approx, da_dyn_exact, da_p3m, da_late = np.loadtxt(log_path, delimiter=",", unpack=True, skiprows=0)

    if aiDrift is not None:
        assert np.allclose(a, aiDrift), "aiDrift does not match the time steps in the log file."

    fig, ax = plt.subplots(figsize=(5.5, 4))
    ax.set_ylim(ymin, 0.35)
    ax.set_xlim(np.min(a), np.max(a))

    # ax.loglog(a, np.ones_like(a)*1/50 / a, label=r"$\Delta a=1/50$")
    ax.loglog(a, da_H_approx / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta a/a\approx 0.031$ (Hubble, approx)")
    ax.loglog(a, da_dyn_approx / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta a\propto aH(a)\left(G\rho\right)^{-1/2}$ (Dynamical, approx)")
    ax.loglog(a, da_p3m / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta t\propto x_s/\sqrt{\langle v^2\rangle}$ (P3M, untuned)")
    if exact:
        ax.loglog(a, da_dyn_exact/a, marker="o", markersize=4, linestyle="-", label=r"$\Delta t=\eta_\textrm{d}\left(G\rho\right)^{-1/2}$ (Dynamical, exact)")
        ax.loglog(a, da_H_exact / a, marker="x", markersize=4, linestyle="-", label=r"$\Delta t=0.031/H(t)$ (Hubble, exact)")
    ax.loglog(a, da_late / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta a_\textrm{max}=0.05$ (late)")
    ls = "dotted" if TimeStepDistribution in [1, 3] else ""
    ax.loglog(a, da_max / a, color="black", marker="o", markersize=4, linestyle=ls, label="CONCEPT global step")
    ax.loglog(a, da / a, marker="x", markersize=4, linestyle="", label=(
        rf"$\Delta a=1/{nsteps}$ (baseline)" if TimeStepDistribution == 0 else
        f"Baseline ({nsteps} log-spaced steps)" if TimeStepDistribution == 1 else
        f"Actual {nsteps} steps"))

    ax.set_xlabel(r"Scale factor $a$", fontsize=12)
    ax.set_ylabel(r"$\Delta a/a$", fontsize=12)
    ax.xaxis.grid(True, which="both", linestyle=":", linewidth=0.6)
    ax.tick_params(labelsize=10)
    fig.subplots_adjust(bottom=0.37)
    fig.legend(fontsize=10, loc='lower center', ncol=2, frameon=False)

    if save_path:
        plt.savefig(save_path)
        logger.info(f"Figure saved to: {save_path}")
    UNINDENT()
    logger.info(f"Plotting timestep limiters from {log_path} done.")
    if show:
        plt.show()
    # plt.close(fig)


def plot_custom_timestepping_diagnostics(
    log_path,
    aiDrift=None,
    TimeStepDistribution=None,
    nsteps=None,
    fac_hubble=None,
    fac_bend=None,
    fac_p3m_fit=None,
    da_max_early=None,
    da_max_late=None,
    plot_bend=True,
    ymin=2e-3,
    ymax=0.35,
    save_path=None,
    show=True,
):
    """
    Plot and compare timestep limiters from log file, as given by the
    `compute_time_step` function in p3m.c.

    Parameters
    ----------
    log_path : str
        Path to the timestep log file.
    aiDrift : array_like or None
        Reference array of scale factors.
    TimeStepDistribution : int or None
        Baseline timestep distribution mode (0, 1, or 2).
    nsteps : int or None
        Total number of baseline steps (used for labelling).
    fac_hubble : float, optional
        Factor for the Hubble timestep limiter (default: None).
    fac_bend : float, optional
        Factor for the bend timestep limiter (default: None).
    fac_p3m_fit : float, optional
        Factor for the P3M fit timestep limiter (default: None).
    da_max_early : float, optional
        Maximum timestep allowed for early stages (default: None).
    da_max_late : float, optional
        Maximum timestep allowed for late stages (default: None).
    plot_bend : bool, optional
        Whether to plot the bend limiter (default: True).
    ymin : float, optional
        Minimum y-axis limit for the plot (default: 5e-3).
    save_path : str or None
        If provided, path to save the figure.
    show : bool, optional
        Whether to display the plot (default: True).
    """
    log_path_custom = log_path[:-4] + "_custom.txt"
    logger.info(f"Plotting timestep limiters from {log_path} and {log_path_custom}...")
    INDENT()
    a, da, da_max, da_dyn, da_bend, da_p3m, da_p3m_fit, da_hubble, da_hubble_lenient, da_late = np.loadtxt(log_path_custom, delimiter=",", unpack=True, skiprows=0)

    if aiDrift is not None:
        assert np.allclose(a, aiDrift), "aiDrift does not match the time steps in the log file."

    fig, ax = plt.subplots(figsize=(5.5, 4))
    ax.set_ylim(ymin, ymax)
    ax.set_xlim(np.min(a), np.max(a))
    bottom_pad = 0.37

    if fac_hubble is not None:
        lab = rf"$\Delta a/a= {fac_hubble:.3f}$ (Hubble)"
    else:
        lab = r"$\Delta a/a$ constant (Hubble)"
    ax.loglog(a, da_hubble / a, marker="o", markersize=4, linestyle="-", label=lab)
    if da_max_early is not None:
        lab = rf"Hubble, lenient (early $\Delta a_\textrm{{max}}={da_max_early:.3f}$)"
    else:
        lab = "Hubble, lenient"
    ax.loglog(a, da_hubble_lenient / a, marker="o", markersize=4, linestyle="-", label=lab)
    ax.loglog(a, da_dyn / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta a\propto aH(a)\left(G\rho\right)^{-1/2}$ (Dynamical)")
    if plot_bend:
        if fac_bend is not None:
            lab = rf"Bend-approx. P3M ($\eta_\textrm{{P}}={fac_bend:.5f}$)"
        else:
            lab = "Bend-approx. P3M"
        ax.loglog(a, da_bend / a, marker="o", markersize=4, linestyle="-", label=lab)
    if fac_p3m_fit is not None:
        lab = rf"P3M fit ($\eta_\textrm{{Pf}}={fac_p3m_fit:.3f}$)"
        ax.loglog(a, da_p3m_fit / a, marker="o", markersize=4, linestyle="-", label=lab)
        bottom_pad += 0.06
    # _, _, _, _, _, _, _, da_p3m_concept, _ = np.loadtxt(log_path, delimiter=",", unpack=True, skiprows=0)
    # ax.loglog(a, da_p3m_concept / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta t\propto x_s/\sqrt{\langle v^2\rangle}$ (P3M, untuned)")
    ax.loglog(a, da_p3m / a, marker="o", markersize=4, linestyle="-", label=r"$\Delta t\propto x_s/\sqrt{\langle v^2\rangle}$ (P3M)")
    if da_max_late is not None:
        lab = rf"$\Delta a_\textrm{{max}}={da_max_late:.3f}$ (late)"
    else:
        lab = r"$\Delta a_\textrm{max}$ (late)"
    ax.loglog(a, da_late / a, marker="o", markersize=4, linestyle="-", label=lab)
    ls = "dotted" if TimeStepDistribution in [1, 3] else ""
    ax.loglog(a, da_max / a, color="black", marker="o", markersize=4, linestyle=ls, label="Max time step allowed (custom)")
    ax.loglog(a, da / a, marker=".", markersize=4, linestyle="", label=(
        rf"$\Delta a=1/{nsteps}$ (baseline)" if TimeStepDistribution == 0 else
        f"Baseline ({nsteps} log-spaced steps)" if TimeStepDistribution == 1 else
        f"Actual {nsteps} steps"))

    ax.set_xlabel(r"Scale factor $a$", fontsize=12)
    ax.set_ylabel(r"$\Delta a/a$", fontsize=12)
    ax.xaxis.grid(True, which="both", linestyle=":", linewidth=0.6)
    ax.tick_params(labelsize=10)
    fig.subplots_adjust(bottom=bottom_pad)
    fig.legend(fontsize=10, loc='lower center', ncol=2, frameon=False)

    if save_path:
        plt.savefig(save_path, dpi=300)
        logger.info(f"Figure saved to: {save_path}")
    UNINDENT()
    logger.info(f"Plotting timestep limiters from {log_path} and {log_path_custom} done.")
    if show:
        plt.show()
    # plt.close(fig)


def log_poly_model(log_a, *coeffs):
    """Polynomial model. log(y) = c0 + c1*log(a) + c2*log(a)^2 + ..."""
    return sum(c * log_a**i for i, c in enumerate(coeffs))


def fit_p3m(a, da_p3m, degree=3, amin=1e-1):
    """
    Fit log(da_p3m / a) = poly(log(a)) for a > 1e-1.

    Parameters
    ----------
    a : ndarray
        Scale factor.
    da_p3m : ndarray
        P3M limiter values.
    degree : int
        Polynomial degree.

    Returns
    -------
    coeffs : ndarray
        Polynomial coefficients [c0, ..., c_degree].
    """
    from scipy.optimize import curve_fit

    mask = a > amin
    if not np.any(mask):
        raise ValueError(f"No values of 'a' greater than {amin} found.")
    log_a = np.log(a[mask])
    log_y = np.log(da_p3m[mask] / a[mask])
    if degree == 3:
        initial_guess = [-3.0, 0.2, 2.5, 0.6]
    elif degree == 5:
        initial_guess = [-3.5, -0.3, -0.5, -4.0, -3.0, -0.5]
    else:
        initial_guess = np.ones(degree + 1)
    coeffs, _ = curve_fit(log_poly_model, log_a, log_y, p0=initial_guess)
    return coeffs


def evaluate_polynomial_log(coeffs, a):
    log_a = np.log(a)
    return np.exp(log_poly_model(log_a, *coeffs))


def plot_p3m(a, da_p3m, coeffs, amin_fit=1e-1, amin_plot=5e-2):
    """
    Plot da_p3m/a, fit, and residuals.

    Parameters
    ----------
    a : ndarray
        Scale factor values.
    da_p3m : ndarray
        P3M limiter values.
    coeffs : ndarray
        Polynomial coefficients.
    amin_fit : float
        Minimum a value used in the fit (for legend).
    amin_plot: float
        Minimum a value to include in plot.
    """
    mask = a > amin_plot  # Only plot for a > amin_plot
    a = a[mask]
    da_p3m = da_p3m[mask]
    ratio = da_p3m / a
    a_dense = np.logspace(np.log10(a.min()), np.log10(a.max()), 500)
    fit_vals = evaluate_polynomial_log(coeffs, a_dense)
    fit_interp = evaluate_polynomial_log(coeffs, a)
    # ratio = da_p3m/a
    # ratio_fit = exp(fit(log(da_p3m)/log(a))) and we're interested in
    # the residuals, res, of the da_p3m fit, that is,
    # res = log(da_p3m) - log_da_p3m_fit
    #     = log(da_p3m) - log(exp(log_da_p3m_fit))
    #     = log(da_p3m) - log(a*fit_interp)
    #     = log(da_p3m / (a*fit_interp))
    #     = log(da_p3m/a / fit_interp)
    #     = log(ratio / fit_interp)
    # res_exp = ratio / fit_interp
    residuals = ratio / fit_interp

    fig = plt.figure(figsize=(6, 3.5))
    gs = GridSpec(2, 1, height_ratios=[3, 1], hspace=0.05)
    ax1 = fig.add_subplot(gs[0])
    ax2 = fig.add_subplot(gs[1], sharex=ax1)

    # P3M limiter and fit
    ax1.plot(a, ratio, ".", label=r"$\Delta t \propto x_s/\sqrt{\langle v^2 \rangle}$", alpha=0.7)
    ax1.plot(a_dense, fit_vals, "-", label=rf"Fit for $a > {amin_fit}$", lw=1.8)
    ax1.set_yscale("log")
    ax1.set_xscale("log")
    ax1.set_ylabel(r"$\Delta a/a$")
    ax1.legend(loc="upper right", frameon=False)
    ax1.tick_params(which="both", direction="in", top=True, right=True)
    ax1.grid(True, which="major", ls=":", lw=0.5)

    # Residuals
    ax2.plot(a, residuals, "o", color="grey", alpha=0.6, markersize=3)
    ax2.axhline(1.0, color="black", lw=1, ls="--")
    ax2.set_ylabel("res", labelpad=2)
    ax2.set_xlabel(r"Scale factor $a$")
    ax2.set_xscale("log")
    ax2.set_yscale("log")

    ax2.tick_params(which="both", direction="in", top=True, right=True)
    ax2.grid(True, which="major", ls=":", lw=0.5)
    ax1.tick_params(labelbottom=False)
    ax2.get_yaxis().set_major_formatter(plt.FuncFormatter(lambda y, _: f"{y:g}"))

    plt.tight_layout()
    plt.show()
