from pathlib import Path

import matplotlib.pyplot as plt
import sklearn

from geocif import geocif_runner as gc
from geocif import logger as log
from geocif import utils as ut

plt.style.use("default")
sklearn.set_config(transform_output="pandas")

import warnings

warnings.simplefilter(action="ignore", category=FutureWarning)


def main(inputs, logger, parser, section, item, type, values):
    # Set experiment_name
    experiment_name = f"{section}_{item}"
    parser.set("DEFAULT", "experiment_name", experiment_name)

    if type == "str":
        original_value = parser.get(section, item)
    elif type == "bool":
        original_value = parser.getboolean(section, item)
    elif type == "int":
        original_value = parser.getint(section, item)
    elif type == "float":
        original_value = parser.getfloat(section, item)

    for value in values:
        if type == "str":
            parser.set(section, item, value)
        elif type == "bool":
            parser.set(section, item, str(value))
        elif type == "int":
            parser.set(section, item, str(value))
        elif type == "float":
            parser.set(section, item, str(value))

        gc.execute_models(inputs, logger, parser)

    parser.set(section, item, str(original_value))

    return parser


def run(path_config_files=[Path("../config/geocif.txt")]):
    logger, parser = log.setup_logger_parser(path_config_files)
    inputs = gc.gather_inputs(parser)

    # Build list of experiments for summary
    experiments = [
        ("Models", "DEFAULT", "model", "str", ["catboost", "merf", "linear"]),
        ("feature_selection", "ML", "feature_selection", "str",
         ["SelectKBest", "BorutaPy", "Leshy", "RFECV", "RFE", "none"]),
        ("lag_years", "ML", "lag_years", "int", [1, 2, 3, 4, 5]),
        ("lag_yield_as_feature", "ML", "lag_yield_as_feature", "bool", [True, False]),
        ("median_years", "ML", "median_years", "int", [2, 3, 4, 5]),
        ("median_yield_as_feature", "ML", "median_yield_as_feature", "bool", [True, False]),
        ("analogous_year_yield_as_feature", "ML", "analogous_year_yield_as_feature", "bool", [True, False]),
    ]

    params = gc._build_summary_params(parser, inputs)
    for i, (name, _sec, _item, _, values) in enumerate(experiments):
        params.append((f"  Exp {i}: {name}", ", ".join(str(v) for v in values)))
    params.append(("Total experiments", str(len(experiments))))
    ut.display_run_summary("GeoCIF Experiments Runner", params, wait=20)

    # Experiment: Models
    logger.info("Experiment 0: Models")
    parser = main(
        inputs,
        logger,
        parser,
        "DEFAULT",
        "model",
        "str",
        ["catboost", "merf", "linear"],
    )

    # Experiment: feature_selection
    logger.info("Experiment 1: feature_selection")
    parser = main(
        inputs,
        logger,
        parser,
        "ML",
        "feature_selection",
        "str",
        ["SelectKBest", "BorutaPy", "Leshy", "RFECV", "RFE", "none"],
    )

    # Experiment: lag_years
    logger.info("Experiment 2: lag_years")
    parser = main(inputs, logger, parser, "ML", "lag_years", "int", [1, 2, 3, 4, 5])

    # Experiment: lag_yield_as_feature
    logger.info("Experiment 3: lag_yield_as_feature")
    parser = main(
        inputs,
        logger,
        parser,
        "ML",
        "lag_yield_as_feature",
        "bool",
        [True, False],
    )

    # Experiment: median_years
    logger.info("Experiment 4: median_years")
    parser = main(inputs, logger, parser, "ML", "median_years", "int", [2, 3, 4, 5])

    # Experiment: median_yield_as_feature
    logger.info("Experiment 5: median_yield_as_feature")
    parser = main(
        inputs,
        logger,
        parser,
        "ML",
        "median_yield_as_feature",
        "bool",
        [True, False],
    )

    # Experiment: analogous_year_yield_as_feature
    logger.info("Experiment 6: analogous_year_yield_as_feature")
    parser = main(
        inputs,
        logger,
        parser,
        "ML",
        "analogous_year_yield_as_feature",
        "bool",
        [True, False],
    )

    # Experiment: optimize
    # logger.info("Experiment 7: optimize")
    # parser = main(inputs, logger, parser, "DEFAULT", "optimize", "bool", [True, False])


if __name__ == "__main__":
    run()
