#!/usr/bin/env bash
#/******************************************************************************
#* Copyright (c) 2024 Lawrence Livermore National Security, LLC and other
#* HYPRE Project Developers. See the top-level COPYRIGHT file for details.
#*
#* SPDX-License-Identifier: MIT
#******************************************************************************/

set -euo pipefail

MODE="run-and-plot"
STUDY="presets"
PLOT_TIMES=0

for arg in "$@"; do
    case "${arg}" in
        --plot-only)
            MODE="plot-only"
            ;;
        --two-material)
            STUDY="two-material"
            ;;
        --times)
            PLOT_TIMES=1
            ;;
        -h|--help)
            MODE="help"
            ;;
        *)
            echo "Unknown option: ${arg}"
            echo "Use -h or --help for usage information."
            exit 1
            ;;
    esac
done

if [[ "${MODE}" == "help" ]]; then
    echo "Usage: $0"
    echo ""
    echo "Runs three elasticity solver preset configurations -- plain unknown-based"
    echo "AMG (elasticity_3D), function-filtered AMG (elasticity_sdc_3D), and the"
    echo "rigid-body-aware nodal/GM2 AMG (elasticity_nodal_3D) -- across 8 mesh-size"
    echo "variants (~1e4 to ~1e6 DOFs) and produces iterations and total (setup+solve)"
    echo "time plots using scripts/analyze_statistics.py. This is the rigid-body-modes"
    echo "(RBM) study: it shows the nodal/GM2 preset converging in the fewest iterations"
    echo "and the least time."
    echo ""
    echo "Modes:"
    echo "  (default)       Run the single-material preset sweep and generate plots"
    echo "  --two-material  Two-material bar: the two locking-free discretizations,"
    echo "                  B-bar (Q1-P0) and mixed u-p, sweeping nu_top in"
    echo "                  {0.49,0.499,0.4999} over 5 mesh sizes (to ~1e6 DOFs);"
    echo "                  writes two_material_iters.csv and the two-panel grouped-bar"
    echo "                  figure iters_two_material_bars.png"
    echo "  --times         With --two-material, additionally render the stacked"
    echo "                  setup/solve time figure for the mixed u-p discretization"
    echo "                  (times_two_material_mixed.png). Combine with --plot-only"
    echo "                  to render it from an existing two_material_iters.csv."
    echo "  --plot-only     Skip runs and only generate plots from existing output logs"
    echo ""
    echo "Environment overrides:"
    echo "  MPIEXEC_EXECUTABLE   MPI launcher executable (default: mpirun)"
    echo "  MPIEXEC_NUMPROC_FLAG MPI rank flag (default: -np)"
    echo "  MPI_RANKS            Number of MPI ranks (default: 1)"
    echo "  PGRID                Processor grid 'Px Py Pz' (default: 1 1 1)"
    echo "  EXEC                 Elasticity executable path (default: ./elasticity)"
    echo "  STATS                Statistics script path (default: ../../../scripts/analyze_statistics.py)"
    echo "  ARGS                 Common elasticity arguments (without -n/-P)"
    echo "  MAX_VARIANTS         Limit number of size variants (default: 8)"
    echo ""
    echo "  Two-material study (--two-material) overrides:"
    echo "  MPI_RANKS_TM         MPI ranks for the two-material study (default: 16)"
    echo "  PGRID_TM             Processor grid 'Px Py Pz' (default: 4 2 2)"
    echo "  NU_BOTTOM            Bottom-material Poisson ratio (default: 0.3)"
    echo "  CONFIG_DIR           Directory holding the .yml configs + plot script (default: .)"
    echo ""
    echo "Outputs:"
    echo "  two_material_iters.csv           (--two-material)"
    echo "  iters_two_material_bars.png      (--two-material)"
    echo "  times_two_material_mixed.png     (--two-material --times)"
    echo "  elasticity_builtin.out"
    echo "  elasticity_sdc.out"
    echo "  elasticity_nodal.out"
    echo "  iters_elasticity_presets_dofs.png"
    echo "  total_elasticity_presets_dofs.png"
    exit 0
fi

MPIEXEC_EXECUTABLE="${MPIEXEC_EXECUTABLE:-mpirun}"
MPIEXEC_NUMPROC_FLAG="${MPIEXEC_NUMPROC_FLAG:--np}"
MPI_RANKS="${MPI_RANKS:-8}"
PGRID="${PGRID:-2 2 2}"
EXEC="${EXEC:-./elasticity}"
STATS="${STATS:-../../../scripts/analyze_statistics.py}"
ARGS="${ARGS:--L 3 1 1 -ns 1 -v 1}"
MAX_VARIANTS="${MAX_VARIANTS:-8}"

# 8 size variants spanning roughly 1e4 to 1e6 DOFs.
# DOFs = 3 * Nx * Ny * Nz with (Nx, Ny, Nz) listed below.
SIZE_VARIANTS=(
    "33 11 11"   # 11,979
    "39 13 13"   # 19,773
    "48 16 16"   # 36,864
    "60 20 20"   # 72,000
    "75 25 25"   # 140,625
    "93 31 31"   # 268,119
    "114 38 38"  # 493,848
    "144 48 48"  # 995,328
)

PRESETS=(
    "elasticity_3D"
    "elasticity_sdc_3D"
    "elasticity_nodal_3D"
)

OUT_FILES=(
    "elasticity_builtin.out"
    "elasticity_sdc.out"
    "elasticity_nodal.out"
)

METHODS=(
    "elasticity_3D"
    "elasticity_sdc_3D"
    "elasticity_nodal_3D"
)

# Each variant is a separate executable run, so a raw '.out' log holds one
# standalone STATISTICS SUMMARY table (single entry) per size variant. The
# statistics script treats every table as a distinct data series, which would
# turn a single preset's size sweep into many one-point series. Collapse each
# raw log into one combined STATISTICS SUMMARY table whose entries are the
# per-variant rows, so the sweep plots as a single curve per preset. The
# 'Solving linear system #N' markers (carrying rows/nonzeros) are preserved so
# the script can map each entry to its DOF count for the '-t rows' abscissa.
build_combined_log() {
    local src="$1"
    local dst="$2"
    awk '
        BEGIN { nmark = 0; ndata = 0 }
        /^Solving linear system #[0-9]+ with [0-9]+ rows and [0-9]+ nonzeros/ {
            markers[nmark++] = $0
            next
        }
        # Data rows inside a STATISTICS SUMMARY table start with "| <entry> |".
        /^\|[[:space:]]*[0-9]+[[:space:]]*\|/ {
            n = split($0, c, "|")
            # c[2]=entry c[3]=build c[4]=setup c[5]=solve c[6]=r0 c[7]=relres c[8]=iters
            build[ndata]  = c[3]; setup[ndata] = c[4]; solve[ndata] = c[5]
            r0[ndata]     = c[6]; rr[ndata]    = c[7]; it[ndata]    = c[8]
            ndata++
            next
        }
        END {
            print "Using HYPRE_RELEASE_VERSION: combined"
            for (i = 0; i < nmark; i++) print markers[i]
            print ""
            print "STATISTICS SUMMARY:"
            print ""
            div = "+------------+-------------+-------------+-------------+------------+------------+--------+"
            print div
            print "|            |    LS build |       setup |       solve |    initial |   relative |        |"
            print "|      Entry |   times [s] |   times [s] |   times [s] |  res. norm |  res. norm |  iters |"
            print div
            for (i = 0; i < ndata; i++) {
                printf "|%11d |%s|%s|%s|%s|%s|%s|\n", i, build[i], setup[i], solve[i], r0[i], rr[i], it[i]
            }
            print div
        }
    ' "${src}" > "${dst}"
}

plot_results() {
    for f in "${OUT_FILES[@]}"; do
        if [[ ! -f "${f}" ]]; then
            echo "Missing required output log '${f}'. Run without --plot-only first."
            exit 1
        fi
    done

    # Sanity-check that each source has enough size points.
    local min_points=2
    for f in "${OUT_FILES[@]}"; do
        local variant_points
        variant_points=$(awk '/Solving linear system #[0-9]+ with [0-9]+ rows and [0-9]+ nonzeros\.\.\./{c++} END{print c+0}' "${f}")
        if (( variant_points < min_points )); then
            echo "Error: '${f}' has only ${variant_points} size point(s)."
            echo "       Need at least ${min_points} points for a meaningful curve."
            echo "       Re-run full sweep (default) or increase MAX_VARIANTS."
            exit 1
        fi
    done

    # Collapse each raw per-variant log into a single combined table so each
    # preset plots as one curve over the DOF sweep.
    local combined_files=()
    local cf
    for f in "${OUT_FILES[@]}"; do
        cf="${f%.out}.combined.out"
        build_combined_log "${f}" "${cf}"
        combined_files+=("${cf}")
    done

    echo "Generating plots (log-scale X axis, DOFs)..."

    echo "  Plotting: iterations vs DOFs -> iters_elasticity_presets_dofs.png"
    PYTHONWARNINGS=ignore::UserWarning MPLBACKEND=Agg python3 "${STATS}" -f "${combined_files[@]}" -ln "${METHODS[@]}" -m "iters" \
        --style docs -t rows -l "DOFs (rows)" -s "elasticity_presets_dofs.png" --log-x -T "Linear solver iterations"

    echo "  Plotting: total (setup+solve) time vs DOFs -> total_elasticity_presets_dofs.png"
    PYTHONWARNINGS=ignore::UserWarning MPLBACKEND=Agg python3 "${STATS}" -f "${combined_files[@]}" -ln "${METHODS[@]}" -m "total" \
        --style docs -t rows -l "DOFs (rows)" -s "elasticity_presets_dofs.png" --log-x --log-y -T "Linear solver total (setup + solve) time [s]"

    echo "Done. Plots saved in current directory."
}

# Two-material bar: the two LOCKING-FREE discretizations of the SAME problem
# (bottom NU_BOTTOM, top nu_top) on the same meshes and rank grid, sweeping nu_top
# across {0.49, 0.499, 0.4999} and each mesh size:
#   bbar  = Q1-P0 mean-dilatation (PCG + AMG, amg-pcg.yml) -- locking-free but the
#           condensed operator keeps the lambda penalty, so it stays
#           ill-conditioned and PCG iterations grow / cap out as nu -> 1/2.
#   mixed = u-p saddle point (FGMRES + MGR, mgr2-gmres.yml + --coarse-schur) --
#           locking-free AND, with the scaled pressure-mass Schur as MGR's coarse
#           operator, the iteration count stays bounded as nu -> 1/2 (lambda-robust);
#           the displacement F-block uses a rigid-body-aware GM2 AMG (2 V-cycles),
#           so the count is also nearly mesh-independent (h-robust).
# (Standard CG Q1 is omitted: it locks, so its solution is wrong -- the comparison
# of interest is between the two locking-free options.) Iterations are collected
# into two_material_iters.csv and rendered as two side-by-side grouped bar charts
# (x = mesh resolution, one bar per nu_top) by plot_two_material.py. Configs and
# the plot script live in this directory; set CONFIG_DIR to relocate.
run_two_material_study() {
    local ranks="${MPI_RANKS_TM:-16}"
    local pgrid="${PGRID_TM:-4 2 2}"
    local cfgdir="${CONFIG_DIR:-.}"
    local nu_bottom="${NU_BOTTOM:-0.3}"
    local csv="two_material_iters.csv"
    # (Nx, Ny, Nz) with Ny odd (interface node layer); ~4e4 .. ~1e6 displacement DOFs.
    local sizes=("31 21 21" "45 31 31" "61 41 41" "81 51 51" "91 61 61")
    # Top-material Poisson ratios: mild -> strongly near-incompressible.
    local nus=("0.49" "0.499" "0.4999")

    if [[ "${MODE}" != "plot-only" ]]; then
        echo "disc,nu,nx,ny,nz,dofs,iters,relres,setup,solve" > "${csv}"
        local s nx ny nz dofs nu spec disc yml row
        for s in "${sizes[@]}"; do
            read -r nx ny nz <<< "${s}"
            dofs=$((3 * nx * ny * nz))
            for nu in "${nus[@]}"; do
                # bbar = displacement-only PCG+AMG; mixed = FGMRES+MGR with a
                # rigid-body-aware GM2 displacement F-block and the mass-Schur coarse
                # operator (needs --coarse-schur so the driver assembles/supplies it).
                for spec in "bbar amg-pcg.yml" "mixed mgr2-gmres.yml"; do
                    read -r disc yml <<< "${spec}"
                    local extra=""
                    if [[ "${disc}" == "mixed" ]]; then extra="--coarse-schur"; fi
                    echo "N=${nx}x${ny}x${nz} (DOFs=${dofs}) nu_top=${nu} ${disc} ..."
                    # Pull the single STATISTICS SUMMARY data row. Columns are
                    # $2 entry, $3 LS build, $4 setup, $5 solve, $6 initial res,
                    # $7 relative res, $8 iters.
                    if ! row=$(${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} ${ranks} ${EXEC} \
                                  --problem two-material --discretization "${disc}" \
                                  -n "${nx}" "${ny}" "${nz}" -P ${pgrid} -nu "${nu_bottom}" \
                                  --nu-top "${nu}" -ns 1 -v 1 ${extra} \
                                  -i "${cfgdir}/${yml}" 2>&1 \
                               | awk -F'|' '/^\|[[:space:]]*[0-9]+[[:space:]]*\|/ \
                                            { for (i = 4; i <= 8; i++) gsub(/ /, "", $i); \
                                              su=$4; so=$5; rr=$7; it=$8; found=1 } \
                                            END { if (found) print it","rr","su","so }'); then
                        echo "    -> run failed; skipping this CSV row" >&2
                        continue
                    fi
                    if [[ -z "${row}" ]]; then
                        echo "    -> no statistics summary; skipping this CSV row" >&2
                        continue
                    fi
                    echo "${disc},${nu},${nx},${ny},${nz},${dofs},${row}" >> "${csv}"
                    echo "    -> iters,relres,setup,solve = ${row}"
                done
            done
        done
    fi

    echo "Plotting grouped iteration bars -> iters_two_material_bars.png"
    PYTHONWARNINGS=ignore::UserWarning MPLBACKEND=Agg python3 \
        "${cfgdir}/plot_two_material.py" "${csv}" "iters_two_material_bars.png"
    echo "Done. Plot: iters_two_material_bars.png (data: ${csv})"

    if (( PLOT_TIMES )); then
        echo "Plotting stacked setup/solve times -> times_two_material_mixed.png"
        PYTHONWARNINGS=ignore::UserWarning MPLBACKEND=Agg python3 \
            "${cfgdir}/plot_two_material_times.py" "${csv}" "times_two_material_mixed.png"
        echo "Done. Plot: times_two_material_mixed.png (data: ${csv})"
    fi
}

if [[ "${STUDY}" == "two-material" ]]; then
    run_two_material_study
    exit 0
fi

if [[ "${MODE}" == "plot-only" ]]; then
    echo "Plot-only mode: using existing output logs."
    plot_results
    exit 0
fi

for i in "${!PRESETS[@]}"; do
    preset="${PRESETS[$i]}"
    outfile="${OUT_FILES[$i]}"
    : > "${outfile}"
    echo "Running ${preset} -> ${outfile}"

    for j in "${!SIZE_VARIANTS[@]}"; do
        if (( j >= MAX_VARIANTS )); then
            break
        fi

        read -r nx ny nz <<< "${SIZE_VARIANTS[$j]}"
        dofs=$((3 * nx * ny * nz))

        echo "  Variant $((j + 1)): N=${nx}x${ny}x${nz} (DOFs=${dofs})"
        {
            echo ""
            echo "===== Variant $((j + 1)): N=${nx}x${ny}x${nz}, DOFs=${dofs} ====="
            # Provide rows/nonzeros marker expected by analyze_statistics.py for -t rows mode.
            echo "Solving linear system #${j} with ${dofs} rows and 1 nonzeros..."
        } >> "${outfile}"

        ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} ${MPI_RANKS} ${EXEC} ${ARGS} \
            -P ${PGRID} -n "${nx}" "${ny}" "${nz}" --solver-preset "${preset}" \
            2>&1 | tee -a "${outfile}"
    done
done

plot_results
