#!/usr/bin/env julia
# -----------------------------------------------------------------------------
# Solid (2D axisymmetric r–z) + 1D plug-flow fluid.
# - Builds a Gmsh mesh with named edges: PIPE, OUTER, ZMIN, ZMAX.
# - Couples via Robin BC:  -k ∂T/∂n = h (T_f(z) - T_s).
#   * Solid RHS uses T_f(z) as a boundary CellField (no fragile point sampling).
#   * Fluid uses wall temperature T_w(z) by sampling the solid just inside the wall.
# - Writes VTK with FLUID aligned to the SOLID automatically (axes detected).
# -----------------------------------------------------------------------------

# push!(LOAD_PATH, "/home/heikki/0Pilvi/Lampoakku")
include("Kappyrat.jl")

using LinearAlgebra
using SparseArrays
using Printf
using Dates
using Logging
using Gmsh
using Gmsh: gmsh
using Gridap
using GridapGmsh
using WriteVTK
using Gridap: VectorValue
using JSON
using YAML
using Base.Threads
using DelimitedFiles

include("controller.jl")

const AXIAL  = 1  # x=z
const RADIAL = 2  # y=r

# VTK config
const VTK_NZVIS  = 256
const VTK_NRVIS  = 256

sci3e(x) = Printf.@sprintf("%.3e", x)
fix3f(x) = Printf.@sprintf("%.3f", x)

# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ Utilities                                                                ║
# ╚══════════════════════════════════════════════════════════════════════════╝

macro banner(msg)
    :( @info "$(string($msg))" )
end

using Gridap: Triangulation, BoundaryTriangulation, Measure, CellField, integrate

#=
# ───────────────────────────── Measures (axisymmetric, full 3D) ─────────────────────────────
"Build measures and axisymmetric weights. x=z, y=r ⇒ radial index = 2."

"""
    build_measures_axisym(model, params) -> (Ω, Γ_pipe, dΩ, dΓ_pipe, wΩ, wΓ_pipe)

Convenience builder for axisymmetric measures and weights: `wΩ(x) = 2π * x[RADIAL]`, `wΓ_pipe = 2π * R_pipe`.
Used by audits and plotting.
"""
function build_measures_axisym(model, params)
    Ω       = Triangulation(model)
    Γ_pipe  = BoundaryTriangulation(model; tags="PIPE")
    dΩ      = Measure(Ω, 2)
    dΓ_pipe = Measure(Γ_pipe, 1)

    wΩ      = CellField(x -> 2π * x[RADIAL], Ω)               # 2π r   (domain weight)
    wΓ_pipe = CellField(_ -> 2π * params.R_pipe, Γ_pipe)       # 2π R   (pipe boundary)

    return (Ω=Ω, Γ_pipe=Γ_pipe, dΩ=dΩ, dΓ_pipe=dΓ_pipe, wΩ=wΩ, wΓ_pipe=wΓ_pipe)
end
=#

"Solid stored energy (J) for the full 3D axisymmetric body relative to Tref."

"""
    energy_solid_axisym(uh, S_or_M, params; Tref=T_initial) -> Float64 (J)

Full 3‑D solid energy via axisymmetric weight `2π r`: `∫ ρ_s c_s (T_s - Tref) 2π r dΩ`.
Accepts either the `SolidCoupler` or a measures tuple `M`.
"""
function energy_solid_axisym(uh, M, params, Vax)
    Tref = params.T_initial
    ρc = params.rho_s * params.cp_s
    return sum(integrate(ρc * M.wΩ * uh, M.dΩ)) - ρc * Tref * Vax
    # return sum(integrate( ρc * M.wΩ * (uh - Tref), M.dΩ ))
end

"Oil stored energy (J) in the pipe (axisymmetric) relative to Tref."

"""
    energy_oil_axisym(T_f, z, params; Tref=T_initial) -> Float64 (J)

Oil energy in the pipe: trapezoidal integral of `ρ_f c_f (T_f - Tref)` over `z`,
times cross‑section `A = π R^2`. Handles non‑uniform `z`.
"""
function energy_oil_axisym(T_f::AbstractVector{<:Real}, z::AbstractVector{<:Real}, wz, params;
                           Tref::Float64 = params.T_initial)
    @assert length(T_f) == length(z)
    A  = π * params.R_pipe^2
    ρc = params.rho_oil * params.cp_oil
    return ρc * A * sum(wz .* (T_f .- Tref))
end


"Inlet thermal power (W) based on bulk inlet/outlet temperatures (axisymmetric)."

"""
    inlet_power_axisym(T_in, params) -> Float64 (W)

Plug‑flow inlet power relative to `Tref`: `ρ_f c_f v A (T_in - Tref)`.
"""
function inlet_power_axisym(T_in::Float64, T_out::Float64, params::Params)
    A  = π * params.R_pipe^2
    ṁ  = params.rho_oil * A * params.v
    return ṁ * params.cp_oil * (T_in - T_out)
end

"Wall power on the solid side (W): Q_wall = ∫_PIPE h · 2πR · (T_s − T_f*) dΓ."

"""
    wall_power_axisym(uh, Tf_b, S_or_M, params) -> Float64 (W)

Instantaneous wall power on the solid side: `∫_PIPE h (T_s - T_f*) 2π R dΓ`.
Accepts either the `SolidCoupler` or a measures tuple `M`.
"""
function wall_power_axisym(uh, Tf_b, M, params)
    return sum(integrate( params.h * M.wΓ_pipe * (uh - Tf_b), M.dΓ_pipe ))
end

# Trapz weights on a 1D grid (z increasing). If zero_inlet=true, w[1]=0.
function trapezoid_weights(z::AbstractVector{<:Real}; zero_inlet::Bool=false)
    n = length(z); @assert n ≥ 2
    w = zeros(Float64, n)
    for i in 2:n-1
        w[i] = 0.5*(z[i+1]-z[i-1])
    end
    w[1]  = (z[2]-z[1])*(zero_inlet ? 0.0 : 0.5)
    w[n]  = 0.5*(z[n]-z[n-1])
    return w
end

"""
    build_model_gmsh(params; Nr, Nz, grad, bump, quads, mshfile) -> model, mshfile

Build the 2‑D r–z Gmsh geometry and mesh with named edges: `PIPE`, `OUTER`, `ZMIN`, `ZMAX`.
Supports transfinite meshing with optional radial grading (`grad`) and axial bump (`bump`).
Writes `mshfile` and returns a `GmshDiscreteModel` for Gridap.
"""

function build_model_gmsh(params::Params, mshfile::String)

    Nr=params.Nr_solid
    Nz=params.Nz_solid
    grad=params.grad
    bump=params.bump
    quads=true

    @assert grad > 0 "grad must be > 0"
    gmsh.initialize()
    try
        gmsh.model.add("store2D")
        L, Rin, Rout = params.L, params.R_pipe, params.R_solid

        # Points: x=z, y=r
        p1 = gmsh.model.geo.addPoint(0.0, Rin,  0.0)   # (0, Rin)
        p2 = gmsh.model.geo.addPoint(L,   Rin,  0.0)   # (L, Rin)
        p3 = gmsh.model.geo.addPoint(L,   Rout, 0.0)   # (L, Rout)
        p4 = gmsh.model.geo.addPoint(0.0, Rout, 0.0)   # (0, Rout)

        # Edges with desired parameter directions
        l1 = gmsh.model.geo.addLine(p1, p2)  # z @ Rin:   0 → L
        l2 = gmsh.model.geo.addLine(p2, p3)  # r @ z=L:   Rin → Rout
        l3 = gmsh.model.geo.addLine(p4, p3)  # z @ Rout:  0 → L  (p4→p3)
        l4 = gmsh.model.geo.addLine(p1, p4)  # r @ z=0:   Rin → Rout

        # CCW loop: reverse edges in the loop (NOT in their transfinite definition)
        cl = gmsh.model.geo.addCurveLoop([ l1, l2, -l3, -l4 ])
        s  = gmsh.model.geo.addPlaneSurface([cl])

        # Structured (transfinite) setup
        gmsh.model.geo.mesh.setTransfiniteCurve(l1, Nz, "Bump", bump)
        gmsh.model.geo.mesh.setTransfiniteCurve(l3, Nz, "Bump", bump)
        # gmsh.model.geo.mesh.setTransfiniteCurve(l1, Nz)
        # gmsh.model.geo.mesh.setTransfiniteCurve(l3, Nz)
        gmsh.model.geo.mesh.setTransfiniteCurve(l2, Nr, "Progression", grad)
        gmsh.model.geo.mesh.setTransfiniteCurve(l4, Nr, "Progression", grad)
        gmsh.model.geo.mesh.setTransfiniteSurface(s, "Left", [p1,p2,p3,p4])

        if quads
            gmsh.model.geo.mesh.setRecombine(2, s)
            gmsh.option.setNumber("Mesh.RecombineAll", 1)             # safety
            gmsh.option.setNumber("Mesh.RecombinationAlgorithm", 1)
        end

        gmsh.model.geo.synchronize()

        # Physical groups
        gmsh.model.addPhysicalGroup(2, [s], 1); gmsh.model.setPhysicalName(2, 1, "SOLID")
        gmsh.model.addPhysicalGroup(1, [l1], 2); gmsh.model.setPhysicalName(1, 2, "PIPE")
        gmsh.model.addPhysicalGroup(1, [l3], 3); gmsh.model.setPhysicalName(1, 3, "OUTER")
        gmsh.model.addPhysicalGroup(1, [l4], 4); gmsh.model.setPhysicalName(1, 4, "ZMIN")
        gmsh.model.addPhysicalGroup(1, [l2], 5); gmsh.model.setPhysicalName(1, 5, "ZMAX")


        gmsh.model.mesh.generate(2)
        gmsh.write(mshfile)
    finally
        gmsh.finalize()
    end
    return GmshDiscreteModel(mshfile)
end

# === MATCHED FLUID GRID FROM PIPE ============================================

"Return z nodes along the PIPE boundary (sorted, unique)."

"""
    fluid_z_from_PIPE(model) -> Vector{Float64}

Build the 1‑D fluid axial grid `z` from the PIPE boundary segments (segment endpoints → nodes).
Honors any axial bump used in the solid mesh.
"""
function fluid_z_from_PIPE(model)::Vector{Float64}
    Γ = BoundaryTriangulation(model; tags="PIPE")
    XΓ = get_cell_coordinates(Γ)                # vector of edge-vertex arrays
    z_nodes = Float64[]
    for cell in XΓ, p in cell
        push!(z_nodes, p[AXIAL])                    # x is the axial (z) coordinate
    end
    return unique(sort(z_nodes))
end

"Length-weighted conversion of segment averages (N-1) to node values (N)."

"""
    segment_to_node_average(Tw_seg, z) -> Tw_node

Map per‑segment wall averages (length `length(z)-1`) to nodal values (length `length(z)`)
by arithmetic averaging of adjacent segments; endpoints copy their only neighbor.
"""
function segment_to_node_average(Tw_seg::Vector{Float64}, z::Vector{Float64})
    N = length(z); @assert length(Tw_seg) == N-1
    Tw_node = similar(z, Float64)
    Tw_node[1] = Tw_seg[1]
    for i in 2:N-1
        Δzm = z[i]-z[i-1]; Δzp = z[i+1]-z[i]
        Tw_node[i] = (Δzm*Tw_seg[i-1] + Δzp*Tw_seg[i])/(Δzm+Δzp)
    end
    Tw_node[N] = Tw_seg[end]
    return Tw_node
end

# === STABLE TIME STEP (tightest of advection / diffusion / exchange) =========
"Return (dt, dbg) for a non-uniform z; α_eff = α_phys + α_num."
function stable_dt_nonuniform(z::Vector{Float64}, params::Params)
    cfl=params.cfl_target
    alpha_f=params.alpha_f
    Δzmin = minimum(diff(z))
    v  = abs(params.v)
    ρ  = params.rho_oil
    cp = params.cp_oil
    R  = params.R_pipe
    h  = params.h
    k  = params.k_oil

    α_phys = k/(ρ*cp)                     # physical thermal diffusivity
    α_num  = 0.5 * v * Δzmin * max(alpha_f, 0.0)   # SUPG-like artificial part
    α_eff  = α_phys + α_num

    β      = (h * (2/R)) / (ρ * cp)       # exchange rate [1/s]

    dt_adv  = cfl * Δzmin / max(v, eps())
    dt_diff = (α_eff>0) ? Δzmin^2/(2α_eff) : Inf
    dt_xchg = 1/max(β, eps())
    dt      = 0.9 * min(dt_adv, dt_diff, dt_xchg)
    return dt, (dt_adv=dt_adv, dt_diff=dt_diff, dt_xchg=dt_xchg,
                alpha_eff=α_eff, alpha_phys=α_phys, alpha_num=α_num,
                beta=β, dzmin=Δzmin)
end

function step_fluid_nonuniform!(T_f::Vector{Float64}, Tw_node::Vector{Float64},
        params, dt::Float64, z::Vector{Float64}; alpha_eff::Float64)

N = length(z)
@assert length(T_f) == N == length(Tw_node)
@assert N ≥ 3 "Need at least 3 fluid nodes for this stencil"

v  = params.v
ρ  = params.rho_oil; cp = params.cp_oil
R  = params.R_pipe;  h  = params.h
β  = (2h)/(ρ*cp*R)                       # exchange coefficient

Δ   = diff(z)                            # face spacings: Δ[i] = z[i+1]-z[i], i=1..N-1

# control-volume widths (½ at the ends, average inside)
Δi  = similar(T_f)
Δi[1] = Δ[1]/2
@inbounds for i=2:N-1
Δi[i] = 0.5*(Δ[i-1] + Δ[i])
end
Δi[N] = Δ[end]/2

# Face advection flux Fadv[i] at face i+1/2 (i=1..N-1), upwind
Fadv = zeros(Float64, N-1)
if v ≥ 0
@inbounds for i=1:N-1
Fadv[i] = v*T_f[i]              # upwind = left cell
end
else
@inbounds for i=1:N-1
Fadv[i] = v*T_f[i+1]            # upwind = right cell
end
end

# Face diffusion flux Fdiff[i] = -alpha_eff * (ΔT/Δz_face)
Fdiff = zeros(Float64, N-1)
if alpha_eff > 0
@inbounds for i=1:N-1
Fdiff[i] = -alpha_eff * (T_f[i+1] - T_f[i]) / Δ[i]
end
end

Tnew = similar(T_f)
Tnew[1] = T_f[1]                         # Dirichlet inlet handled by controller

# Interior updates: flux divergence + exchange
@inbounds for i=2:N-1
div_adv = -(Fadv[i]  - Fadv[i-1])  / Δi[i]
div_dif = -(Fdiff[i] - Fdiff[i-1]) / Δi[i]
xchg    =  β*(Tw_node[i] - T_f[i])
Tnew[i] =  T_f[i] + dt*(div_adv + div_dif + xchg)
end

# Outlet (Neumann diffusion ⇒ right diffusive flux = 0)
div_advN = -( Fadv[N-1] - Fadv[N-2] ) / Δi[N]
div_difN = -( 0.0        - Fdiff[N-2] ) / Δi[N]
xchgN    =  β*(Tw_node[N] - T_f[N])
Tnew[N]  =  T_f[N] + dt*(div_advN + div_difN + xchgN)
Tnew[N] = Tnew[N-1]
copy!(T_f, Tnew)
return
end


# Build a boundary CellField of T_f(z) on Γ: interpolates from 1D fluid grid

"""
    make_Tf_boundary_field(z, T_f, Γ, axial_axis) -> CellField

Interpolate the 1‑D fluid temperature array `(z, T_f)` along the boundary triangulation `Γ`
(tag `"PIPE"`), returning a `CellField` T_f^\ast(x) evaluated by axial coordinate.
"""
function make_Tf_boundary_field(z::AbstractVector{<:Real},
                                T_f::AbstractVector{<:Real},
                                Γ)
    @assert length(z) == length(T_f) "z and T_f must have same length"
    @assert all(diff(z) .> 0) "z must be strictly increasing"

    function Tf_of_z(zz::Float64)
        if zz <= z[1]
            return float(T_f[1])
        elseif zz >= z[end]
            return float(T_f[end])
        else
            k = clamp(searchsortedlast(z, zz), 1, length(z)-1)
            zk, zk1 = z[k], z[k+1]
            denom = max(zk1 - zk, eps(Float64))
            λ = (zz - zk) / denom
            return (1-λ)*T_f[k] + λ*T_f[k+1]
        end
    end

    return CellField(x -> Tf_of_z(Float64(x[AXIAL])), Γ)
end

# Robust wall temperature field for Γ_pipe by sampling uh just inside the solid
function wall_T_cellfield(uh, Γ_pipe, params, z; frac::Float64 = 0.25)
    Rin, Rout = params.R_pipe, params.R_solid
    Nr_guess  = hasproperty(params, :Nr_solid) ? params.Nr_solid :
                (hasproperty(params, :Nr) ? params.Nr : 60)
    dr_est = max((Rout - Rin)/Nr_guess, 1e-6*(Rout - Rin))
    δr     = max(frac * dr_est, 1e-5*(Rout - Rin))

    dzmin = (length(z) > 1) ? minimum(diff(z)) : params.L
    δz    = max(frac * dzmin, 1e-5*params.L)

    return CellField(x -> begin
        zz = Float64(x[AXIAL])
        rr = Float64(x[RADIAL])
        # step a hair **into** the solid
        rr = min(rr + δr, Rout - 1e-12)
        # pull off the axial caps to avoid corner-eval misses
        if zz < δz
            zz += δz
        elseif zz > params.L - δz
            zz -= δz
        end
        uh(VectorValue(zz, rr))
    end, Γ_pipe)
end


# VTK writer. Exports:
#  1) solid volume with uh
#  2) PIPE boundary surface with Tf_b (fluid-on-wall) and T_wall (solid-on-wall)
#  3) fluid rect r–z grid with T_f and Tw_node repeated across radius

"""
    write_vtk_step(step, t, model, uh, z, T_f, params; kwargs...)

Writes solid field as unstructured VTK and fluid field as rectilinear r–z grid
(x=z, y=r). The 1‑D fluid temperature is replicated across radius for visualization only.
"""
function write_vtk_step(step, t::Float64, model, uh, z::Vector{Float64}, T_f::Vector{Float64}, params::Params, outdir::String; Nz_vis::Int=256, Nr_vis::Int=256, Tw_node::Union{Nothing,Vector{Float64}}=nothing)

                        # isdir(outdir) || mkpath(outdir)

    # --- 1) Solid volume ------------------------------------------------------
    base_solid = joinpath(outdir, @sprintf("solid_%06d", step))
    Ω = Triangulation(model)
    writevtk(Ω, base_solid; cellfields=["T"=>uh])

    # --- 2) PIPE boundary: Tf_b and T_wall -----------------------------------
    Γ_pipe = BoundaryTriangulation(model; tags="PIPE")

    # fluid temperature traced to wall (CellField on Γ)
    Tf_b = make_Tf_boundary_field(z, T_f, Γ_pipe)

    # solid wall temperature as a boundary CellField
    T_wall = wall_T_cellfield(uh, Γ_pipe, params, z)        # safe sample inside solid

    base_pipe = joinpath(outdir, @sprintf("pipe_%06d", step))
    writevtk(Γ_pipe, base_pipe; cellfields=["Tf_b"=>Tf_b, "T_wall"=>T_wall])

    # --- 3) Fluid rect grid in (z,r) with T_f (and optional Tw_node) ---------
    # We use the solver’s matched nodes for z; r just spans [0,R_pipe]
    z_vis = z
    r_vis = collect(range(0.0, params.R_pipe; length=Nr_vis))

    base_fluid = joinpath(outdir, @sprintf("fluid_%06d", step))
    vtkf = (AXIAL==1) ? vtk_grid(base_fluid, z_vis, r_vis) : vtk_grid(base_fluid, r_vis, z_vis)

    # T_f defined on z → repeat across radius
    Tf_grid = repeat(reshape(T_f, 1, length(z_vis)), Nr_vis, 1)  # (Nr, Nz)
    vtkf["T_fluid", VTKPointData()] = (AXIAL==1) ? Tf_grid' : Tf_grid

    if Tw_node !== nothing
        @assert length(Tw_node) == length(z_vis)
        Tw_grid = repeat(reshape(Tw_node, 1, length(z_vis)), Nr_vis, 1)
        vtkf["T_wall_node", VTKPointData()] = (AXIAL==1) ? Tw_grid' : Tw_grid
    end

    vtk_save(vtkf)
end

"Build measures & weights once (axisymmetric audit if/when you need it)."
function build_measures_axisym(model, params)
    Ω       = Triangulation(model)
    Γ_pipe  = BoundaryTriangulation(model; tags="PIPE")
    dΩ      = Measure(Ω, 2)
    dΓ_pipe = Measure(Γ_pipe, 1)
    wΩ      = CellField(x -> 2π * x[RADIAL], Ω)            # 2π r
    wΓ_pipe = CellField(_ -> 2π * params.R_pipe, Γ_pipe)   # 2π R (constant on PIPE)
    return (Ω=Ω, Γ_pipe=Γ_pipe, dΩ=dΩ, dΓ_pipe=dΓ_pipe, wΩ=wΩ, wΓ_pipe=wΓ_pipe)
end

function gnielinski(params::Params)
    # Compute HTC if not given
    D  = 2*params.R_pipe
    Re = params.v * D / params.nu_oil                  # ν [m^2/s]
    Pr = (params.rho_oil * params.nu_oil) * params.cp_oil / params.k_oil  # = μ*cp/k

    f  = (0.79*log(Re) - 1.64)^(-2)                    # Petukhov (Darcy), log = natural log in Julia
    Nu = (f/8) * (Re - 1000) * Pr / (1 + 12.7*sqrt(f/8)*(Pr^(2/3) - 1))
    h  = Nu * params.k_oil / D
    @info "Computed HTC via Gnielinski: $(h) Re: $(Re) Pr: $(Pr)"
    return h
end

function check_turbulent_flow(params::Params)
    D = 2 * params.R_pipe
    Re = params.v * D / params.nu_oil
    re_threshold = 4000.0  # Lower bound for turbulent flow (Gnielinski validity)
    if Re < re_threshold
        min_v = re_threshold * params.nu_oil / D
        println("Error: Flow speed $(params.v) m/s results in Re = $Re < $re_threshold, which may not ensure turbulent flow (Gnielinski correlation assumes turbulent regime). Minimum required speed: $min_v m/s.")
        exit(1)  # Stop the simulation
        #=
        you could use throw(ErrorException(...)) instead for more graceful error handling.
        =#
    end
end

"""
State for robust wall-temperature segment averages via L² projection on Γ=PIPE.
- Builds a P0 space on Γ, assembles/factorizes its mass matrix once.
- Stores segment midpoints & axial order for fast extraction each step.
"""

struct SolidCoupler
    # FE spaces & measures
    U
    V
    Γ_pipe
    dΩ
    dΓ
    wΩ
    wΓ
    # solid factor (CHOLMOD)
    F_solid                # untyped: SparseArrays.CHOLMOD.Factor
    # boundary P0 projector bits
    VΓ
    UΓ
    diagAΓ::Vector{Float64}        # diagonal of P0 mass matrix on Γ
    mids::Vector{VectorValue{2,Float64}}
    order::Vector{Int}
    # constants
    ρ::Float64
    c::Float64
    κ::Float64
    h::Float64
    θ::Float64
    L::Float64
    dt::Float64
end

"""
    prepare_coupled_solid_axisym(model, params, dt; θ=1.0) -> S::SolidCoupler

Assemble the time‑invariant axisymmetric solid system (implicit θ‑method), factorize it once
with CHOLMOD, and prepare the P0 boundary projector on `"PIPE"` (stores the diagonal only).
Returns a `SolidCoupler` with cached data for fast timesteps.
"""

function prepare_coupled_solid_axisym(model, params, dt; θ::Float64=1.0)
    # FE spaces
    reffe = ReferenceFE(lagrangian, Float64, 1)
    V = TestFESpace(model, reffe; conformity=:H1)
    U = TrialFESpace(V)

    # Measures & weights
    Ω = Triangulation(model)
    Γ_pipe = BoundaryTriangulation(model; tags="PIPE")
    dΩ = Measure(Ω, 2)
    dΓ = Measure(Γ_pipe, 1)
    wΩ = CellField(x -> 2π * x[RADIAL], Ω)
    wΓ = CellField(_ -> 2π * params.R_pipe, Γ_pipe)

    ρ, c, κ, h, L = params.rho_s, params.cp_s, params.k_s, params.h, params.L

    # Solid LHS (axisymmetric) + factorization
    a(u, v) = integrate(wΩ * ((ρ * c / dt) * v * u + θ * κ * (∇(v) ⋅ ∇(u))), dΩ) +
              integrate(θ * h * wΓ * v * u, dΓ)
    A = assemble_matrix(a, U, V)
    F_solid = cholesky(Symmetric(A))   # CHOLMOD

    # Boundary P0 (segment averages)
    reffeΓ = ReferenceFE(lagrangian, Float64, 0)
    VΓ = TestFESpace(Γ_pipe, reffeΓ; conformity=:L2)
    UΓ = TrialFESpace(VΓ)
    aΓ(u, v) = integrate(v * u, dΓ)
    AΓ = assemble_matrix(aΓ, UΓ, VΓ)
    diagAΓ = diag(AΓ)

    # Midpoints & axial order
    XΓ = get_cell_coordinates(Γ_pipe)
    mids = map(c -> 0.5 * (c[1] + c[2]), XΓ)
    order = sortperm(eachindex(mids); by=i -> mids[i][AXIAL])

    return SolidCoupler(U, V, Γ_pipe, dΩ, dΓ, wΩ, wΓ, F_solid, VΓ, UΓ, diagAΓ, mids, order, ρ, c, κ, h, θ, L, dt)
end

# Per-step solid solve (uses S.dt)
function solve_solid!(S::SolidCoupler, uh_prev::FEFunction, Tf_b)
    l(v) = integrate(S.wΩ * (S.ρ * S.c / S.dt) * v * uh_prev, S.dΩ) +
           integrate(S.θ * S.h * S.wΓ * v * Tf_b, S.dΓ)
    b = assemble_vector(l, S.V)
    x = S.F_solid \ b
    return FEFunction(S.U, x)
end

# Exact per-segment wall averages (L² on Γ), ordered along axis

"""
    wall_T_segment_average(S::SolidCoupler, uh) -> Tw_seg::Vector{Float64}

Exact per‑segment **L² averages** of the solid wall temperature on `"PIPE"`, ordered by axial
position to match the fluid segments. Implemented by assembling the linear form
`ℓ_Γ(v) = ∫ v * tr(uh) dΓ`, dividing elementwise by the diagonal of the P0 mass matrix on Γ,
then sampling at segment midpoints.
"""
function wall_T_segment_average(S::SolidCoupler, uh::FEFunction)::Vector{Float64}
    lΓ(v) = integrate(v * uh, S.dΓ)                # trace(uh) handled by Gridap
    bΓ = Vector(assemble_vector(lΓ, S.VΓ))
    xΓ = bΓ ./ S.diagAΓ                           # elementwise divide
    TwΓ = FEFunction(S.UΓ, xΓ)                    # 2-arg ctor (unconstrained)
    vals = [TwΓ(S.mids[i]) for i in 1:length(S.mids)]
    return vals[S.order]
end

make_Tf_b(S::SolidCoupler, z::AbstractVector, T_f::AbstractVector) =
    make_Tf_boundary_field(z, T_f, S.Γ_pipe)

# (optional) audits that reuse same measures/weights
energy_solid_axisym(uh, S::SolidCoupler, params; Tref=params.T_initial) =
    sum(integrate((params.rho_s * params.cp_s) * S.wΩ * (uh - Tref), S.dΩ))

wall_power_axisym(uh, Tf_b, S::SolidCoupler, params) =
    sum(integrate(params.h * S.wΓ * (uh - Tf_b), S.dΓ))

# ╔══════════════════════════════════════════════════════════════════════════╗
# ║ Main driver                                                              ║
# ╚══════════════════════════════════════════════════════════════════════════╝

"""
    run_simulation(config_file_or_params, args...) -> nothing

Top‑level driver: reads config, builds mesh/model, constructs measures and couplers,
matches fluid grid to PIPE, computes stable `dt`, time‑steps (controller → solid → wall‑avg → fluid),
runs energy/power audits, and writes outputs.
"""

function run_simulation(params, tulosdir::String, conf_id::String)
    # --- params, model, measures -----------------------------------------------------
    BLAS.set_num_threads(Threads.nthreads())
    omp_threads = get(ENV, "OMP_NUM_THREADS", "")
    @info "Threads $(Threads.nthreads()) $omp_threads"

    vtk_dir = joinpath(tulosdir, "vtk_$conf_id")
    rm(vtk_dir, recursive=true, force=true)
    mkdir(vtk_dir)

    if length(params.phases) > 1
        t_phase_change = params.phases[1][2]
    else
        t_phase_change = 0.0
    end

    hc = HeaterCooler(params)
    history = Tuple{Float64,Float64,Float64,Float64,Float64,Float64,Float64,Float64}[]

    if params.h === nothing
        params.h = gnielinski(params)
    end

    check_turbulent_flow(params)
    msh_file =joinpath(tulosdir, "$conf_id.msh")
    model = build_model_gmsh(params, msh_file)
 
    M = build_measures_axisym(model, params)      # (Ω, Γ_pipe, dΩ, dΓ_pipe, wΩ, wΓ_pipe)
    Γ_pipe = M.Γ_pipe                             # we need the boundary here

    # --- matched fluid grid (fixed) --------------------------------------------------
    z = fluid_z_from_PIPE(model)                  # matches solid PIPE edge (handles bump)
    w_trap = trapezoid_weights(z; zero_inlet=false)
    w_state = trapezoid_weights(z; zero_inlet=true)

    # --- time step from tightest limit (fixed) ---------------------------------------
    dt, dbg = stable_dt_nonuniform(z, params)
    @info "Fluid dt limits" (; dt, dbg...)

    T_f = fill(params.T_initial, length(z))                # fluid
    Vax = sum(integrate(M.wΩ, M.dΩ))  # build once

    # after model, z, dt are built:
    S = prepare_coupled_solid_axisym(model, params, dt; θ=1.0)
    uh = interpolate_everywhere(x -> params.T_initial, S.U)

    E_solid = 0.0
    E_solid_prev = E_solid
    E_oil = 0.0
    P_in = 0.0
    Q_wall = 0.0
    Energy_in = 0.0
    # Q_wall_cumul = 0.0

    # --- time loop -------------------------------------------------------------------
    t = 0.0; step = 0
    t_next_vtk = 0

    while t < params.t_end - 1e-12
        step += 1

        # Outlet for controller
        T_out = T_f[end]
        # Controller provides inlet; auditing omitted per your note
        T_in = control(hc, T_out, t, dt, params)
        T_f[1] = T_in

        # controller sets T_f[1] ...
        Tf_b = make_Tf_b(S, z, T_f)
        uh = solve_solid!(S, uh, Tf_b)
        Tw_seg = wall_T_segment_average(S, uh)
        Tw_node = segment_to_node_average(Tw_seg, z)
        # --- FLUID: non-uniform step (upwind ±, non-uniform Laplacian, wall exchange)
        step_fluid_nonuniform!(T_f, Tw_node, params, dt, z; alpha_eff=dbg.alpha_eff)

        E_solid = energy_solid_axisym(uh, M, params, Vax)
        E_oil   = energy_oil_axisym(T_f, z, w_state, params)
        P_in    = inlet_power_axisym(T_in, T_out, params)

        Q_wall = -wall_power_axisym(uh, Tf_b, S, params)
        
        # sum(integrate(params.h * (CellField(_ -> 2π * params.R_pipe, Wwall.Γ))
        #                       * (uh - Tf_b), Wwall.dΓ))
        Energy_in += P_in*dt
        # Q_wall_cumul += Q_wall*dt

        push!(history,(t, T_in, T_out, P_in, Q_wall, E_solid, E_oil, Energy_in))
        if t + 1e-12 >= t_next_vtk
            write_vtk_step(step, t, model, uh, z, T_f, params, vtk_dir; Nz_vis=256, Nr_vis=128)
            #=
            """
            Seuraavassa logiikka meni vähän solmun :-)
            Tarkoitus on hidastaa simuloinnin ensimmäinen puoli tuntia ja puoli tuntia kahden puolen
            ensimmäistä ajotavan vaihtoa.
            """
            if t < 1800 
                t_next_vtk += floor(Int, params.save_every/5.0)
            elseif t_phase_change > 1800
                if t > t_phase_change - 1800 && t < t_phase_change + 1800
                    t_next_vtk += floor(Int, params.save_every/5.0)
                else
                    t_next_vtk += params.save_every
                end
            else
                t_next_vtk += params.save_every
            end
            =#
            t_next_vtk += params.save_every
        end

        # Advance time
        t += dt
    end

    save_history(history, joinpath(tulosdir, "$(conf_id)_history.dat"))
end

function max_storage(params)
    V_solid = π*params.R_solid^2*params.L
    V_pipe = π * params.R_pipe^2 * params.L
    M = params.rho_s*(V_solid - V_pipe)
    Qmax = M*params.cp_s*(params.T_hot - params.T_initial)
    return Qmax
end

#####################################################################
# CLI
if abspath(PROGRAM_FILE) == @__FILE__

    run_id = length(ARGS) >= 1 ? ARGS[1] : "ajo"
    tulosdir = joinpath("tulokset", run_id)
    if !ispath(tulosdir)
        mkpath(tulosdir)
    end

    conf_id = length(ARGS) >= 2 ? ARGS[2] : "x"

    log_file = joinpath(tulosdir, "$(conf_id)_log_simulation.log")

    # Set up file logging
    log_io = open(log_file, "w")
    logger = SimpleLogger(log_io)
    global_logger(logger)
    @info Dates.format(now(), "yyyy-mm-dd_HH:MM")
    config_file = "config_$conf_id.yaml"

    cp(config_file, joinpath(tulosdir, config_file), force=true)
    @info config_file, tulosdir

    params = Params(config_file)
    log_params(params)
    @info "maximum storage $(max_storage(params))"  
    simulation_time = @elapsed run_simulation(params, tulosdir, conf_id)
    simulation_time = format_time(simulation_time, "simulation time")
    @info "$simulation_time"
    plot_all(params, tulosdir, conf_id)
end