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

__all__ = [
    "ROOT_PATH",
    "GPU_ID",
    "use_coco_test",
    "config",
    "MAX_EPOCHS",
    "TRAIN_MORE_THAN_ONE_EPOCH",
    "pipeline",
    "module",
    "trainer",
    "recorder",
    "before_train_module_weights",
    "after_train_module_weights",
    "load_trainer_and_module_with_checkpoint_path_from_pipeline",
    "load_current_checkpoint_path_from_pipeline",
    "checkpoint_path",
    "inference_pipeline",
    "before_load_inference_module_weights",
    "pl_checkpoint",
    "after_load_inference_module_weights",
    "checkpoint_path",
    "inference_pipeline",
    "samples",
    "test_dataloader",
    "recorder",
    "results",
    "checkpoint_callback",
    "mlflow_save_path",
    "test_samples_json",
    "images_df",
    "mlflow_predictions",
    "test_samples_with_paths_json",
    "images_df",
    "source_path",
    "model_image",
    "port",
    "container",
]

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

import cloudpickle
import mlflow
import numpy as np
import pandas as pd
import torch
import torch.utils
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.logger import set_debug_logs, set_no_logs
from roheboam.engine.pipeline import Pipeline
from roheboam.engine.pipeline.config_loader import dump_config_to_path, 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_bounding_box_detection.data import create_image_bboxes_detection_sample
from roheboam.engine.vision.tasks.image_bounding_box_detection.deployment import MLFlowPyfuncImageBoundingBoxDetectionModel
from roheboam.engine.vision.utils.image.plot import plot_image_grid

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

# Cell
def use_coco_test(config):
    config[VARIABLE_KEY]["train_image_path"] = str(ROOT_PATH / "data" / "images" / "train2017")
    config[VARIABLE_KEY]["train_image_glob_string"] = "*.jpg"

    config[VARIABLE_KEY]["train_label_path"] = str(ROOT_PATH / "data" / "labels" / "train2017")
    config[VARIABLE_KEY]["train_label_glob_string"] = "*.txt"

    config[VARIABLE_KEY]["test_image_path"] = str(ROOT_PATH / "data" / "images" / "test2017")
    config[VARIABLE_KEY]["test_image_glob_string"] = "*.jpg"

    config[VARIABLE_KEY]["test_label_path"] = str(ROOT_PATH / "data" / "labels" / "test2017")
    config[VARIABLE_KEY]["test_label_glob_string"] = "*.txt"
    config[VARIABLE_KEY]["max_epochs"] = 1
    config[VARIABLE_KEY]["train_num_workers"] = 8
    config[VARIABLE_KEY]["val_num_workers"] = 8
    config[VARIABLE_KEY]["test_num_workers"] = 8
    config[VARIABLE_KEY]["labels_to_keep"] = None
    config[VARIABLE_KEY]["gpus"] = GPU_ID if torch.cuda.is_available() else None
    return config, (12, 8)


# Cell
config = load_config_from_paths(
    [
        ROOT_PATH / "configs/callbacks.yml",
        ROOT_PATH / "configs/dataset.yml",
        ROOT_PATH / "configs/data_setup.yml",
        ROOT_PATH / "configs/loss.yml",
        ROOT_PATH / "configs/metric.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_coco_test(config)

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

# Cell
set_no_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=np.array(pipeline.get_node_output("TrainSamples")[train_sample_idx].tolist()),
    val_samples=pipeline.get_node_output("TrainSamples")[val_sample_idx],
    test_samples=pipeline.get_node_output("TestSamples"),
)
trainer = pipeline.get_node_output("TrainerCreator")(checkpoint_callback=pipeline.get_node_output("PyTorchLightningModelCheckpoint"))
recorder = pipeline.get_node_output("ImageBoundingBoxDetectionResultRecorder")

# 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
import numpy as np

from roheboam.engine.vision.tasks.image_bounding_box_detection.data import ImageBoundingBoxDetectionSample
from roheboam.engine.vision.utils.bbox import nms_with_scores


def load_trainer_and_module_with_checkpoint_path_from_pipeline(pipeline, checkpoint_path, attach_recorder=True):

    pipeline.run()
    module = pipeline.get_node_output("ModuleCreator")()
    module.load_state_dict(pl_load(checkpoint_path, map_location=lambda storage, loc: storage)["state_dict"])
    trainer = pipeline.get_node_output("TrainerCreator")(checkpoint_callback=pipeline.get_node_output("PyTorchLightningModelCheckpoint"))
    recorder_exists = module.callback_handler.get_callback_by_name("ImageBoundingBoxDetectionInferenceRecorder") is not None
    if attach_recorder and not recorder_exists:
        pipeline.get_node("ImageBoundingBoxDetectionInferenceRecorder").should_remove = False
        pipeline.get_node("ImageBoundingBoxDetectionInferenceRecorder").create_output()
        recorder = pipeline.get_node_output("ImageBoundingBoxDetectionInferenceRecorder")
        module.callback_handler.callbacks.append(recorder)
    return trainer, module


def load_current_checkpoint_path_from_pipeline(pipeline):
    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)
    return checkpoint_path


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

before_load_inference_module_weights = copy_and_flatten_model_weights(inference_pipeline.get_node_output("ModuleCreator")().model)

(
    inference_trainer,
    inference_module,
) = load_trainer_and_module_with_checkpoint_path_from_pipeline(inference_pipeline, checkpoint_path)

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 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)
    ]
)

# Cell
checkpoint_path = load_current_checkpoint_path_from_pipeline(pipeline)

# Cell
inference_pipeline = Pipeline.create_from_config(
    config,
    get_toolbox_lookup(),
)
(
    inference_trainer,
    inference_module,
) = load_trainer_and_module_with_checkpoint_path_from_pipeline(inference_pipeline, checkpoint_path)

samples = inference_pipeline.get_node_output("TestSamples")
test_dataloader = inference_pipeline.get_node_output("TestDataLoaderCreator")(inference_pipeline.get_node_output("TestDatasetCreator")(samples))
recorder = inference_pipeline.get_node_output("ImageBoundingBoxDetectionInferenceRecorder")
results = recorder.create_results_from_dataloader(test_dataloader, inference_trainer, inference_module)

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

checkpoint_callback = pipeline.get_node_output("PyTorchLightningModelCheckpoint")
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(samples))],
    "data": [[s.image.data.tolist()] for s in 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(samples))],
    "data": [[str(s.image.path)] for s in samples],
}
images_df = pd.read_json(json.dumps(test_samples_with_paths_json), orient="split")

# Cell
import tempfile
import time

import requests
from urllib3.exceptions import NewConnectionError

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"coco-image-bbox-detection:{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()
