import os
import sys
import joblib
import numpy as np
import pandas as pd

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
INPUT_PATH = os.path.join(os.path.dirname(BASE_DIR), "corrected_csvs", "final_test.csv.gz")
OUTPUT_PATH = os.path.join(BASE_DIR, "baseline_output.csv")
MODEL_PATH = os.path.join(BASE_DIR, "baseline_model.joblib")

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 = 3000) -> 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 levenshtein_distance(seq1: list[int], seq2: list[int]) -> int:
    n, m = len(seq1), len(seq2)
    dp = np.zeros((n + 1, m + 1), dtype=int)
    for i in range(n + 1):
        dp[i, 0] = i
    for j in range(m + 1):
        dp[0, j] = j
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            cost = 0 if seq1[i - 1] == seq2[j - 1] else 1
            dp[i, j] = min(
                dp[i - 1, j] + 1,
                dp[i, j - 1] + 1,
                dp[i - 1, j - 1] + cost
            )
    return int(dp[n, m])


def predict_baseline(input_csv_path: str, model_path: str = MODEL_PATH) -> list[int]:
    bundle = joblib.load(model_path)
    window_size = bundle["window_size"]
    scaler = bundle["scaler"]
    clf = bundle["classifier"]

    df = pd.read_csv(input_csv_path)
    segments = segment_telemetry_naive(df, window_size=window_size)

    X_test = np.array([extract_baseline_features(s) for s in segments])
    X_test_scaled = scaler.transform(X_test)

    predictions = [int(p) for p in clf.predict(X_test_scaled)]
    return predictions


def main():
    input_path = sys.argv[1] if len(sys.argv) > 1 else INPUT_PATH
    output_path = sys.argv[2] if len(sys.argv) > 2 else OUTPUT_PATH

    predictions = predict_baseline(input_path, MODEL_PATH)

    with open(output_path, "w", encoding="utf-8") as f:
        for p in predictions:
            f.write(f"{p}\n")


if __name__ == "__main__":
    main()