import pandas as pd
from training.types import TransformationsOverTime
from utils.helpers import get_first_valid_return_index
from tqdm import tqdm
from transformations.base import Transformation
from typing import Optional
from data_loader.types import ForwardReturnSeries, XDataFrame, ySeries
import ray
from utils.parallel import parallel_compute_with_bar


def walk_forward_process_transformations(
    X: XDataFrame,
    y: ySeries,
    forward_returns: ForwardReturnSeries,
    window_size: int,
    retrain_every: int,
    from_index: Optional[pd.Timestamp],
    transformations: list[Transformation],
) -> TransformationsOverTime:
    transformations_over_time = [
        pd.Series(index=y.index).rename(t.get_name()) for t in transformations
    ]

    first_nonzero_return = max(
        get_first_valid_return_index(forward_returns),
        get_first_valid_return_index(X.iloc[:, 0]),
        get_first_valid_return_index(y),
    )
    train_from = (
        first_nonzero_return + window_size + 1
        if from_index is None
        else X.index.to_list().index(from_index)
    )
    train_till = len(y)

    processed_transformations = parallel_compute_with_bar(
        [
            preprocess_transformations_window.remote(
                X,
                y,
                transformations,
                first_nonzero_return,
                index,
            )
            for index in range(train_from, train_till, retrain_every)
        ]
    )

    for transformation, index_time in processed_transformations:
        for transformation_index, transformation in enumerate(transformation):
            transformations_over_time[transformation_index][
                X.index[index_time]
            ] = transformation

    return transformations_over_time


@ray.remote
def preprocess_transformations_window(
    X: XDataFrame,
    y: ySeries,
    transformations: list[Transformation],
    first_nonzero_return: int,
    index: int,
) -> tuple[list[Transformation], int]:
    train_window_start = X.index[first_nonzero_return]
    train_window_end = X.index[index - 1]

    X_expanding_window = X[train_window_start:train_window_end]
    y_expanding_window = y[train_window_start:train_window_end]

    current_transformations = [t.clone() for t in transformations]
    for transformation in current_transformations:
        X_expanding_window = transformation.fit_transform(
            X_expanding_window, y_expanding_window
        )

    return (current_transformations, index)