import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report
 
IDLE_CSV  = "raw_data_label0.csv"
SHAKE_CSV = "raw_data_label1.csv"
 
WINDOW_SIZE = 50   # 0.5s window @ ~100 Hz
STEP_SIZE   = 25   # 50% overlap
 
MODEL_HEADER_OUT = "include/model.h"
 
def load_and_merge():
    df0 = pd.read_csv(IDLE_CSV)
    df1 = pd.read_csv(SHAKE_CSV)
 
    # Just sanity: enforce int label for safety
    df0["label"] = 0
    df1["label"] = 1
 
    df = pd.concat([df0, df1], ignore_index=True)
    # Sort by timestamp is not strictly required, but helps consistent windows
    df = df.sort_values(by="timestamp_ms").reset_index(drop=True)
    return df
 
def extract_features(df):
    """
    df columns: timestamp_ms, ax_g, ay_g, az_g, label
    We'll slide windows, compute features, and assign the majority label.
    """
    feats = []
    ax = df["ax_g"].to_numpy()
    ay = df["ay_g"].to_numpy()
    az = df["az_g"].to_numpy()
    lbl = df["label"].to_numpy()
 
    mag = np.sqrt(ax**2 + ay**2 + az**2)
 
    n = len(df)
    idx = 0
    while idx + WINDOW_SIZE <= n:
        sl = slice(idx, idx+WINDOW_SIZE)
        win_mag = mag[sl]
        win_lbl = lbl[sl]
 
        mean_mag = np.mean(win_mag)
        std_mag  = np.std(win_mag)
        p2p_mag  = np.max(win_mag) - np.min(win_mag)
 
        # majority/avg label in this window
        label_window = int(round(np.mean(win_lbl)))
 
        feats.append({
            "mean_mag": mean_mag,
            "std_mag":  std_mag,
            "p2p_mag":  p2p_mag,
            "label":    label_window
        })
 
        idx += STEP_SIZE
 
    return pd.DataFrame(feats)
 
def main():
    df = load_and_merge()
    feats = extract_features(df)
 
    X = feats[["mean_mag", "std_mag", "p2p_mag"]].to_numpy()
    y = feats["label"].to_numpy()
 
    clf = LogisticRegression()
    clf.fit(X, y)
 
    y_pred = clf.predict(X)
    print(classification_report(y, y_pred))
 
    w = clf.coef_[0]      # [w1, w2, w3]
    b = clf.intercept_[0] # bias
 
    print("Weights:", w)
    print("Bias:", b)
 
    with open(MODEL_HEADER_OUT, "w") as f:
        f.write("// Auto-generated by train_model.py\n")
        f.write("#pragma once\n\n")
        f.write("struct ShakeModel {\n")
        f.write("    static constexpr float W0 = %.8ff;\n" % b)
        f.write("    static constexpr float W1 = %.8ff;\n" % w[0])
        f.write("    static constexpr float W2 = %.8ff;\n" % w[1])
        f.write("    static constexpr float W3 = %.8ff;\n" % w[2])
        f.write("};\n\n")
        f.write("// Feature extraction params\n")
        f.write("#define WINDOW_SIZE %d\n" % WINDOW_SIZE)
        f.write("#define STEP_SIZE   %d\n" % STEP_SIZE)
 
if __name__ == "__main__":
    main()