"""
STEP 2: Scaling
STEP 3: Prepare Time Series for LSTM Autoencoder
"""

import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from typing import List, Tuple


def build_sequences_with_labels(
    df: pd.DataFrame,
    feature_cols: List[str],
    label_col: str = 'DysbiosisLabel',
    seq_len: int = 14,
) -> Tuple[np.ndarray, np.ndarray, List[Tuple]]:
    """
    Slice a dataframe into rolling time windows for each patient,
    assigning each sequence a label based on the window.

    Returns:
        X_sequences: np.ndarray of shape (num_sequences, seq_len, num_features)
        y_labels: np.ndarray of shape (num_sequences,)
        time_indices: list of (PatientID, DayRelativeToNearestHCT) tuples
    """
    X_sequences = []
    y_labels = []
    time_indices = []

    for pid, group in df.groupby('PatientID'):
        group = group.sort_values('DayRelativeToNearestHCT')
        values = group[feature_cols].values
        labels = group[label_col].values
        days = group['DayRelativeToNearestHCT'].values

        for i in range(len(values) - seq_len + 1):
            seq = values[i:i+seq_len]
            label_window = labels[i:i+seq_len]
            # 1 if any point in window is dysbiotic
            label = int(label_window.max())
            time = days[i+seq_len-1]        # label by last day of window
            X_sequences.append(seq)
            y_labels.append(label)
            time_indices.append((pid, time))

    return np.array(X_sequences), np.array(y_labels), time_indices


if __name__ == "__main__":
    # --- [1] Scaling ---
    scaler = MinMaxScaler()
    merged_df[feature_cols] = scaler.fit_transform(merged_df[feature_cols])

    # --- [2] Define features ---
    seq_len = 14
    feature_cols = [
        col for col in merged_df.columns
        if col not in [
            'PatientID', 'SampleID', 'DayRelativeToNearestHCT',
            'DysbiosisLabel', 'MaxTemperature', 'NeutrophilCount'
        ] and not col.startswith('Consistency_')
    ]

    # --- [3] Build sequences ---
    X_seq, y_seq, time_idx = build_sequences_with_labels(
        merged_df, feature_cols, seq_len=seq_len
    )

    # --- [4] Create DataFrame for temporal splitting ---
    split_df = pd.DataFrame(
        time_idx, columns=['PatientID', 'DayRelativeToNearestHCT'])
    split_df['y'] = y_seq
    split_df['X_idx'] = range(len(X_seq))

    # --- [5] Sort by patient and time ---
    split_df = split_df.sort_values(['PatientID', 'DayRelativeToNearestHCT'])

    # --- [6] Determine split indices (70/15/15 by number of sequences) ---
    n_total = len(split_df)
    n_train = int(0.7 * n_total)
    n_val = int(0.15 * n_total)

    train_idx = split_df.iloc[:n_train]['X_idx'].values
    val_idx = split_df.iloc[n_train:n_train+n_val]['X_idx'].values
    test_idx = split_df.iloc[n_train+n_val:]['X_idx'].values

    # --- [7] Final splits ---
    X_train, y_train = X_seq[train_idx], y_seq[train_idx]
    X_val, y_val = X_seq[val_idx], y_seq[val_idx]
    X_test, y_test = X_seq[test_idx], y_seq[test_idx]

    print(
        f"Train: {X_train.shape}, Validation: {X_val.shape}, Test: {X_test.shape}")

    # Filter X_train to only include normal samples (where y_train is 0)
    X_train_normal = X_train[y_train == 0]

    timesteps = X_train_normal.shape[1]
    n_features = X_train_normal.shape[2]

    print(f"Timesteps: {timesteps}, Features: {n_features}")