from pathlib import Path

import numpy as np
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.preprocessing import LabelEncoder
from tabpfn import TabPFNClassifier

from dcat_ap_hub import Dataset

# %% [markdown]
# Notebook-style script:
# 1) Download DCAT-AP industrial dataset
# 2) Run dataset processor
# 3) Load processed train/test arrays
# 4) Train TabPFN and evaluate

# %% 1) Config
# Change this URL if you want another dataset from your hub.
DATASET_URL = (
    "http://localhost:8081/datasets/b4286742-f993-4ec7-b869-109e98f00a58.jsonld"
)

# Where downloaded files go.
DATA_DIR = "./data"

# Name of the processed folder inside the downloaded dataset directory.
PROCESSED_DIR = "processed"

# Set True if you want to re-download / re-process.
FORCE_DOWNLOAD = False
FORCE_PROCESS = False

# %% 2) Download + process with dcat-ap-hub
ds = Dataset.from_url(DATASET_URL)
ds.download(data_dir=DATA_DIR, force=FORCE_DOWNLOAD)
ds.process(processed_dir=PROCESSED_DIR, force=FORCE_PROCESS)

processed_path = ds.processed_path
if processed_path is None:
    raise RuntimeError("Processed path is unavailable after ds.process().")

print(f"Processed directory: {processed_path}")

# %% 3) Load processed arrays generated by the processor
required_files = ["X_train.npy", "X_test.npy", "y_train.npy", "y_test.npy"]
missing_files = [
    name for name in required_files if not (processed_path / name).exists()
]
if missing_files:
    raise FileNotFoundError(
        f"Missing processed arrays in '{processed_path}': {', '.join(missing_files)}"
    )

X_train = np.load(Path(processed_path) / "X_train.npy")
X_test = np.load(Path(processed_path) / "X_test.npy")
y_train = np.ravel(np.load(Path(processed_path) / "y_train.npy"))
y_test = np.ravel(np.load(Path(processed_path) / "y_test.npy"))

print(f"Train shape: X={X_train.shape}, y={y_train.shape}")
print(f"Test shape: X={X_test.shape}, y={y_test.shape}")

# %% 4) Encode labels (TabPFN classifier expects class ids)
label_encoder = LabelEncoder()
y_train_encoded = label_encoder.fit_transform(y_train)
y_test_encoded = label_encoder.transform(y_test)

print("Classes:", list(label_encoder.classes_))

# %% 5) Fit TabPFN and run inference
# Default TabPFNClassifier() uses the current default model weights.
clf = TabPFNClassifier()

clf.fit(X_train, y_train_encoded)
prediction_probabilities = clf.predict_proba(X_test)
predictions = clf.predict(X_test)

# %% 6) Metrics
accuracy = accuracy_score(y_test_encoded, predictions)
print("Accuracy:", accuracy)

n_classes = len(label_encoder.classes_)
if n_classes == 2:
    auc = roc_auc_score(y_test_encoded, prediction_probabilities[:, 1])
    print("ROC AUC:", auc)
elif n_classes > 2:
    auc = roc_auc_score(
        y_test_encoded,
        prediction_probabilities,
        multi_class="ovr",
        average="weighted",
    )
    print("ROC AUC (ovr, weighted):", auc)
