import os
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import LocalOutlierFactor
from sklearn.preprocessing import StandardScaler
from sklearn.svm import OneClassSVM

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, "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 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(df_t1)
    segs_t2 = segment_telemetry(df_t2)

    X_kin = np.array([extract_kinematic_features(s) for s in segs_t1 + segs_t2])
    y_kin = np.array([1] * len(segs_t1) + [2] * len(segs_t2))

    clf_task = RandomForestClassifier(n_estimators=200, random_state=42)
    clf_task.fit(X_kin, y_kin)

    complete_t1_segs = [s for s in segs_t1 if is_complete_task1_cycle(s)]

    X_dyn = np.array([extract_task1_phase_features(s) for s in complete_t1_segs])
    scaler_dyn = StandardScaler().fit(X_dyn)
    Z_dyn = scaler_dyn.transform(X_dyn)

    lof_model = LocalOutlierFactor(n_neighbors=5, novelty=True).fit(Z_dyn)
    ocsvm_model = OneClassSVM(kernel="rbf", nu=0.05, gamma="scale").fit(Z_dyn)

    loo_rms_scores = []
    for i in range(len(X_dyn)):
        X_sub = np.delete(X_dyn, i, axis=0)
        sc_sub = StandardScaler().fit(X_sub)
        z_i = sc_sub.transform(X_dyn[i : i + 1])
        loo_rms_scores.append(float(np.sqrt(np.mean(z_i ** 2))))
    loo_rms_scores = np.array(loo_rms_scores)

    rms_threshold = float(np.max(loo_rms_scores) * 1.02)

    model_bundle = {
        "dt_threshold": 0.5,
        "clf_task": clf_task,
        "scaler_dyn": scaler_dyn,
        "lof_model": lof_model,
        "ocsvm_model": ocsvm_model,
        "rms_threshold": rms_threshold,
        "feature_names_dyn": [
            "A_e1_ext",
            "A_e1_late",
            "B_e1_ext",
            "B_minus_A_e1",
            "ascent_de2_max",
            "B_v5_mean",
        ],
    }

    joblib.dump(model_bundle, MODEL_OUTPUT_PATH)


if __name__ == "__main__":
    main()