package net.heartshield.signal;

import android.util.Log;
import android.util.Pair;

import java.util.ArrayList;
import java.util.List;

/**
 * Running signal quality indicator.
 *
 * @author David Madl (git@abanbytes.eu)
 * @date 2017-04-23
 */
public class RunningQuality {
    private static final String TAG = "RunningQuality";

    /** template beat is resampled to this #samples */
    private static final int BEAT_LEN = 30;

    /** threshold for accepting initial beats */
    private static final double BEAT_CORR_THR_1 = 0.8;

    /** threshold for accepting subsequent beats */
    public static final double BEAT_CORR_THR_2 = 0.6;

    /** for second measurement after failed first one */
    public static final double BEAT_CORR_THR_2_LOWER = 0.45;

    private double mBeatCorrThr2 = BEAT_CORR_THR_2;


    private Double startTime = null;
    private List<Double> x = new ArrayList<>();
    private List<Integer> edges = new ArrayList<>();
    private double fps;
    private List<double[]> beatTemplates = new ArrayList<>();
    private double[] beatTemplate = null;
    private List<IRange> badBeatRanges = new ArrayList<>();
    private List<Pair<Integer, Double>> labels = new ArrayList<>();
    private List<Double> runningCorrs = new ArrayList<>();
    private Integer startIdx = null;
    private int numNoisy = 0;
    //private int numTotal = 0;

    /** Pythonic range [start:end) */
    public static class IRange {
        public int start;
        public int end;

        public IRange(int start, int end) {
            this.start = start;
            this.end = end;
        }
    }

    public RunningQuality(double fps) {
        this.fps = fps;
    }

    public void clear() {
        x = new ArrayList<>();
        edges = new ArrayList<>();
        beatTemplates = new ArrayList<>();
        beatTemplate = null;
        badBeatRanges = new ArrayList<>();
        labels = new ArrayList<>();
        runningCorrs = new ArrayList<>();
        startIdx = null;
        numNoisy = 0;
        //numTotal = 0;
        startTime = null;

        Log.i(TAG, "clear()");
    }

    /*
     * edge: detected downslope
     * beat: timespan between two edges
     */

    /**
     * @param t   time of sample 0 in this slice (but only used for very first call)
     * @param sl  recent slice to append
     */
    public void append(double t, double[] sl) {
        //Log.i(TAG, "append(sl.length=" + sl.length + ")");

        if(startTime == null)
            startTime = t;

        final int beatdet_hwin = 30;
        int prev_ix = Math.max(x.size() - beatdet_hwin, 0);
        x.addAll(DVec.toList(sl));

        // nb. CPU waste. nice-to: incremental beatdetect()
        int[] idx = IVec.where(Signal.beatdetect(fps, DVec.toArray(x)));

        Log.i(TAG, "append(sl.length=" + sl.length + "). all idx: " + IVec.toString(idx));

        // ignore new beats unless closer within - at least 1 sec (half-window of savgol in getrr())
        int[] recent = IVec.ge(idx, prev_ix);
        int[] oldEnough = IVec.lt(idx, x.size() - beatdet_hwin);
        int[] new_idx = IVec.index(idx, IVec.where(IVec.mul(recent, oldEnough)));
        int prev_iedge = edges.size();
        edges.addAll(IVec.toList(new_idx));

        if(edges.size() > 1 && new_idx.length > 0)
            appendBeats(prev_ix, prev_iedge);
    }

    public interface BeatListener {
        /**
         * Listener callback for first-time locking.
         *
         * @param startIdx  index of sample where this beat starts
         */
        void onLocked(int startIdx);

        /**
         * Listener callback for all detected beats, good and bad.
         * Note that this may be dispatched several times in succession.
         *
         * @param startIdx      index of sample where this beat starts
         * @param goodBeat      whether this beat is good
         * @param posCorrCoeff  positive-clipped correlation coefficient in range [0:1]
         */
        void onBeat(int startIdx, boolean goodBeat, double posCorrCoeff);
    }
    private BeatListener mListener;

    public void setListener(BeatListener listener) { mListener = listener; }
    private void dispatchBeat(int startIdx, boolean goodBeat, double posCorrCoeff) {
        if(mListener != null)
            mListener.onBeat(startIdx, goodBeat, posCorrCoeff);
    }
    private void dispatchLocked(int startIdx) {
        if(mListener != null)
            mListener.onLocked(startIdx);
    }

    public double getTimeAt(int idx) {
        return startTime + idx / fps;
    }

    public List<IRange> getBadBeatRanges() { return badBeatRanges; }

    private void appendBeats(int prev_ix, int prev_iedge) {
        int start = Math.max(prev_iedge-1, 0);
        int end = edges.size()-1;
        for(int i = start; i < end; i++) {
            int s = edges.get(i);
            int e = edges.get(i+1);

            // TODO: should ignore crazy-long and very short beats here. (filter up on beat detector)

            //Log.i(TAG, "appendBeat(" + s + ", " + e + ")");
            double[] beat = resample(DVec.toArray(x.subList(s, e)));

            double corr = Double.NaN;
            double posCorr = Double.NaN;
            boolean goodBeat = false;
            if(beatTemplates.size() > 0) {
                corr = Signal.crossCorr(beat, beatTemplate);
                posCorr = Math.min(Math.max(corr, 0.0), 1.0);
                double corrThreshold;
                if(beatTemplates.size() > 2)
                    corrThreshold = mBeatCorrThr2;
                else
                    corrThreshold = BEAT_CORR_THR_1;
                goodBeat = (corr > corrThreshold);
            }

            boolean justLocked = false;
            if(beatTemplates.size() == 0) {
                // cannot correlate the first beat, no template yet
                Log.i(TAG, "appendBeat(" + s + ", " + e + ") first beat");
                addTemplate(beat);
            } else if(beatTemplates.size() <= 2) {
                Log.i(TAG, "appendBeat(" + s + ", " + e + ") init corr=" + corr);
                // initialization: restart if there is no clear correlation between beats
                // nb.: this may fail for low BP with double-peaks of current getrr()
                if(goodBeat) {
                    addTemplate(beat);
                    if(beatTemplates.size() > 2)
                        justLocked = true;
                } else {
                    replaceTemplate(beat);
                    badBeatRanges.clear();
                }
            } else {
                Log.i(TAG, "appendBeat(" + s + ", " + e + ") run corr=" + corr);
                // running mode: collect bad beats, but may be OK not to restart immediately
                if(startIdx == null)
                    startIdx = s;

                if(goodBeat) {
                    // keep updating beat template (nb. this gets O(n**2) currently)
                    addTemplate(beat);
                } else {
                    badBeatRanges.add(new IRange(s, e));
                    numNoisy++;
                }

                runningCorrs.add(posCorr);

                if(justLocked)
                    dispatchLocked(s);
                dispatchBeat(s, goodBeat, posCorr);
            }

            if(!Double.isNaN(corr))
                labels.add(new Pair<>(s, corr));
        }
    }

    public boolean isRunning() {
        return beatTemplates.size() > 2;
    }

    public void setBeatCorrThr2(double beatCorrThr2) {
        mBeatCorrThr2 = beatCorrThr2;
    }

    private void addTemplate(double[] beat) {
        beatTemplates.add(beat);
        // nb. this gets O(n**2) currently
        beatTemplate = Signal.mean(beatTemplates);
    }

    private void replaceTemplate(double[] beat) {
        beatTemplates.clear();
        beatTemplates.add(beat);
        beatTemplate = Signal.mean(beatTemplates);
    }

    /** resample to BEAT_LEN */
    private double[] resample(double[] x) {
        double[] t = DVec.linspace(0, x.length, BEAT_LEN, false);
        return Signal.interp(t, IVec.toDouble(IVec.arange(0, x.length)), x);
    }
}
