__all__ = ["encode_data"]

from pathlib import Path

import numpy as np

from move.conf.schema import DataConfig
from move.core.logging import get_logger
from move.data import io, preprocessing
from move.visualization.dataset_distributions import plot_value_distributions


def encode_data(config: DataConfig):
    """Encodes categorical and continuous datasets specified in configuration.
    Categorical data is one-hot encoded, whereas continuous data is z-score
    normalized.

    Args:
        config: data configuration
    """
    logger = get_logger(__name__)
    logger.info("Beginning task: encode data")

    raw_data_path = Path(config.raw_data_path)
    raw_data_path.mkdir(exist_ok=True)
    interim_data_path = Path(config.interim_data_path)
    interim_data_path.mkdir(exist_ok=True, parents=True)
    output_path = Path(config.results_path) / "encoded_datasets"
    output_path.mkdir(exist_ok=True, parents=True)

    sample_names = io.read_names(raw_data_path / f"{config.sample_names}.txt")

    mappings = {}
    for dataset_name in config.categorical_names:
        # ! the string representation seems to be the variable value (here a string)
        logger.info(f"Encoding '{dataset_name}'")
        filepath = raw_data_path / f"{dataset_name}.tsv"
        names, values = io.read_tsv(filepath, sample_names)
        values, mapping = preprocessing.one_hot_encode(values)
        mappings[dataset_name] = mapping
        io.dump_names(interim_data_path / f"{dataset_name}.txt", names)
        np.save(interim_data_path / f"{dataset_name}.npy", values)
    if mappings:
        io.dump_mappings(interim_data_path / "mappings.json", mappings)

    for input_config in config.continuous_inputs:
        scale = not hasattr(input_config, "scale") or input_config.scale
        action_name = "Encoding" if scale else "Reading"
        dataset_name = input_config.name
        logger.info(f"{action_name} '{dataset_name}'")
        filepath = raw_data_path / f"{dataset_name}.tsv"
        names, values = io.read_tsv(filepath, sample_names)

        # Plotting the value distribution for all continuous datasets
        # before preprocessing:
        fig = plot_value_distributions(values)
        fig_path = str(
            output_path / f"Value_distribution_{dataset_name}_unprocessed.png"
        )
        fig.savefig(fig_path)

        if scale:
            logger.debug(
                f"Scaling dataset: {dataset_name}, log2 transform: {input_config.log2}"
            )
            values, mask_1d = preprocessing.scale(values, input_config.log2)
            names = names[mask_1d]
            logger.debug(f"Columns with zero variance: {np.sum(~mask_1d)}")
            # Plotting the value distribution for all continuous datasets:
            fig = plot_value_distributions(values)
            fig_path = str(output_path / f"Value_distribution_{dataset_name}.png")
            fig.savefig(fig_path)

        io.dump_names(interim_data_path / f"{dataset_name}.txt", names)
        np.save(interim_data_path / f"{dataset_name}.npy", values)