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

Copyright (C) 2004, 2006.

A. Ronald Gallant
Post Office Box 659
Chapel Hill NC 27514-0659
USA   

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.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

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

#include <sstream>
#include "snp.h"
#include "squash.h"

using namespace scl;
using namespace libsnp;
using namespace snp;
using namespace std;

snp::control::control(std::istream& ctrl_stream) 
: control_stream(ctrl_stream) { }

bool snp::control::read_line()
{
  control_stream >> input_file;
  control_stream >> output_file;
  control_stream >> fnew;
  control_stream >> fold;
  control_stream >> nstart;
  control_stream >> jseed;
  return control_stream.good();
}

string snp::control::get_line()
{
  string ctrl_line;
  ctrl_line += input_file;
  while (ctrl_line.size()<20) ctrl_line = " " + ctrl_line;
  ctrl_line += " ";
  while (output_file.size()+ctrl_line.size()<41) ctrl_line += " ";
  ctrl_line += output_file;
  ctrl_line += fmt('e',10,2,fnew).get_ostr();
  ctrl_line += fmt('e',10,2,fold).get_ostr();
  ctrl_line += fmt('d',7,nstart).get_ostr();
  ctrl_line += fmt('d',9,jseed).get_ostr();
  return ctrl_line;
}

string snp::control::get_input_file()
{
  return input_file;
}

string snp::control::get_output_file()
{
  return output_file;
}

REAL snp::control::get_fnew()
{
  return fnew;
}

REAL snp::control::get_fold()
{
  return fold;
}

INTEGER snp::control::get_nstart()
{
  return nstart;
}

INT_32BIT snp::control::get_jseed()
{
  return jseed;
}

snp::trnfrm::trnfrm(const datparms& dp,const tranparms& tp,const realmat& data)
: dpm(dp), tpm(tp), R(realmat(dp.M,dp.M,0.0)), P(realmat(dp.M,dp.M,0.0))
{
  if (!tpm.useold) {
    fill(tpm.mean);
    for (INTEGER t=1 ; t<=data.ncol(); ++t) {
      for (INTEGER i=1; i<=dpm.M; ++i) {
        tpm.mean[i] += data(i,t);
      }
    }
    tpm.mean = tpm.mean/data.ncol();

    fill(tpm.variance);
    for (INTEGER t=1 ; t<=data.ncol(); ++t) {
      for (INTEGER j=1; j<=dpm.M; ++j) {
        for (INTEGER i=1; i<=dpm.M; ++i) {
          tpm.variance(i,j) += (data(i,t)-tpm.mean[i])*(data(j,t)-tpm.mean[j]);
        }
      }
    }
    tpm.variance = tpm.variance/data.ncol();

    if (tpm.diag) {
      for (INTEGER j=1; j<=dpm.M; ++j) {
        for (INTEGER i=1; i<=dpm.M; ++i) {
          if (i!=j) tpm.variance(i,j) = 0.0;
        }
      }
    }
  }

  if (dpm.M == 1) {
    R[1] = sqrt(tpm.variance[1]);
    P[1] = 1.0/R[1];
    detP = P[1];
  }
  else if (tpm.diag) {
    detP = 1.0;
    for (INTEGER i=1; i<=dpm.M; ++i) {
      R(i,i) = sqrt(tpm.variance(i,i));
      P(i,i) = 1.0/R(i,i);
      detP *= P(i,i);
    }
  }
  else {
    detP = 1.0;
    realmat U,S,V;
    INTEGER rank = svd(tpm.variance,U,S,V);
    if (rank != dpm.M) warn("Error, tran, data appear to be singular");
    for (INTEGER i=1; i<=dpm.M; ++i) {
      S[i] = sqrt(S[i]);
    }
    R = U*diag(S);

    for (INTEGER i=1; i<=dpm.M; ++i) {
      S[i] = 1.0/S[i];
      detP *= S[i];
    }
    P = diag(S)*T(U);
  }
}

snp::trnfrm::trnfrm() 
: dpm(), tpm(), R(), P(), detP() 
{ }

snp::trnfrm::trnfrm(const trnfrm& tr)
: dpm(tr.dpm), tpm(tr.tpm), R(tr.R), P(tr.P), detP(tr.detP)
{ }

trnfrm& snp::trnfrm::operator=(const trnfrm& tr)
{
  if (this != &tr) {
    dpm=tr.dpm; tpm=tr.tpm; R=tr.R; P=tr.P; detP=tr.detP;
  }
  return *this;
}

void snp::trnfrm::normalize(realmat& y) const
{
  for (INTEGER t=1; t<=y.ncol(); ++t) {
    realmat z = P*(y("",t)-tpm.mean);
    for (INTEGER i=1; i<=dpm.M; ++i) {
      y(i,t) = z[i];
    }
  }
}    

void snp::trnfrm::unnormalize(realmat& y) const
{
  for (INTEGER t=1; t<=y.ncol(); ++t) {
    realmat z = R*y("",t) + tpm.mean;
    for (INTEGER i=1; i<=dpm.M; ++i) {
      y(i,t) = z[i];
    }
  }
}    

void snp::trnfrm::unscale(realmat& sigma) const
{
  sigma = R*sigma*T(R);
}

void snp::trnfrm::spline(realmat& x) const
{
  for (INTEGER i=1; i<=x.size(); ++i) {
     ave0(x[i],tpm.inflec);
  }
}

void snp::trnfrm::logistic(realmat& x) const
{
  for (INTEGER i=1; i<=x.size(); ++i) {
     gst0(x[i],tpm.inflec);
  }
}

snp::sumry::sumry(ostream& sumry_os)
: os(sumry_os), best_sn(REAL_MAX)
{ }

void snp::sumry::set_filename(const std::string& fn)
{
  filename = fn;
}

void snp::sumry::set_pname(const string pn)
{
  pname = pn;
}

void snp::sumry::set_ll_ptr(const snpll* llptr)
{
  ll_ptr = llptr;
}

void snp::sumry::set_n(INTEGER in)
{
  n = in;
}

void snp::sumry::set_sn(REAL isn)
{
  sn = isn;
}

void snp::sumry::set_iter(INTEGER itr)
{
  iter = itr;
}

void snp::sumry::set_termination_code(INTEGER code)
{
   termination_code = code;
}

bool snp::sumry::write_header()
{
  os << "SNP Restart\n"
     << "--------------------------------------------------"
     << "-----------------------------\n"
     << "             output_file      fitted_model   p opt iter"
     << "   obj_func        bic  \n"
     << "--------------------------------------------------"
     << "-----------------------------\n";
  return os.good();
}

bool snp::sumry::write_line()
{
  INTEGER p = (*ll_ptr).get_rho().size();
  REAL bic = sn + (REAL(p)/REAL(n))*log(REAL(n))/2.0;
  if (filename.size()< 24) {
    for (size_t i=1; i<=24-filename.size(); ++i) os << " ";
  } 
  else {
    os << " ";
  }
  os << filename << "  "
     << (*ll_ptr).get_ufunc().get_lag()
     << (*ll_ptr).get_rfunc().get_q()
     << (*ll_ptr).get_rfunc().get_Qtype()
     << (*ll_ptr).get_rfunc().get_p()
     << (*ll_ptr).get_rfunc().get_Ptype()
     << (*ll_ptr).get_rfunc().get_v()
     << (*ll_ptr).get_rfunc().get_Vtype()
     << (*ll_ptr).get_rfunc().get_w()
     << (*ll_ptr).get_rfunc().get_Wtype()
     << (*ll_ptr).get_afunc().get_lag()
     << (*ll_ptr).get_snpden().get_Kz()
     << (*ll_ptr).get_snpden().get_Iz()
     << (*ll_ptr).get_afunc().get_maxKz()
     << (*ll_ptr).get_afunc().get_maxIz()
     << (*ll_ptr).get_afunc().get_Kx()
     << (*ll_ptr).get_afunc().get_Ix()
     << fmt('d',4,p)
     << fmt('d',4,termination_code)
     << fmt('d',5,iter)
     << fmt('f',11,5,sn) 
     << fmt('f',11,5,bic);
  if (sn < best_sn) {
     best_sn = sn; 
     os << "*\n";
  }
  else {
     os << '\n';
  }
  os.flush();
  return os.good();
}

bool snp::sumry::write_partial_line()
{
  if (filename.size()< 24) {
    for (size_t i=1; i<=24-filename.size(); ++i) os << " ";
  }
  else {
    os << " ";
  }
  os << filename << '\n';
  os.flush();
  return os.good();
}

