/*
 * Philip Crotwell
 * University of South Carolina, 2019
 * https://www.seis.sc.edu
 */
import { Duration, DateTime } from "luxon";
import { Seismogram } from "./seismogram";
import { InstrumentSensitivity } from "./stationxml";
import {
  Butterworth,
  ChebyshevI,
  ChebyshevII,
  PassbandType,
  IIRFilter,
  CenteredHilbertTransform,
  LOWPASS,
  BANDPASS,
  HIGHPASS,
} from "./oregondsputil";
import { isDef } from "./util";

/**
 * Constant for bandpass OregonDSP filter creation.
 */
export const BAND_PASS = "BANDPASS";

/**
 * Constant for lowpass OregonDSP filter creation.
 */
export const LOW_PASS = "LOWPASS";

/**
 * Constant for highpass OregonDSP filter creation.
 */
export const HIGH_PASS = "HIGHPASS";

export function amplitude(real: number, imag: number): number {
  return Math.hypot(real, imag);
}

/**
 * Remove the mean from a seismogram. Subtract the mean from each data point.
 *
 * @param   seis input seismogram
 * @returns       seismogram with mean of zero
 */
export function rMean(seis: Seismogram): Seismogram {
  if (seis instanceof Seismogram) {
    const meanVal = seis.mean();
    const rmeanSeismogram = new Seismogram(
      seis.segments.map((s) => {
        const demeanY = s.y.map(function (d) {
          return d - meanVal;
        });
        const out = s.cloneWithNewData(demeanY);
        return out;
      }),
    );
    return rmeanSeismogram;
  } else {
    throw new Error("rMean arg not a Seismogram");
  }
}

export type LineFitType = {
  slope: number;
  intercept: number;
  reference_time: DateTime;
  sigma: number;
  sigma_a: number;
  sigma_b: number;
  correlation: number;
};

/**
 * Calculate best fit line to seismogram. Limited to contiguous data currently.
 * Code derived from scm/lifite.c in SAC.
 * Original version from Steve Taylor.
 *
 * @param  seis                        [description]
 * @param  referenceTime               [description]
 * @returns               best fit line
 */
export function lineFit(
  seis: Seismogram,
  referenceTime?: DateTime,
): LineFitType {
  if (seis.numPoints === 0) {
    throw new Error(`cannot lineFit a seismogram with no points`);
  }
  /* - Initialize accumulators. */
  const rn = seis.numPoints;
  let sumx = 0;
  let sumy = 0;
  let sumxy = 0;
  let sumx2 = 0;
  let sumy2 = 0;

  /* - Loop on each data point. */

  referenceTime = referenceTime ? referenceTime : seis.start;
  const x1 = referenceTime.toMillis() / 1000; // seconds
  seis.segments.forEach((seg) => {
    const seg_start_x = seg.start.toMillis() / 1000 - x1;
    const dx = 1 / seg.sampleRate; // seconds
    const Y = seg.y;
    for (let i = 0; i < Y.length; i++) {
      const yi = Y[i];
      const xi = seg_start_x + dx * i;
      sumx = sumx + xi;
      sumy = sumy + yi;
      sumxy = sumxy + xi * yi;
      sumx2 = sumx2 + xi * xi;
      sumy2 = sumy2 + yi * yi;
    }
  });

  /* - Calculate linear fit. */

  const d = rn * sumx2 - sumx * sumx;
  // zero denominator would cause NaN, assume zero slope intercept
  const b = d !== 0 ? (sumx2 * sumy - sumx * sumxy) / d : 0;
  const a = d !== 0 ? (rn * sumxy - sumx * sumy) / d : 0;

  /* - Estimate standard deviation in data. */

  const sig2 =
    (sumy2 +
      rn * b * b +
      a * a * sumx2 -
      2 * b * sumy -
      2 * a * sumxy +
      2 * b * a * sumx) /
    (seis.numPoints - 1);
  const sig = Math.sqrt(sig2);

  /* - Estimate errors in linear fit. */

  const siga2 = (rn * sig2) / d;
  const sigb2 = (sig2 * sumx2) / d;
  const siga = Math.sqrt(siga2);
  const sigb = Math.sqrt(sigb2);

  /* - Calculate correlation coefficient between data and model. */

  let cc =
    (rn * sumxy - sumx * sumy) / Math.sqrt(d * (rn * sumy2 - sumy * sumy));
  cc = Math.abs(cc);

  return {
    slope: a,
    intercept: b,
    reference_time: referenceTime,
    sigma: sig,
    sigma_a: siga,
    sigma_b: sigb,
    correlation: cc,
  };
}

/**
 * Returns a new Seismogram with the trend removed by
 * subtracting the trend line from each data point.
 *
 * @param   seis input seismogram
 * @param fitLine optional fit type
 * @returns       seismogram with mean of zero and best fit line horizontal
 */
export function removeTrend(
  seis: Seismogram,
  fitLine?: LineFitType,
): Seismogram {
  if (seis instanceof Seismogram) {
    const linfit = fitLine ? fitLine : lineFit(seis);
    if (Number.isNaN(linfit.slope) || Number.isNaN(linfit.intercept)) {
      throw new Error(
        `Can't remove trend with NaN, slope: ${linfit.slope} int: ${linfit.intercept}`,
      );
    }
    const ref_secs = linfit.reference_time.toMillis() / 1000; // seconds
    const rtr_segments = seis.segments.map((seg) => {
      const start_secs = seg.start.toMillis() / 1000; // seconds
      const start_offset = start_secs - ref_secs;
      const dx = 1 / seg.sampleRate; // seconds
      const rtr_y = seg.y.map((y, idx) => {
        const out =
          y - (start_offset + dx * idx) * linfit.slope - linfit.intercept;
        return out;
      });
      const rtr_seg = seg.cloneWithNewData(rtr_y);
      return rtr_seg;
    });
    return new Seismogram(rtr_segments);
  } else {
    throw new Error("removeTrend arg not a Seismogram");
  }
}

/**
 * Apply the frequency independent overall gain to a seismogram. This does not
 * do a full transfer using poles and zero, this only applies the scalar conversion
 * factor to convert counts back to original real world units and update the units.
 *
 * @param   seis                  the seismogram to correct
 * @param   instrumentSensitivity overall gain object, usually pulled from stationxml
 * @returns                        new seismogram with original units, like m/s and gain applied.
 */
export function gainCorrect(
  seis: Seismogram,
  instrumentSensitivity: InstrumentSensitivity,
): Seismogram {
  const gain = instrumentSensitivity.sensitivity;
  const out = mul(seis, 1 / gain);
  out.segments.forEach((s) => (s.yUnit = instrumentSensitivity.inputUnits));
  return out;
}

export function mul(seis: Seismogram, factor: number): Seismogram {
  if (seis instanceof Seismogram) {
    const gainSeismogram = new Seismogram(
      seis.segments.map((s) => {
        let gainY;

        if (s.y instanceof Int32Array || s.y instanceof Float32Array) {
          gainY = Float32Array.from(s.y);
        } else {
          gainY = Float64Array.from(s.y);
        }

        gainY = gainY.map(function (d) {
          return d * factor;
        });
        const outS = s.cloneWithNewData(gainY);
        return outS;
      }),
    );
    return gainSeismogram;
  } else {
    throw new Error(`Expected Seismogram but was ${typeof seis}`);
  }
}

export function add(seis: Seismogram, factor: number): Seismogram {
  if (seis instanceof Seismogram) {
    const gainSeismogram = new Seismogram(
      seis.segments.map((s) => {
        let gainY;

        if (s.y instanceof Int32Array || s.y instanceof Float32Array) {
          gainY = Float32Array.from(s.y);
        } else {
          gainY = Float64Array.from(s.y);
        }

        gainY = gainY.map(function (d) {
          return d + factor;
        });
        const outS = s.cloneWithNewData(gainY);
        return outS;
      }),
    );
    return gainSeismogram;
  } else {
    throw new Error(`Expected Seismogram but was ${typeof seis}`);
  }
}

export function getPassband(
  type: string,
): typeof LOWPASS | typeof BANDPASS | typeof HIGHPASS {
  if (type === LOW_PASS) {
    return LOWPASS;
  } else if (type === BAND_PASS) {
    return PassbandType.BANDPASS;
  } else if (type === HIGH_PASS) {
    return PassbandType.HIGHPASS;
  } else {
    throw new Error(`unknown pass band: ${type}`);
  }
}
/**
 * Creates a Butterworth IIR filter using the OregonDSP library.
 *
 * @param   numPoles       number of poles
 * @param   passband       type, use constants of BAND_PASS, LOW_PASS, HIGH_PASS
 * @param   lowFreqCorner  low corner frequency
 * @param   highFreqCorner high corner frequency
 * @param   delta          delta, period, of timeseries
 * @returns                 Butterworth IIR filter
 */
export function createButterworth(
  numPoles: number,
  passband: string,
  lowFreqCorner: number,
  highFreqCorner: number,
  delta: number,
): InstanceType<typeof Butterworth> {
  const passbandtype = getPassband(passband);
  return new Butterworth(
    numPoles,
    passbandtype,
    lowFreqCorner,
    highFreqCorner,
    delta,
  );
}

/**
 * Creates a Chebyshev I IIR filter using the OregonDSP library.
 *
 * @param   numPoles       number of poles
 * @param   epsilon        Chebyshev epsilon value
 * @param   passband       type, use constants of BAND_PASS, LOW_PASS, HIGH_PASS
 * @param   lowFreqCorner  low corner frequency
 * @param   highFreqCorner high corner frequency
 * @param   delta          delta, period, of timeseries
 * @returns                 Chebyshev I IIR filter
 */
export function createChebyshevI(
  numPoles: number,
  epsilon: number,
  passband: string,
  lowFreqCorner: number,
  highFreqCorner: number,
  delta: number,
): InstanceType<typeof ChebyshevI> {
  const passbandtype = getPassband(passband);
  return new ChebyshevI(
    numPoles,
    epsilon,
    passbandtype,
    lowFreqCorner,
    highFreqCorner,
    delta,
  );
}

/**
 * Creates a Chebyshev II IIR filter using the OregonDSP library.
 *
 * @param   numPoles       number of poles
 * @param   epsilon        Chebyshev epsilon value
 * @param   passband       type, use constants of BAND_PASS, LOW_PASS, HIGH_PASS
 * @param   lowFreqCorner  low corner frequency
 * @param   highFreqCorner high corner frequency
 * @param   delta          delta, period, of timeseries
 * @returns                 Chebyshev II IIR filter
 */
export function createChebyshevII(
  numPoles: number,
  epsilon: number,
  passband: string,
  lowFreqCorner: number,
  highFreqCorner: number,
  delta: number,
): InstanceType<typeof ChebyshevII> {
  const passbandtype = getPassband(passband);
  return new ChebyshevII(
    numPoles,
    epsilon,
    passbandtype,
    lowFreqCorner,
    highFreqCorner,
    delta,
  );
}

/**
 * Applies the filter to the given seismogram.
 *
 * @param   iirFilter filter to apply
 * @param   seis      seismogram to apply filter to
 * @returns            filtered seismogram
 */
export function applyFilter(
  iirFilter: InstanceType<typeof IIRFilter>,
  seis: Seismogram,
): Seismogram {
  // check delta and samplePeriod with 0.1% of each other
  if (
    Math.abs(iirFilter.getDelta() - seis.samplePeriod) / seis.samplePeriod >
    0.001
  ) {
    throw new Error(
      `Filter, delta=${iirFilter.getDelta()}, has different delta from seis, ${1 / seis.sampleRate}`,
    );
  }
  const filteredSegments = [];

  for (let i = 0; i < seis.segments.length; i++) {
    const outData = Float32Array.from(seis.segments[i].y);
    iirFilter.filterInPlace(outData);
    filteredSegments.push(seis.segments[i].cloneWithNewData(outData));
  }

  return new Seismogram(filteredSegments);
}

/**
 * Calculates the envelope, y_i = sqrt( y_i * y_i + h_i * h_i)
 *  where h is the hilber transform of y. The default configuration
 *  for the hilbet transform is n=100, lowEdge=.05 and highEdge = 0.95
 *
 * @param seis seismogram to apply envelope to
 * @returns seismogram cloned but with data as the envelope
 */
export function envelope(seis: Seismogram): Seismogram {
  if (seis.isContiguous()) {
    const seisY = seis.y;
    const s = hilbert(seis);
    const hilbertY = s.y;
    let outY;

    if (seis.y instanceof Int32Array || seis.y instanceof Float32Array) {
      outY = new Float32Array(seisY.length);
    } else {
      outY = new Float64Array(seisY.length);
    }

    for (let n = 0; n < seisY.length; n++) {
      outY[n] = Math.sqrt(hilbertY[n] * hilbertY[n] + seisY[n] * seisY[n]);
    }

    return seis.cloneWithNewData(outY);
  } else {
    throw new Error("Cannot take envelope of non-contiguous seismogram");
  }
}

/**
 * Calculates the hilbert transform using the OregonDSP library
 *  with default number of points, n=10 (to yield a 21 pt FIR transform)
 *  and default low and high edge of 0.05 and 0.95. Low and high edge are
 *  given normalized 0 to 1.
 *
 * Note this uses Float32Array, other array types will be converted,
 * possibly losing precision.
 *
 * @param seis seismogram to calculate from
 * @param n optional number of points in transform, default is 10
 * @param lowEdge low edge of filter, normailized to 0-1, default is 0.05
 * @param highEdge high edge of filter, normailized to 0-1, default is 0.95
 * @returns hilbert transformed data
 */
export function hilbert(
  seis: Seismogram,
  n?: number,
  lowEdge?: number,
  highEdge?: number,
): Seismogram {
  if (seis.isContiguous()) {
    let seisY: Float32Array;
    if (seis.y instanceof Float32Array) {
      seisY = seis.y;
    } else {
      seisY = Float32Array.from(seis.y);
    }

    if (!isDef(n)) {
      n = 10;
    }

    if (!isDef(lowEdge)) {
      lowEdge = 0.05;
    }

    if (!isDef(highEdge)) {
      highEdge = 0.95;
    }

    const hilbert = new CenteredHilbertTransform(n, lowEdge, highEdge);

    const coeff = hilbert.getCoefficients();

    for (const c of coeff) {
      if (Number.isNaN(c)) {
        throw new Error(`Hilbert FIR coeff includes NaN: ${coeff.join()}`);
      }
    }

    const hilbertY = hilbert.filter(seisY);
    const s = seis.cloneWithNewData(hilbertY);
    return s;
  } else {
    throw new Error("Cannot take hilbert of non-contiguous seismogram");
  }
}

/**
 * Differentiate a seismogram.
 *
 * @param   seis input seismogram
 * @returns       differentiated seismogram
 */
export function differentiate(seis: Seismogram): Seismogram {
  if (seis instanceof Seismogram) {
    const diffSeismogram = new Seismogram(
      seis.segments.map((s) => {
        const origY = s.y;
        const sampRate = 1.0 * s.sampleRate; // same as 1/delta

        const diffY = new Float32Array(origY.length - 1);

        for (let i = 0; i < diffY.length; i++) {
          diffY[i] = (origY[i + 1] - origY[i]) * sampRate;
        }

        const out = s.cloneWithNewData(diffY);
        out.startTime = out.startTime.plus(
          Duration.fromMillis(1000 / out.sampleRate / 2),
        ); // second
        out.yUnit = out.yUnit + "/s";
        return out;
      }),
    );
    return diffSeismogram;
  } else {
    throw new Error("diff arg not a Seismogram");
  }
}

/**
 * Integrate a seismogram.
 *
 * @param   seis input seismogram
 * @param   integrationConst integration constant, optional, defaults to 0
 * @returns       integrated seismogram
 */
export function integrate(seis: Seismogram, integrationConst=0): Seismogram {
  let prior = integrationConst;
  if (seis instanceof Seismogram) {
    const intSeismogram = new Seismogram(
      seis.segments.map((s) => {
        const origY = s.y;
        const sampPeriod = s.samplePeriod; // same as 1/delta

        const intY = new Float32Array(origY.length +1);
        intY[0] = prior;
        for (let i = 1; i < intY.length; i++) {
          prior += sampPeriod * origY[i-1];
          intY[i] = prior;
        }

        const out = s.cloneWithNewData(intY);
        out.startTime = out.startTime.minus(
          Duration.fromMillis(1000 / out.sampleRate / 2),
        ); // second
        if (out.yUnit.endsWith("/s")) {
          out.yUnit = out.yUnit.slice(0, out.yUnit.length-2);
        } else {
          out.yUnit += "s";
        }
        return out;
      }),
    );
    return intSeismogram;
  } else {
    throw new Error("integrate arg not a Seismogram");
  }
}