import json
from pathlib import Path

from codecarbon import EmissionsTracker
from datasets import load_dataset
from sklearn.metrics import accuracy_score

from model import FastModel, save_pipeline

dataset = load_dataset("rfcx/frugalai")
train_dataset = dataset["train"]
test_dataset = dataset["test"]
tracker = EmissionsTracker(allow_multiple_runs=True)
with open("../config.json", "r") as file:
    config = json.load(file)
model = FastModel(
    config["audio_processing_params"],
    config["features_params"],
    config["lgbm_params"],
)
model.fit(dataset["train"])

# INFERENCE
tracker.start()
tracker.start_task("inference")
true_label = dataset["test"]["label"]
predictions = model.predict(dataset["test"])

emissions_data = tracker.stop_task()

print(accuracy_score(true_label, predictions))
print("energy_consumed_wh", emissions_data.energy_consumed * 1000)
print("emissions_gco2eq", emissions_data.emissions * 1000)

save_pipeline(model, Path("../"))