# Databricks notebook source
# MAGIC %md ## Import Run
# MAGIC
# MAGIC Import run from the folder that was created by the [Export_Run]($Export_Run) notebook.
# MAGIC
# MAGIC #### Widgets
# MAGIC * `1. Destination experiment name` - Import run into this experiment. Will create if it doesn't exist.
# MAGIC * `2. Input directory` - Input directory containing an exported run.
# MAGIC * `3. Import source tags`

# COMMAND ----------

# MAGIC %md ### Include setup

# COMMAND ----------

# MAGIC %run ./Common

# COMMAND ----------

# MAGIC %md ### Widget setup

# COMMAND ----------

dbutils.widgets.text("1. Destination experiment name", "")
experiment_name = dbutils.widgets.get("1. Destination experiment name")

dbutils.widgets.text("2. Input directory", "")
input_dir = dbutils.widgets.get("2. Input directory")

dbutils.widgets.dropdown("3. Import source tags", "no", ["yes", "no"])
import_source_tags = dbutils.widgets.get("3. Import source tags") == "yes"

print("input_dir:", input_dir)
print("experiment_name:", experiment_name)
print("import_source_tags:", import_source_tags)

# COMMAND ----------

assert_widget(experiment_name, "1. Destination experiment name")
assert_widget(input_dir, "2. Input base directory")

# COMMAND ----------

# MAGIC %md ### Import Run

# COMMAND ----------

from mlflow_migration.run.import_run import import_run

run, _ = import_run(
    experiment_name=experiment_name,
    input_dir=input_dir,
    import_source_tags=import_source_tags,
)
print("Run ID:", run.info.run_id)

# COMMAND ----------

# MAGIC %md ### Display run UI link

# COMMAND ----------

display_run_uri(run.info.run_id)

# COMMAND ----------

# MAGIC %md ### Check imported source tags

# COMMAND ----------

if import_source_tags:
    import pandas as pd

    run = mlflow_client.get_run(run.info.run_id)
    data = [(k, v) for k, v in run.data.tags.items() if k.startswith("mlflow_exim")]
    df = pd.DataFrame(data, columns=["Key", "Value"])
    display(df)
