"""Fit preprocessing assets and the endpoint-specific TabICLv2 estimators."""

import os
import json
import random
import logging
import argparse

import numpy as np

from sklearn.preprocessing import StandardScaler
from datetime import datetime

from src.model import Tox21TabICL
from src.utils import TASKS, KNOWN_DESCR, create_dir, write_pickle, normalize_config

parser = argparse.ArgumentParser(description="TabICLv2 training script for Tox21")


parser.add_argument(
    "--config",
    type=str,
    default="config/config.json",
)

parser.add_argument(
    "--seed",
    type=int,
    default=None,
)


def main(config):
    print("Preprocess train molecules")

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

    # setup logger
    logger = logging.getLogger(__name__)
    script_name = os.path.splitext(os.path.basename(__file__))[0]
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=[
            logging.FileHandler(
                os.path.join(
                    config["log_folder"],
                    f"{script_name}_{timestamp}.log",
                )
            ),
            logging.StreamHandler(),
        ],
    )

    logger.info(config)

    # seeding
    random.seed(config["seed"])
    np.random.seed(config["seed"])

    train_data = np.load(
        os.path.join(
            config["data_folder"], f"tox21_train_cv{config['cvfold']}.npz"
        )
    )
    train_X = np.concatenate([train_data[descr] for descr in KNOWN_DESCR], axis=1)
    train_y = train_data["labels"]

    val_data = np.load(
        os.path.join(
            config["data_folder"], f"tox21_validation_cv{config['cvfold']}.npz"
        )
    )
    val_X = np.concatenate([val_data[descr] for descr in KNOWN_DESCR], axis=1)
    val_y = val_data["labels"]

    train_X = np.concatenate([train_X, val_X], axis=0)
    train_y = np.concatenate([train_y, val_y], axis=0)

    # remove molecules that couldn't be sanitized
    mask = ~np.isnan(train_X).any(axis=1)
    train_X = train_X[mask]
    train_y = train_y[mask]

    # select 500 features with most variance
    variances_ = np.nanvar(train_X, axis=0)
    order = np.argsort(variances_)[
        : -(config["feature_selection"]["max_features"] + 1) : -1
    ]
    logger.info(
        f"Save used feature order as: {os.path.join(config['data_folder'], 'feature_order.npy')}"
    )
    np.save(os.path.join(config["data_folder"], "feature_order.npy"), order)

    train_X = train_X[:, order]

    max_samples = (
        config["max_samples"] if config["max_samples"] > 0 else train_X.shape[0]
    )
    resample_idxs = np.random.choice(
        np.arange(train_X.shape[0]), size=(max_samples,), replace=True
    )
    train_X = train_X[resample_idxs]
    train_y = train_y[resample_idxs]

    # create and save scaler
    scaler = StandardScaler()
    train_X = scaler.fit_transform(train_X)
    logger.info(
        f"Save fitted scaler as: {os.path.join(config['data_folder'], 'scaler.pkl')}"
    )
    write_pickle(os.path.join(config["data_folder"], "scaler.pkl"), scaler)
    if config["ckpt_path"]:
        logger.info(
            f"Fitted TabICLClassifier state will be saved in: {config['ckpt_path']}"
        )
    else:
        logger.info("Fitted TabICLClassifier state will NOT be saved.")

    model = Tox21TabICL(
        seed=config["seed"],
        device=config["device"],
        model_config=config["model_config"],
    )

    logger.info("Start training.")
    for i, task in enumerate(TASKS):
        logger.info(f"Fitting task: {task}")

        # select task-specific data
        task_train_y = train_y[:, i]
        task_mask = ~np.isnan(task_train_y)
        task_train_y = task_train_y[task_mask].astype(int)
        task_train_X = train_X[task_mask]
        print(f"Fit task {task} using {sum(task_mask)} samples")

        model.fit(task, task_train_X, task_train_y)

    log_text = f"Finished training."

    logger.info(log_text)

    if config["ckpt_path"]:
        model.save_model(config["ckpt_path"])
        logger.info(f"Save model under: {config['ckpt_path']}")


if __name__ == "__main__":
    args = parser.parse_args()

    with open(args.config, "r") as f:
        config = json.load(f)
    config = normalize_config(config)

    create_dir(config["log_folder"])

    if args.seed is not None:
        config["seed"] = args.seed

    main(config)