/*
Copyright (C) 1994-1998 Evgenii Rudnyi, http://Evgenii.Rudnyi.Ru/

Implementation of the algorithm described in  
E.B. Rudnyi. Statistical model of systematic errors: linear error model. Chemometrics and Intelligent Laboratory Systems. 1996, V. 34, N 1, p. 41-54. 

This software is a copyrighted work licensed under the terms, described
in the file "FREE_LICENSE". 
*/

#include <iomanip.h>
#include <math.h>
#include "imsl.h"
#include "sumsqr.h"

static int itmp;

void series::init()
{
  hide = false;
  fl_sr2 = same;
  fl_ga = same;
  fl_gb = same;
  isr2 = 0;
  igat = 0;
  igbt = 0;
  adjust = 0.;
  sr2 = 1.;
  sr2old = 1.;
  ga = 0.;
  gaold = 0.;
  gb = 0.;
  gbold = 0.;
  f = NULL;
  mask = 0l;
}

void anal_series()
{
  unsigned long mask;
  int i;
  for (i = 0; i < nparam; i++)
    par[i].use = false;
  for (i = 0; i < nvar; i++)
  {
    Nst[i] = 0;
    Ngat[i] = 0;
    Ngbt[i] = 0;
  }
  iser = ser.begin();
  while(iser != ser.end())
  {
    series& t = *iser;
    if (!t.Ns())
      t.hide = true;
    else if (!t.hide)
    {
      if (!t.Ps())
      {
        t.fl_gb = series::fixed;
        t.gb = 0.;
      }
      if (t.fl_sr2 == series::same)
        Nst[t.isr2]++;
      if (t.fl_ga == series::same)
        Ngat[t.igat]++;
      if (t.fl_gb == series::same)
        Ngbt[t.igbt]++;
      for (i = 0; i < nparam; i++)
      {
        mask = 1l<<i;
        if (t.mask & mask)
          par[i].use = true;
      }
    }
    iser++;
  }
  for (i = 0; i < nparam; i++)
    if (!par[i].use)
      par[i].atr = false;
  for (i = 0; i < nvar; i++)
  {
    if (Nst[i] == 1)
    {
      Nst[i] = 0;
      warning("number of the same sr2[i] is one - set to own");
      iser = ser.begin();
      while(iser != ser.end())
      {
        series& t = *iser;
        if (!t.hide && t.isr2 == i && t.fl_sr2 == series::same)
        {
          t.fl_sr2 = series::own;
          break;
        }
        iser++;
      }
    }
    if (Ngat[i] == 1)
    {
      Ngat[i] = 0;
      warning("number of the same ga[i] is one - set to own");
      iser = ser.begin();
      while(iser != ser.end())
      {
        series& t = *iser;
        if (!t.hide && t.igat == i && t.fl_ga == series::same)
        {
          t.fl_ga = series::own;
          break;
        }
        iser++;
      }
    }
    if (Ngbt[i] == 1)
    {
      Ngbt[i] = 0;
      warning("number of the same gb[i] is one - set to own");
      iser = ser.begin();
      while(iser != ser.end())
      {
        series& t = *iser;
        if (!t.hide && t.igbt == i && t.fl_gb == series::same)
        {
          t.fl_gb = series::own;
          break;
        }
        iser++;
      }
    }
  }
}

void set_sums_1(series& t)
{
  double f;
  int k;
  t.sumi = 0.;
  t.suma = 0.;
  t.sumb = 0.;
  for (k = 0; k < t.Ntot(); k++)
  {
    if (t.atr(k))
    {
      f = (*t.f)(t(k));
      t.sumi += f*f;
      t.suma += f;
      if (t.Ps())
        t.sumb += f*(t(k, 1) - t.xs());
    }
  }
  if (t.adjust)
  {
    if (t.Ps())
      t.sumi += t.adjust*(t.Ns()-2);
    else if (t.Ns() > 0)
      t.sumi += t.adjust*(t.Ns()-1);
  }
}

void set_sums(series& t, double* f, int& i)
{
  int is = i;
  double da, db;
  int k;
  t.sumi = 0.;
  t.suma = 0.;
  t.sumb = 0.;
  for (k = 0; k < t.Ntot(); k++)
  {
    if (t.atr(k))
    {
      f[i] = (*t.f)(t(k));
      t.sumi += f[i]*f[i];
      t.suma += f[i];
      if (t.Ps())
        t.sumb += f[i]*(t(k, 1) - t.xs());
      i++;
    }
  }
  da = (1. - 1./sqrt(1. + t.Ns()*t.ga))/t.Ns();
  if (t.Ps())
    db = (1. - 1./sqrt(1. + t.Ps()*t.gb))/t.Ps();
  else
    db = 0.;
  t.SS = 0.;
  for (k = 0; k < t.Ntot(); k++)
  {
    if (t.atr(k))
    {
      f[is] -= da*t.suma;
      if (t.Ps())
        f[is] -= db*(t(k, 1) - t.xs())*t.sumb;
      f[is] /= sqrt(t.sr2);
      t.SS += f[is]*f[is];
      is++;
    }
  }
  if (t.adjust)
  {
    if (t.Ps())
    {
      t.sumi += t.adjust*(t.Ns()-2);
      t.SS += t.adjust*(t.Ns()-2);
    }
    else if (t.Ns() > 0)
    {
      t.sumi += t.adjust*(t.Ns()-1);
      t.SS += t.adjust*(t.Ns()-1);
    }
  }
}

extern "C" {
void ss(double* x, int *m_, int *n_, double* f)
{
	int &m = *m_;
	int &n = *n_;
  int i = 0, j;
  for (j = 0; j < nparam; j++)
  {
    if (par[j].atr)
    {
      par[j].x = x[i];
      i++;
    }
  }
  init_est();
  i = 0;
  iser = ser.begin();
  while(iser != ser.end())
  {
    series& t = *iser;
    if (!t.hide)
      set_sums(t, f, i);
    iser++;
  }
  iss++;
  if (print == all_iter)
  {
    of << "ZXSSQ function eval. " << setw(3) << iss << ", ";
    print_ss_par_1();
  }
  cout << "\b\b\b\b" << setw(4) << iss;
}
}

void eval_L(double& L, double& SS)
{
  L = 0.;
  SS = 0.;
  iser = ser.begin();
  while(iser != ser.end())
  {
    series& t = *iser;
    if (!t.hide)
    {
      L -= t.Ns()*log(t.sr2) + log(1. + t.Ns()*t.ga) + log(1. + t.Ps()*t.gb);
      SS += (t.sumi - t.ga/(1. + t.Ns()*t.ga)*t.suma*t.suma
                    - t.gb/(1. + t.Ps()*t.gb)*t.sumb*t.sumb)/t.sr2;
    }
    iser++;
  }
  L -= SS;
}

extern "C" {

double gat_solve(double *x_)
{
	double &x = *x_;
  if (x < 0.) return 0.;
  double f = 0.;
  int i = itmp;
  iser = ser.begin();
  while (iser != ser.end())
  {
    series& t = *iser;
    if (t.fl_ga == series::same && !t.hide && t.igat == i)
    {
      if (t.fl_sr2 == series::same)
        f += t.suma*t.suma/sr2t[t.isr2]/pow((1. + t.Ns()*x), 2.)
                      - t.Ns()/(1. + t.Ns()*x);
      else
        f += t.suma*t.suma/t.sr2/pow((1. + t.Ns()*x), 2.)
                      - t.Ns()/(1. + t.Ns()*x);
    }
    iser++;
  }
  if (print == all_iter)
    of << "gat_solve " << i << setw(wdth) << x
       << setw(wdth) << f << endl;
  return f;
}

double gbt_solve(double *x_)
{
	double &x = *x_;
  if (x < 0.) return 0.;
  double f = 0.;
  int i = itmp;
  iser = ser.begin();
  while (iser != ser.end())
  {
    series& t = *iser;
    if (t.fl_gb == series::same && !t.hide && t.igbt == i)
    {
      if (t.fl_sr2 == series::same)
        f += t.sumb*t.sumb/sr2t[t.isr2]/pow((1. + t.Ps()*x), 2.)
                      - t.Ps()/(1. + t.Ps()*x);
      else
        f += t.sumb*t.sumb/t.sr2/pow((1. + t.Ps()*x), 2.)
                      - t.Ps()/(1. + t.Ps()*x);
    }
    iser++;
  }
  if (print == all_iter)
    of << "gbt_solve " << i << setw(wdth) << x
       << setw(wdth) << f << endl;
  return f;
}
}

void eval_var_comp()
{
  double L = 1.;
  double Lold = 0.;
  double SS;
  double epsL = 1e-7;
// for ZREAL2 and ZBRENT
  double eps = 0., eps2 = 0., eta = 0.;
  double x;
  int nsig = 6, n = 1, itmax = 50, ier;
  double a, b;
  int i;

  int iter = 0;

  while (fabs((L - Lold)/L) > epsL && iter < 50)
  {
    iter++;
    Lold = L;
    for (i = 0; i < nvar; i++)
    {
      sr2t[i] = 0.;
      Nst[i] = 0;
    }
    iser = ser.begin();
    while (iser != ser.end())
    {
      series& t = *iser;
      if (!t.hide)
      {
        switch (t.fl_sr2)
        {
          case series::fixed:
            break;
          case series::same:
            sr2t[t.isr2] += t.sumi - t.ga/(1. + t.Ns()*t.ga)*t.suma*t.suma
                            - t.gb/(1. + t.Ps()*t.gb)*t.sumb*t.sumb;
            Nst[t.isr2] += t.Ns();
            break;
          case series::own:
            t.sr2 = (t.sumi - t.ga/(1. + t.Ns()*t.ga)*t.suma*t.suma
                            - t.gb/(1. + t.Ps()*t.gb)*t.sumb*t.sumb)/t.Ns();
            break;
          default:
            warning("check fl._sr2");
        }
      }
      iser++;
    }
    for (i = 0; i < nvar; i++)
      if (Nst[i])
        sr2t[i] /= Nst[i];

    for (i = 0; i < nvar; i++)
    {
      itmp = i;
      if (Ngat[i] > 1)
      {
        if (gat[i])
        {
          itmax = 50;
          x = gat[i];
          zreal2_(gat_solve, &eps, &eps2, &eta, &nsig, &n, &x, &itmax, &ier);
          if (ier)
          {
            if (print == all_iter)
              warning("can't find solution for gat[i]");
            gat[i] = 0.;
          }
          else
            gat[i] = x;
          if (gat[i] < 0. || gat[i] > 100000.) gat[i] = 0.;
        }
        if (!gat[i])
        {
          itmax = 50;
          a = 0.;
          b = 100000.;
          zbrent_(gat_solve, &eps, &nsig, &a, &b, &itmax, &ier);
          if (ier == 130)
            gat[i] = 0.;
          else
            gat[i] = b;
        }
      }

      if (Ngbt[i] > 1)
      {
        if (gbt[i])
        {
          itmax = 50;
          x = gbt[i];
          zreal2_(gbt_solve, &eps, &eps2, &eta, &nsig, &n, &x, &itmax, &ier);
          if (ier)
          {
            if (print == all_iter)
              warning("can't find solution for gbt[i]");
            gbt[i] = 0.;
          }
          else
            gbt[i] = x;
          if (gbt[i] < 0. || gbt[i] > 1e9) gbt[i] = 0.;
        }
        if (!gbt[i])
        {
          itmax = 50;
          a = 0.;
          b = 1e9;
          zbrent_(gbt_solve, &eps, &nsig, &a, &b, &itmax, &ier);
          if (ier == 130)
            gbt[i] = 0.;
          else
            gbt[i] = b;
        }
      }
    }

    iser = ser.begin();
    while (iser != ser.end())
    {
      series& t = *iser;
      if (!t.hide)
      {
        if (t.fl_sr2 == series::same)
          t.sr2 = sr2t[t.isr2];
        if (t.Ps())
        {
          switch (t.fl_gb)
          {
            case series::fixed:
              break;
            case series::same:
              t.gb = gbt[t.igbt];
              break;
            case series::own:
              t.gb = (t.sumb*t.sumb/t.Ps()/t.sr2 - 1.)/t.Ps();
              break;
            default:
              warning("check fl_gb");
          }
        }
        switch (t.fl_ga)
        {
          case series::fixed:
            break;
          case series::same:
            t.ga = gat[t.igat];
            break;
          case series::own:
            t.ga = (t.suma*t.suma/t.Ns()/t.sr2 - 1.)/t.Ns();
            break;
          default:
            warning("check fl_ga");
        }
        if (t.ga < 0.) t.ga = 0;
        if (t.gb < 0.) t.gb = 0;
      }
      iser++;
    }
    eval_L(L, SS);
    if (print == all_iter)
    {
      of << "var_comp - iteration " << setw(4) << iter << endl;
      of << "value of L is " << setw(15) << L
         << "  and of SS is " << setw(15)
         << SS << endl;
    }
    cout << "\b\b\b\b" << setw(4) << iter;
  }
  if (print != results)
  {
    of << "New variance components after " << setw(4) << iter
       << "  iterations" << endl;
    if ( print != all_iter)
      of << "value of L is " << setw(15) << L
         << "  and of SS is " << setw(15)
         << SS << endl;
    init_est();
    print_all_points();
    post_analysis();
  }
}