## Example for surrogate model with a custom pretrained model

# This example shows how to pre-train a model and use it as a surrogate.
# Please note that the model is not designed to be useful but to demonstrate the workflow.

# This example assumes some basic familiarity with using BayBE.
# We thus refer to [`campaign`](./../Basics/campaign.md) for a basic example.

### Necessary imports

import numpy as np
import torch
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx.operator_converters.linear_regressor import convert_sklearn_bayesian_ridge
from sklearn.linear_model import BayesianRidge

from baybe.campaign import Campaign
from baybe.objective import Objective
from baybe.parameters import NumericalDiscreteParameter
from baybe.recommenders import FPSRecommender, SequentialGreedyRecommender
from baybe.searchspace import SearchSpace
from baybe.strategies import TwoPhaseStrategy
from baybe.surrogates import CustomONNXSurrogate
from baybe.targets import NumericalTarget
from baybe.utils.dataframe import add_fake_results, to_tensor

### Experiment Setup

parameters = [
    NumericalDiscreteParameter(
        name="Pressure[bar]",
        values=[1, 5, 10],
        tolerance=0.2,
    ),
    NumericalDiscreteParameter(
        name="Temperature[degree_C]",
        values=np.linspace(100, 200, 10),
    ),
]


### "Pre-training" stage

# Note that this example trains with several helpers built-in to BayBE.
# This can be done independently (and elsewhere).
# The only requirement that BayBE needs is that the model is in an onnx format.
# And The format should return both the mean and standard deviation.
# This example is based on a `BayesianRidge` regressor from `sklearn`.
# Its native conversion to onnx is supported via `skl2onnx`.
# Please also note that this example does not give a useful model.
# Its purpose is to show the workflow for using pre-trained surrogates in BayBE.

searchspace = SearchSpace.from_product(parameters=parameters, constraints=None)
train_x = to_tensor(searchspace.discrete.comp_rep)
train_y = torch.rand(train_x.size(dim=0))  # train with a random y vector

# Define model and fit

model = BayesianRidge()
model.fit(train_x, train_y)


### Convert model to onnx

# Need the option to return standard deviation

options = {type(model): {"return_std": True}}

# Specify what the input name is

ONNX_INPUT_NAME = "example_input_name"

# input dimensions and input type (should always be a float)

input_dim = train_x.size(dim=1)
initial_type = [(ONNX_INPUT_NAME, FloatTensorType([None, input_dim]))]

# Conversion

onnx_str = convert_sklearn(
    model,
    initial_types=initial_type,
    options=options,
    custom_conversion_functions={type(model): convert_sklearn_bayesian_ridge},
).SerializeToString()  # serialize to string to save in file


### Create a surrogate model with a pretrained model

surrogate_model = CustomONNXSurrogate(
    onnx_str=onnx_str,
    onnx_input_name=ONNX_INPUT_NAME,  # specify input name
)

### Create campaign

campaign = Campaign(
    searchspace=SearchSpace.from_product(parameters=parameters, constraints=None),
    objective=Objective(
        mode="SINGLE", targets=[NumericalTarget(name="Yield", mode="MAX")]
    ),
    strategy=TwoPhaseStrategy(
        recommender=SequentialGreedyRecommender(surrogate_model=surrogate_model),
        initial_recommender=FPSRecommender(),
    ),
)

### Iterate with recommendations and measurements

# Let's do a first round of recommendation
recommendation = campaign.recommend(batch_size=2)

print("Recommendation from campaign:")
print(recommendation)

# Add some fake results

add_fake_results(recommendation, campaign)
campaign.add_measurements(recommendation)

### Model Outputs

# Do another round of recommendations
recommendation = campaign.recommend(batch_size=2)

# Print second round of recommendations

print("Recommendation from campaign:")
print(recommendation)


### Using configuration instead

# Note that this can be placed inside an overall baybe config
# Refer to [`create_from_config`](./../Serialization/create_from_config.md) for an example

CONFIG = {
    "type": "CustomONNXSurrogate",
    "onnx_str": onnx_str,
    "onnx_input_name": ONNX_INPUT_NAME,
}

### Model creation from dict (or json if string)
model_from_python = CustomONNXSurrogate(
    onnx_str=onnx_str, onnx_input_name=ONNX_INPUT_NAME
)
model_from_configs = CustomONNXSurrogate.from_dict(CONFIG)

# This configuration creates the same model

assert model_from_python == model_from_configs

# JSON configuration (expects onnx_str to be decoded with `ISO-8859-1`)

model_json = model_from_python.to_json()
assert model_from_python == CustomONNXSurrogate.from_json(model_json)
