import json
from joblib import dump
from sklearn.datasets import load_breast_cancer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import recall_score
from src.utils import MODEL_PATH, META_PATH


def load_data():
    data = load_breast_cancer()
    return data.data, data.target, list(data.feature_names), list(data.target_names)


def train(smoke=False):
    X, y, feature_names, target_names = load_data()
    Xtr, Xte, ytr, yte = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    pipe = Pipeline([
        ("scaler", MinMaxScaler()),
        ("rf", RandomForestClassifier(random_state=42))
    ])

    if smoke:
        params = {"rf__n_estimators": [50]}
        cv = 2
    else:
        params = {"rf__n_estimators": [100, 200], "rf__max_depth": [None, 8, 12]}
        cv = 5

    grid = GridSearchCV(pipe, params, scoring="recall", cv=cv, n_jobs=-1)
    grid.fit(Xtr, ytr)

    best = grid.best_estimator_
    yhat = best.predict(Xte)
    rec = recall_score(yte, yhat)

    dump(best, MODEL_PATH)
    META_PATH.write_text(json.dumps({
        "best_params": grid.best_params_,
        "cv": cv,
        "recall_test": rec,
        "feature_names": feature_names,
        "target_names": target_names
    }, indent=2))

    return rec


if __name__ == "__main__":
    r = train(smoke=False)
    print(f"Recall (test): {r:.4f}")