package io.pslab.communication;

import android.util.Log;

import org.apache.commons.math3.analysis.ParametricUnivariateFunction;
import org.apache.commons.math3.complex.Complex;
import org.apache.commons.math3.fitting.CurveFitter;
import org.apache.commons.math3.optim.nonlinear.vector.jacobian.LevenbergMarquardtOptimizer;
import org.apache.commons.math3.stat.descriptive.DescriptiveStatistics;
import org.apache.commons.math3.transform.DftNormalization;
import org.apache.commons.math3.transform.FastFourierTransformer;
import org.apache.commons.math3.transform.TransformType;
import io.pslab.filters.BandstopFilter;
import io.pslab.filters.Lfilter;

import java.util.ArrayList;
import java.util.Arrays;

import static java.lang.Math.cos;
import static org.apache.commons.lang3.math.NumberUtils.min;
import static org.apache.commons.lang3.math.NumberUtils.max;
import static org.apache.commons.math3.util.FastMath.abs;
import static org.apache.commons.math3.util.FastMath.exp;
import static org.apache.commons.math3.util.FastMath.sin;

/**
 * Created by akarshan on 5/13/17.
 * <p>
 * Unlike python, the curve fitting and the fourier transorm methods only accepts the arrays
 * of length of power of 2.
 * <p>
 * JAVA equivalent of scipy.fftpack.rfft, scipy.fftpack.rfftfreq and scipy.fftpack.fftfreq functions are'nt
 * available so they are been added in this class by the name fftToRfft, rfftFrequency and fftFrequency
 * respectively.
 * <p>
 * The fit functions have passed some dry tests but required to be tested with data points generated by PSLab.
 * </p>
 */

@SuppressWarnings("ALL")
public class AnalyticsClass {

    // todo : check the accuracy of the all the curve fitting methods with data points generated by PSLab
    //-------------------------- Exponential Fit ----------------------------------------//

    private ParametricUnivariateFunction exponentialParametricUnivariateFunction = new ParametricUnivariateFunction() {
        @Override
        public double value(double x, double... parameters) {
            double a = parameters[0];
            double b = parameters[1];
            double c = parameters[2];
            return a * exp(-x / b) + c;
        }

        @Override
        public double[] gradient(double x, double... parameters) {
            double a = parameters[0];
            double b = parameters[1];
            double c = parameters[2];

            return new double[]{
                    exp(-x / b),
                    (a * exp(-x / b) * x) / (b * b),
                    1
            };                                                      //partial derivatives
        }
    };

    public ArrayList<double[]> fitExponential(double time[], double voltage[]) {
        //length of time and voltage arrays should be in the power of 2
        double size = time.length;
        double v80 = voltage[0] * 0.8;
        double rc = 0;
        double[] vf = new double[time.length];
        for (int k = 0; k < size - 1; k++) {
            if (voltage[k] < v80) {
                rc = time[k] / .223;
                break;
            }
        }
        double[] initialGuess = new double[]{voltage[0], rc, 0};
        //initialize the optimizer and curve fitter.
        LevenbergMarquardtOptimizer optimizer = new LevenbergMarquardtOptimizer();
        CurveFitter fitter = new CurveFitter(optimizer);
        for (int i = 0; i < time.length; i++)
            fitter.addObservedPoint(time[i], voltage[i]);
        double[] result = fitter.fit(exponentialParametricUnivariateFunction, initialGuess);    //result -> the fitted parameters.
        for (int i = 0; i < time.length; i++)
            vf[i] = result[0] * exp(-time[i] / result[1]) + result[2];
        return new ArrayList<double[]>(Arrays.asList(result, vf));
    }

    //-------------------------- Sine Fit ----------------------------------------//

    private ParametricUnivariateFunction sineParametricUnivariateFunction = new ParametricUnivariateFunction() {
        @Override
        public double value(double x, double... parameters) {
            double a1 = parameters[0];
            double a2 = parameters[1];
            double a3 = parameters[2];
            double a4 = parameters[3];
            return a4 + a1 * sin(abs(a2 * (2 * Math.PI)) * x + a3);
        }

        @Override
        public double[] gradient(double x, double... parameters) {
            double a = parameters[0];
            double b = parameters[1];
            double c = parameters[2];
            double d = parameters[3];
            return new double[]{
                    Math.sin(c + 2 * Math.PI * b * x),
                    2 * a * Math.PI * x * cos(c + 2 * b * Math.PI * x),
                    a * cos(c + 2 * b * Math.PI * x),
                    1
            };
        }
    };

    public double[] sineFit(double[] xReal, double[] yReal) {
        int n = xReal.length;
        int index = 0;
        double[] frequencyArray = new double[]{};
        double[] yReal2 = new double[yReal.length];
        double yHat[];
        double yHatSquare[];
        double max = 0.;
        double returnOffset = 0;
        double frequency;
        double returnFrequency;
        double amplitude;
        double returnAmplitude;
        double phase;
        double returnPhase;
        double[] guess;
        Complex complex[];

        double offset = ((double) max(yReal) + (double) min(yReal)) / 2;
        for (int i = 0; i < yReal.length; i++)
            yReal2[i] = yReal[i] - offset;

        FastFourierTransformer fastFourierTransformer = new FastFourierTransformer(DftNormalization.STANDARD);
        complex = fastFourierTransformer.transform(yReal2, TransformType.FORWARD);
        yHat = fftToRfft(complex);      //yHat is an array of Discrete Fourier transform of a real sequence
        yHatSquare = new double[yHat.length];
        for (int i = 0; i < yHat.length; i++) {
            yHatSquare[i] = Math.pow(yHat[i], 2);
            if (yHatSquare[i] > max) {        //yHatSquare is an array square of each element of yHat
                max = yHatSquare[i];
                index = i;
            }
        }

        frequencyArray = rfftFrequency(n, (xReal[1] - xReal[0]) / (2 * Math.PI));
        frequency = frequencyArray[index];
        frequency /= (2 * Math.PI);

        amplitude = (max(yReal) - min(yReal)) / 2.0;

        phase = 0;

        LevenbergMarquardtOptimizer optimizer = new LevenbergMarquardtOptimizer();
        CurveFitter fitter = new CurveFitter(optimizer);
        for (int i = 0; i < n; i++)
            fitter.addObservedPoint(xReal[i], yReal2[i]);

        guess = new double[]{amplitude, frequency, phase, 0};
        double[] result = fitter.fit(sineParametricUnivariateFunction, guess);

        amplitude = result[0];
        frequency = result[1];
        phase = result[2];
        returnOffset = result[3];

        if (frequency < 0)
            Log.v("sineFit", "negative frequency");

        returnOffset += offset;
        returnPhase = ((phase) * 180 / (3.14));
        if (amplitude < 0)
            returnPhase -= 180;
        if (returnPhase < 0)
            returnPhase = (returnPhase + 720) % 360;
        returnFrequency = 1e6 * abs(frequency);
        returnAmplitude = abs(amplitude);
        return new double[]{returnAmplitude, returnFrequency, returnOffset, returnPhase};
    }

    ParametricUnivariateFunction squareParametricUnivariateFunction = new ParametricUnivariateFunction() {
        @Override
        public double value(double x, double... parameters) {
            double amp = parameters[0];
            double freq = parameters[1];
            double phase = parameters[2];
            double dc = parameters[3]; //dc - duty cycle
            double offset = parameters[4];
            return offset + amp * signalSquare(2 * Math.PI * freq * (x - phase), freq, dc);
        }

        @Override
        public double[] gradient(double x, double... parameters) {
            /*partial derivatives w.r.t all the five variables. Square functions are not differentiable
            at finitely many points, still we have used it anyway since the curve fitter uses the value of
            gradients. The values are true except at points where the transition from high to low or low to high
            takes place.*/
            double a = parameters[0]; //unused for partial derivative
            double b = parameters[1];
            double c = parameters[2];
            double d = parameters[3];
            double e = parameters[4]; //unused for partial derivative

            return new double[]{
                    signalSquare(2 * Math.PI * b * (x - c), b, d), 0, 0, 0, 1
            };
        }
    };

    public double[] squareFit(double[] xReal, double[] yReal) {
        int n = xReal.length;
        double mx = max(yReal);
        double mn = min(xReal);
        double offset = (mx + mn) / 2;
        double sumGreaterThanOffset = 0;
        double sumLesserThanOffset = 0;
        double n1 = 0; // count of numbers less than offset
        double n2 = 0; // count of numbers less than offset
        double[] yTmp = new double[yReal.length];
        double[] yReal2 = new double[yReal.length];
        double[] guess;
        double returnOffset;
        double returnFrequency;
        double returnAmplitude;
        double returnPhase;
        double returnDC;

        for (int i = 0; i < yReal.length; i++)
            yReal2[i] = yReal[i] - offset;

        for (int i = 0; i < yReal.length; i++) {
            if (yReal[i] > offset) {
                sumGreaterThanOffset += yReal[i];
                yTmp[i] = 2;
                n1++;
            } else if (yReal[i] < offset) {
                sumLesserThanOffset += yReal[i];
                yTmp[i] = 0;
                n2++;
            }
        }

        double amplitude = (sumGreaterThanOffset / n1) - (sumLesserThanOffset / n2);
        boolean[] bools = new boolean[yTmp.length];
        double tmp;
        for (int i = 0; i < yTmp.length - 1; i++) {
            tmp = yTmp[i + 1] - yTmp[i];
            tmp = abs(tmp);
            bools[i] = tmp > 1;
        }
        double[] edges = new double[bools.length];
        double[] levels = new double[bools.length];
        int j = 0;
        for (int i = 0; i < bools.length; i++) {
            if (bools[i]) {
                edges[j] = xReal[i];
                levels[j] = yTmp[i];
                j++;
            }
        }

        double frequency = 1 / (edges[2] - edges[0]);
        double phase = edges[0];
        double dc = 0.5;

        if (edges.length >= 4) {
            if (levels[0] == 0)
                dc = (edges[1] - edges[0]) / (edges[2] - edges[0]);
            else {
                dc = (edges[2] - edges[1]) / (edges[3] - edges[1]);
                phase = edges[1];
            }
        }

        LevenbergMarquardtOptimizer optimizer = new LevenbergMarquardtOptimizer();
        CurveFitter fitter = new CurveFitter(optimizer);
        for (int i = 0; i < n; i++)
            fitter.addObservedPoint(xReal[i], yReal2[i]);

        guess = new double[]{amplitude, frequency, phase, dc, 0};
        double[] result = fitter.fit(squareParametricUnivariateFunction, guess);

        amplitude = result[0];
        frequency = result[1];
        phase = result[2];
        dc = result[3];
        returnOffset = result[4];

        if (frequency < 0)
            Log.v("squareFit", "negative frequency");

        returnOffset += offset;
        returnFrequency = 1e6 * abs(frequency);
        returnAmplitude = abs(amplitude);
        returnPhase = phase;
        returnDC = dc;

        return new double[]{returnAmplitude, returnFrequency, returnPhase, returnDC, returnOffset};
    }

    public double findFrequency(double[] voltage, double samplingInterval) {
        int voltageLength = voltage.length;
        double[] frequency;
        double[] amplitude;
        int index = 0;
        double max = 0;
        Complex[] complex;
        DescriptiveStatistics stats = new DescriptiveStatistics();
        for (int i = 0; i < voltageLength; i++)
            stats.addValue(voltage[i]);
        double voltageMean = stats.getMean();
        for (int i = 0; i < voltageLength; i++)
            voltage[i] = voltage[i] - voltageMean;                                                                  // remove DC component
        frequency = Arrays.copyOfRange(fftFrequency(voltageLength, samplingInterval), 0, voltageLength / 2);        // take only the +ive half of the frequncy array
        FastFourierTransformer fastFourierTransformer = new FastFourierTransformer(DftNormalization.STANDARD);
        complex = fastFourierTransformer.transform(voltage, TransformType.FORWARD);
        amplitude = new double[complex.length / 2];
        for (int i = 0; i < complex.length / 2; i++) {                                                                 // take only the +ive half of the fft result
            amplitude[i] = complex[i].abs() / voltageLength;
            if (amplitude[i] > max) {                                                                                // search for the tallest peak, the fundamental
                max = amplitude[i];
                index = i;
            }
        }
        return frequency[index];
    }

    public ArrayList<double[]> amplitudeSpectrum(double[] voltage, int samplingInterval, int nHarmonics) {
        int voltageLength = voltage.length;
        Complex[] complex;
        double[] amplitude;
        double[] newAmplitude;
        int index = 0;
        double max = 0;
        double[] frequency = Arrays.copyOfRange(fftFrequency(voltageLength, samplingInterval), 0, voltageLength / 2);   // take only the +ive half of the frequncy array
        FastFourierTransformer fastFourierTransformer = new FastFourierTransformer(DftNormalization.STANDARD);
        complex = fastFourierTransformer.transform(voltage, TransformType.FORWARD);
        amplitude = new double[complex.length / 2];
        for (int i = 0; i < complex.length / 2; i++) {                                                                     // take only the +ive half of the fft result
            amplitude[i] = complex[i].abs() / voltageLength;
            if (amplitude[i] > max) {                                                                                    // search for the tallest peak, the fundamental
                max = amplitude[i];
                index = i;
            }
        }
        if (index == 0) {                                                                                               // DC component is dominating
            max = 0;
            newAmplitude = Arrays.copyOfRange(amplitude, 4, amplitude.length);                                          // skip frequencies close to zero
            for (int i = 0; i < newAmplitude.length; i++) {
                if (newAmplitude[i] > max) {
                    max = newAmplitude[i];
                    index = i;
                }
            }
        }
        return new ArrayList<double[]>(Arrays.asList(                                                                   // restrict to 'nHarmonics' harmonics
                Arrays.copyOfRange(frequency, index * nHarmonics, frequency.length),
                Arrays.copyOfRange(amplitude, index * nHarmonics, amplitude.length)
        ));
    }

    //-------------------------- Damped Sine Fit ----------------------------------------//

    public ParametricUnivariateFunction dampedSineParametricUnivariateFunction = new ParametricUnivariateFunction() {
        @Override
        public double value(double x, double... parameters) {
            double amplitude = parameters[0];
            double frequency = parameters[1];
            double phase = parameters[2];
            double offset = parameters[3];
            double damp = parameters[4];
            return offset + amplitude * exp(-damp * x) * sin(Math.abs(frequency) * x + phase);
        }

        @Override
        public double[] gradient(double x, double... parameters) {
            double amplitude = parameters[0];
            double frequency = parameters[1];
            double phase = parameters[2];
            double offset = parameters[3];
            double damp = parameters[4];
            return new double[]{
                    exp(-damp * x) * sin(x * frequency + phase),
                    amplitude * exp(-damp * x) * x * cos(x * frequency + phase),
                    amplitude * exp(-damp * x) * cos(x * frequency + phase),
                    -amplitude * exp(-damp * x) * x * sin(x * frequency + phase),
                    1
            };
        }
    };


    public double[] getGuessValues(double xReal[], double yReal[], String function) {
        if (function.equals("sine") || function.equals("damped sine")) {
            int n = xReal.length;
            Complex[] complex;
            double[] yReal2 = new double[yReal.length];
            double[] yHatSquare;
            double[] yHat;
            double[] frequencyArray;
            double frequency;
            double amplitude;
            double phase;
            double max = 0;
            int index = 0;
            DescriptiveStatistics stats = new DescriptiveStatistics();
            for (int i = 0; i < yReal.length; i++)
                stats.addValue(yReal[i]);
            double offset = stats.getMean();
            for (int i = 0; i < yReal.length; i++)
                yReal2[i] = yReal[i] - offset;
            FastFourierTransformer fastFourierTransformer = new FastFourierTransformer(DftNormalization.STANDARD);
            complex = fastFourierTransformer.transform(yReal, TransformType.FORWARD);
            yHat = fftToRfft(complex);
            yHatSquare = new double[yHat.length];
            for (int i = 0; i < yHat.length; i++) {
                yHatSquare[i] = Math.pow(yHat[i], 2);
                if (yHatSquare[i] > max) {        //yHatSquare is an array square of each element of yHat
                    max = yHatSquare[i];
                    index = i;
                }
            }
            frequencyArray = rfftFrequency(n, (xReal[1] - xReal[0]) / 2);
            frequency = frequencyArray[index];
            amplitude = (max(yReal) - min(yReal)) / 2.0;
            phase = 0;
            if (function.equals("sine"))
                return new double[]{amplitude, frequency, phase, offset};
            if (function.equals("damped sine"))
                return new double[]{amplitude, frequency, phase, offset, 0};
        }
        return null;
    }

    public double[] arbitFit(double[] xReal, double[] yReal, ParametricUnivariateFunction function, double[] guessValues) {
        int n = xReal.length;
        LevenbergMarquardtOptimizer optimizer = new LevenbergMarquardtOptimizer();
        CurveFitter fitter = new CurveFitter(optimizer);
        for (int i = 0; i < n; i++)
            fitter.addObservedPoint(xReal[i], yReal[i]);
        double[] result = fitter.fit(function, guessValues);
        return result;
    }

    ArrayList<double[]> butterNotch(double lowCut, double highCut, double fs, int order) {
        double nyq = 0.5 * fs;
        double low = lowCut / nyq;
        double high = highCut / nyq;
        BandstopFilter bandstopFilter = new BandstopFilter(order, new double[]{low, high});
        return bandstopFilter.abGetter();
    }

    double[] butterNotchFilter(double[] data, double lowCut, double highCut, double fs, int order) {
        ArrayList<double[]> arrayList = butterNotch(lowCut, highCut, fs, order);
        double[] b = arrayList.get(0);
        double[] a = arrayList.get(1);
        Lfilter lfilter = new Lfilter();
        return lfilter.filter(b, a, data);
    }

    public ArrayList<double[]> fft(double[] signal, double samplingInterval) {
        /*
         Returns positive half of the Fourier transform of the signal.
         Sampling interval 'samplingInterval', in milliseconds
         */
        int ns = signal.length;
        Complex[] complex;
        double[] terms;
        double[] frequencyArray;
        double[] x;
        double[] y;
        if (ns % 2 == 1) {
            ns = ns - 1;
            signal = Arrays.copyOfRange(signal, 0, ns);
        }
        FastFourierTransformer fastFourierTransformer = new FastFourierTransformer(DftNormalization.STANDARD);
        complex = fastFourierTransformer.transform(signal, TransformType.FORWARD);
        terms = new double[complex.length];
        for (int i = 0; i < complex.length; i++)
            terms[i] = complex[i].abs() / ns;
        frequencyArray = fftFrequency(ns, samplingInterval);
        x = Arrays.copyOfRange(frequencyArray, 0, ns / 2);
        y = Arrays.copyOfRange(terms, 0, ns / 2);
        return new ArrayList<double[]>(Arrays.asList(x, y));
    }

    public double[] rfftFrequency(int n, double space) {
        /*
         The returned array contains the frequency bins in
    	cycles/unit (with zero at the start) given a window length `n` and a
   		sample spacing `space`
		 */
        double[] returnArray = new double[n + 1];
        for (int i = 0; i < n + 1; i++) {
            returnArray[i] = Math.floor(i / 2) / (n * space);
        }
        return Arrays.copyOfRange(returnArray, 1, returnArray.length);
    }

    public double[] fftFrequency(int n, double space) {
        /*
        Return the Discrete Fourier Transform sample frequencies.
        The returned array contains the frequency bin centers in cycles
        per unit of the sample spacing (with zero at the start).  For instance, if
        the sample spacing is in seconds, then the frequency unit is cycles/second.
        Given a window length `n` and a sample spacing `spacing`.
         */
        double value = 1.0 / (n * space);
        int N = Math.floorDiv(n - 1, 2) + 1;
        double[] results = new double[n];
        for (int i = 0; i < N; i++) {
            results[i] = i;
            results[i] = results[i] * value;
        }
        int j = N;
        for (int i = -Math.floorDiv(n, 2); i < 0; i++) {

            results[j] = i;
            results[j] = results[j] * value;
            j++;
        }
        return results;
    }

    public double[] fftToRfft(Complex[] complex) {
        //The returned array contains Discrete Fourier transform of a real sequence.
        double[] real = new double[complex.length];
        double[] imaginary = new double[complex.length];
        double[] result = new double[complex.length];
        int j = 0;
        int k = 0;
        int l = 0;
        for (int i = 0; i < complex.length / 2 + 1; i++) {
            real[i] = complex[i].getReal();
            imaginary[i] = complex[i].getImaginary();
        }

        for (int i = 0; i < complex.length / 2 + 1; i++) {
            if (real[j] == 0.0 && imaginary[k] == 0) {
                result[l++] = 0.0;
                j++;
                k++;
            } else if (real[j] != 0 && imaginary[k] == 0) {
                result[l++] = real[j++];
                k++;
            } else {
                result[l++] = real[j++];
                result[l++] = imaginary[k++];
            }
        }
        return result;
    }

    public double signalSquare(double xAxisValue, double freq, double dc) {
        //This method determines whether at a given x value, the value of y is in the upper half or lower half
        if (xAxisValue % (2 * Math.PI * freq) <= dc)
            return 1;
        else
            return -1;
    }
}