import numpy as np
import pandas as pd
from pathlib import Path
from sklearn.preprocessing import MinMaxScaler
from scipy.signal import welch
from pygazeanalyser.detectors import fixation_detection, saccade_detection


##############################################################################
# KONFIGURATION
##############################################################################

SAMPLING_RATE = 25          # Hz
MIN_DUR_BLINKS = 2          # x * 40ms


##############################################################################
# EYE-TRACKING FUNKTIONEN
##############################################################################

def clean_eye_df(df):
    """Extrahiert nur Eye-Tracking Spalten und entfernt leere Zeilen."""
    eye_cols = [c for c in df.columns if c.startswith("EYE_")]

    if not eye_cols:
        return pd.DataFrame()
    
    df_eye = df[eye_cols].copy()
    df_eye = df_eye.replace([np.inf, -np.inf], np.nan)
    df_eye = df_eye.dropna(subset=eye_cols, how="all")
    
    return df_eye.reset_index(drop=True)


def extract_gaze_signal(df):
    """Extrahiert 2D-Gaze-Positionen, maskiert ungültige Samples und interpoliert."""
    gx_L = df["EYE_LEFT_GAZE_POINT_ON_DISPLAY_AREA_X"].astype(float).copy()
    gy_L = df["EYE_LEFT_GAZE_POINT_ON_DISPLAY_AREA_Y"].astype(float).copy()
    gx_R = df["EYE_RIGHT_GAZE_POINT_ON_DISPLAY_AREA_X"].astype(float).copy()
    gy_R = df["EYE_RIGHT_GAZE_POINT_ON_DISPLAY_AREA_Y"].astype(float).copy()

    val_L = (df["EYE_LEFT_PUPIL_VALIDITY"] == 1)
    val_R = (df["EYE_RIGHT_PUPIL_VALIDITY"] == 1)

    # Inf → NaN
    for arr in [gx_L, gy_L, gx_R, gy_R]:
        arr.replace([np.inf, -np.inf], np.nan, inplace=True)

    # Ungültige maskieren
    gx_L[~val_L] = np.nan
    gy_L[~val_L] = np.nan
    gx_R[~val_R] = np.nan
    gy_R[~val_R] = np.nan

    
    # Mittelwert beider Augen
    gx = np.mean(np.column_stack([gx_L, gx_R]), axis=1)
    gy = np.mean(np.column_stack([gy_L, gy_R]), axis=1)

    # Interpolation
    gx = pd.Series(gx).interpolate(limit=None, limit_direction="both").bfill().ffill()
    gy = pd.Series(gy).interpolate(limit=None, limit_direction="both").bfill().ffill()
    
    # MinMax Skalierung
    xscaler = MinMaxScaler()
    gxscale = xscaler.fit_transform(gx.values.reshape(-1, 1))

    yscaler = MinMaxScaler()
    gyscale = yscaler.fit_transform(gy.values.reshape(-1, 1))
    
    return np.column_stack((gxscale, gyscale))


def extract_pupil(df):
    """Extrahiert Pupillengröße (beide Augen gemittelt)."""
    pl = df["EYE_LEFT_PUPIL_DIAMETER"].replace([np.inf, -np.inf], np.nan)
    pr = df["EYE_RIGHT_PUPIL_DIAMETER"].replace([np.inf, -np.inf], np.nan)

    vl = df.get("EYE_LEFT_PUPIL_VALIDITY")
    vr = df.get("EYE_RIGHT_PUPIL_VALIDITY")

    if vl is None or vr is None:
        validity = (~pl.isna() | ~pr.isna()).astype(int).to_numpy()
    else:
        validity = ((vl == 1) | (vr == 1)).astype(int).to_numpy()

    p = np.mean(np.column_stack([pl, pr]), axis=1)
    p = pd.Series(p).interpolate(limit=50, limit_direction="both").bfill().ffill()
    
    return p.to_numpy(), validity


def detect_blinks(pupil_validity, min_duration=5):
    """Erkennt Blinks: Validity=0 → Blink."""
    blinks = []
    start = None

    for i, v in enumerate(pupil_validity):
        if v == 0 and start is None:
            start = i
        elif v == 1 and start is not None:
            if i - start >= min_duration:
                blinks.append([start, i])
            start = None

    return blinks


def compute_IPA(pupil, fs=25):
    """Index of Pupillary Activity (Duchowski 2018)."""
    f, Pxx = welch(pupil, fs=fs, nperseg=int(fs*2))
    hf_band = (f >= 0.6) & (f <= 2.0)
    return np.sum(Pxx[hf_band])


def extract_eye_features(df_eye, fs=25, min_dur_blinks=2):
    """
    Extrahiert Eye-Tracking Features für ein einzelnes Window.
    Gibt Dictionary mit allen Eye-Features zurück.
    """
    # Gaze
    gaze = extract_gaze_signal(df_eye)
    
    # Pupille
    pupil, pupil_validity = extract_pupil(df_eye)

    
    # ----------------------------
    # FIXATIONS
    # ----------------------------
    time_ms = np.arange(len(df_eye)) * 1000.0 / fs

    fix, efix = fixation_detection(
        x=gaze[:, 0], y=gaze[:, 1], time=time_ms,
        missing=0.0, maxdist=0.003, mindur=10
    )

    fixation_durations = [f[2] for f in efix if np.isfinite(f[2]) and f[2] > 0]

    # Kategorien
    F_short = sum(66 <= d <= 150 for d in fixation_durations)
    F_medium = sum(300 <= d <= 500 for d in fixation_durations)
    F_long = sum(d >= 1000 for d in fixation_durations)
    F_hundred = sum(d > 100 for d in fixation_durations)

    # ----------------------------
    # SACCADES
    # ----------------------------
    sac, esac = saccade_detection(
        x=gaze[:, 0], y=gaze[:, 1], time=time_ms, 
        missing=0, minlen=12, maxvel=0.2, maxacc=1
    )

    sac_durations = [s[2] for s in esac]
    sac_amplitudes = [((s[5]-s[3])**2 + (s[6]-s[4])**2)**0.5 for s in esac]

    # ----------------------------
    # BLINKS
    # ----------------------------
    blinks = detect_blinks(pupil_validity, min_duration=min_dur_blinks)
    blink_durations = [(b[1] - b[0]) / fs for b in blinks]

    # ----------------------------
    # PUPIL
    # ----------------------------
    if np.all(np.isnan(pupil)):
        mean_pupil = np.nan
        ipa = np.nan
    else:
        mean_pupil = np.nanmean(pupil)
        ipa = compute_IPA(pupil, fs=fs)

    # Feature Dictionary
    return {
        "Fix_count_short_66_150": F_short,
        "Fix_count_medium_300_500": F_medium,
        "Fix_count_long_gt_1000": F_long,
        "Fix_count_100": F_hundred,
        "Fix_mean_duration": np.mean(fixation_durations) if fixation_durations else 0,
        "Fix_median_duration": np.median(fixation_durations) if fixation_durations else 0,
        "Sac_count": len(sac),
        "Sac_mean_amp": np.mean(sac_amplitudes) if sac_amplitudes else 0,
        "Sac_mean_dur": np.mean(sac_durations) if sac_durations else 0,
        "Sac_median_dur": np.median(sac_durations) if sac_durations else 0,
        "Blink_count": len(blinks),
        "Blink_mean_dur": np.mean(blink_durations) if blink_durations else 0,
        "Blink_median_dur": np.median(blink_durations) if blink_durations else 0,
        "Pupil_mean": mean_pupil,
        "Pupil_IPA": ipa
    }

def compute_features_from_parquet(parquet_path):
    df = pd.read_parquet(parquet_path)
    df_eye = clean_eye_df(df)

    if df_eye.empty:
        return None

    features = extract_eye_features(
        df_eye,
        fs=SAMPLING_RATE,
        min_dur_blinks=MIN_DUR_BLINKS
    )

    return features