import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import fclusterdata


def is_above_quantile(series: pd.Series, q: float) -> np.array:
    x = series.to_numpy()
    threshold = np.quantile(x, q=q)
    return series > threshold


def add_significant(x: pd.DataFrame, q: float = 0.25) -> pd.DataFrame:
    x["significant"] = is_above_quantile(x["photons"], q=q)
    return x


def substract_frameoffset(x: pd.DataFrame, firstframe: int) -> pd.DataFrame:
    x["frame"] -= firstframe
    return x


def flip_z(x: pd.DataFrame) -> pd.DataFrame:
    x["z"] *= -1
    return x

def flip_x(x: pd.DataFrame, mirrorx: float) -> pd.DataFrame:
    x["x"] = mirrorx - x["x"]
    return x


# config/dataset/EPFL_MT0N1HDAS.yaml
def bin(x: pd.DataFrame, binsize: int, threshold: float = 1.0) -> pd.DataFrame:
    x["frame"] = x["frame"] // binsize

    xmerged = []
    for i in x["frame"].unique():
        xi = x[x["frame"] == i].copy()
        coords = xi[["x", "y", "z"]].values
        clusters = fclusterdata(coords, t=threshold, criterion="distance")
        xi["cluster"] = clusters
        xi = xi.groupby("cluster")
        xi = xi.agg(
            {"frame": "max", "x": "max", "y": "max", "z": "max", "photons": "mean"}
        )
        xmerged.append(xi.reset_index(drop=True))
    xmerged = pd.concat(xmerged, ignore_index=True)
    return xmerged
