codegen.rs

// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2026 Jonathan D.A. Jewell <j.d.a.jewell@open.ac.uk>
//
// Code generation module — Generates Betlang source code from parsed
// manifest variables. The generated code declares probabilistic variables,
// sets up distribution sampling, defines ternary logic propagation rules,
// and configures the simulation harness.

use crate::abi::SimulationConfig as AbiSimConfig;
use crate::codegen::distribution::{betlang_type, distribution_info, propagation_rule};
use crate::codegen::parser::{ParsedVariable, parse_simulation_config, parse_variables};
use crate::manifest::Manifest;
use std::fmt::Write as FmtWrite;

/// Generate complete Betlang source code from a validated manifest.
///
/// The generated code includes:
/// 1. A header comment with project metadata and generation timestamp.
/// 2. Distribution declarations for each probabilistic variable.
/// 3. Propagation rules mapping distributions to arithmetic/logic modes.
/// 4. Ternary logic helper declarations for Bernoulli variables.
/// 5. A simulation configuration block.
/// 6. An entry-point `simulate` function that orchestrates the Monte Carlo run.
pub fn generate_betlang_code(manifest: &Manifest) -> Result<String, String> {
    let variables = parse_variables(manifest)?;
    let sim_config = parse_simulation_config(&manifest.simulation);

    let mut code = String::with_capacity(4096);

    // --- Header ---
    emit_header(&mut code, &manifest.project.name);

    // --- Distribution declarations ---
    emit_distribution_declarations(&mut code, &variables);

    // --- Propagation rules ---
    emit_propagation_rules(&mut code, &variables);

    // --- Ternary logic helpers (only if Bernoulli variables exist) ---
    let has_bernoulli = variables
        .iter()
        .any(|v| v.distribution.kind() == "bernoulli");
    if has_bernoulli {
        emit_ternary_helpers(&mut code, &variables);
    }

    // --- Simulation configuration ---
    emit_simulation_config(&mut code, &sim_config);

    // --- Entry point ---
    emit_entry_point(&mut code, &variables, &sim_config);

    Ok(code)
}

/// Emit the Betlang file header with project metadata.
fn emit_header(code: &mut String, project_name: &str) {
    writeln!(
        code,
        "// Generated by betlangiser — ternary probabilistic modelling"
    )
    .expect("TODO: handle error");
    writeln!(code, "// Project: {}", project_name).expect("TODO: handle error");
    writeln!(code, "// SPDX-License-Identifier: MPL-2.0").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");
    writeln!(code, "module {} where", sanitize_identifier(project_name))
        .expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");
}

/// Emit distribution declarations — each variable gets a `let` binding
/// with its Betlang distribution constructor and type annotation.
fn emit_distribution_declarations(code: &mut String, variables: &[ParsedVariable]) {
    writeln!(code, "// --- Distribution Declarations ---").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    for var in variables {
        let info = distribution_info(&var.distribution);
        let btype = betlang_type(&var.distribution);

        writeln!(code, "/// {} — {}", var.name, info.description).expect("TODO: handle error");
        writeln!(code, "/// Sampling method: {}", info.sampling_method)
            .expect("TODO: handle error");
        writeln!(code, "/// Support: {}", info.support).expect("TODO: handle error");
        writeln!(
            code,
            "let {} : {} = {}",
            var.name, btype, info.betlang_constructor
        )
        .expect("TODO: handle error");
        writeln!(code).expect("TODO: handle error");
    }
}

/// Emit propagation rule declarations — each variable is annotated with
/// how arithmetic operations on it should propagate uncertainty.
fn emit_propagation_rules(code: &mut String, variables: &[ParsedVariable]) {
    writeln!(code, "// --- Propagation Rules ---").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    for var in variables {
        let rule = propagation_rule(&var.distribution);
        writeln!(code, "@propagation({}, \"{}\")", var.name, rule).expect("TODO: handle error");
    }
    writeln!(code).expect("TODO: handle error");
}

/// Emit ternary logic helper functions for Bernoulli variables.
///
/// For each Bernoulli variable, generates:
/// - A ternary evaluation function that maps sampled proportions to Kleene truth values
/// - AND/OR/NOT combinators for ternary logic chains
fn emit_ternary_helpers(code: &mut String, variables: &[ParsedVariable]) {
    writeln!(code, "// --- Ternary Logic (Kleene Algebra) ---").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    writeln!(code, "/// Kleene strong three-valued AND").expect("TODO: handle error");
    writeln!(code, "let ternary_and(a: Ternary, b: Ternary) -> Ternary =")
        .expect("TODO: handle error");
    writeln!(code, "  match (a, b) with").expect("TODO: handle error");
    writeln!(code, "  | (False, _) | (_, False) -> False").expect("TODO: handle error");
    writeln!(code, "  | (True, True)            -> True").expect("TODO: handle error");
    writeln!(code, "  | _                       -> Unknown").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    writeln!(code, "/// Kleene strong three-valued OR").expect("TODO: handle error");
    writeln!(code, "let ternary_or(a: Ternary, b: Ternary) -> Ternary =")
        .expect("TODO: handle error");
    writeln!(code, "  match (a, b) with").expect("TODO: handle error");
    writeln!(code, "  | (True, _) | (_, True)   -> True").expect("TODO: handle error");
    writeln!(code, "  | (False, False)           -> False").expect("TODO: handle error");
    writeln!(code, "  | _                        -> Unknown").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    writeln!(code, "/// Kleene strong three-valued NOT").expect("TODO: handle error");
    writeln!(code, "let ternary_not(a: Ternary) -> Ternary =").expect("TODO: handle error");
    writeln!(code, "  match a with").expect("TODO: handle error");
    writeln!(code, "  | True    -> False").expect("TODO: handle error");
    writeln!(code, "  | False   -> True").expect("TODO: handle error");
    writeln!(code, "  | Unknown -> Unknown").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");

    // Per-variable ternary evaluators.
    for var in variables {
        if var.distribution.kind() == "bernoulli" {
            writeln!(
                code,
                "/// Evaluate '{}' as ternary: True if p >= 0.9, False if p <= 0.1, else Unknown",
                var.name
            )
            .expect("TODO: handle error");
            writeln!(
                code,
                "let {}_ternary(observed: Float) -> Ternary =",
                var.name
            )
            .expect("TODO: handle error");
            writeln!(code, "  if observed >= 0.9 then True").expect("TODO: handle error");
            writeln!(code, "  else if observed <= 0.1 then False").expect("TODO: handle error");
            writeln!(code, "  else Unknown").expect("TODO: handle error");
            writeln!(code).expect("TODO: handle error");
        }
    }
}

/// Emit the simulation configuration block.
fn emit_simulation_config(code: &mut String, config: &AbiSimConfig) {
    writeln!(code, "// --- Simulation Configuration ---").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");
    writeln!(code, "@config(samples = {})", config.samples).expect("TODO: handle error");
    writeln!(code, "@config(confidence = {})", config.confidence).expect("TODO: handle error");
    if let Some(seed) = config.seed {
        writeln!(code, "@config(seed = {})", seed).expect("TODO: handle error");
    }
    writeln!(code, "@config(output = \"{}\")", config.output_format).expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");
}

/// Emit the main `simulate` entry point that runs the Monte Carlo simulation.
fn emit_entry_point(code: &mut String, variables: &[ParsedVariable], config: &AbiSimConfig) {
    writeln!(code, "// --- Entry Point ---").expect("TODO: handle error");
    writeln!(code).expect("TODO: handle error");
    writeln!(code, "/// Run the Monte Carlo simulation.").expect("TODO: handle error");
    writeln!(code, "let simulate() -> SimulationResults =").expect("TODO: handle error");
    writeln!(code, "  let results = {{").expect("TODO: handle error");

    for var in variables {
        let info = distribution_info(&var.distribution);
        writeln!(
            code,
            "    {}: sample({}, {}),",
            var.name, info.betlang_constructor, config.samples
        )
        .expect("TODO: handle error");
    }

    writeln!(code, "  }}").expect("TODO: handle error");
    writeln!(code, "  report(results, confidence={})", config.confidence)
        .expect("TODO: handle error");
}

/// Sanitize a project name into a valid Betlang module identifier.
///
/// Replaces hyphens and spaces with underscores, removes non-alphanumeric
/// characters, and ensures the identifier starts with a letter.
fn sanitize_identifier(name: &str) -> String {
    let sanitized: String = name
        .chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect();

    // Ensure it starts with a letter.
    if sanitized.starts_with(|c: char| c.is_ascii_digit()) {
        format!("m_{}", sanitized)
    } else if sanitized.is_empty() {
        "unnamed_module".to_string()
    } else {
        sanitized
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::manifest::{Manifest, ProjectConfig, SimulationConfig, VariableDecl};

    /// Helper: create a manifest with the given variables.
    fn make_manifest(name: &str, vars: Vec<VariableDecl>) -> Manifest {
        Manifest {
            project: ProjectConfig {
                name: name.to_string(),
                description: None,
            },
            variables: vars,
            simulation: SimulationConfig::default(),
        }
    }

    #[test]
    fn test_generate_produces_betlang_code() {
        let vars = vec![
            VariableDecl {
                name: "price".to_string(),
                distribution: "normal".to_string(),
                mean: Some(100.0),
                std_dev: Some(5.0),
                min: None,
                max: None,
                alpha: None,
                beta_param: None,
                probability: None,
                expression: None,
            },
            VariableDecl {
                name: "will_buy".to_string(),
                distribution: "bernoulli".to_string(),
                mean: None,
                std_dev: None,
                min: None,
                max: None,
                alpha: None,
                beta_param: None,
                probability: Some(0.7),
                expression: None,
            },
        ];
        let m = make_manifest("risk-model", vars);
        let code = generate_betlang_code(&m).expect("TODO: handle error");

        // Header present.
        assert!(code.contains("Generated by betlangiser"));
        assert!(code.contains("module risk_model where"));

        // Distribution declarations present.
        assert!(code.contains("let price : Prob<Float> = Normal(100, 5)"));
        assert!(code.contains("let will_buy : Prob<Ternary> = Bernoulli(0.7)"));

        // Propagation rules present.
        assert!(code.contains("@propagation(price, \"linear-error-propagation\")"));
        assert!(code.contains("@propagation(will_buy, \"kleene-ternary\")"));

        // Ternary helpers present (because there is a Bernoulli variable).
        assert!(code.contains("let ternary_and"));
        assert!(code.contains("let ternary_or"));
        assert!(code.contains("let ternary_not"));
        assert!(code.contains("let will_buy_ternary"));

        // Simulation config present.
        assert!(code.contains("@config(samples = 10000)"));
        assert!(code.contains("@config(confidence = 0.95)"));

        // Entry point present.
        assert!(code.contains("let simulate()"));
    }

    #[test]
    fn test_generate_without_bernoulli_omits_ternary() {
        let vars = vec![VariableDecl {
            name: "x".to_string(),
            distribution: "normal".to_string(),
            mean: Some(0.0),
            std_dev: Some(1.0),
            min: None,
            max: None,
            alpha: None,
            beta_param: None,
            probability: None,
            expression: None,
        }];
        let m = make_manifest("no-bernoulli", vars);
        let code = generate_betlang_code(&m).expect("TODO: handle error");

        // No ternary helpers when there are no Bernoulli variables.
        assert!(!code.contains("ternary_and"));
        assert!(!code.contains("Kleene"));
    }

    #[test]
    fn test_sanitize_identifier() {
        assert_eq!(sanitize_identifier("risk-model"), "risk_model");
        assert_eq!(sanitize_identifier("my model 2"), "my_model_2");
        assert_eq!(sanitize_identifier("123abc"), "m_123abc");
        assert_eq!(sanitize_identifier(""), "unnamed_module");
        assert_eq!(sanitize_identifier("valid_name"), "valid_name");
    }

    #[test]
    fn test_generate_all_distributions() {
        let vars = vec![
            VariableDecl {
                name: "a".to_string(),
                distribution: "normal".to_string(),
                mean: Some(0.0),
                std_dev: Some(1.0),
                min: None,
                max: None,
                alpha: None,
                beta_param: None,
                probability: None,
                expression: None,
            },
            VariableDecl {
                name: "b".to_string(),
                distribution: "uniform".to_string(),
                mean: None,
                std_dev: None,
                min: Some(0.0),
                max: Some(10.0),
                alpha: None,
                beta_param: None,
                probability: None,
                expression: None,
            },
            VariableDecl {
                name: "c".to_string(),
                distribution: "beta".to_string(),
                mean: None,
                std_dev: None,
                min: None,
                max: None,
                alpha: Some(2.0),
                beta_param: Some(5.0),
                probability: None,
                expression: None,
            },
            VariableDecl {
                name: "d".to_string(),
                distribution: "bernoulli".to_string(),
                mean: None,
                std_dev: None,
                min: None,
                max: None,
                alpha: None,
                beta_param: None,
                probability: Some(0.5),
                expression: None,
            },
            VariableDecl {
                name: "e".to_string(),
                distribution: "custom".to_string(),
                mean: None,
                std_dev: None,
                min: None,
                max: None,
                alpha: None,
                beta_param: None,
                probability: None,
                expression: Some("test()".to_string()),
            },
        ];
        let m = make_manifest("all-dists", vars);
        let code = generate_betlang_code(&m).expect("TODO: handle error");

        // All five variables declared.
        assert!(code.contains("let a :"));
        assert!(code.contains("let b :"));
        assert!(code.contains("let c :"));
        assert!(code.contains("let d :"));
        assert!(code.contains("let e :"));
    }

    #[test]
    fn test_generate_with_seed() {
        let vars = vec![VariableDecl {
            name: "x".to_string(),
            distribution: "normal".to_string(),
            mean: Some(0.0),
            std_dev: Some(1.0),
            min: None,
            max: None,
            alpha: None,
            beta_param: None,
            probability: None,
            expression: None,
        }];
        let mut m = make_manifest("seeded", vars);
        m.simulation.seed = Some(42);
        let code = generate_betlang_code(&m).expect("TODO: handle error");
        assert!(code.contains("@config(seed = 42)"));
    }
}