/***************************************************************************
 *   Copyright (C) 2007 by Faubet Pierre   *
 *   pierre.faubet@e.ujf-grenoble.fr   *
 *                                                                         *
 *   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 <iostream>
#include <fstream>
#include <sstream>
#include <cmath>

#include "environmental_factors.h"
#include "BIMrErrors.h"

environmental_factors::environmental_factors(const string& filename,bool interactions)
{
  /* Consider interactions ? */
  with_interaction = interactions;
  int info = 1;
  bool readfactors=false;
  int factorsread=0;

  clog << "Read environmental factor data from file " << filename.data() << endl;

  ifstream infile;
  infile.open(filename.data());
  if (infile.fail())
    {
      EXCEPTION_INFOS();
      throw OpenFileError();
    }

  string line;
  while (!infile.eof())
    {
      getline(infile,line,'=');

      /* Read number of populations */
      if ((line.length() >= 13) && (line.substr(line.length()-13,line.length())=="[populations]"))
        {
          getline(infile,line);
          istringstream dat(line);
          if (dat >> nb_pop)
            info++;
        }
      /* Read number of factors */
      else if ((line.length() >= 9) && (line.substr(line.length()-9,line.length())=="[factors]"))
        {

          getline(infile,line);
          istringstream dat(line);
          if (dat >> nb_factors)
            info++;
        }

      /* Get factor values */
      else if ((line.length() >= 3) && (line.substr(line.length()-3,line.length())=="[G]"))
        {
          if (info != 3)
            {
              cerr << "Sample sizes" << endl;
              EXCEPTION_INFOS();
              throw FileFormatError();
            }

          getline(infile,line);
          istringstream factor(line);
          int r;
          if ((factor >> r) && (r<=nb_factors))
            {
              int q=0;
              while (!infile.eof() && q<nb_pop)
                {
                  getline(infile,line);
                  if (line.length())
                    {
                      if ((line.find('=') == string::npos))
                        {
                          istringstream fct(line);
                          int l=0;
                          while ((fct >> G[r][q][l]) && (l<nb_pop))
                            l++;
                          if (l!=nb_pop)
                            {
                              cerr << "Factor " << r << '/' << nb_factors << endl;
                              cerr << "Population " << q+1 << ',' << l << endl;
                              EXCEPTION_INFOS();
                              throw FileFormatError();
                            }
                          q++;
                        }
                      else
                        break;
                    }
                }
              if (q!=nb_pop)
                {
                  cerr << "Factor " << r << '/' << nb_factors << endl;
                  cerr << "Population " << q+1 << endl;
                  EXCEPTION_INFOS();
                  throw FileFormatError();
                }
              factorsread++;
            }
          else
            {
              cerr << "Factor number" << endl;
              EXCEPTION_INFOS();
              throw FileFormatError();
            }
        }

      if (info==3 && !readfactors)
        {
          createarrays();
          readfactors=true;
        }
    }
  if ((factorsread!=nb_factors) || factorsread==0)
    {
      if (factorsread)
        {
          cerr << "Environmental factors missing." << endl;
          cerr << factorsread << " read, " << nb_factors << " expected" << endl;
        }
      EXCEPTION_INFOS();
      throw FileFormatError();
    }

  infile.close();

  normalize();
  if (with_interaction)
    addinteractions();

  nb_model = 1 << ( size - 1 );

  inputfilename=filename;

  clog << "Read environmental factor data ... (done)" << endl;
}


environmental_factors::~environmental_factors()
{
  clog << "Drop environmental data" << endl;

  for (int r=0;r<size;r++)
    {
      for (int q=0;q<nb_pop;q++)
        delete [] G[r][q];
      delete [] G[r];
    }
  delete [] G;

  clog << "Drop environmental data ... (done)" << endl;
}

inline void environmental_factors::createarrays()
{
  clog << "Allocate memory for environmental factor data" << endl;
//   clog << "Number of populations: " << nb_pop  << endl;
//   clog << "Number of factors: " << nb_factors  << endl;

  /* Size of data */
  size = 1 + nb_factors + (with_interaction ? nb_factors*(nb_factors-1)/2 : 0);
  G = new double **[size];
  if (!G)
    {
      cerr << "Environmental factors" << endl;
      EXCEPTION_INFOS();
      throw MemError();
    }

  for (int r=0;r<size;r++)
    {
      G[r] = new double *[nb_pop];
      if (!G[r])
        {
          cerr << "Environmental factors " << r << endl;
          EXCEPTION_INFOS();
          throw MemError();
        }

      for (int q=0;q<nb_pop;q++)
        {
          G[r][q] = new double [nb_pop];
          if (!G[r][q])
            {
              cerr << "Environmental factors " << r << endl;
              cerr << "Population " << q+1 << endl;
              EXCEPTION_INFOS();
              throw MemError();
            }
        }
    }
  for (int q=0;q<nb_pop;q++)
    for (int l=0;l<nb_pop;l++)
      G[0][q][l] = 1;
}

void environmental_factors::write(const string& outfil)
{
  /* Open output file */
  ofstream out;
  out.open(outfil.data(),ios::app);
  if (out.fail())
    {
      EXCEPTION_INFOS();
      throw OpenFileError();
    }

  out << endl;
  out << "Environmental factors: " << nb_factors <<endl;
  out << "Pop/Factor";
  for (int r=1;r<size;r++)
    out << '\t' << r;
  out << endl;

  /* Write data */
  for (int q=0;q<nb_pop;q++)
    {
      for (int l=0;l<nb_pop;l++)
        {
          out << q+1 << ',' << l+1;
          for (int r=1;r<size;r++)
            out << '\t' << G[r][q][l];
          out << endl;
        }
      out << endl;
    }

  out << "Interactions: " << (with_interaction ? "yes" : "no") << endl;

  if (with_interaction)
    {
      int r=nb_factors+1;
      for (int r1=1;r1<=nb_factors;r1++)
        for (int r2=r1+1;r2<=nb_factors;r2++)
          {
            out << "G" << r << "= " << "G" << r1 << "*G" << r2 << endl;
            r++;
          }
    }

  /* Print alternative models */
  out << "Model " << 0 << ": no factors" << endl ;
  for (int model = 1; model < nb_model; model++ )
    if ( is_a_valid_model (model) )
      {
        out << "Model " << model << ":";
        for ( int r = 1; r < size; r++ )
          if ( ( model >> ( r - 1 ) ) % 2 )
            out << " G" << r;
        out << endl;
      }
  out.close();
}

bool environmental_factors::is_a_valid_model(int M)
{
  int r = nb_factors;
  if ( with_interaction && ( M > 1 << ( nb_factors - 1 ) ) )
    for ( int r1 = 0; r1 < nb_factors; r1++ )
      for ( int r2 = r1 + 1; r2 < nb_factors; r2++ )
        {
          if ( ( ( M >> r ) % 2 )
               && ! ( ( ( M >> r1 ) % 2 ) && ( ( M >> r2 ) % 2 ) ) )
            return false;
          r++;
        }
  return true;
}

inline void environmental_factors::normalize()
{
  for (int r=1;r<=nb_factors;r++)
    {
      double mean = 0;
      double var = 0;
      for (int q=0;q<nb_pop;q++)
        for (int l=0;l<nb_pop;l++)
          if (q != l)
            {
              mean += G[r][q][l];
              var += G[r][q][l]*G[r][q][l];
            }
      mean /= (nb_pop)*(nb_pop-1);
      var = var/((nb_pop)*(nb_pop-1)) - mean*mean;

      for (int q=0;q<nb_pop;q++)
        for (int l=0;l<nb_pop;l++)
          if (q != l)
            G[r][q][l] = (G[r][q][l]-mean)/ (var!=0 ? sqrt(var) : 1);
    }
}

inline void environmental_factors::addinteractions()
{
  clog << "Add first order interactions between environmental factors" << endl;
  int k=1;
  for (int r1=1;r1<=nb_factors;r1++)
    for (int r2=r1+1;r2<=nb_factors;r2++)
      {
        for (int q=0;q<nb_pop;q++)
          for (int l=0;l<nb_pop;l++)
            G[nb_factors+k][q][l] = (q != l ? G[r1][q][l]*G[r2][q][l] : 0);
        k++;
      }
  clog << "Add first order interactions between environmental factors ... (done)" << endl;
}
