import os
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(os.path.dirname(BASE_DIR), "corrected_csvs")
if not os.path.exists(DATA_DIR):
    DATA_DIR = os.path.dirname(BASE_DIR)

MODEL_OUTPUT_PATH = os.path.join(BASE_DIR, "baseline_model.joblib")

WINDOW_SIZE = 3000

SENSOR_COLS = (
    [f"pos_joint{i}" for i in range(1, 7)]
    + [f"vel_joint{i}" for i in range(1, 7)]
    + [f"eff_joint{i}" for i in range(1, 7)]
)


def segment_telemetry_naive(df: pd.DataFrame, window_size: int = WINDOW_SIZE) -> list[pd.DataFrame]:
    segments = []
    for start in range(0, len(df), window_size):
        chunk = df.iloc[start : start + window_size]
        if len(chunk) >= window_size // 2:
            segments.append(chunk.reset_index(drop=True))
    return segments


def extract_baseline_features(seg: pd.DataFrame) -> np.ndarray:
    means = seg[SENSOR_COLS].mean().values
    stds = seg[SENSOR_COLS].std().values
    return np.concatenate([means, stds])


def main():
    t1_path = os.path.join(DATA_DIR, "task1_training.csv.gz")
    t2_path = os.path.join(DATA_DIR, "task2_training.csv.gz")

    df_t1 = pd.read_csv(t1_path)
    df_t2 = pd.read_csv(t2_path)

    segs_t1 = segment_telemetry_naive(df_t1, WINDOW_SIZE)
    segs_t2 = segment_telemetry_naive(df_t2, WINDOW_SIZE)

    X_t1 = np.array([extract_baseline_features(s) for s in segs_t1])
    X_t2 = np.array([extract_baseline_features(s) for s in segs_t2])

    X_train = np.vstack([X_t1, X_t2])
    y_train = np.array([1] * len(X_t1) + [2] * len(X_t2))

    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)

    clf = RandomForestClassifier(n_estimators=50, random_state=42)
    clf.fit(X_train_scaled, y_train)

    model_bundle = {
        "window_size": WINDOW_SIZE,
        "scaler": scaler,
        "classifier": clf,
    }

    joblib.dump(model_bundle, MODEL_OUTPUT_PATH)


if __name__ == "__main__":
    main()