package net.heartshield.signal;

import android.util.Log;

import java.util.Arrays;
import java.util.List;

/**
 * @author David Madl (git@abanbytes.eu)
 * @date 2017-03-12
 */
public class Signal {
    private static final String TAG = "MeasureActivity"; // TODO

    /**
     * Calculate the local maxima of a (heartbeat) signal vector.
     * @return one-hot 'boolean' vector of same length as `dv`.
     */
    public static int[] heartbeat_localmax(double[] dv) {
        if(dv.length < 3) // if cannot even diff twice, this input is too short
            return IVec.zeros(dv.length);

        // note: I ported this straight from Python, to get a feel for how numpy functions
        // would look like in Java. There is no good reason not to just write a single loop.
        //
        // (the lack of operators in Java makes formula notation very painful;
        // Scala may be preferred)

        double[] minf = new double[]{Double.NEGATIVE_INFINITY};
        double[] diffvec = DVec.stack(minf, dv, minf);

        double[] diffScore = DVec.diff(DVec.sign(DVec.diff(diffvec)));

        //System.out.println(DVec.toString(diffScore));

        int[] e = new int[]{dv.length-1};
        double[] a = DVec.get(dv, 0, dv.length-1);
        double[] b = DVec.get(dv, 1, dv.length);
        int[] ep = IVec.where(DVec.eq(a, b));
        int[] runEndingPositions = IVec.stack(ep, e);

        int[] m1 = new int[]{-1};
        int[] runLengths = IVec.diff(IVec.stack(m1, runEndingPositions));
        int[] runStarts = IVec.add(IVec.sub(runEndingPositions, runLengths), 1);

        int[] greaterRuns = IVec.where(IVec.gt(runLengths, 1));
        int[] realRunStarts = IVec.index(runStarts, greaterRuns);
        int[] realRunStops = IVec.index(runEndingPositions, greaterRuns);
        int[] realRunLengths = IVec.index(runLengths, greaterRuns);

        int[] sta = DVec.eq(DVec.index(diffScore, realRunStarts), -1.0);
        int[] sto = DVec.eq(DVec.index(diffScore, realRunStops), -1.0);
        int[] maxRuns = IVec.mul(sta, sto);

        double[] maxRunMiddles = DVec.sub(DVec.round(DVec.add(IVec.toDouble(IVec.index(realRunStarts, maxRuns)), DVec.mul(IVec.toDouble(IVec.index(realRunLengths, maxRuns)), 0.5))), 1.0);
        int[] maxRunMiddlesIntDir = IVec.dirac(diffScore.length, DVec.toInteger(maxRunMiddles));

        int[] maxima = DVec.eq(diffScore, -2.0);
        int[] maximaPlus = IVec.add(maxima, maxRunMiddlesIntDir);
        int[] imaxima = IVec.ge(maxima, 1);

        return imaxima;
    }

    /** moving average filter of length N */
    public static double[] running_mean(double[] x, int N) {
        if(x.length < N)
            throw new IllegalArgumentException("sequence too short to filter");

        double[] p = new double[]{0.0};
        double[] cs = DVec.cumsum(DVec.stack(p, x));
        double[] a = DVec.get(cs, N, cs.length);
        double[] b = DVec.get(cs, 0, cs.length - N);
        double[] y = DVec.mul(DVec.sub(a, b), 1.0 / (double) N);
        return y;
    }

    /**
     * q-th percentile of x, interpolates linearly between values if necessary
     * @param x  input array
     * @param q  percentile in range [0:100]
     */
    public static double percentile(double[] x, double q) {
        if(q < 0.0 || q > 100.0)
            throw new IllegalArgumentException("percentile q must be between [0:100]");

        double[] y = DVec.sort(x);

        double idx = q / 100.0 * (y.length - 1);

        // linearly interpolate between neighboring data points y[i] and y[j]
        int i = (int) idx;
        double frac = idx - i;
        int j = (i != y.length - 1) ? (i + 1) : i;

        return y[i] + (y[j] - y[i]) * frac;
    }

    private static void throwIfNotAscending(double[] xp) {
        for(int i = 0; i < xp.length - 1; i++)
            if((xp[i+1] - xp[i]) <= 0.0)
                throw new IllegalArgumentException("xp is not strictly monotonically ascending");
    }

    /** return np.interp(x, xp, fp) - boundary values are 0th order (constant) */
    public static double[] interp(double[] x, double[] xp, double[] fp) {
        if(xp.length != fp.length)
            throw new IllegalArgumentException("xp.length != fp.length");
        if(xp.length < 2)
            throw new IllegalArgumentException("cannot interpolate without at least 2 data points");
        throwIfNotAscending(xp);

        double[] y = new double[x.length];
        //double[] dxp = DVec.diff(xp);
        //double[] dfp = DVec.diff(fp);
        double[] slope = new double[xp.length-1];
        double[] intercept = new double[xp.length-1];

        for(int i = 0; i < xp.length-1; i++) {
            double dxp = xp[i+1] - xp[i];
            double dfp = fp[i+1] - fp[i];
            slope[i] = dfp / dxp;
            intercept[i] = fp[i] - slope[i] * xp[i];
        }

        for(int i = 0; i < x.length; i++) {
            if(x[i] < xp[0]) {
                y[i] = fp[0];
            } else if(x[i] > xp[xp.length-1]) {
                y[i] = fp[xp.length-1];
            } else {
                int idx = Arrays.binarySearch(xp, x[i]);
                if (idx <= -1) {
                    // not found precisely. interpolate
                    idx = -idx - 1; // convert back to array index (insertion point)
                    // idx==0 or idx==x.length cannot happen, handled above.
                    //assert(idx > 0 && idx < xp.length);
                    idx = idx - 1; // one below insertion point (left point)
                    y[i] = slope[idx] * x[i] + intercept[idx];
                } else {
                    // found precisely
                    y[i] = fp[idx];
                }
            }
        }

        return y;
    }

    /** return ibis[np.where((ibis>0.4)&(ibis<1.4))[0]] */
    public static double[] goodIbis(double[] ibis) {
        final double MIN_IBI = 0.4;
        final double MAX_IBI = 1.4;

        int[] largeEnough = DVec.lt(DVec.mul(DVec.ones(ibis.length), MIN_IBI), ibis);
        int[] smallEnough = DVec.lt(ibis, DVec.mul(DVec.ones(ibis.length), MAX_IBI));
        int[] goodIndices = IVec.where(IVec.mul(largeEnough, smallEnough));
        return DVec.index(ibis, goodIndices);
    }

    /**
     * @return one-hot 'boolean' vector of same length as `dv`.
     */
    public static int[] beatdetect(double fps, double[] x) {
        final double DIFF_THRESHOLD = 0.35;

        final int N1 = (int) (0.2 * fps);
        final int N2 = (int) (0.2334 * fps);
        if(x.length < N1 + N2) {
            //Log.i(TAG, "calculateBpm() signal too short");
            return new int[0]; // signal too short
        }
        double[] ds = DVec.mul(Signal.running_mean(Signal.running_mean(DVec.diff(x), N1), N2), -1);

        // make threshold adaptive to max slope
        double T = DIFF_THRESHOLD * Signal.percentile(ds, 95); // outlier-robust max()

        //Log.i(TAG, "calculateBpm() threshold T=" + T + " 95th=" + Signal.percentile(ds, 95) + " 50th=" + Signal.percentile(ds, 50) + " max=" + Signal.percentile(ds, 100));

        for(int i = 0; i < ds.length; i++)
            if(ds[i] < T)
                ds[i] = 0;

        return Signal.heartbeat_localmax(ds);
    }

    // TODO(david): testbed for calculateBpm() accuracy, on existing signals having a reference ECG
    // (or even the PlotFilter variant and slope localmax() beatdetector, vs. getrr2()

    /**
     * detect beats and compute heart rate in a recent slice of detrended intensity signal.
     * @param fps  frames per second
     * @param x    intensity samples
     * @return heart rate in BPM, or null if none could be computed
     */
    public static Double calculateBpm(double fps, double[] x) {
        final int NUM_RECENT_BEATS = 20;
        final int MIN_NUM_BEATS = 6;
        final double MIN_BPM = 30;
        final double MAX_BPM = 180;

        double[] t = DVec.linspace(0.0, x.length / fps, x.length, false);

        // peak detection and bpm calculation (on smoothed signal; final bpm comes from server)
        // ds = -1*running_mean(running_mean(np.diff(self.pf.detrended()), 6), 7)

        int[] beats = beatdetect(fps, x);

        int[] beatLoc = IVec.where(beats);
        Log.i(TAG, "calculateBpm() beatLoc.length=" + beatLoc.length);
        final int skipFront = 3;
        final int skipBack = 1;
        if(beatLoc.length < skipFront + skipBack + 1) {
            //Log.i(TAG, "calculateBpm() not enough beats to skip");
            return null; // not enough beats
        }
        // skip some early beats (noisy), and the most recent beat (may be spurious)
        beatLoc = IVec.get(beatLoc, skipFront, beatLoc.length - skipBack);

        double[] ibis = DVec.diff(DVec.index(t, beatLoc));
        Log.i(TAG, "calculateBpm() ibis=" + DVec.toString(ibis));
        ibis = Signal.goodIbis(ibis);
        // only use recent beats to provide some visual sense of heart rate trends
        ibis = DVec.get(ibis, Math.max(0, ibis.length - NUM_RECENT_BEATS), ibis.length);

        if(ibis.length < MIN_NUM_BEATS) {
            Log.i(TAG, "calculateBpm() not enough beats < MIN_NUM_BEATS");
            return null; // not enough beats
        }

        double bpm = 60.0 / DVec.mean(ibis);
        if(bpm > MAX_BPM || bpm < MIN_BPM) {
            Log.i(TAG, "calculateBpm() computed invalid BPM = " + Double.toString(bpm));
            return null; // invalid BPM
        }

        return bpm;
    }

    /** normalized cross-correlation of the two signals of same length */
    public static double crossCorr(double[] x, double[] y) {
        return DVec.sum(DVec.mul(x, y)) / Math.sqrt(DVec.sum(DVec.mul(x, x)) * DVec.sum(DVec.mul(y, y)));
    }

    /** two-dimensional mean of a collection of signals */
    public static double[] mean(List<double[]> sigs) {
        final int nsigs = sigs.size();
        if(nsigs == 0)
            throw new IllegalArgumentException("cannot compute mean of empty stack");
        double[] y = new double[sigs.get(0).length];
        for(int i = 0; i < nsigs; i++) {
            double[] x = sigs.get(i);
            if(x.length != y.length)
                throw new IllegalArgumentException("stack entry " + i + " has different length from the rest");
            for(int j = 0; j < y.length; j++)
                y[j] += x[j];
        }

        for(int j = 0; j < y.length; j++)
            y[j] /= nsigs;
        return y;
    }

    public static double[] triangle(int ntaps) {
        int htaps = ntaps/2;
        double[] taps = new double[ntaps];

        for(int i = 0; i <= htaps; i++)
            taps[i] = ((double) i) / htaps;
        for(int i = htaps+1; i < ntaps; i++)
            taps[i] = 1.0 - ((double) (i - htaps)) / htaps;

        taps = DVec.mul(taps, 1.0 / DVec.sum(taps));
        return taps;
    }
}
