// Uncertainty propagation through a conservative one-dimensional heat model.
// Serial configuration -> 1000 independent parallel realizations -> serial statistics.
// Requires four admitted CPU slots. Run with --max-steps 5000000.
import aner.parallel

class HeatExperiment {
    public let points: Int64
    public let length_m: Float64
    public let dt_s: Float64
    public let steps: Int64
    public let pulse_steps: Int64
    public let ambient_k: Float64
    public let diffusivity_m2_s: Float64
    public let diffusivity_spread: Float64
    public let source_k_s: Float64
    public let source_spread: Float64
    public let threshold_k: Float64
}

fn finite_value(value: Float64) -> Bool {
    return value - value == 0.0
}

fn validate_experiment(c: HeatExperiment) -> Result<HeatExperiment, String> {
    if !finite_value(c.length_m) || !finite_value(c.dt_s) || !finite_value(c.ambient_k) ||
       !finite_value(c.diffusivity_m2_s) || !finite_value(c.diffusivity_spread) ||
       !finite_value(c.source_k_s) || !finite_value(c.source_spread) || !finite_value(c.threshold_k) {
        return Result.err<HeatExperiment, String>("Every Float64 configuration field must be finite.")
    }
    if c.points < 3 || c.points % 2 == 0 {
        return Result.err<HeatExperiment, String>("Use an odd grid of at least three points.")
    }
    if c.length_m <= 0.0 || c.dt_s <= 0.0 || c.steps < 1 {
        return Result.err<HeatExperiment, String>("Length, time step and step count must be positive.")
    }
    if c.pulse_steps < 0 || c.pulse_steps > c.steps || c.ambient_k <= 0.0 {
        return Result.err<HeatExperiment, String>("Pulse duration must fit the horizon; initial Kelvin temperature must be positive.")
    }
    if c.diffusivity_m2_s <= 0.0 || c.diffusivity_spread < 0.0 || c.diffusivity_spread >= 1.0 {
        return Result.err<HeatExperiment, String>("Diffusivity must be positive with relative spread in [0, 1).")
    }
    if c.source_k_s < 0.0 || c.source_spread < 0.0 || c.source_spread > 1.0 || c.threshold_k <= 0.0 {
        return Result.err<HeatExperiment, String>("Use nonnegative heating, source spread in [0, 1], and a positive Kelvin threshold.")
    }
    let dx = c.length_m / Float64(c.points - 1)
    let dx2 = dx * dx
    let alpha_max = c.diffusivity_m2_s * (1.0 + c.diffusivity_spread)
    if !finite_value(dx2) || dx2 <= 0.0 || !finite_value(alpha_max) || alpha_max <= 0.0 {
        return Result.err<HeatExperiment, String>("Grid spacing and maximum diffusivity must be positive and representable.")
    }
    let dt_over_dx2 = c.dt_s / dx2
    let courant_sum = 2.0 * alpha_max * dt_over_dx2
    let maximum_heating_increment = c.dt_s * c.source_k_s * (1.0 + c.source_spread)
    if !finite_value(dt_over_dx2) || dt_over_dx2 <= 0.0 || !finite_value(courant_sum) ||
       !finite_value(maximum_heating_increment) || courant_sum > 1.0 {
        return Result.err<HeatExperiment, String>("Explicit diffusion parameters must be representable and stable: reduce dt_s or use a coarser grid.")
    }
    return Result.ok<HeatExperiment, String>(c)
}

fn heat_realization(parameters: List, iteration: Int64, seed: Int64) -> List {
    let c = validate_experiment(cast<HeatExperiment>(parameters.get(0))).unwrap()
    let dx = c.length_m / Float64(c.points - 1)
    let dt_over_dx2 = c.dt_s / (dx * dx)
    let alpha = Array<Float64>(c.points - 1, 0.0)
    let source_shape = Array<Float64>(c.points, 0.0)
    var temperature = Array<Float64>(c.points, c.ambient_k)
    var next_temperature = Array<Float64>(c.points, c.ambient_k)

    // Each face coefficient is drawn once, then shared by both adjacent fluxes.
    var face = 0
    while face < c.points - 1 {
        let perturbation = 2.0 * parallel.random(seed, draw: face) - 1.0
        alpha.set(face, c.diffusivity_m2_s * (1.0 + c.diffusivity_spread * perturbation))
        face = face + 1
    }
    // A distinct draw controls one heating amplitude for this realization.
    let amplitude_draw = parallel.random(seed, draw: c.points - 1)
    let amplitude = 1.0 + c.source_spread * (2.0 * amplitude_draw - 1.0)

    // Average the central-third heater over each node-centered control volume.
    // This gives half strength at nodes 10 and 20 on the default 31-point grid.
    var node = 1
    while node < c.points - 1 {
        let x = Float64(node) * dx
        var left = x - 0.5 * dx
        var right = x + 0.5 * dx
        if left < c.length_m / 3.0 { left = c.length_m / 3.0 }
        if right > 2.0 * c.length_m / 3.0 { right = 2.0 * c.length_m / 3.0 }
        if right > left { source_shape.set(node, (right - left) / dx) }
        node = node + 1
    }

    var peak_k = c.ambient_k
    var step = 0
    while step < c.steps {
        var heating_increment = 0.0
        if step < c.pulse_steps { heating_increment = c.dt_s * c.source_k_s * amplitude }
        node = 1
        while node < c.points - 1 {
            let current = temperature.get(node)
            let left_flux = alpha.get(node - 1) * (temperature.get(node - 1) - current)
            let right_flux = alpha.get(node) * (temperature.get(node + 1) - current)
            let updated = current + dt_over_dx2 * (left_flux + right_flux) +
                heating_increment * source_shape.get(node)
            next_temperature.set(node, updated)
            if updated > peak_k { peak_k = updated }
            node = node + 1
        }
        // Both arrays retain the fixed endpoint temperatures. Swap, do not update in place.
        let previous = temperature
        temperature = next_temperature
        next_temperature = previous
        step = step + 1
    }
    let center_k = temperature.get(c.points / 2)
    let spatial_mean_k = (temperature.sum() - c.ambient_k) / Float64(c.points - 1)
    var exceeded = 0.0
    if peak_k > c.threshold_k { exceeded = 1.0 }
    return List(peak_k, center_k, spatial_mean_k, exceeded)
}

// Serial statistics below need roots only of bounded probabilities divided by n.
// Forty Newton steps suffice for this example's n=1000 (no general sqrt API implied).
fn probability_root(value: Float64) -> Float64 {
    if value == 0.0 { return 0.0 }
    var estimate = 1.0
    var step = 0
    while step < 40 {
        estimate = 0.5 * (estimate + value / estimate)
        step = step + 1
    }
    return estimate
}

// 1. SERIAL SETUP: one declared experiment; all initial/boundary nodes are 293.15 K.
let experiment = HeatExperiment(points: 31, length_m: 0.1, dt_s: 0.1,
    steps: 400, pulse_steps: 200, ambient_k: 293.15,
    diffusivity_m2_s: 0.00001, diffusivity_spread: 0.25,
    source_k_s: 3.0, source_spread: 0.2, threshold_k: 335.0)
let checked = validate_experiment(experiment)
if checked.is_err() { print(checked.unwrap_err()) }
let initial_parameters = List(checked.unwrap())

// BEGIN PARALLEL
// 2. Independent realizations share the same nominal inputs, with private field buffers.
let results = parallel.repeat(iterations: 1000, task: heat_realization,
    parameters: initial_parameters, execution: parallel.local(workers: 4), seed: 2026)
// END PARALLEL

// 3. SERIAL ANALYSIS: simulation spread and Monte Carlo estimator uncertainty differ.
let n = Float64(results.rows())
let probability = results.mean(column: 3)
let mc_standard_error = probability_root(probability * (1.0 - probability) / n)
let z = 1.96
let z2 = z * z
let denominator = 1.0 + z2 / n
let interval_center = (probability + z2 / (2.0 * n)) / denominator
let interval_half = z * probability_root(probability * (1.0 - probability) / n +
    z2 / (4.0 * n * n)) / denominator
print("Independent realizations: " + String(results.rows()))
print("Peak temperature mean (K): " + String(results.mean(column: 0)))
print("Peak temperature sample std (K): " + String(results.std(column: 0, sample: true)))
print("Final center temperature mean (K): " + String(results.mean(column: 1)))
print("Final spatial mean temperature (K): " + String(results.mean(column: 2)))
print("Fraction of peaks above 335 K: " + String(probability))
print("Monte Carlo probability standard error: " + String(mc_standard_error))
print("Approximate 95% Wilson interval: " + String(interval_center - interval_half) +
    " to " + String(interval_center + interval_half))
print("Native executor seconds: " + String(results.seconds()))
// Optional new-file export preserves iteration, seed and the four output columns:
// results.save_csv("heat-realizations.csv")