# AUTOGENERATED! DO NOT EDIT! File to edit: test_end_to_end_resnet_image_classification.ipynb (unless otherwise specified).

__all__ = [
    "ROOT_PATH",
    "GPU_ID",
    "use_hotdog_not_hotdog_test",
    "config",
    "MAX_EPOCHS",
    "TRAIN_MORE_THAN_ONE_EPOCH",
    "module",
    "trainer",
    "recorder",
    "before_train_module_weights",
    "after_train_module_weights",
    "checkpoint_callback",
    "checkpoint_path",
    "inference_pipeline",
    "inference_module",
    "inference_trainer",
    "before_load_inference_module_weights",
    "pl_checkpoint",
    "after_load_inference_module_weights",
    "test_samples",
    "test_dataset",
    "test_dataloader",
    "trainer_predictions",
    "inference_trainer_predictions",
    "mlflow_save_path",
    "test_samples_json",
    "images_df",
    "mlflow_predictions",
    "test_samples_with_paths_json",
    "images_df",
    "mlflow_predictions",
    "source_path",
    "model_image",
    "port",
    "container",
]

# Cell
import json
import os
import random
from pathlib import Path

import cloudpickle
import mlflow
import numpy as np
import pandas as pd
from pytorch_lightning.utilities.cloud_io import load as pl_load

from roheboam.engine import get_toolbox_lookup
from roheboam.engine.integrations.mlflow.loading import load_mlflow_pyfunc_model
from roheboam.engine.integrations.mlflow.saving import save_mlflow_pyfunc_model
from roheboam.engine.pipeline import Pipeline
from roheboam.engine.pipeline.config_loader import load_config_from_paths
from roheboam.engine.pipeline.constants import VARIABLE_KEY
from roheboam.engine.utils.convenience import is_notebook
from roheboam.engine.utils.saving import save_to_temporary_file
from roheboam.engine.utils.testing import all_methods_in_callback_handler_are_called, arrays_are_same, copy_and_flatten_model_weights
from roheboam.engine.vision.tasks.image_classification.data import create_image_classification_sample
from roheboam.engine.vision.tasks.image_classification.deployment import MLFlowPyfuncImageClassificationModel

# Cell
ROOT_PATH = Path(globals()["_dh"][0]) if is_notebook() else Path(__file__).parent.parent
GPU_ID = 1

# Cell
import torch


def use_hotdog_not_hotdog_test(config):
    config[VARIABLE_KEY]["train_data_path"] = str(ROOT_PATH / "data" / "hotdog_not_hotdog_test" / "train")
    config[VARIABLE_KEY]["test_data_path"] = str(ROOT_PATH / "data" / "hotdog_not_hotdog_test" / "test")
    config[VARIABLE_KEY]["max_epochs"] = 1
    config[VARIABLE_KEY]["train_num_workers"] = 0
    config[VARIABLE_KEY]["val_num_workers"] = 0
    config[VARIABLE_KEY]["test_num_workers"] = 0
    config[VARIABLE_KEY]["first_n_samples"] = 5
    config[VARIABLE_KEY]["gpus"] = 1 if torch.cuda.is_available() else None
    return config, (12, 8)


# Cell
config = load_config_from_paths(
    [
        ROOT_PATH / "configs/augmentation.yml",
        ROOT_PATH / "configs/callbacks.yml",
        ROOT_PATH / "configs/dataset.yml",
        ROOT_PATH / "configs/data_setup.yml",
        ROOT_PATH / "configs/loss.yml",
        ROOT_PATH / "configs/model.yml",
        ROOT_PATH / "configs/optimizer.yml",
        ROOT_PATH / "configs/trainer.yml",
        ROOT_PATH / "configs/transformations.yml",
        ROOT_PATH / "configs/variables.yml",
        ROOT_PATH / "configs/splitter.yml",
    ]
)

# Cell
config, FIG_SIZE = use_hotdog_not_hotdog_test(config)

# Cell
MAX_EPOCHS = config[VARIABLE_KEY]["max_epochs"]
TRAIN_MORE_THAN_ONE_EPOCH = MAX_EPOCHS > 1

# Cell
from roheboam.engine.logger import no_engine_logs

with no_engine_logs():
    pipeline = Pipeline.create_from_config(
        config,
        get_toolbox_lookup(),
    )
    pipeline.run()

# Cell
train_sample_idx, val_sample_idx = next(iter(pipeline.get_node_output("DataSplitCreator")().split()))
module = pipeline.get_node_output("ModuleCreator")(
    train_samples=pipeline.get_node_output("TrainSampleDataCreator")[train_sample_idx],
    val_samples=pipeline.get_node_output("TrainSampleDataCreator")[val_sample_idx],
    test_samples=pipeline.get_node_output("TestSampleDataCreator"),
)
trainer = pipeline.get_node_output("TrainerCreator")()
recorder = pipeline.get_node_output("ImageClassificationResultRecorder")

# Cell
before_train_module_weights = copy_and_flatten_model_weights(module.model)
trainer.fit(module)
after_train_module_weights = copy_and_flatten_model_weights(module.model)
trainer.test()

assert not all(
    [
        arrays_are_same(before_parameters, after_parameters)
        for before_parameters, after_parameters in zip(before_train_module_weights, after_train_module_weights)
    ]
)
assert all_methods_in_callback_handler_are_called(pipeline.get_node_output("CallbackHandler"))

# Cell
checkpoint_callback = pipeline.get_node_output("PyTorchLightningModelCheckpoint")
checkpoint_path = str(checkpoint_callback.best_model_path if TRAIN_MORE_THAN_ONE_EPOCH else checkpoint_callback.current_model_path)

# Cell
inference_pipeline = Pipeline.create_from_config(
    config,
    get_toolbox_lookup(),
)
inference_pipeline.run()

inference_module = inference_pipeline.get_node_output("ModuleCreator")(
    train_samples=inference_pipeline.get_node_output("TrainSampleDataCreator")[train_sample_idx],
    val_samples=inference_pipeline.get_node_output("TrainSampleDataCreator")[val_sample_idx],
    test_samples=inference_pipeline.get_node_output("TestSampleDataCreator"),
)
inference_trainer = inference_pipeline.get_node_output("TrainerCreator")(
    checkpoint_callback=inference_pipeline.get_node_output("PyTorchLightningModelCheckpoint")
)

before_load_inference_module_weights = copy_and_flatten_model_weights(inference_module.model)

assert id(trainer) != id(inference_trainer)
assert id(module.model) != id(inference_module.model)
assert id(before_train_module_weights) != id(before_load_inference_module_weights)
assert not all(
    [
        arrays_are_same(module_parameters, checkpoint_parameters)
        for module_parameters, checkpoint_parameters in zip(before_train_module_weights, before_load_inference_module_weights)
    ]
)
pl_checkpoint = pl_load(checkpoint_path, map_location=lambda storage, loc: storage)
inference_module.load_state_dict(pl_checkpoint["state_dict"])

after_load_inference_module_weights = copy_and_flatten_model_weights(inference_module.model)

assert id(before_load_inference_module_weights) != id(after_load_inference_module_weights)
assert not all(
    [
        arrays_are_same(before_parameters, after_parameters)
        for before_parameters, after_parameters in zip(before_load_inference_module_weights, after_load_inference_module_weights)
    ]
)

assert TRAIN_MORE_THAN_ONE_EPOCH ^ all(
    [
        arrays_are_same(after_train_parameters, after_load_parameters)
        for after_train_parameters, after_load_parameters in zip(after_train_module_weights, after_load_inference_module_weights)
    ]
) or str(checkpoint_callback.best_model_path) == str(checkpoint_callback.current_model_path)

# Cell
test_samples = pipeline.get_node_output("TestSampleDataCreator")
test_dataset = pipeline.get_node_output("TestDatasetCreator")(test_samples)
test_dataloader = pipeline.get_node_output("TestDataLoaderCreator")(test_dataset)
trainer_predictions = trainer.test(module, test_dataloaders=[test_dataloader])
inference_trainer_predictions = inference_trainer.test(inference_module, test_dataloaders=[test_dataloader])
assert str(trainer_predictions) == str(inference_trainer_predictions)

# Cell
from roheboam.engine.logger import no_engine_logs
from roheboam.engine.utils.convenience import remove_path

mlflow_save_path = pipeline.get_node_output("TimeStampedSavePath") / "mlflow"
if mlflow_save_path.exists():
    remove_path(mlflow_save_path)

save_mlflow_pyfunc_model(
    mlflow_save_path,
    pipeline.get_node_output("ServingModel"),
    artifacts={
        "config_path": str(checkpoint_callback.best_config_path),
        "model_save_path": str(checkpoint_callback.best_model_path),
    },
)
with no_engine_logs():
    mlflow_model = mlflow.pyfunc.load_model(str(mlflow_save_path))

# Cell
test_samples_json = {
    "columns": ["images"],
    "index": [i for i in range(len(test_samples))],
    "data": [[create_image_classification_sample(**s).image.data.tolist()] for s in test_samples],
}

images_df = pd.read_json(json.dumps(test_samples_json), orient="split")

# Cell
mlflow_predictions = mlflow_model.predict(images_df)

# Cell
test_samples_with_paths_json = {
    "columns": ["image_paths"],
    "index": [i for i in range(len(test_samples))],
    "data": [[str(create_image_classification_sample(**s).image.path)] for s in test_samples],
}
images_df = pd.read_json(json.dumps(test_samples_with_paths_json), orient="split")

# Cell
mlflow_predictions = mlflow_model.predict(images_df)

# Cell
import tempfile
import time

import requests
from requests.exceptions import ConnectionError

from roheboam.engine.integrations.docker import containerise_model, containerise_model_debug, run_docker_container
from roheboam.engine.utils.convenience import get_open_port

# Cell
source_path = str(ROOT_PATH).split("/examples")[0]
model_image = f"hotdog-not-hotdog-image-classification:{pipeline.get_node_output('TimeStampedSavePath').name}"
containerise_model_debug(
    source_path=source_path,
    model_path=mlflow_save_path,
    model_image=model_image,
)

# Cell
port = get_open_port()
container = run_docker_container(model_image, port=port, gpus=str(config["Variables"]["gpus"]))

# Cell
try:
    for _ in range(5):
        try:
            response = requests.post(
                f"http://0.0.0.0:{port}/invocations",
                data=json.dumps(test_samples_json),
                headers={"Content-Type": "application/json"},
            )
            assert response.status_code == 200
            break
        except Exception:
            print("Connection Error")
            time.sleep(5)
finally:
    container.stop()
