import logging

import pandas as pd

from .. import constants

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def main():
    # Entrees : dataset source deja prepare (hors git)
    feature_file_path = constants.SOURCE_DIR / constants.FEATURE_FILENAME
    target_file_path = constants.SOURCE_DIR / constants.TARGET_FILENAME

    logger.info(f"Read dataset from {feature_file_path} and {target_file_path}")
    df_features = pd.read_parquet(feature_file_path)
    df_target = pd.read_parquet(target_file_path)

    df = df_features.join(df_target)
    timestamps = df.index.get_level_values("timestamp")

    # Sorties : splits versionnes par DVC dans le depot
    constants.DATASET_DIR.mkdir(parents=True, exist_ok=True)
    logger.info(f"Split strategy: {constants.CHOSEN_SPLIT_STRATEGY.value}")
    for dataset_name, (start_date, end_date) in constants.DATASET_SPLIT_DATES[constants.CHOSEN_SPLIT_STRATEGY].items():
        file_path = constants.DATASET_DIR / f"{dataset_name}.parquet"
        mask = (timestamps.date >= start_date) & (timestamps.date <= end_date)
        df_split = df[mask]
        logger.info(f"Split dataset into {dataset_name} ({start_date} -> {end_date}) with shape: {df_split.shape}")
        df_split.to_parquet(file_path)


if __name__ == "__main__":
    main()