from __future__ import annotations

import os
from pathlib import Path
from shutil import copy2

import gmsh
import matplotlib.pyplot as plt
import numpy as np
import ogstools as ot
import pyvista as pv
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D

ot.plot.setup.show_region_bounds = False

out_dir = Path(os.environ.get("OGS_TESTRUNNER_OUT_DIR", "_out"))
out_dir.mkdir(parents=True, exist_ok=True)

_dim_map = {"1d": 1, "2d": 2}
dim_norm = str(dimension_str).strip().lower()
if dim_norm not in _dim_map:
    msg = (
        f"Unsupported dimension '{dimension_str}'. Use one of: {list(_dim_map.keys())}"
    )
    raise ValueError(msg)
target_dim = _dim_map[dim_norm]
dim_suffix = dim_norm

MESH_NAME = "rect_band_default"
MESH_ROOT = out_dir / "meshes"
MESH_ROOT.mkdir(parents=True, exist_ok=True)
(MESH_ROOT / dim_suffix).mkdir(parents=True, exist_ok=True)


# geometry / mesh
GEOMETRY = {
    "L": 0.07,  # m
    "H": 0.0099,  # m
    "left_L": 0.024,  # m
    "band_L": 0.022,  # m
    "right_L": 0.024,  # m
    "h": 0.002,  # m
}
MESH = {
    "mesh_name": MESH_NAME,
    "msh_file": MESH_ROOT / dim_suffix / f"{MESH_NAME}.msh",
    "vtu_file": MESH_ROOT / dim_suffix / "domain.vtu",
}

# chemical fields
FIELDS = {
    "C_main": 1.6e-5,
    "Ca_main": 1.0e-8,
    "H_main": 2.5119e-6,  # [H⁺] = 10^(-pH) = 10^(-5.6) = 2.5119e-6 mol/L
    "C_mid": 3.7934e-2,
    "Ca_mid": 7.9567e-3,
    "H_mid": 7.59e-7,  # [H⁺] = 10^(-pH) = 10^(-6.12) ≈ 7.59e-7 mol/L
}

# material properties
KAPPA = {
    "kappa_upstream": 2.95e-12,
    "kappa_reactive_init": 1.52e-12,
    "kappa_downstream": 3.21e-12,
}
FLUID_PHYSICS = {
    "rho_liq": 1000.0,  # kg/m^3
    "mu_liq": 1.0e-3,  # Pa·s
    "alpha_L": 1e-4,  # m
    "alpha_T": 1e-4,  # m
    "pore_diffusion": 1.0e-9,  # m^2/s
    "retardation_factor": 1.0,
    "decay_rate": 0.0,
    "components": ["S(6)", "C", "Ca", "Na", "Cl", "H"],
}
SOLID = {
    "molar_volume": 3.693e-05,  # m^3/mol
    "volume_fraction_default": 0.0,
    "volume_fraction_by_mid": {"1": 0.2185},
}

# initials / boundaries (scalars)
PARAMS = {
    "porosity_upstream": 0.4277,
    "porosity_reactive": 0.3255,
    "porosity_downstream": 0.4685,
    "p_initial": 101325,
    "p_outlet": 101325,
    "q_inlet": 2.16948e-3,  # mass-flux = rho*q if rho=1000 and q≈2.16948e-6 m/s
    "S6_initial": 1.0e-8,  # Initial sulfate in pore water: 10⁻⁵ mmol/L = 1e-8 mol/L
    "S6_inlet": 5.0e-3,  # Inlet sulfate from 0.005M H₂SO₄ = 0.005 mol/L SO₄²⁻
    "H_inlet": 7.4456e-3,  # [H⁺] = 10^(-pH) = 10^(-2.1281) ≈ 0.0074456
    "C_inlet": 1.6e-5,
    "Ca_inlet": 1.0e-8,
    "NaCl_initial": 3.0e-8,
    "NaCl_inlet": 3.0e-8,
}

# project / io paths
PROJECT = {
    "basename": "porosityIncrease_hc",
    "pvd_path": Path(f"{out_dir}/hc_porosity_change_{dim_suffix}.pvd"),
    "chem_db_src": Path.cwd() / "llnl.dat",
    "chem_db_dst": out_dir / "llnl.dat",
    "workdir": Path.cwd(),
    "out_dir": out_dir,
}


#  time stepping values
time_settings = {
    "t_initial": 0,
    "t_end": 360000,
    "initial_dt": 120,
    "minimum_dt": 120,
    "maximum_dt": 120,
}

# copy chemistry database
if not PROJECT["chem_db_src"].is_file():
    msg = f"Source file not found: {PROJECT['chem_db_src']}"
    raise FileNotFoundError(msg)
copy2(PROJECT["chem_db_src"], PROJECT["chem_db_dst"])

def _compute_nx(
    left_L: float, band_L: float, right_L: float, h: float
) -> tuple[int, int, int]:
    """Compute element counts in x for left, band, right regions."""
    nx_left = max(1, round(left_L / h))
    nx_band = max(2, round(band_L / h))
    nx_right = max(1, round(right_L / h))

    return nx_left, nx_band, nx_right


# --------------------------------------------------------------------------------------
# 2D structured rectangle
# --------------------------------------------------------------------------------------
def create_rect_structured(
    filepath: str | Path,
    L: float = 0.07,
    H: float = 0.0099,
    left_L: float = 0.024,
    band_L: float = 0.022,
    right_L: float = 0.024,
    h: float = 0.0005,
    generate_mesh: bool = True,
) -> Path:
    filepath = Path(filepath)
    model_name = filepath.stem

    if abs((left_L + band_L + right_L) - L) > 1e-12:
        msg = f"Invalid lengths: left_L + band_L + right_L = {left_L + band_L + right_L} != L = {L}"
        raise ValueError(msg)
    if not (0.0 < band_L < L):
        msg = f"Invalid band_L={band_L} for L={L}."
        raise ValueError(msg)
    if H <= 0:
        msg = f"Invalid height H={H}."
        raise ValueError(msg)

    # element counts
    nx_left, nx_band, nx_right = _compute_nx(left_L, band_L, right_L, h)
    ny = max(1, round(H / h))

    # geometry coordinates
    xL = -L / 2.0
    x1 = xL + left_L
    x2 = x1 + band_L
    xR = L / 2.0
    yB = -H / 2.0
    yT = H / 2.0

    gmsh.initialize()
    try:
        gmsh.option.setNumber("General.Verbosity", 1)
        gmsh.option.setNumber("Mesh.RecombineAll", 1)  # global recombine

        gmsh.model.add(model_name)
        geo = gmsh.model.geo
        mgeo = gmsh.model.geo.mesh

        # -----------------------
        # Points (counter-clockwise)
        # -----------------------
        geo.addPoint(xL, yB, 0, h, tag=1)
        geo.addPoint(x1, yB, 0, h, tag=2)
        geo.addPoint(x2, yB, 0, h, tag=3)
        geo.addPoint(xR, yB, 0, h, tag=4)
        geo.addPoint(xR, yT, 0, h, tag=5)
        geo.addPoint(x2, yT, 0, h, tag=6)
        geo.addPoint(x1, yT, 0, h, tag=7)
        geo.addPoint(xL, yT, 0, h, tag=8)

        # -----------------------
        # Lines
        # -----------------------
        geo.addLine(1, 2, tag=1)  # bottom left
        geo.addLine(2, 3, tag=2)  # bottom band
        geo.addLine(3, 4, tag=3)  # bottom right
        geo.addLine(4, 5, tag=4)  # right vertical
        geo.addLine(5, 6, tag=5)  # top right
        geo.addLine(6, 7, tag=6)  # top band
        geo.addLine(7, 8, tag=7)  # top left
        geo.addLine(8, 1, tag=8)  # left vertical
        geo.addLine(2, 7, tag=9)  # internal left vertical
        geo.addLine(3, 6, tag=10)  # internal right vertical

        # -----------------------
        # Surfaces
        # -----------------------
        def plane(loop_id: int, edges: list[int], surf_id: int) -> None:
            geo.addCurveLoop(edges, loop_id)
            geo.addPlaneSurface([loop_id], surf_id)

        plane(1, [1, 9, 7, 8], 1)
        plane(2, [2, 10, 6, -9], 2)
        plane(3, [3, 4, 5, -10], 3)

        # -----------------------
        # Transfinite definitions
        # -----------------------
        mgeo.setTransfiniteCurve(1, nx_left + 1)
        mgeo.setTransfiniteCurve(7, nx_left + 1)

        mgeo.setTransfiniteCurve(2, nx_band + 1)
        mgeo.setTransfiniteCurve(6, nx_band + 1)

        mgeo.setTransfiniteCurve(3, nx_right + 1)
        mgeo.setTransfiniteCurve(5, nx_right + 1)

        mgeo.setTransfiniteCurve(8, ny + 1)
        mgeo.setTransfiniteCurve(4, ny + 1)
        mgeo.setTransfiniteCurve(9, ny + 1)
        mgeo.setTransfiniteCurve(10, ny + 1)

        for s in (1, 2, 3):
            mgeo.setTransfiniteSurface(s, "Alternate")
            mgeo.setRecombine(2, s)

        geo.synchronize()

        # -----------------------
        # Physical groups
        # -----------------------
        pg = gmsh.model.addPhysicalGroup
        pn = gmsh.model.setPhysicalName

        # boundaries (lines, dim = 1)
        right = pg(1, [4])
        pn(1, right, "Right")

        left = pg(1, [8])
        pn(1, left, "Left")

        # domains (surfaces, dim = 2)
        dom1 = pg(2, [1])
        pn(2, dom1, "domain1")

        mid = pg(2, [2])
        pn(2, mid, "Reactive_domain")

        dom2 = pg(2, [3])
        pn(2, dom2, "domain2")

        comp = pg(2, [1, 2, 3])
        pn(2, comp, "Computational_domain")
        if generate_mesh:
            gmsh.model.mesh.generate(2)

        out = filepath if filepath.suffix == ".msh" else filepath.with_suffix(".msh")
        gmsh.write(str(out))
        return out

    finally:
        gmsh.finalize()


# --------------------------------------------------------------------------------------
# 1D structured line
# --------------------------------------------------------------------------------------
def create_line_structured(
    filepath: str | Path,
    L: float = 0.07,
    left_L: float = 0.024,
    band_L: float = 0.022,
    right_L: float = 0.024,
    h: float = 0.0005,
    generate_mesh: bool = True,
) -> Path:
    filepath = Path(filepath)
    model_name = filepath.stem

    if abs((left_L + band_L + right_L) - L) > 1e-12:
        msg = f"Invalid lengths: left_L + band_L + right_L = {left_L + band_L + right_L} != L = {L}"
        raise ValueError(msg)
    if not (0.0 < band_L < L):
        msg = f"band_L must be in (0, L). Got band_L={band_L} for L={L}."
        raise ValueError(msg)

    nx_left, nx_band, nx_right = _compute_nx(left_L, band_L, right_L, h)

    xL = -L / 2.0
    x1 = xL + left_L
    x2 = x1 + band_L
    xR = L / 2.0

    gmsh.initialize()
    try:
        gmsh.option.setNumber("General.Verbosity", 1)
        gmsh.model.add(model_name)
        geo = gmsh.model.geo
        mgeo = gmsh.model.geo.mesh

        # points
        geo.addPoint(xL, 0.0, 0.0, h, tag=1)
        geo.addPoint(x1, 0.0, 0.0, h, tag=2)
        geo.addPoint(x2, 0.0, 0.0, h, tag=3)
        geo.addPoint(xR, 0.0, 0.0, h, tag=4)

        # lines
        geo.addLine(1, 2, tag=1)
        geo.addLine(2, 3, tag=2)
        geo.addLine(3, 4, tag=3)

        # transfinite along the line
        mgeo.setTransfiniteCurve(1, nx_left + 1)
        mgeo.setTransfiniteCurve(2, nx_band + 1)
        mgeo.setTransfiniteCurve(3, nx_right + 1)

        geo.synchronize()

        pg = gmsh.model.addPhysicalGroup
        pn = gmsh.model.setPhysicalName

        # boundaries
        right = pg(0, [4])
        pn(0, right, "Right")

        left = pg(0, [1])
        pn(0, left, "Left")

        # domains (lines, dim = 1)
        dom1 = pg(1, [1])
        pn(1, dom1, "domain1")

        mid = pg(1, [2])
        pn(1, mid, "Reactive_domain")

        dom2 = pg(1, [3])
        pn(1, dom2, "domain2")

        comp = pg(1, [1, 2, 3])
        pn(1, comp, "Computational_domain")

        if generate_mesh:
            gmsh.model.mesh.generate(1)

        out = filepath if filepath.suffix == ".msh" else filepath.with_suffix(".msh")
        gmsh.write(str(out))
        return out

    finally:
        gmsh.finalize()


def generate_mesh(dim: str) -> Path:
    d = dim.strip().lower()
    mesh_dir = MESH_ROOT / d
    mesh_dir.mkdir(parents=True, exist_ok=True)
    fp = mesh_dir / f"{MESH_NAME}.msh"

    if d == "2d":
        return create_rect_structured(
            filepath=fp,
            L=GEOMETRY["L"],
            H=GEOMETRY["H"],
            left_L=GEOMETRY["left_L"],
            band_L=GEOMETRY["band_L"],
            right_L=GEOMETRY["right_L"],
            h=GEOMETRY["h"],
            generate_mesh=True,
        )
    if d == "1d":
        return create_line_structured(
            filepath=fp,
            L=GEOMETRY["L"],
            left_L=GEOMETRY["left_L"],
            band_L=GEOMETRY["band_L"],
            right_L=GEOMETRY["right_L"],
            h=GEOMETRY["h"],
            generate_mesh=True,
        )

    msg = "dim must be '1d' or '2d'"
    raise ValueError(msg)

def plot_2d_meshes(meshes_2d):
    figures = []
    if not meshes_2d:
        return figures

    domain_mesh = meshes_2d.get("domain", next(iter(meshes_2d.values())))

    try:
        fig = plt.figure(figsize=(10, 3.0), dpi=150)
        ax = fig.add_subplot(111)

        ot.plot.contourf(
            domain_mesh, "MaterialIDs", ax=ax, cbar=True, linewidth=0.2, cmap="viridis"
        )

        ax.set_aspect("equal")
        ax.set_xlabel("x / m", fontsize=11)
        ax.set_ylabel("y / m", fontsize=11)
        ax.tick_params(labelsize=10)
        ax.set_title("2D Mesh - Material IDs", fontsize=13, pad=15)

        boundary_colors = {"Left": "red", "Right": "blue"}
        for name, mesh in meshes_2d.items():
            for boundary_key, color in boundary_colors.items():
                if boundary_key.lower() in name.lower():
                    try:
                        points = mesh.points
                        if len(points) > 1:
                            if boundary_key in ["Left", "Right"]:
                                sorted_idx = np.argsort(points[:, 1])
                            else:
                                sorted_idx = np.argsort(points[:, 0])
                            ax.plot(
                                points[sorted_idx, 0],
                                points[sorted_idx, 1],
                                color=color,
                                linewidth=5,
                                alpha=0.8,
                                label=name,
                            )
                    except Exception:
                        continue

        if any(
            any(
                boundary_key.lower() in name.lower() for boundary_key in boundary_colors
            )
            for name in meshes_2d
        ):
            ax.legend(fontsize=9, loc="upper right", frameon=True)

        fig.tight_layout()
        figures.append(fig)

    except Exception as e:
        print(f"Could not create 2D contour plot: {e}")

    return figures


def plot_1d_meshes(meshes_1d):
    figures = []
    if not meshes_1d:
        return figures

    mesh = meshes_1d.get("domain", next(iter(meshes_1d.values())))
    fig, ax = plt.subplots(figsize=(10, 2.5), dpi=150)

    points = mesh.points
    if len(points) == 0:
        return figures

    cell_data = None
    for key in ("MaterialIDs", "MaterialID", "MaterialId"):
        if key in mesh.cell_data and len(mesh.cell_data[key]) > 0:
            cell_data = np.asarray(mesh.cell_data[key]).flatten()
            break

    segments = []
    if hasattr(mesh, "cells"):
        cells = np.asarray(mesh.cells)
        if len(cells) > 0:
            if cells.ndim == 1:
                i = 0
                while i < len(cells):
                    n_points = cells[i]
                    cell_points = cells[i + 1 : i + 1 + n_points]
                    for j in range(len(cell_points) - 1):
                        x1 = points[cell_points[j]][0]
                        x2 = points[cell_points[j + 1]][0]
                        segments.append([[x1, 0], [x2, 0]])
                    i += 1 + n_points
            else:
                for cell in cells:
                    if len(cell) >= 2:
                        x1 = points[cell[0]][0]
                        x2 = points[cell[1]][0]
                        segments.append([[x1, 0], [x2, 0]])

    if not segments and len(points) > 1:
        for i in range(len(points) - 1):
            x1 = points[i][0]
            x2 = points[i + 1][0]
            segments.append([[x1, 0], [x2, 0]])
        if cell_data is None:
            cell_data = np.ones(len(segments), dtype=int)

    if not segments:
        return figures

    segments = np.array(segments)

    lc = LineCollection(
        segments,
        linewidth=8.0,
        capstyle="round",
        joinstyle="round",
        alpha=0.9,
        zorder=2,
        cmap="viridis",
    )

    if cell_data is not None and len(cell_data) == len(segments):
        lc.set_array(cell_data)
        cbar = plt.colorbar(lc, ax=ax, pad=0.02, shrink=0.8)
        cbar.set_label("MaterialIDs", fontsize=11, labelpad=8)
        cbar.ax.tick_params(labelsize=9)
        if np.issubdtype(cell_data.dtype, np.integer):
            unique_vals = np.unique(cell_data)
            cbar.set_ticks(unique_vals)
            cbar.set_ticklabels([str(int(v)) for v in unique_vals])

    ax.add_collection(lc)

    x_coords = points[:, 0]
    ax.plot(
        x_coords,
        np.zeros_like(x_coords),
        "o",
        markersize=4,
        color="white",
        markeredgecolor="black",
        markeredgewidth=1.2,
        alpha=0.95,
        zorder=3,
        label="Nodes",
    )

    boundary_colors = {"Left": "red", "Right": "blue"}
    boundary_markers = []

    for name, boundary_mesh in meshes_1d.items():
        for boundary_key, color in boundary_colors.items():
            if boundary_key.lower() in name.lower():
                try:
                    boundary_points = boundary_mesh.points
                    if len(boundary_points) > 0:
                        if boundary_key == "Left":
                            x_pos = boundary_points[0][0]
                            marker = ">"
                        else:
                            x_pos = (
                                boundary_points[0][0]
                                if len(boundary_points) == 1
                                else boundary_points[-1][0]
                            )
                            marker = "<"

                        ax.plot(
                            x_pos,
                            0,
                            marker=marker,
                            markersize=10,
                            color=color,
                            markeredgecolor="black",
                            markeredgewidth=1.0,
                            zorder=4,
                        )
                        boundary_markers.append(True)
                except Exception:
                    continue

    if len(x_coords) > 0:
        x_min, x_max = x_coords.min(), x_coords.max()
        if x_max == x_min:
            x_min -= 0.5
            x_max += 0.5
        else:
            padding = (x_max - x_min) * 0.05
            x_min -= padding
            x_max += padding
        ax.set_xlim(x_min, x_max)

    ax.set_ylim(-0.5, 0.5)
    ax.set_yticks([])

    for spine in ["top", "right", "left"]:
        ax.spines[spine].set_visible(False)
    ax.spines["bottom"].set_visible(True)
    ax.spines["bottom"].set_color("#404040")
    ax.spines["bottom"].set_linewidth(0.8)

    ax.tick_params(axis="x", which="both", bottom=True, labelsize=10)
    ax.set_xlabel("x / m", fontsize=11, labelpad=8)
    ax.set_ylabel("")

    ax.grid(True, axis="x", alpha=0.2, linestyle="-", linewidth=0.5)
    ax.set_title("1D Mesh - Material IDs", fontsize=13, pad=15, weight="normal")

    legend_handles = []
    legend_handles.append(
        Line2D(
            [0],
            [0],
            marker="o",
            color="white",
            markeredgecolor="black",
            markersize=6,
            linestyle="None",
            label="Nodes",
        )
    )

    if boundary_markers:
        legend_handles.append(
            Line2D(
                [0],
                [0],
                marker=">",
                color="red",
                markeredgecolor="black",
                markersize=8,
                linestyle="None",
                label="Left",
            )
        )
        legend_handles.append(
            Line2D(
                [0],
                [0],
                marker="<",
                color="blue",
                markeredgecolor="black",
                markersize=8,
                linestyle="None",
                label="Right",
            )
        )

    if legend_handles:
        ax.legend(
            handles=legend_handles,
            loc="upper right",
            frameon=True,
            framealpha=0.9,
            edgecolor="#CCCCCC",
            fontsize=9,
            handletextpad=0.5,
        )

    ax.set_facecolor("#FAFAFA")
    fig.patch.set_facecolor("white")
    fig.set_constrained_layout(True)

    figures.append(fig)
    return figures


def generate_vtu_meshes_from_Gmsh(msh_file: Path, target_dim: int):
    if not msh_file.is_file():
        raise FileNotFoundError(msh_file)
    out_dir = msh_file.parent
    figures = []
    if target_dim == 2:
        meshes = ot.Meshes.from_gmsh(msh_file, dim=[2], log=False)
        for name, mesh in meshes.items():
            pv.save_meshio(out_dir / f"{name}.vtu", mesh)
        if meshes:
            figures.extend(plot_2d_meshes(meshes))
    elif target_dim == 1:
        meshes = ot.Meshes.from_gmsh(msh_file, dim=[1], log=False)
        for name, mesh in meshes.items():
            pv.save_meshio(out_dir / f"{name}.vtu", mesh)
        if meshes:
            figures.extend(plot_1d_meshes(meshes))
    if figures:
        plt.show()
    return figures


def run_mesh(dim: str):
    d = dim.strip().lower()
    td = {"1d": 1, "2d": 2}[d]
    base_msh = MESH["msh_file"]
    base_vtu = MESH["vtu_file"]
    msh = base_msh.parent.parent / d / base_msh.name
    vtu = base_vtu.parent.parent / d / base_vtu.name
    if not msh.is_file():
        raise FileNotFoundError(msh)
    generate_vtu_meshes_from_Gmsh(msh, target_dim=td)
    return vtu


def assert_structured_mesh_size(mesh_path: Path, target_dim: int) -> None:
    mesh = pv.read(str(mesh_path))

    nx_left, nx_band, nx_right = _compute_nx(
        GEOMETRY["left_L"],
        GEOMETRY["band_L"],
        GEOMETRY["right_L"],
        GEOMETRY["h"],
    )

    if target_dim == 1:
        expected_cells = nx_left + nx_band + nx_right
        expected_points = expected_cells + 1
    elif target_dim == 2:
        ny = max(1, round(GEOMETRY["H"] / GEOMETRY["h"]))
        nx_total = nx_left + nx_band + nx_right
        expected_cells = nx_total * ny
        expected_points = (nx_total + 1) * (ny + 1)
    else:
        msg = f"Unsupported target_dim={target_dim!r} (expected 1 or 2)."
        raise ValueError(msg)

    if mesh.n_cells != expected_cells or mesh.n_points != expected_points:
        msg = (
            "Unexpected mesh size:\n"
            f"  cells : {mesh.n_cells} (expected {expected_cells})\n"
            f"  points: {mesh.n_points} (expected {expected_points})"
        )
        raise RuntimeError(msg)

    print(
        f"Mesh check passed: {mesh.n_cells} cells, {mesh.n_points} points "
        f"(dim = {target_dim}D)."
    )

def pre_processing_chemo_by_params(
    mesh_path: str | Path,
    L: float,
    H: float,
    left_L: float,
    band_L: float,
    right_L: float,
    C_main=1.6e-5,
    Ca_main=1.0e-8,
    H_main=2.461e-6,
    C_mid=3.7934e-2,
    Ca_mid=7.9567e-3,
    H_mid=7.485e-7,
    tol: float = 1e-12,
    save_reactive_submesh: bool = True,
):
    mesh_path = Path(mesh_path)
    if not mesh_path.is_file():
        raise FileNotFoundError(mesh_path)

    if abs((left_L + band_L + right_L) - L) > 1e-12 or not (0 < band_L < L):
        msg = f"Invalid geometry parameters: left_L={left_L}, band_L={band_L}, right_L={right_L}, L={L}"
        raise ValueError(msg)

    xL = -L / 2.0
    x1 = xL + left_L
    x2 = x1 + band_L

    m = pv.read(str(mesh_path))
    if m.n_points == 0:
        msg = f"No points in mesh: {mesh_path}"
        raise RuntimeError(msg)

    x = m.points[:, 0]
    is_mid = (x >= x1 - tol) & (x <= x2 + tol)

    n = m.n_points
    C = np.full(n, float(C_main))
    Ca = np.full(n, float(Ca_main))
    H = np.full(n, float(H_main))

    C[is_mid] = float(C_mid)
    Ca[is_mid] = float(Ca_mid)
    H[is_mid] = float(H_mid)

    m.point_data["C"] = C
    m.point_data["Ca"] = Ca
    m.point_data["H"] = H
    m.save(str(mesh_path))

    if save_reactive_submesh:
        reactive_mesh_path = mesh_path.with_name("Reactive_domain.vtu")
        if reactive_mesh_path.exists():
            r = pv.read(str(reactive_mesh_path))
            if r.n_points > 0:
                nr = r.n_points
                r.point_data["C"] = np.full(nr, float(C_mid))
                r.point_data["Ca"] = np.full(nr, float(Ca_mid))
                r.point_data["H"] = np.full(nr, float(H_mid))
                r.save(str(reactive_mesh_path))
                return str(mesh_path), str(reactive_mesh_path)
        else:
            print("Reactive submesh not found; nothing to update.")

    return str(mesh_path)

pre_processing_chemo_by_params(
    mesh_path=vtu_path,
    L=GEOMETRY["L"],
    H=GEOMETRY["H"],
    left_L=GEOMETRY["left_L"],
    band_L=GEOMETRY["band_L"],
    right_L=GEOMETRY["right_L"],
    C_main=FIELDS["C_main"],
    Ca_main=FIELDS["Ca_main"],
    H_main=FIELDS["H_main"],
    C_mid=FIELDS["C_mid"],
    Ca_mid=FIELDS["Ca_mid"],
    H_mid=FIELDS["H_mid"],
)

def identify_subdomains(dim: str):
    d = dim.strip().lower()
    path = MESH_ROOT / d
    bulk = path / "domain.vtu"
    if not bulk.is_file():
        return

    names = [
        "domain.vtu",
        "Computational_domain.vtu",
        "domain1.vtu",
        "domain2.vtu",
        "Left.vtu",
        "Reactive_domain.vtu",
        "Right.vtu",
    ]

    cwd = Path.cwd()
    os.chdir(path)
    try:
        ot.cli().identifySubdomains("-m", "domain.vtu", "--", *names)
    finally:
        os.chdir(cwd)


identify_subdomains(dim_suffix)

prj_in = Path(f"{PROJECT['basename']}.prj")
prj_out = out_dir / f"{PROJECT['basename']}_updated_{dim_suffix}.prj"
prj = ot.Project(input_file=prj_in, output_file=prj_out)

def _set_values(name: str, vec4):
    """Write a 4-vector into <values> (or <value>) for parameter `name`"""
    txt = " ".join(f"{v:.16g}" for v in vec4)
    n = prj.replace_text(txt, xpath=f".//parameters/parameter[name='{name}']/values")
    if n == 0:
        n = prj.replace_text(txt, xpath=f".//parameters/parameter[name='{name}']/value")
    if n == 0:
        msg = f"Parameter '{name}' not found (neither <values> nor <value>)."
        raise KeyError(msg)


def _set_value(name: str, val):
    """Write a scalar into <value> (or <values>) for parameter `name`"""
    sval = f"{val:.16g}" if isinstance(val, int | float) else str(val)
    n = prj.replace_text(sval, xpath=f".//parameters/parameter[name='{name}']/value")
    if n == 0:
        n = prj.replace_text(
            sval, xpath=f".//parameters/parameter[name='{name}']/values"
        )
    if n == 0:
        msg = f"Parameter '{name}' not found (neither <value> nor <values>)."
        raise KeyError(msg)


def _set_all(xpath: str, val):
    """Replace all matches at `xpath` with `val`"""
    sval = f"{val:.16g}" if isinstance(val, int | float) else str(val)
    n = prj.replace_text(sval, xpath=xpath)
    if n == 0:
        msg = f"No nodes matched xpath: {xpath}"
        raise KeyError(msg)


def add_kappa_parameters(prj, dim: int, kappa: dict):
    if dim == 1:
        for name, K in kappa.items():
            prj.parameters.add_parameter(
                name=name, type="Constant", value=f"{float(K):.16g}"
            )
    elif dim == 2:
        for name, K in kappa.items():
            prj.parameters.add_parameter(
                name=name,
                type="Constant",
                values=f"{float(K):.16g} 0 0 {float(K):.16g}",
            )
    elif dim == 3:
        for name, K in kappa.items():
            prj.parameters.add_parameter(
                name=name,
                type="Constant",
                values=" ".join(f"{v:.16g}" for v in [K, 0, 0, 0, K, 0, 0, 0, K]),
            )
    else:
        msg = f"Invalid dim: {dim}. Must be 1, 2, or 3."
        raise ValueError(msg)

add_kappa_parameters(prj, dim=target_dim, kappa=KAPPA)

for k, v in PARAMS.items():
    _set_value(k, v)

for param, value in time_settings.items():
    _set_all(f".//time_stepping/{param}", value)

# dispersion (all media)
_set_all(
    ".//media/medium/properties/property[name='longitudinal_dispersivity']/value",
    FLUID_PHYSICS["alpha_L"],
)
_set_all(
    ".//media/medium/properties/property[name='transversal_dispersivity']/value",
    FLUID_PHYSICS["alpha_T"],
)

# fluid properties (AqueousLiquid in all media)
_set_all(
    ".//phases/phase[type='AqueousLiquid']/properties/property[name='density']/value",
    FLUID_PHYSICS["rho_liq"],
)
_set_all(
    ".//phases/phase[type='AqueousLiquid']/properties/property[name='viscosity']/value",
    FLUID_PHYSICS["mu_liq"],
)

# component transport (each listed component in all media)
for comp in FLUID_PHYSICS["components"]:
    base = f".//phases/phase[type='AqueousLiquid']/components/component[name='{comp}']/properties/property"
    _set_all(base + "[name='pore_diffusion']/value", FLUID_PHYSICS["pore_diffusion"])
    _set_all(
        base + "[name='retardation_factor']/value", FLUID_PHYSICS["retardation_factor"]
    )
    _set_all(base + "[name='decay_rate']/value", FLUID_PHYSICS["decay_rate"])

# solid Calcite (all media)
_set_all(
    ".//phases/phase[type='Solid']"
    "/components/component[name='Calcite']"
    "/properties/property[name='molar_volume']/value",
    SOLID["molar_volume"],
)

_set_all(
    ".//phases/phase[type='Solid']"
    "/components/component[name='Calcite']"
    "/properties/property[name='volume_fraction']/value",
    SOLID["volume_fraction_default"],
)

for mid, vf in SOLID["volume_fraction_by_mid"].items():
    n = prj.replace_text(
        f"{vf:.16g}",
        xpath=(
            f".//media/medium[@id='{mid}']"
            f"/phases/phase[type='Solid']"
            f"/components/component[name='Calcite']"
            f"/properties/property[name='volume_fraction']/value"
        ),
    )
    if n == 0:
        msg = f"Calcite volume_fraction node not found for medium @id='{mid}. Check your PRJ structure and ids."
        raise KeyError(msg)

force_vec = {1: "0", 2: "0 0", 3: "0 0 0"}[target_dim]


prj.replace_text(force_vec, xpath=".//processes/process/specific_body_force")
prj.replace_text(f"hc_porosity_change_{dim_suffix}", xpath=".//output/prefix")

prj.write_input()
print("Wrote:", prj_out)