/*! \file sigfit.cc
 * \brief fit signal by trial-signals
 * 
 * ----------------------------------------------------------------------------
 * 
 * $Id: sigfit.cc,v 1.19 2004-09-24 15:36:33 tforb Exp $
 * \author Thomas Forbriger
 * \date 28/01/2004
 * 
 * fit signal by trial-signals
 * 
 * Copyright (c) 2004 by Thomas Forbriger (BFO Schiltach) 
 * 
 * REVISIONS and CHANGES 
 *  - 28/01/2004   V1.0   Thomas Forbriger
 *  - 24/02/2004   V1.1   introduced search range stabilization
 *                        more output values
 *                 V1.2   correction: searchfac is a factor!
 *  - 16/03/2004   V1.3   add trend and offset fit
 *  - 24/09/2004   V1.4   
 *                        - allow to skip samples if this is helpful to remove
 *                          transient filter loading response
 *                        - remove correction from synthetics and signal
 *                          and make this process transparent
 *                        - provide exponential decay
 * 
 * ============================================================================
 */
#define SIGFIT_VERSION \
  "SIGFIT   V1.4   fit signal by trial-signals"
#define SIGFIT_CVSID \
  "$Id: sigfit.cc,v 1.19 2004-09-24 15:36:33 tforb Exp $"

#include <fstream>
#include <iostream>
#include <list>
#include <vector>
#include <cmath>
#include <tfxx/commandline.h>
#include <tfxx/error.h>
#include <sffxx.h>
#include <sffostream.h>
#include <tsxx/tsxx.h>
#include <tsxx/innerproduct.h>
#include <aff/array.h>
#include <aff/dump.h>
#include <aff/seriesoperators.h>
#include <linearxx/lapackxx.h>

typedef std::list<std::string> Tnamelist;
typedef ts::TDsfftimeseries Tbundle;
typedef Tbundle::Tseries Tseries;
typedef std::vector<Tbundle> Tbundlevec;
typedef aff::Array<Tbundle::Tseries::Tvalue> Tmatrix;

using std::cout;
using std::cerr;
using std::endl;

// structure to store commandline options
struct Options {
  bool verbose;
  double Tdate;
  bool truncate;
  std::string residualname;
  bool writeresidual;
  double searchrange;
  bool usesearchrange;
  bool equalsearch;
  bool fittrend, fitoffset, doskip, fitexp;
  double amptrend, ampoffset, skip, tcexp;
}; // struct Options

/*======================================================================*/
/* functions */

// remove average
void removeavg(Tseries& s)
{
  double avg=0.;
  for (int i=s.first(); i<=s.last(); ++i) { avg += s(i); }
  avg /= s.size();
  for (int i=s.first(); i<=s.last(); ++i) { s(i) -= avg; }
}

// remove trend
void removetre(Tseries& s)
{
  double sum=0.;
  double jser=0.;
  double jqser=0.;
  for (int i=s.first(); i<=s.last(); ++i) 
  { 
    sum += s(i); 
    jser += i;
    jqser += i*i;
  }
  //double a,b;
  sum /= s.size();
  for (int i=s.first(); i<=s.last(); ++i) { s(i) -= sum; }
}

// formatted number output
std::string formatfloat(const double &v)
{
  char seq[20];
  std::sprintf(seq, "%8.5f ", v);
  return(std::string(seq));
}

/*======================================================================*/
/* main progran */

int main(int iargc, char* argv[])
{

  // define usage information
  char usage_text[]=
  {
    SIGFIT_VERSION "\n"
    "usage: sigfit [-v] [-Tdate v] [-truncate]" "\n"
    "              [-Sramp[=v]] [-Sconst[=v]] [-Sexp[=v]]" "\n"
    "              [-residual f]" "\n"
    "              [-searchrange[=r]] [-equalsearch] [-skip n]" "\n"
    "              signal trial [trial ...]" "\n"
    "   or: sigfit --help|-h" "\n"
  };

  // define full help text
  char help_text[]=
  {
    SIGFIT_CVSID "\n"
    "\n"
    "-v           be verbose" "\n"
    "-Sramp[=v]   add a ramp to the set of trial signals with amplitude v" "\n"
    "-Sconst[=v]  add a constant (of value v) to the set of trial signals" "\n"
    "-Sexp[=v]    add a exponential decay (with time constant v relative" "\n"
    "             to the length of the time series) to the set of trial" "\n"
    "             signals" "\n"
    "-Tdate v     tolerance for comparison of date of first sample" "\n"
    "             v give the tolerance in units of the sampling interval" "\n"
    "-truncate    truncate all series to the same number of samples" "\n"
    "-residual f  write waveforms containing residual to file \"f\"" "\n"   
    "-searchrange[=r] set search range to stabilize the fit in cases" "\n"
    "             where a trial series does not (or almost not)" "\n"
    "             contribute to the signal" "\n"
    "             the search range r is given in units of" "\n"
    "             rms(signal)/rms(trial signals)" "\n"
    "-equalsearch use same search range for each trial signal" "\n"
    "-skip n      skip n seconds at beginning of each trace" "\n"
    "\n"
    "signal       signal to by explained by a linear combination" "\n"
    "             of trial signals" "\n"
    "trial        any set of trial signals"
  };

  // define commandline options
  using namespace tfxx::cmdline;
  static Declare options[]= 
  {
    // 0: print help
    {"help",arg_no,"-"},
    // 1: verbose mode
    {"v",arg_no,"-"},
    // 2: date tolerance
    {"Tdate",arg_yes,"0."},
    // 3: truncate to a common number of samples
    {"truncate",arg_no,"-"},
    // 4: output SFF file
    {"residual",arg_yes,"-"},
    // 5: search range
    {"searchrange",arg_opt,"100."},
    // 6: equal search ranges
    {"equalsearch",arg_no,"-"},
    // 7: fit a trend
    {"Sramp",arg_opt,"1."},
    // 8: fit an offset
    {"Sconst",arg_opt,"1."},
    // 9: fit an offset
    {"skip",arg_yes,"0"},
    // 10: fit an offset
    {"Sexp",arg_opt,"1."},
    {NULL}
  };

  // no arguments? print usage...
  if (iargc<2) 
  {
    cerr << usage_text << endl;
    exit(0);
  }

  // collect options from commandline
  Commandline cmdline(iargc, argv, options);

  // help requested? print full help text...
  if (cmdline.optset(0))
  {
    cerr << usage_text << endl;
    cerr << help_text << endl;
    exit(0);
  }

  // read options
  Options opt;
  opt.verbose=cmdline.optset(1);
  opt.Tdate=cmdline.double_arg(2);
  opt.truncate=cmdline.optset(3);
  opt.writeresidual=cmdline.optset(4);
  opt.residualname=cmdline.string_arg(4);
  opt.usesearchrange=cmdline.optset(5);
  opt.searchrange=cmdline.double_arg(5);
  opt.equalsearch=cmdline.optset(6);
  opt.fittrend=cmdline.optset(7);
  opt.fitoffset=cmdline.optset(8);
  opt.amptrend=cmdline.double_arg(7);
  opt.ampoffset=cmdline.double_arg(8);
  opt.skip=cmdline.double_arg(9);
  opt.doskip=cmdline.optset(9);
  opt.fitexp=cmdline.optset(10);
  opt.tcexp=cmdline.double_arg(10);

  /*----------------------------------------------------------------------*/

  // read file names
  TFXX_assert(cmdline.extra(),"ERROR: missing signal file name!");
  std::string signalname=cmdline.next();

  TFXX_assert((cmdline.extra()||opt.fittrend||opt.fitoffset),
              "ERROR: missing trial signal file name!");
  Tnamelist namelist;
  while (cmdline.extra()) { namelist.push_back(cmdline.next()); }

  // read signal file
  Tbundle signal;
  {
    if (opt.verbose)
    { cout << "read signal file \"" << signalname << "\"" << endl; }
    std::ifstream is(signalname.c_str());
    sff::FileHeader fileheader(is);
    if (opt.verbose)
    { cout << "  read trace" << endl; }
    sff::InputWaveform<Tseries> inputwaveform;
    try {
      inputwaveform.read(is);
    } 
    catch(...) {
      cerr << "ERROR (sigfit): read failed!" << endl;
      throw;
    }
    sff::TraceHeader traceheader=inputwaveform.header();
    signal.header=traceheader.wid2();
    signal=inputwaveform.series();
    bool last=traceheader.last();
    std::string wid2line=signal.header.line();
    if (opt.verbose)
    { 
      cout << "  " << wid2line.substr(0,69) << endl; 
      if (!last)
      {
        cout << "there are more traces in that files" << endl 
          << "ignoring them..." << endl;
      }
    }
  }
  
  // read trial files
  if (opt.verbose)
  { cout << namelist.size() << " trial data files to read" << endl; }
  Tbundlevec bundlevec;
  for (Tnamelist::const_iterator i=namelist.begin(); i!=namelist.end(); i++)
  {
    if (opt.verbose)
    { cout << "read trial data file \"" << *i << "\"" << endl; }
    std::ifstream is(i->c_str());
    sff::FileHeader fileheader(is);
    bool last=false;
    while(!last)
    {
      if (opt.verbose)
      { cout << "  read trial signal #" << bundlevec.size()+1 << endl; }
      sff::InputWaveform<Tseries> inputwaveform(is);
      sff::TraceHeader traceheader=inputwaveform.header();
      Tbundle bundle;
      bundle.header=traceheader.wid2();
      bundle=inputwaveform.series();
      bundlevec.push_back(bundle);
      last=traceheader.last();
      std::string wid2line=bundle.header.line();
      if (opt.verbose)
      { cout << "  " << wid2line.substr(0,69) << endl; }
    }
  }
  unsigned int ntracesread=bundlevec.size();

  // add synthetic trial files
  if (opt.fitoffset)
  {
    if (opt.verbose) { cout << "add offset to trial signals" << endl; }
    Tbundle synsignal;
    synsignal.header=signal.header;
    synsignal=Tseries(signal.shape());
    synsignal.series()=opt.ampoffset;
    synsignal.header.station="NSP";
    synsignal.header.channel="off";
    synsignal.header.auxid="NSP";
    synsignal.header.instype="NSP";
    bundlevec.push_back(synsignal);
  }

  if (opt.fittrend)
  {
    if (opt.verbose) { cout << "add trend to trial signals" << endl; }
    Tbundle synsignal;
    synsignal.header=signal.header;
    synsignal=Tseries(signal.shape());
    for (int i=synsignal.first(); i<=synsignal.last(); ++i)
    {
      synsignal(i)=2.*opt.amptrend*
        (i-(0.5*(synsignal.first()+synsignal.last())))/
        synsignal.size();
    }
    synsignal.header.station="NSP";
    synsignal.header.channel="ramp";
    synsignal.header.auxid="NSP";
    synsignal.header.instype="NSP";
    bundlevec.push_back(synsignal);
  }

  if (opt.fitexp)
  {
    if (opt.verbose) 
    {
      cout << "add exponential decay to trial signals" << endl; 
    }
    Tbundle synsignal;
    synsignal.header=signal.header;
    synsignal=Tseries(signal.shape());
    for (int i=synsignal.first(); i<=synsignal.last(); ++i)
    {
      double argval=-1.*double(i-synsignal.first())/
        (double(synsignal.size())*opt.tcexp);
      synsignal(i)=std::exp(argval);
    }
    synsignal.header.station="NSP";
    synsignal.header.channel="exp";
    synsignal.header.auxid="NSP";
    synsignal.header.instype="NSP";
    bundlevec.push_back(synsignal);
  }

  /*----------------------------------------------------------------------*/

  // truncate length if requested
  if (opt.truncate)
  {
    if (opt.verbose) { cout << "truncate signals if necessary..." << endl; }
    long int n=signal.last();
    for (Tbundlevec::const_iterator i=bundlevec.begin(); 
         i!=bundlevec.end(); i++)
    { n=n<i->last() ? n : i->last(); }
    if (signal.last()>n)
    {
      signal.setlastindex(n);
      signal.header.nsamples=signal.size();
      if (opt.verbose) { cout << "truncate signal to "
        << signal.header.nsamples << " samples" << endl; }
    }
    for (Tbundlevec::iterator i=bundlevec.begin(); 
         i!=bundlevec.end(); i++)
    { 
      if (i->last()>n)
      {
        i->setlastindex(n);
        i->header.nsamples=signal.size();
        if (opt.verbose) { cout << "truncate trial signal" << endl
          << i->header.line() << endl; }
      }
    }
  }

  /*----------------------------------------------------------------------*/

  // check header consistency
  sff::WID2compare compare(sff::Fnsamples | sff::Fdt | sff::Fdate);
  compare.setdatetolerance(opt.Tdate);
  if (opt.verbose) { cout << "checking consistency..." << endl; }
  for (Tbundlevec::const_iterator i=bundlevec.begin(); i!=bundlevec.end(); i++)
  {
    if (!compare (i->header,signal.header))
    {
      cerr << "ERROR: header signature mismatch:" << endl;
      cerr << "signal:" << endl;
      cerr << signal.header.line();
      cerr << "trial signal:" << endl;
      cerr << i->header.line();
      TFXX_abort("baling out...");
    }
  }

  /*----------------------------------------------------------------------*/

  // skip samples if requested
  if (opt.doskip) 
  {
    if (opt.verbose)
    {
      cout << "skip " << opt.skip << " seconds" << endl;
    }
    int nskip=int(floor(opt.skip/signal.header.dt));
    if (nskip>0)
    {
      libtime::TRelativeTime add=nskip*libtime::double2time(signal.header.dt);
      signal.setfirstindex(signal.first()+nskip);
      signal.header.date+=add;
      for (Tbundlevec::iterator i=bundlevec.begin(); 
           i!=bundlevec.end(); i++)
      {
        i->setfirstindex(i->first() + nskip);
        i->header.date += add;
      }
      if (opt.verbose)
      {
        cout << "skipped " << nskip << " samples ("
          << add.timestring() << ")" << endl;
      }
    }
    else
    {
      if (opt.verbose)
      {
        cout << "NOTICE: nothing to skip..." << endl;
        cout << "  seems like " 
          << opt.skip << " seconds to skip < " 
          << signal.header.dt << " seconds sampling interval" << endl;
      }
    }
  }

  /*----------------------------------------------------------------------*/
  double signalrms=ts::rms(signal);
  
  // set up system of linear equations
  if (opt.verbose) { cout << "set up system of linear equations..." << endl; }
  int N=bundlevec.size();
  if (opt.verbose) { cout << "system is of size " << N << "x" << N << endl; }
  Tmatrix Matrix(N,N), rhs(N), coeff(N), normproduct(N), trialrms(N);
  double fulltrialrms=0.;
  for (int i=1; i<=N; ++i)
  {
    for (int k=i; k<=N; ++k)
    {
      Matrix(i,k)=ts::innerproduct(bundlevec[i-1], bundlevec[k-1]);
      Matrix(k,i)=Matrix(i,k);
    }
    rhs(i)=ts::innerproduct(bundlevec[i-1], signal);
    trialrms(i)=ts::rms(bundlevec[i-1]);
    fulltrialrms+=(trialrms(i)*trialrms(i));
    normproduct(i)=rhs(i)/(signalrms*trialrms(i)*signal.size());
  }
  fulltrialrms=sqrt(fulltrialrms/double(N));

  // add stabilization
  if (opt.usesearchrange)
  {
    if (opt.verbose) 
    {
      cout << "stabilize fit:" << endl; 
      cout << "  relative search range: " << opt.searchrange << endl;
    }
    double searchfac=signal.size()*signalrms;
    for (int i=1; i<=N; ++i)
    {
      double range;
      if (opt.equalsearch)
      { range=opt.searchrange*signalrms/fulltrialrms; }
      else
      { range=opt.searchrange*signalrms/trialrms(i); }
      if (opt.verbose)
      {
        cout << "  search range for trial series " << i 
          << " is " << range << endl;
      }
      Matrix(i,i)+=pow((searchfac/range),2);
    }
  }

  // solve
  coeff=linear::lapack::dposv(Matrix,rhs);

  // set up synthetics
  Tbundle synthetics;
  synthetics=Tseries(signal.shape());
  synthetics.series()=0.;
  for (int i=1; i<=N; ++i)
  { synthetics += coeff(i) * bundlevec[i-1]; }
  synthetics.header=signal.header;
  synthetics.header.channel="synt";

  // set up correction
  Tbundle correction;
  correction=Tseries(signal.shape());
  correction.series()=0.;
  if (ntracesread<bundlevec.size())
  {
    for (int i=ntracesread; i<N; ++i)
    { correction += coeff(i+1) * bundlevec[i]; }
  }
  correction.header=signal.header;
  correction.header.channel="corr";
  Tbundle corrsignal;
  corrsignal.header=signal.header;
  corrsignal.header.channel="csig";
  corrsignal=signal-correction;
  Tbundle corrsynthetics;
  corrsynthetics.header=signal.header;
  corrsynthetics.header.channel="csyn";
  corrsynthetics=synthetics-correction;

  // set up residual
  Tbundle residual;
  residual=Tseries(signal.shape());
  residual=signal-synthetics;
  residual.header=signal.header;
  residual.header.channel="diff";

  // write waveforms
  if (opt.writeresidual)
  {
    if (opt.verbose) 
    {
      cout << "write residual waveform to " 
      << opt.residualname << endl;
    }
    std::ofstream ofs(opt.residualname.c_str());
    sff::SFFostream<Tseries> os(ofs);
    os << signal;
    os << synthetics;
    os << residual;
    for (Tbundlevec::const_iterator i=bundlevec.begin(); 
         i!=bundlevec.end(); i++)
    { os << *i; }
    for (int i=1; i<=N; ++i)
    {
      Tbundle psyn=bundlevec[i-1];
      psyn=bundlevec[i-1].series()*coeff(i);
      char chan[6];
      std::sprintf(chan,"s%1.1d", i);
      psyn.header.channel=std::string(chan);
      os << psyn;
    }
    os << corrsignal;
    os << corrsynthetics;
    os << correction;
  }
  
  // calculate rms values
  double residualrms=ts::rms(residual);
  double explained=1.-(residualrms/signalrms);

  // calculate direction cosines
  Tmatrix normcoeff(N);
  double length=0.;
  for (int i=1; i<=N; ++i)
  { length += coeff(i)*coeff(i); }
  length=sqrt(length);
  for (int i=1; i<=N; ++i)
  { normcoeff(i) = coeff(i)/length; }
  
  std::string timestring(signal.header.date.timestring());
  cout << "         time: " << timestring.substr(4,21) << endl;
  cout << "      channel: " << signal.header.channel << endl;
  cout << "      station: " << signal.header.station << endl;
  cout << "   instrument: " << signal.header.instype << endl;
  cout << "    signalrms: " << signalrms << endl;
  cout << "  residualrms: " << residualrms << endl;
  cout << "explained rms: " << explained << endl;
  cout << " coefficients: ";
  for (int i=1; i<=N; ++i)
  { cout << coeff(i) << "  "; }
  cout << endl;
  cout << " rms coeffic.: " << length << endl;
  cout << " normalized coefficients: ";
  for (int i=1; i<=N; ++i)
  { cout << normcoeff(i) << "  "; }
  cout << endl;
  cout << " coefficients x rms: ";
  for (int i=1; i<=N; ++i)
  { 
    cout << coeff(i)*trialrms(i) << "  "; 
  }
  cout << endl;
  cout << " normalized scalar product: ";
  for (int i=1; i<=N; ++i)
  { 
    cout << normproduct(i) << "  "; 
  }
  cout << endl;

  // reportline
  cout << signal.header.station << " "
    << signal.header.channel << " "
    << signal.header.instype << " "
    << timestring << " ";
  cout << formatfloat(explained);
  cout << formatfloat(signalrms);
  cout << formatfloat(residualrms);
  for (int i=1; i<=N; ++i)
  { 
    cout << formatfloat(coeff(i));
  }
  cout << formatfloat(length);
  for (int i=1; i<=N; ++i)
  { 
    cout << formatfloat(normcoeff(i));
  }
  for (int i=1; i<=N; ++i)
  { 
    cout << formatfloat(coeff(i)*trialrms(i)); 
  }
  for (int i=1; i<=N; ++i)
  { 
    cout << formatfloat(normproduct(i)); 
  }
  cout << endl;
}

/* ----- END OF sigfit.cc ----- */