// Synthetic building observations: energy = 50 + 3*temperature + 2*occupancy.
// Whole buildings stay in one declared population. CPU Float64 regression.
import aner.data
import aner.ml
import aner.tensor

fn energy_schema() -> DataSchema {
    var schema = data.schema()
    schema = data.column(schema, "row_id", "Int64", false)
    schema = data.column(schema, "building_id", "String", false)
    schema = data.column(schema, "temperature", "Float64", false)
    schema = data.column(schema, "occupancy", "Float64", false)
    schema = data.column(schema, "energy", "Float64", false)
    return schema
}

let table = data.collect(
    data.scan_csv("examples/fixtures/contextual/energy.csv", energy_schema(), "NA"),
    max_rows: 16, max_bytes: 1048576, batch_rows: 4)
let observations = data.frame(table, "row_id", "building_id", "energy",
    List("", "", "degC", "1", "kWh"))
let populations = data.partition_groups(observations,
    List("B0", "B1", "B2", "B3"), List("B4", "B5"), List("B6", "B7"))
let preprocessing = ml.fit_standardizer(populations.train())
let training = data.batch(ml.transform(preprocessing, populations.train()), 1048576, "allow")
let validation = data.batch(ml.transform(preprocessing, populations.validation()), 1048576, "allow")
let testing = data.batch(ml.transform(preprocessing, populations.test()), 1048576, "allow")
let model = ml.fit_linear(training)
let prediction = ml.predict(model, validation)

print("Training, validation and test rows")
print(populations.train().rows())
print(populations.validation().rows())
print(populations.test().rows())
print("Input features exclude identity, group and target")
print(training.inputs().cols())
print("First validation observation identity")
print(prediction.row_id(0))
print("Prediction retains declared target unit")
print(prediction.unit())
print("Validation and test reproduce the synthetic rule")
print(ml.rmse(prediction, validation) < 0.000000001)
print(ml.rmse(ml.predict(model, testing), testing) < 0.000000001)
print("First validation prediction is 108 kWh within tolerance")
print(prediction.values().value(0, 0) > 107.999999999 &&
    prediction.values().value(0, 0) < 108.000000001)

// A separate three-row fixture isolates training-only standardization.
var oracle_schema = data.schema()
oracle_schema = data.column(oracle_schema, "row_id", "Int64", false)
oracle_schema = data.column(oracle_schema, "building_id", "String", false)
oracle_schema = data.column(oracle_schema, "temperature", "Float64", false)
oracle_schema = data.column(oracle_schema, "energy", "Float64", false)
let oracle_table = data.collect(
    data.scan_csv("examples/fixtures/contextual/standardization.csv", oracle_schema, "NA"),
    3, 1048576, 2)
let oracle = data.frame(oracle_table, "row_id", "building_id", "energy",
    List("", "", "degC", "kWh"))
let oracle_parts = data.partition_groups(oracle,
    List("TrainA", "TrainB"), List(), List("TestC"))
let oracle_preprocessing = ml.fit_standardizer(oracle_parts.train())
let held_out = data.batch(ml.transform(oracle_preprocessing, oracle_parts.test()), 1024, "allow")
print("Training center, scale and held-out transformed value")
print(oracle_preprocessing.center().value(0, 0))
print(oracle_preprocessing.scale().value(0, 0))
print(held_out.inputs().value(0, 0))

// Save new artifacts once; restoration requires this live trusted training anchor.
preprocessing.save("energy.scaler.anermodel")
model.save("energy.linear.anermodel")
let restored_preprocessing = ml.load_standardizer("energy.scaler.anermodel", populations.train())
let restored_training = data.batch(ml.transform(restored_preprocessing, populations.train()), 1048576)
let restored_model = ml.load_linear("energy.linear.anermodel", restored_training)
let restored_validation = data.batch(ml.transform(restored_preprocessing, populations.validation()), 1048576)
print("Restored workflow preserves the verified prediction contract")
print(ml.rmse(ml.predict(restored_model, restored_validation), restored_validation) < 0.000000001)