// Heat-diffusion benchmark: 10000 realizations on one and four CPU workers.
// Same task, configuration and seed as repeat_heat_diffusion.aner.
// Requires four admitted CPU slots. Run with --max-steps 5000000.
// The two ensembles run in sequence; this benchmark creates no file.
// Set four_first to true to reverse the run order without changing the model.
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)
}

// 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())

// Each run evolves 31 points through 400 dependent time steps per realization.
// Initial parameters and physical assumptions are the same for both ensembles.
let four_first = false
var first_workers = 1
var second_workers = 4
if four_first { first_workers = 4; second_workers = 1 }
let first_execution = parallel.local(workers: first_workers)
let second_execution = parallel.local(workers: second_workers)
print("Starting first ensemble with workers: " + String(first_workers))
let first = parallel.repeat(iterations: 10000, task: heat_realization,
    parameters: initial_parameters, execution: first_execution, seed: 2026)
print("First ensemble executor seconds: " + String(first.seconds()))
print("Starting second ensemble with workers: " + String(second_workers))
let second = parallel.repeat(iterations: 10000, task: heat_realization,
    parameters: initial_parameters, execution: second_execution, seed: 2026)
print("Second ensemble executor seconds: " + String(second.seconds()))
var one_seconds = first.seconds()
var four_seconds = second.seconds()
if four_first { one_seconds = second.seconds(); four_seconds = first.seconds() }
print("One-worker executor seconds: " + String(one_seconds))
print("Four-worker executor seconds: " + String(four_seconds))

// Serial verification checks shapes, every ordered cell and every derived seed.
var identical = first.rows() == 10000 && second.rows() == 10000 &&
    first.cols() == 4 && second.cols() == 4 && first.workers() == first_workers && second.workers() == second_workers
var mismatched_cells = 0
var mismatched_seeds = 0
var iteration = 0
while iteration < first.rows() {
    if first.seed(iteration: iteration) != second.seed(iteration: iteration) {
        mismatched_seeds = mismatched_seeds + 1
    }
    var column = 0
    while column < first.cols() {
        if first.value(iteration: iteration, column: column) !=
           second.value(iteration: iteration, column: column) {
            mismatched_cells = mismatched_cells + 1
        }
        column = column + 1
    }
    iteration = iteration + 1
}
identical = identical && mismatched_cells == 0 && mismatched_seeds == 0
var identical_statistics = true
var column = 0
while column < first.cols() {
    if first.mean(column: column) != second.mean(column: column) ||
       first.variance(column: column, sample: true) != second.variance(column: column, sample: true) ||
       first.std(column: column, sample: true) != second.std(column: column, sample: true) {
        identical_statistics = false
    }
    column = column + 1
}
print("Identical output cells and seeds: " + String(identical))
print("Identical means, sample variances and standard deviations: " + String(identical_statistics))
print("Output cells checked: " + String(first.rows() * first.cols()))
print("Derived seeds checked: " + String(first.rows()))
print("Mismatched output cells: " + String(mismatched_cells))
print("Mismatched derived seeds: " + String(mismatched_seeds))
print("Measured executor speedup: " + String(one_seconds / four_seconds))
if four_first {
    print("One-worker summary: " + second.summary())
    print("Four-worker summary: " + first.summary())
} else {
    print("One-worker summary: " + first.summary())
    print("Four-worker summary: " + second.summary())
}
column = 0
while column < second.cols() {
    print("Column " + String(column) + " mean: " + String(second.mean(column: column)))
    print("Column " + String(column) + " sample variance: " + String(second.variance(column: column, sample: true)))
    print("Column " + String(column) + " sample std: " + String(second.std(column: column, sample: true)))
    column = column + 1
}
print("Peak temperature mean (K): " + String(second.mean(column: 0)))
print("Peak temperature sample std (K): " + String(second.std(column: 0, sample: true)))
print("Final center temperature mean (K): " + String(second.mean(column: 1)))
print("Final spatial mean temperature (K): " + String(second.mean(column: 2)))
print("Fraction of peaks above 335 K: " + String(second.mean(column: 3)))