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, "output.csv")
MODEL_PATH = os.path.join(BASE_DIR, "model.joblib")


def segment_telemetry(df: pd.DataFrame, dt_threshold: float = 0.5) -> list[pd.DataFrame]:
    timestamps = df["timestamp"].values
    dt = np.diff(timestamps)
    jump_indices = np.where(np.abs(dt) > dt_threshold)[0]
    boundaries = [0] + [int(idx + 1) for idx in jump_indices] + [len(df)]

    segments = []
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        if end - start > 50:
            segments.append(df.iloc[start:end].reset_index(drop=True))
    return segments


def extract_kinematic_features(seg: pd.DataFrame) -> np.ndarray:
    feats = []
    duration = float(seg["timestamp"].iloc[-1] - seg["timestamp"].iloc[0])
    feats.append(duration)

    for j in range(1, 7):
        p = seg[f"pos_joint{j}"].values
        v = seg[f"vel_joint{j}"].values
        e = seg[f"eff_joint{j}"].values
        feats.extend([
            np.mean(p), np.std(p), np.min(p), np.max(p), np.max(p) - np.min(p),
            np.std(v), np.max(np.abs(v)),
            np.mean(e), np.std(e), np.min(e), np.max(e),
        ])

    d = seg["pos_drive_joint"].values
    feats.extend([np.min(d), np.max(d), np.max(d) - np.min(d), np.std(d)])
    return np.array(feats, dtype=np.float64)


def is_complete_task1_cycle(seg: pd.DataFrame) -> bool:
    d = seg["pos_drive_joint"].values
    active_idx = np.where(d > 0.15)[0]
    if len(active_idx) < 200:
        return False
    d_max = float(np.max(d))
    post_deact_samples = len(seg) - int(active_idx[-1])
    return (0.15 < d_max < 0.35) and (post_deact_samples > 800)


def extract_task1_phase_features(seg: pd.DataFrame) -> np.ndarray:
    d = seg["pos_drive_joint"].values
    active_idx = np.where(d > 0.15)[0]
    i_act = int(active_idx[0])
    i_deact = int(active_idx[-1])
    i_apex = i_act + int(np.argmin(seg["pos_joint2"].values[i_act:i_deact]))

    phase_ascent = seg.iloc[i_act:i_apex]
    forward_sub = seg.iloc[i_apex:i_deact]
    moving_forward = np.where(forward_sub["pos_joint2"].values > 0.01)[0]
    i_start_A = i_apex + int(moving_forward[0])
    subtraj_A = seg.iloc[i_start_A : i_start_A + 506].reset_index(drop=True)

    return_sub = seg.iloc[i_deact:]
    moving_return = np.where(return_sub["pos_joint2"].values < 0.86)[0]
    i_start_B = i_deact + int(moving_return[0])
    subtraj_B = seg.iloc[i_start_B : i_start_B + 506].reset_index(drop=True)

    a_e1_ext = float(np.mean(subtraj_A["eff_joint1"].values[250:506]))
    a_e1_late = float(np.mean(subtraj_A["eff_joint1"].values[400:500]))
    b_e1_ext = float(np.mean(subtraj_B["eff_joint1"].values[250:506]))
    b_minus_a_e1 = b_e1_ext - a_e1_ext
    ascent_de2_max = float(np.max(np.abs(np.diff(phase_ascent["eff_joint2"].values))))
    b_v5_mean = float(np.mean(subtraj_B["vel_joint5"].values))

    return np.array([
        a_e1_ext,
        a_e1_late,
        b_e1_ext,
        b_minus_a_e1,
        ascent_de2_max,
        b_v5_mean,
    ], dtype=np.float64)


def predict_sequence(input_csv_path: str, model_path: str = MODEL_PATH) -> list[int]:
    bundle = joblib.load(model_path)
    dt_threshold = bundle["dt_threshold"]
    clf_task = bundle["clf_task"]
    scaler_dyn = bundle["scaler_dyn"]
    lof_model = bundle["lof_model"]
    ocsvm_model = bundle["ocsvm_model"]
    rms_threshold = bundle["rms_threshold"]

    df = pd.read_csv(input_csv_path)
    segments = segment_telemetry(df, dt_threshold=dt_threshold)

    predictions = []
    for seg in segments:
        x_kin = extract_kinematic_features(seg).reshape(1, -1)
        pred_task = int(clf_task.predict(x_kin)[0])

        if pred_task == 1 and is_complete_task1_cycle(seg):
            x_dyn = extract_task1_phase_features(seg).reshape(1, -1)
            z_dyn = scaler_dyn.transform(x_dyn)

            rms_z = float(np.sqrt(np.mean(z_dyn ** 2)))
            is_rms_anom = rms_z > rms_threshold
            is_lof_anom = int(lof_model.predict(z_dyn)[0]) == -1
            is_svm_anom = float(ocsvm_model.decision_function(z_dyn)[0]) < -0.05

            if (int(is_rms_anom) + int(is_lof_anom) + int(is_svm_anom)) >= 2:
                pred_task = 0

        predictions.append(pred_task)

    return predictions


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 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_sequence(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()