/* Copyright (C) 2017 Massachusetts Institute of Technology.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#include <stdlib.h>
#include <stdio.h>
#include <ctype.h>
#include <math.h>

#include "harminv-int.h"
#include "check.h"
#include "copyright.h"

#ifdef HAVE_UNISTD_H
#include <unistd.h>
#endif
#ifdef HAVE_GETOPT_H
#include <getopt.h>
#endif

/* eat whitespace, including #... comments, from the file.  Returns the
   number of newlines read (so that a line count can be maintained).  If
   echo_comments != 0, then echo #... comments to stdout.  Commas count
   as whitespace, so that we can read comma-delimited text. */
static int eat_whitespace(FILE *f, int echo_comments) {
  int c, newlines = 0;
  do {
    do {
      c = getc(f);
      newlines += c == '\n';
    } while (isspace(c) || c == ',');

    if (c == '#') { /* # begins comments that extend to the newline */
      if (echo_comments) putc(c, stdout);
      do {
        c = getc(f);
        if (echo_comments) {
          if (c != EOF)
            putc(c, stdout);
          else /* terminate line if we hit EOF */
            putc('\n', stdout);
        }
        newlines += c == '\n';
      } while (c != EOF && c != '\n');
    }
  } while (isspace(c));

  ungetc(c, f); /* put back the last character read */
  newlines -= c == '\n';

  return newlines;
}

static int eat_plus(FILE *f) {
  int c = getc(f);
  if (c != EOF && c != '+') ungetc(c, f);
  return (c == '+' || c == '-');
}

static int eat_i(FILE *f) {
  int c = getc(f);
  if (c != EOF && tolower(c) != 'i') ungetc(c, f);
  return (tolower(c) == 'i');
}

static cmplx *read_input_data(FILE *f, int *n, int verbose) {
  cmplx *data = NULL;
  int line = 1, n_alloc = 0;
  *n = 0;

  do {
    double re = 0.0, im = 0.0;
    int nread;
    line += eat_whitespace(f, verbose);
    nread = fscanf(f, "%lg", &re);
    if (nread == 1 && eat_plus(f)) {
      nread = fscanf(f, "%lg", &im);
      if (nread == 1) nread = eat_i(f);
    }
    else
      im = 0.0;
    if (nread != EOF) {
      if (nread < 1) {
        fprintf(stderr, "harminv: invalid input on line %d.\n", line);
        free(data);
        *n = 0;
        return NULL;
      }
      if (*n >= n_alloc) {
        n_alloc = (n_alloc + 1) * 2;
        data = (cmplx *)realloc(data, sizeof(cmplx) * n_alloc);
        CHECK(data, "out of memory");
      }
      data[*n] = re + I * im;
      ++*n;
    }
  } while (!feof(f));

  data = (cmplx *)realloc(data, sizeof(cmplx) * *n);
  return data;
}

static const double inf = HUGE_VAL;

#define DENSITY 0.0
#define NFMIN 100
#define NFMAX 300
#define ERR_THRESH 0.1
#define REL_ERR_THRESH inf
#define AMP_THRESH 0.0
#define REL_AMP_THRESH -1.0
#define Q_THRESH 10.0

static void usage(FILE *f) {
  fprintf(f,
          "Usage: harminv [options] <freq-min>-<freq-max>...\n"
          "Options: \n"
          "         -h : this help message\n"
          "         -V : print version number and copyright\n"
          "         -v : verbose output\n"
          "         -T : specify periods instead of frequencies\n"
          "         -w : specify/output angular frequency, not frequency\n"
          "         -n : flip the sign of the frequency convention\n"
          "    -t <dt> : specify sampling interval dt [default: 1]\n"
          "     -d <d> : specify spectral density [default: %g]\n"
          "    -f <nf> : use at least <nf> basis functions [default: %d]\n"
          "  -s <sort> : sort by <sort> = freq/err/decay/amp [default: freq]\n"
          "        -F : discard frequencies outside of specified range\n"
          "    -a <a> : discard amplitudes < max * <a> [default: %e]\n"
          "    -A <A> : discard amplitudes < <A> [default: %g]\n"
          "    -e <e> : discard relative errors > min * <e> [default: %e]\n"
          "    -E <E> : discard relative errors > <E> [default: %e]\n"
          "    -Q <Q> : discard Q > <Q> [default: %g]\n",
          DENSITY, NFMIN, AMP_THRESH, REL_AMP_THRESH, ERR_THRESH, REL_ERR_THRESH, Q_THRESH);
}

#define TWOPI 6.2831853071795864769252867665590057683943388

harminv_data hd;

enum { SORT_FREQUENCY, SORT_DECAY, SORT_ERROR, SORT_AMPLITUDE, SORT_Q } sortby = SORT_FREQUENCY;

static int cmp(double a, double b) { return a > b ? 1 : (a < b ? -1 : 0); }

static int compar(const void *a, const void *b) {
  const int *ia = (const int *)a;
  const int *ib = (const int *)b;
  cmplx aa, ab;

  switch (sortby) {
    case SORT_FREQUENCY: return cmp(harminv_get_freq(hd, *ia), harminv_get_freq(hd, *ib));
    case SORT_DECAY: return cmp(harminv_get_decay(hd, *ia), harminv_get_decay(hd, *ib));
    case SORT_ERROR: return cmp(harminv_get_freq_error(hd, *ia), harminv_get_freq_error(hd, *ib));
    case SORT_AMPLITUDE:
      harminv_get_amplitude(&aa, hd, *ia);
      harminv_get_amplitude(&ab, hd, *ia);
      return cmp(cabs(aa), cabs(ab));
    case SORT_Q:
      return cmp(harminv_get_freq(hd, *ia) / harminv_get_decay(hd, *ia),
                 harminv_get_freq(hd, *ib) / harminv_get_decay(hd, *ib));
  }
  return 0;
}

typedef struct {
  int verbose;
  double fmin, fmax;
  int only_f_inrange;
  double err_thresh, rel_err_thresh, amp_thresh, rel_amp_thresh, Q_thresh;
  double min_err, max_amp;
  int num_ok;
} mode_ok_data;

static int mode_ok(harminv_data d, int k, void *ok_d_) {
  mode_ok_data *ok_d = (mode_ok_data *)ok_d_;
  double errk, ampk, f;
  int ok;
  cmplx aa;

  if (k == -1) { /* initialize */
    int i;
    ok_d->num_ok = 0;
    if (!harminv_get_num_freqs(d)) return 0;
    ok_d->min_err = harminv_get_freq_error(d, 0);
    ;
    harminv_get_amplitude(&aa, d, 0);
    ok_d->max_amp = cabs(aa);
    for (i = 1; i < harminv_get_num_freqs(d); ++i) {
      double err, amp;
      if ((err = harminv_get_freq_error(d, i)) < ok_d->min_err) ok_d->min_err = err;
      harminv_get_amplitude(&aa, d, i);
      if ((amp = cabs(aa)) > ok_d->max_amp) ok_d->max_amp = amp;
    }
    return 0;
  }
  else if (k == -2) { /* finish */
    if (ok_d->verbose && harminv_get_num_freqs(d))
      printf("# harminv: %d/%d modes are ok: "
             "errs <= %e and %e * %e\n, "
             "amps >= %g, %e * %g, "
             "|Q| >= %g\n",
             ok_d->num_ok, harminv_get_num_freqs(d), ok_d->err_thresh, ok_d->rel_err_thresh,
             ok_d->min_err, ok_d->amp_thresh, ok_d->rel_amp_thresh, ok_d->max_amp, ok_d->Q_thresh);
    return 0;
  }

  f = fabs(harminv_get_freq(d, k));
  errk = harminv_get_freq_error(d, k);
  harminv_get_amplitude(&aa, d, k);
  ampk = cabs(aa);

  ok = ((!ok_d->only_f_inrange || (f >= ok_d->fmin && f <= ok_d->fmax)) &&
        errk <= ok_d->err_thresh &&
        (!isfinite(ok_d->rel_err_thresh) || errk <= ok_d->min_err * ok_d->rel_err_thresh) &&
        ampk >= ok_d->amp_thresh && ampk >= ok_d->rel_amp_thresh * ok_d->max_amp &&
        fabs(harminv_get_Q(d, k)) >= ok_d->Q_thresh);

  ok_d->num_ok += ok;

  return ok;
}

#define SOLVE_ONCE_ONLY 0 /* 1 to use harminv_solve_once */
#define SOLVE_OK_ONLY 0   /* 1 for experimental solver */

int main(int argc, char **argv) {
  int verbose = 0;
  int c;
  extern char *optarg;
  extern int optind;
  int specify_periods = 0;
  int specify_omega = 0;
  int negate_omega = 0;
  double dt = 1.0;
  mode_ok_data ok_d;
  int n, nf, nfmin = NFMIN;
  double density = DENSITY;
  int iarg;
  cmplx *data;

  ok_d.only_f_inrange = 0;
  ok_d.err_thresh = ERR_THRESH;
  ok_d.rel_err_thresh = REL_ERR_THRESH;
  ok_d.amp_thresh = AMP_THRESH;
  ok_d.rel_amp_thresh = REL_AMP_THRESH;
  ok_d.Q_thresh = Q_THRESH;

  while ((c = getopt(argc, argv, "hvVTFwnt:d:f:s:e:E:a:A:Q:")) != -1)
    switch (c) {
      case 'h': usage(stdout); return EXIT_SUCCESS;
      case 'V':
        printf("harminv " PACKAGE_VERSION " by Steven G. Johnson\n" COPYRIGHT);
        return EXIT_SUCCESS;
      case 'v': verbose = 1; break;
      case 'T': specify_periods = 1; break;
      case 'w': specify_omega = 1; break;
      case 'n': negate_omega = 1; break;
      case 'F': ok_d.only_f_inrange = 1; break;
      case 'a': ok_d.rel_amp_thresh = atof(optarg); break;
      case 'A': ok_d.amp_thresh = atof(optarg); break;
      case 'E': ok_d.err_thresh = atof(optarg); break;
      case 'e': ok_d.rel_err_thresh = atof(optarg); break;
      case 'Q': ok_d.Q_thresh = atof(optarg); break;
      case 't': dt = atof(optarg); break;
      case 'f': nfmin = atoi(optarg); break;
      case 'd':
        density = atof(optarg);
        if (density < 0) {
          fprintf(stderr, "harminv: -d argument must be >= 0\n");
          return EXIT_FAILURE;
        }
        break;
      case 's':
        switch (tolower(optarg[0])) {
          case 'f': sortby = SORT_FREQUENCY; break;
          case 'd': sortby = SORT_DECAY; break;
          case 'e': sortby = SORT_ERROR; break;
          case 'a': sortby = SORT_AMPLITUDE; break;
          case 'q': sortby = SORT_Q; break;
          default:
            fprintf(stderr, "harminv: invalid sort type -s %c\n", tolower(optarg[0]));
            usage(stderr);
            return EXIT_FAILURE;
        }
        break;
      default:
        fprintf(stderr, "harminv: invalid argument -%c\n", c);
        usage(stderr);
        return EXIT_FAILURE;
    }
  if (optind == argc) { /* no parameters left */
    fprintf(stderr, "harminv: missing required frequency range(s)\n");
    usage(stderr);
    return EXIT_FAILURE;
  }

  /* harminv requires nf > 1 */
  if (nfmin < 2) nfmin = 2;

  data = read_input_data(stdin, &n, verbose);

  if (n < 1) {
    fprintf(stderr, "harminv: no data read\n");
    return EXIT_FAILURE;
  }

  if (verbose) printf("# harminv: %d inputs, dt = %g\n", n, dt);

  printf("frequency, decay constant, Q, amplitude, phase, error\n");

  ok_d.verbose = verbose;

  for (iarg = optind; iarg < argc; ++iarg) {
    double fmin, fmax;
    int i;
    int *isort = NULL;

    if (sscanf(argv[iarg], "%lf-%lf", &fmin, &fmax) != 2) {
      fprintf(stderr, "harminv: invalid argument \"%s\"\n", argv[iarg]);
      return EXIT_FAILURE;
    }
    if (specify_periods) {
      if (fmin == 0 || fmax == 0) {
        fprintf(stderr,
                "harminv: invalid argument \"%s\""
                ": 0 not a valid period\n",
                argv[iarg]);
        return EXIT_FAILURE;
      }
      fmin = 1 / fmin;
      fmax = 1 / fmax;
    }
    if (specify_omega) {
      fmin /= TWOPI;
      fmax /= TWOPI;
    }
    if (negate_omega) dt *= -1;
    if ((fmin > fmax && dt > 0) || (fmin < fmax && dt < 0)) {
      double dummy = fmin;
      fmin = fmax;
      fmax = dummy;
    }
    if (verbose) printf("# searching frequency range %g - %g\n", fmin, fmax);

    ok_d.fmin = fabs(fmin * dt);
    ok_d.fmax = fabs(fmax * dt);
    if (ok_d.fmin > ok_d.fmax) {
      double dummy = ok_d.fmin;
      ok_d.fmin = ok_d.fmax;
      ok_d.fmax = dummy;
    }

    nf = (fmax - fmin) * dt * n * density;
    if (nf > NFMAX) nf = NFMAX;
    if (nf < nfmin) nf = nfmin;
    if (verbose)
      printf("# using %d spectral basis functions, density %g\n", nf,
             nf / ((fmax - fmin) * dt * n));

    hd = harminv_data_create(n, data, fmin * dt, fmax * dt, nf);

#if SOLVE_OK_ONLY
    harminv_solve_ok_modes(hd, mode_ok, &ok_d);
#elif SOLVE_ONCE_ONLY
    harminv_solve_once(hd);
#else
    harminv_solve(hd);
#endif

    mode_ok(hd, -1, &ok_d); /* initialize ok_d */

    CHK_MALLOC(isort, int, harminv_get_num_freqs(hd));
    for (i = 0; i < harminv_get_num_freqs(hd); ++i)
      isort[i] = i;
    qsort(isort, harminv_get_num_freqs(hd), sizeof(int), compar);

    for (i = 0; i < harminv_get_num_freqs(hd); ++i) {
      double freq, decay, err;
      cmplx amp;
      int j = isort[i];
#if SOLVE_OK_ONLY
      CHECK(mode_ok(hd, j, &ok_d), "bug: invalid mode");
#else
      if (!mode_ok(hd, j, &ok_d)) continue;
#endif
      freq = harminv_get_freq(hd, j) / dt;
      decay = harminv_get_decay(hd, j) / fabs(dt);
      harminv_get_amplitude(&amp, hd, j);
      err = harminv_get_freq_error(hd, j);
      printf("%g, %e, %g, %g, %g, %e\n", freq * (specify_omega ? TWOPI : 1.0), decay,
             harminv_get_Q(hd, j), cabs(amp), carg(amp) * (dt < 0 ? -1 : 1), err);
    }

#if !SOLVE_OK_ONLY
    mode_ok(hd, -2, &ok_d);
#endif

    harminv_data_destroy(hd);
  }

  free(data);
  return EXIT_SUCCESS;
}

#ifdef F77_DUMMY_MAIN
#ifdef __cplusplus
extern "C"
#endif
    int
    F77_DUMMY_MAIN() {
  return 1;
}
#endif
