"""Untitled0.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1W3XSMmzxThy87iCQTJmPuUZ-3-o3J5tL
"""

import pickle
import pandas as pd
import numpy as np
from scipy.signal import butter, filtfilt, find_peaks, welch
from scipy import stats
import antropy as ant

FS_ACC, FS_GYR, FS_MAG = 204.0, 204.0, 197.0

def butter_filter(cutoff, fs, btype, order=4):
    nyq = 0.5 * fs
    b, a = butter(order, cutoff / nyq, btype=btype, analog=False)
    return b, a

def preprocess_signal(arr, fs, lp=20.0, hp=0.1, clip_s=5.0):
    sig = arr.copy().astype(float)
    for i in range(sig.shape[1]):
        mu, sd = np.nanmean(sig[:, i]), np.nanstd(sig[:, i])
        sig[:, i] = np.clip(sig[:, i], mu - clip_s*sd, mu + clip_s*sd)
        nans = np.isnan(sig[:, i])
        if nans.any():
            x = np.arange(len(sig[:, i]))
            sig[nans, i] = np.interp(x[nans], x[~nans], sig[~nans, i])
    b_lp, a_lp = butter_filter(min(lp, fs/2-1), fs, 'low')
    b_hp, a_hp = butter_filter(hp, fs, 'high')
    for i in range(sig.shape[1]):
        if len(sig) > 30:
            sig[:, i] = filtfilt(b_lp, a_lp, sig[:, i])
            sig[:, i] = filtfilt(b_hp, a_hp, sig[:, i])
        sd = np.std(sig[:, i])
        if sd > 1e-8:
            sig[:, i] = (sig[:, i] - np.mean(sig[:, i])) / sd
    return sig

def safe_entropy(sig, func, **kwargs):
    try:
        return float(func(sig, **kwargs))
    except Exception:
        return np.nan

def time_domain_features(sig, fs):
    f = {}
    f['mean']     = np.mean(sig)
    f['std']      = np.std(sig)
    f['rms']      = np.sqrt(np.mean(sig**2))
    f['range']    = np.max(sig) - np.min(sig)
    f['iqr']      = float(np.percentile(sig, 75) - np.percentile(sig, 25))
    f['skewness'] = float(stats.skew(sig))
    f['kurtosis'] = float(stats.kurtosis(sig))
    zc = np.where(np.diff(np.sign(sig)))[0]
    f['zcr'] = len(zc) / (len(sig) / fs) if len(sig) > 0 else 0
    jerk = np.diff(sig) * fs
    f['jerk_rms'] = np.sqrt(np.mean(jerk**2)) if len(jerk) > 0 else 0
    f['jerk_std'] = np.std(jerk) if len(jerk) > 0 else 0
    f['jerk_max'] = np.max(np.abs(jerk)) if len(jerk) > 0 else 0
    thresh = np.percentile(np.abs(sig), 60)
    min_d  = int(0.3 * fs)
    peaks, props = find_peaks(sig, height=thresh, distance=min_d)
    n_steps = len(peaks)
    dur = len(sig) / fs
    f['step_count'] = n_steps
    f['cadence_spm'] = (n_steps / dur) * 60 if dur > 0 else 0
    if n_steps >= 3:
        ivls = np.diff(peaks) / fs
        f['step_ivl_mean']  = np.mean(ivls)
        f['step_ivl_std']   = np.std(ivls)
        f['step_ivl_cv']    = f['step_ivl_std'] / f['step_ivl_mean'] if f['step_ivl_mean'] > 0 else np.nan
        f['step_ivl_range'] = np.max(ivls) - np.min(ivls)
        amps = props['peak_heights']
        f['peak_amp_cv'] = np.std(amps) / np.mean(amps) if np.mean(amps) > 0 else np.nan
    else:
        for k in ['step_ivl_mean','step_ivl_std','step_ivl_cv','step_ivl_range','peak_amp_cv']:
            f[k] = np.nan
    return f

def freq_domain_features(sig, fs):
    f = {}
    if len(sig) < 256:
        f['total_power'] = f['dominant_freq'] = f['gait_band_power'] = f['gait_power_ratio'] = np.nan
        f['tremor_band_power'] = f['tremor_power_ratio'] = np.nan
        f['spectral_entropy'] = f['spectral_centroid'] = np.nan
        return f

    nperseg = min(256, len(sig) // 2)
    freqs, psd = welch(sig, fs=fs, nperseg=nperseg)
    total_pwr = np.trapezoid(psd, freqs)
    f['total_power'] = total_pwr
    gm = (freqs >= 0.5) & (freqs <= 3.0)
    if gm.any():
        f['dominant_freq']    = freqs[gm][np.argmax(psd[gm])]
        f['gait_band_power']  = np.trapezoid(psd[gm], freqs[gm])
        f['gait_power_ratio'] = f['gait_band_power'] / total_pwr if total_pwr > 0 else 0
    else:
        f['dominant_freq'] = f['gait_band_power'] = f['gait_power_ratio'] = np.nan
    tm = (freqs >= 3.0) & (freqs <= 8.0)
    if tm.any():
        f['tremor_band_power']  = np.trapezoid(psd[tm], freqs[tm])
        f['tremor_power_ratio'] = f['tremor_band_power'] / total_pwr if total_pwr > 0 else 0
    else:
        f['tremor_band_power'] = f['tremor_power_ratio'] = np.nan
    psd_n = psd / (psd.sum() + 1e-12)
    f['spectral_entropy']   = -np.sum(psd_n * np.log2(psd_n + 1e-12)) if len(psd_n) > 0 else np.nan
    f['spectral_centroid']  = np.sum(freqs * psd) / (np.sum(psd) + 1e-12) if np.sum(psd) > 0 else np.nan
    return f

def entropy_features(sig):
    s = sig[:2000]
    if len(s) < 5:
        return {'app_entropy': np.nan, 'samp_entropy': np.nan, 'perm_entropy': np.nan, 'svd_entropy': np.nan}
    return {
        'app_entropy':  safe_entropy(s, ant.app_entropy,    order=2, metric='chebyshev'),
        'samp_entropy': safe_entropy(s, ant.sample_entropy, order=2, metric='chebyshev'),
        'perm_entropy': safe_entropy(s, ant.perm_entropy,   order=3, normalize=True),
        'svd_entropy':  safe_entropy(s, ant.svd_entropy,    order=3, delay=1, normalize=True),
    }


def extract_all_features(subj):
    f = {}
    dur_s = subj['duration_ms'] / 1000.0 if subj['duration_ms'] > 0 else np.nan
    f['walking_speed_ms'] = 10.0 / dur_s if dur_s and dur_s > 0 else np.nan
    f['duration_s']       = dur_s
    f['walking_aids']     = subj['walking_aids']
    f['age']              = subj['age'] if subj['age'] else np.nan

    min_len_acc = len(subj['acc']) if 'acc' in subj and subj['acc'] is not None else 0
    min_len_gyr = len(subj['gyr']) if 'gyr' in subj and subj['gyr'] is not None else 0
    min_len_mag = len(subj['mag']) if 'mag' in subj and subj['mag'] is not None else 0

    for sensor_name, sig_mag_data, sig_3d_data, fs_val in [
        ('acc', subj.get('acc_mag'), subj.get('acc'), FS_ACC),
        ('gyr', subj.get('gyr_mag'), subj.get('gyr'), FS_GYR),
        ('mag', subj.get('mag_mag'), subj.get('mag'), FS_MAG),
    ]:
        if sig_mag_data is not None and len(sig_mag_data) > 0:
            td = time_domain_features(sig_mag_data, fs_val)
            f.update({f'{sensor_name}_mag_{k}': v for k, v in td.items()})
            fd = freq_domain_features(sig_mag_data, fs_val)
            f.update({f'{sensor_name}_mag_{k}': v for k, v in fd.items()})
            if sensor_name in ('acc', 'gyr'):
                ent = entropy_features(sig_mag_data)
                f.update({f'{sensor_name}_mag_{k}': v for k, v in ent.items()})

        if sig_3d_data is not None and len(sig_3d_data) > 0:
            for i, ax in enumerate(['x', 'y', 'z']):
                if sig_3d_data.shape[1] > i:
                    td_ax = time_domain_features(sig_3d_data[:, i], fs_val)
                    f.update({f'{sensor_name}_{ax}_{k}': v for k, v in td_ax.items()})

    mn = min(min_len_acc, min_len_gyr)
    if mn > 10 and subj.get('acc_mag') is not None and subj.get('gyr_mag') is not None:
        f['acc_gyr_corr'] = float(np.corrcoef(subj['acc_mag'][:mn], subj['gyr_mag'][:mn])[0, 1])
    else:
        f['acc_gyr_corr'] = np.nan

    acc_e = np.mean(subj['acc_mag']**2) if subj.get('acc_mag') is not None and len(subj['acc_mag']) > 0 else 0
    gyr_e = np.mean(subj['gyr_mag']**2) if subj.get('gyr_mag') is not None and len(subj['gyr_mag']) > 0 else 0
    f['gyr_acc_energy_ratio'] = gyr_e / (acc_e + 1e-12)

    return f


def predict_neurological_risk_from_raw_data(acc_df, gyr_df, mag_df, user_data, model_path='neurological_gait_model.pkl'):
    """
    Predicts the neurological risk from raw sensor data (accelerometer, gyroscope, magnetometer)
    and user demographic/test data using the first LightGBM model.

    Args:
        acc_df (pd.DataFrame): Raw accelerometer data with 'x', 'y', 'z' columns.
        gyr_df (pd.DataFrame): Raw gyroscope data with 'x', 'y', 'z' columns.
        mag_df (pd.DataFrame): Raw magnetometer data with 'x', 'y', 'z' columns.
        user_data (dict): Dictionary containing user demographic and test data, e.g.,
                          {'duration': 25000, 'age': 75, 'walking_aids': 0}.
        model_path (str): Path to the saved neurological_gait_model.pkl file.

    Returns:
        dict: A dictionary containing the predicted risk score and risk level (LOW, MODERATE, HIGH).
              Returns an error dictionary if the model file is not found or data is insufficient.
    """
    try:
        with open(model_path, 'rb') as f:
            bundle = pickle.load(f)
    except FileNotFoundError:
        return {'error': f"Model file not found at {model_path}. Please ensure it's uploaded."}

    lgb_model = bundle['model']
    scaler = bundle['scaler']
    imputer = bundle['imputer']
    feature_names = bundle['feature_names']
    THRESHOLD_LOW = bundle['threshold_low']
    THRESHOLD_HIGH = bundle['threshold_high']

    processed_acc = preprocess_signal(acc_df[['x','y','z']].values, FS_ACC)
    processed_gyr = preprocess_signal(gyr_df[['x','y','z']].values, FS_GYR)
    processed_mag = preprocess_signal(mag_df[['x','y','z']].values, FS_MAG)

    processed_acc_mag = np.sqrt(np.sum(processed_acc**2, axis=1))
    processed_gyr_mag = np.sqrt(np.sum(processed_gyr**2, axis=1))
    processed_mag_mag = np.sqrt(np.sum(processed_mag**2, axis=1))

    min_len = min(len(processed_acc), len(processed_gyr), len(processed_mag))
    if min_len < 30:
        return {'error': 'Insufficient sensor data length for feature extraction.'}

    subject_data_processed = {
        'subject_id': user_data.get('subject_id', -1),
        'user_data': user_data,
        'acc': processed_acc[:min_len],
        'gyr': processed_gyr[:min_len],
        'mag': processed_mag[:min_len],
        'acc_mag': processed_acc_mag[:min_len],
        'gyr_mag': processed_gyr_mag[:min_len],
        'mag_mag': processed_mag_mag[:min_len],
        'label': 0,
        'duration_ms': user_data.get('duration', 0),
        'age': user_data.get('age', None),
        'walking_aids': int(user_data.get('walking_aids', False)),
    }

    extracted_features = extract_all_features(subject_data_processed)
    X_single_df = pd.DataFrame([extracted_features])

    X_single_df_aligned = X_single_df.reindex(columns=imputer.feature_names_in_, fill_value=np.nan)
    X_imp_single = imputer.transform(X_single_df_aligned)
    X_df_c_single = pd.DataFrame(X_imp_single, columns=imputer.feature_names_in_)

    X_filtered_single = X_df_c_single[feature_names]

    X_scaled_single = scaler.transform(X_filtered_single)

    score = np.pow(lgb_model.predict_proba(X_scaled_single)[0, 1],1.5)

    if score < THRESHOLD_LOW:
        level = 'LOW'
    elif score < THRESHOLD_HIGH:
        level = 'MODERATE'
    else:
        level = 'HIGH'

    return {'score': float(score), 'risk_level': level}

def predict_neurological_subtype_from_raw_data(acc_df, gyr_df, mag_df, user_data, model_path='neurological_subtype_model.pkl'):
    """
    Predicts the neurological subtype from raw sensor data for a subject
    already identified as neurological.

    Args:
        acc_df (pd.DataFrame): Raw accelerometer data with 'x', 'y', 'z' columns.
        gyr_df (pd.DataFrame): Raw gyroscope data with 'x', 'y', 'z' columns.
        mag_df (pd.DataFrame): Raw magnetometer data with 'x', 'y', 'z' columns.
        user_data (dict): Dictionary containing user demographic and test data, e.g.,
                          {'duration': 25000, 'age': 75, 'walking_aids': 0}.
        model_path (str): Path to the saved neurological_subtype_model.pkl file.

    Returns:
        dict: A dictionary containing the predicted subtype label, the raw probability scores
              for each class, and the highest probability.
    """

    try:
        with open(model_path, 'rb') as f:
            bundle_subtype = pickle.load(f)
    except FileNotFoundError:
        return {'error': f"Model file not found at {model_path}. Please ensure it's uploaded."}

    lgb_model_second = bundle_subtype['model']
    scaler_second = bundle_subtype['scaler']
    imputer_second = bundle_subtype['imputer']
    feature_names_second = bundle_subtype['feature_names']
    label_encoder_second = bundle_subtype['label_encoder']

    processed_acc = preprocess_signal(acc_df[['x','y','z']].values, FS_ACC)
    processed_gyr = preprocess_signal(gyr_df[['x','y','z']].values, FS_GYR)
    processed_mag = preprocess_signal(mag_df[['x','y','z']].values, FS_MAG)

    processed_acc_mag = np.sqrt(np.sum(processed_acc**2, axis=1))
    processed_gyr_mag = np.sqrt(np.sum(processed_gyr**2, axis=1))
    processed_mag_mag = np.sqrt(np.sum(processed_mag**2, axis=1))

    min_len = min(len(processed_acc), len(processed_gyr), len(processed_mag))
    if min_len < 30:
        return {'error': 'Insufficient sensor data length for feature extraction.'}

    subject_data_processed = {
        'subject_id': user_data.get('subject_id', -1),
        'user_data': user_data,
        'acc': processed_acc[:min_len],
        'gyr': processed_gyr[:min_len],
        'mag': processed_mag[:min_len],
        'acc_mag': processed_acc_mag[:min_len],
        'gyr_mag': processed_gyr_mag[:min_len],
        'mag_mag': processed_mag_mag[:min_len],
        'label': 1,
        'duration_ms': user_data.get('duration', 0),
        'age': user_data.get('age', None),
        'walking_aids': int(user_data.get('walking_aids', False)),
    }

    extracted_features = extract_all_features(subject_data_processed)
    X_single_df = pd.DataFrame([extracted_features])
    X_single_df_aligned = X_single_df.reindex(columns=imputer_second.feature_names_in_, fill_value=np.nan)
    X_imp_single = imputer_second.transform(X_single_df_aligned)
    X_df_c_single = pd.DataFrame(X_imp_single, columns=imputer_second.feature_names_in_)

    X_filtered_single = X_df_c_single[feature_names_second]

    X_scaled_single = scaler_second.transform(X_filtered_single)

    proba_scores = lgb_model_second.predict_proba(X_scaled_single)[0]

    predicted_class_idx = np.argmax(proba_scores)
    predicted_prob = proba_scores[predicted_class_idx]

    predicted_label = label_encoder_second.inverse_transform([predicted_class_idx])[0]

    return {
        'predicted_subtype': predicted_label,
        'probability_scores': dict(zip([c.replace("_","/") for c in label_encoder_second.classes_], proba_scores)),
        'highest_probability': float(predicted_prob)
    }

import os
import pandas as pd
from fastapi import FastAPI, Body
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def read_root():
    return {"message": "Neurological Risk API is Online on Hugging Face"}

@app.post("/predict_risk")
def predict_risk(dict_in: dict = Body(...)):
    acc_df = pd.DataFrame(dict_in['acc']).T
    acc_df.columns = ['x', 'y', 'z']
    gyr_df = pd.DataFrame(dict_in['gyr']).T
    gyr_df.columns = ['x', 'y', 'z']
    mag_df = pd.DataFrame(dict_in['mag']).T
    mag_df.columns = ['x', 'y', 'z']

    user_data = dict_in['user']

    prediction_result = predict_neurological_risk_from_raw_data(
        acc_df, gyr_df, mag_df, user_data
    )

    return {
        'error': 'error' in prediction_result,
        'score': prediction_result.get('score'),
        'risk_level': prediction_result.get('risk_level', 'UNKNOWN')
    }

@app.post("/predict_subtype")
def predict_subtype(dict_in: dict = Body(...)):
    acc_df = pd.DataFrame(dict_in['acc']).T
    acc_df.columns = ['x', 'y', 'z']
    gyr_df = pd.DataFrame(dict_in['gyr']).T
    gyr_df.columns = ['x', 'y', 'z']
    mag_df = pd.DataFrame(dict_in['mag']).T
    mag_df.columns = ['x', 'y', 'z']

    user_data = dict_in['user']

    prediction_result_standalone = predict_neurological_subtype_from_raw_data(
        acc_df, gyr_df, mag_df, user_data
    )

    if 'error' in prediction_result_standalone:
        return {"error": True, "Predicted subtype": None, "Probability scores": [], "Highest probability": None}

    return {
        "error": False,
        "Predicted subtype": str(prediction_result_standalone['predicted_subtype']),
        "Probability scores": {str(k): float(v) for k, v in prediction_result_standalone['probability_scores'].items()},
        "Highest probability": float(prediction_result_standalone['highest_probability'])
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)