"""Learn classification or regression with real Welt research data.

Run after welt login, or approve the printed browser connection link.
Requires welt-client[parquet]. Each run retains its hosted training dataset/predictor.
"""
import argparse


def run(task="classification"):
    if task == "classification":
        from pathlib import Path
        from tempfile import mkdtemp
        import pandas as pd
        from sklearn.model_selection import train_test_split
        from sklearn.metrics import accuracy_score, classification_report
        from welt import Client, Classifier

        # 1. Get a real table from Welt's research inventory.
        dataset_id = "asset-d839644b6de36fc5cafdda18c5daeca3"
        path = Path(mkdtemp(prefix="welt-iris-")) / "iris.parquet"
        with Client() as client:
            info = client.research_dataset(dataset_id)
            print("Dataset:", info["name"])
            client.download_research_dataset(dataset_id, path)
        data = pd.read_parquet(path)
        target = "target"  # The outcome column in this table.

        # 2. Look at the rows before modelling.
        print(data.head())
        print("Rows and columns:", data.shape)
        data.info()
        print(data[target].value_counts())

        # 3. Set aside test rows before fitting; never train on their answers.
        train, test = train_test_split(
            data, test_size=0.2, random_state=9, stratify=data[target])
        query = test.drop(columns=target)
        y_test = test[target]
        print("Training rows:", len(train), "— Test rows:", len(test))

        # 4. Create a model and train it on the training table.
        model = Classifier(model="tabicl-v2", random_state=9)
        model.fit(train, target=target)

        # 5. Predict the held-out rows, without passing their target column.
        predictions = model.predict(query)
        print(pd.DataFrame({"actual": y_test.to_numpy(), "predicted": predictions}).head())

        # 6. Compare predictions with the answers we kept aside.
        print("Test accuracy:", accuracy_score(y_test, predictions))
        print(classification_report(y_test, predictions, zero_division=0))
        print("Keep this predictor ID:", model.predictor_id_)
    else:
        from pathlib import Path
        from tempfile import mkdtemp
        import pandas as pd
        from sklearn.model_selection import train_test_split
        from sklearn.metrics import mean_absolute_error, root_mean_squared_error, r2_score
        from welt import Client, Regressor

        # 1. Get a real table from Welt's research inventory.
        dataset_id = "asset-bc7dddd884c5e798c1b2abfc559a36ab"
        path = Path(mkdtemp(prefix="welt-yacht-")) / "yacht.parquet"
        with Client() as client:
            info = client.research_dataset(dataset_id)
            print("Dataset:", info["name"])
            client.download_research_dataset(dataset_id, path)
        data = pd.read_parquet(path)
        target = "Residuary.resistance"  # The outcome column in this table.

        # 2. Look at the rows before modelling.
        print(data.head())
        print("Rows and columns:", data.shape)
        data.info()
        print(data[target].describe())

        # 3. Set aside test rows before fitting; never train on their answers.
        train, test = train_test_split(
            data, test_size=0.2, random_state=9)
        query = test.drop(columns=target)
        y_test = test[target]
        print("Training rows:", len(train), "— Test rows:", len(test))

        # 4. Create a model and train it on the training table.
        model = Regressor(model="tabicl-v2", random_state=9)
        model.fit(train, target=target)

        # 5. Predict the held-out rows, without passing their target column.
        predictions = model.predict(query)
        print(pd.DataFrame({"actual": y_test.to_numpy(), "predicted": predictions}).head())

        # 6. Compare predictions with the answers we kept aside.
        print("Test MAE:", mean_absolute_error(y_test, predictions))
        print("Test RMSE:", root_mean_squared_error(y_test, predictions))
        print("Test R²:", r2_score(y_test, predictions))
        print("Keep this predictor ID:", model.predictor_id_)
    return model, predictions


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--task", choices=("classification", "regression"), default="classification")
    args = parser.parse_args()
    from welt import Client
    with Client() as client:
        client.connect()
    run(args.task)


if __name__ == "__main__":
    main()