"""Compile a small model pipeline; upload the YAML through Kubeflow's UI."""

from kfp import compiler, dsl
from kfp.dsl import Dataset, Input, Metrics, Model, Output

MODEL_PACKAGES = ["scikit-learn==1.7.2", "joblib==1.5.2"]


@dsl.component(base_image="python:3.11-slim")
def prepare_data(training_data: Output[Dataset], test_data: Output[Dataset]):
    import json
    from pathlib import Path

    training = {"x": [[1], [2], [3], [4]], "y": [3, 5, 7, 9]}
    testing = {"x": [[5], [6], [7]], "y": [11, 13, 15]}
    Path(training_data.path).write_text(json.dumps(training), encoding="utf-8")
    Path(test_data.path).write_text(json.dumps(testing), encoding="utf-8")
    print("Prepared 4 training pairs and 3 separate test pairs.")


@dsl.component(base_image="python:3.11-slim", packages_to_install=MODEL_PACKAGES)
def train_model(
    training_data: Input[Dataset], fit_intercept: bool, model: Output[Model]
):
    import json
    from pathlib import Path

    import joblib
    import sklearn
    from sklearn.linear_model import LinearRegression

    data = json.loads(Path(training_data.path).read_text(encoding="utf-8"))
    estimator = LinearRegression(fit_intercept=fit_intercept)
    estimator.fit(data["x"], data["y"])
    joblib.dump(estimator, model.path)
    model.metadata.update({
        "fit_intercept": fit_intercept,
        "coefficient": float(estimator.coef_[0]),
        "intercept": float(estimator.intercept_),
        "scikit_learn_version": sklearn.__version__,
    })
    print(f"fit_intercept={fit_intercept}")
    print(f"Learned line: y = {estimator.coef_[0]:.6f} * x + {estimator.intercept_:.6f}")
    print("Saved the fitted model as a Model artifact.")


@dsl.component(base_image="python:3.11-slim", packages_to_install=MODEL_PACKAGES)
def evaluate_model(
    model: Input[Model],
    test_data: Input[Dataset],
    metrics: Output[Metrics],
    predictions: Output[Dataset],
) -> float:
    import json
    from pathlib import Path

    import joblib
    from sklearn.metrics import mean_absolute_error

    # This task reloads the model written by train_model in a separate task.
    estimator = joblib.load(model.path)
    data = json.loads(Path(test_data.path).read_text(encoding="utf-8"))
    predicted = estimator.predict(data["x"])
    mae = float(mean_absolute_error(data["y"], predicted))
    metrics.log_metric("mae", mae)
    rows = [
        {"x": x[0], "expected": y, "predicted": float(p), "absolute_error": abs(y - float(p))}
        for x, y, p in zip(data["x"], data["y"], predicted)
    ]
    Path(predictions.path).write_text(json.dumps(rows, indent=2), encoding="utf-8")
    for row in rows:
        print(f"x={row['x']}: expected={row['expected']}, predicted={row['predicted']:.6f}")
    print(f"Mean absolute error: {mae:.6f}")
    return mae


@dsl.pipeline(name="tiny-model-comparison")
def model_pipeline(fit_intercept: bool = True) -> float:
    data = prepare_data()
    training = train_model(
        training_data=data.outputs["training_data"], fit_intercept=fit_intercept
    )
    evaluation = evaluate_model(
        model=training.outputs["model"], test_data=data.outputs["test_data"]
    )
    for task in (data, training, evaluation):
        task.set_caching_options(False)
    return evaluation.outputs["Output"]


if __name__ == "__main__":
    compiler.Compiler().compile(model_pipeline, "model_pipeline.yaml")
    print("Created model_pipeline.yaml")