#%%
import automatminer as amm
import pandas as pd

import src.utils.logging as logging

#%%
log_dir = logging.get_log_dir("automatminer")
logger = logging.get_logger(path=log_dir)

#%%
train = pd.read_csv("data/gaultois_labels.csv", comment="#")[
    ["T", "formula", "zT"]
].dropna(subset=["zT"])

screen = pd.read_csv("data/screening_formulas.csv", comment="#")

for df in [train, screen]:
    df.rename(columns={"formula": "composition"}, inplace=True)

#%%
# Form Cartesian product between the screening features and the
# 4 temperatures found in Gaultois's db. We'll predict each
# material at all 4 temps.
temps = pd.DataFrame([300, 400, 700, 1000], columns=["T"])
# Create a temporary common key around which to compute the Cartesian product.
screen["key"] = temps["key"] = 0
# Perform the Cartesian product.
screen = temps.merge(screen, how="outer")
for df in [temps, screen]:
    df.drop(columns=["key"], inplace=True)

# #%%

pipe_config = {
    **amm.get_preset_config(),
    "learner": amm.automl.TPOTAdaptor(generations=4),
}

pipe = amm.MatPipe(**pipe_config)
pipe.fit(train, "zT")

# #%%
# pipe.digest(log_dir + "digest.txt")
# pipe.save(log_dir + "mat.pipe")

#%%
pipe = amm.MatPipe.load("logs/automatminer/10-09@09:02:42/mat.pipe")

#%%
pred = pipe.predict(screen.tail(10))


# %%
