/****************************************************************************************/
/*                                                                                      */
/*  Project     :  1D-BRNN                                                              */
/*  Release     :  3.3                                                                  */
/*                                                                                      */
/*  File        :  Model.cpp                                                            */
/*  Description :  BRNN model with three neural networks                                */
/*                                                                                      */
/*  Author(s)   :  Christophe Magnan  (2013) - New generic version                      */
/*              :  Jianlin Cheng      (2003) - New custom version for SCRATCH           */
/*              :  Gianluca Pollastri (2001) - Customized version for SCRATCH           */
/*              :  Paolo Frasconi     (1997) - Initial generic version                  */
/*                                                                                      */
/*  Copyright   :  Institute for Genomics and Bioinformatics                            */
/*                 University of California, Irvine                                     */
/*                                                                                      */
/*  Modified    :  2015/07/01                                                           */
/*                                                                                      */
/****************************************************************************************/

#include "Model.h"

/*******************************************/
/*                Interface                */
/*******************************************/

// Constructor
Model::Model() :
  Num_Features(-1),Num_Classes(-1),Num_Hidden(-1),
  Num_AdjPos_FWD(-1),Num_AdjPos_BWD(-1), Num_Outputs_FWD(-1),
  Num_Outputs_BWD(-1),Num_Hidden_FWD(-1),Num_Hidden_BWD(-1),
  MAINnet(NULL),FWDnet(NULL),BWDnet(NULL),Inputs_MAIN(NULL),
  Inputs_FWD(NULL),Inputs_BWD(NULL){}

// Destructor
Model::~Model()
{
  delete FWDnet; delete BWDnet;
  delete MAINnet; delete[] Inputs_MAIN;
  delete[] Inputs_FWD; delete[] Inputs_BWD;
}

// Write model
void Model::write(char* fileout)
{
  filebuf out_buffer;
  if(out_buffer.open(fileout,ios::out)!=0)
  { ostream os(&out_buffer);
    os << Num_Features << " " << Num_Classes << " " << Num_Hidden << "\n";
    os << Num_AdjPos_FWD << " " << Num_Outputs_FWD << " " << Num_Hidden_FWD << "\n";
    os << Num_AdjPos_BWD << " " << Num_Outputs_BWD << " " << Num_Hidden_BWD << "\n";
    MAINnet->write(os); FWDnet->write(os); BWDnet->write(os); out_buffer.close(); }
  else{ cerr << "Cannot write into file \"" << fileout << "\"\n"; exit(1); }
}

// Load model
void Model::load(char* filein)
{
  filebuf in_buffer;
  if(in_buffer.open(filein,ios::in)!=0)
  { istream is(&in_buffer);
    is >> Num_Features >> Num_Classes >> Num_Hidden;
    is >> Num_AdjPos_FWD >> Num_Outputs_FWD >> Num_Hidden_FWD;
    is >> Num_AdjPos_BWD >> Num_Outputs_BWD >> Num_Hidden_BWD;
    MAINnet=new Network(); MAINnet->load(is);
    FWDnet=new Network(); FWDnet->load(is);
    BWDnet=new Network(); BWDnet->load(is);
    Inputs_MAIN=new float[MAINnet->Num_Inputs];
    Inputs_FWD=new float[FWDnet->Num_Inputs];
    Inputs_BWD=new float[BWDnet->Num_Inputs];
    in_buffer.close(); check_viability(); }
  else{ cerr << "Cannot read file \"" << filein << "\"\n"; exit(1); }
}

// New random model
void Model::initialize(int feat,int cl,int hid,int cF,int cB,int oF,int oB,int hF,int hB)
{
  Num_Features=feat; Num_Classes=cl; Num_Hidden=hid;
  Num_AdjPos_FWD=cF; Num_Outputs_FWD=oF; Num_Hidden_FWD=hF;
  Num_AdjPos_BWD=cB; Num_Outputs_BWD=oB; Num_Hidden_BWD=hB;
  MAINnet=new Network(); MAINnet->initialize(feat+(2*cF+1)*oF+(2*cB+1)*oB,hid,cl);
  FWDnet=new Network();  FWDnet->initialize(feat+oF,hF,oF);
  BWDnet=new Network();  BWDnet->initialize(feat+oB,hB,oB);
  MAINnet->Output_Layer->IsOutputLay=1; Inputs_MAIN=new float[MAINnet->Num_Inputs];
  FWDnet->Output_Layer->IsOutputLay=0;  Inputs_FWD=new float[FWDnet->Num_Inputs];
  BWDnet->Output_Layer->IsOutputLay=0;  Inputs_BWD=new float[BWDnet->Num_Inputs];
}

// Allocate memory for back-propagation
void Model::alloc_backpropagation()
{
  MAINnet->alloc_backpropagation();
  FWDnet->alloc_backpropagation();
  BWDnet->alloc_backpropagation();
}

// Propagate a sequence into the BRNN
void Model::propagate(Sequence* seq)
{
  int length=seq->Num_Positions;
  for(int p=0;p<length;p++){ propagate_FWD(seq,p); }
  for(int p=length-1;p>=0;p--){ propagate_BWD(seq,p); }
  for(int p=0;p<length;p++){ propagate_MAIN(seq,p);
  float max=0; int cl=0; for(int i=0;i<Num_Classes;i++){
  float pb=seq->Outputs_MAIN[p][i]; if(pb>max){
  max=pb; cl=i; } } seq->Predictions[p]=cl; }
}

// Expectation step of the back-propagation method
void Model::expectation(Sequence* seq)
{
  int length=seq->Num_Positions;
  seq->reset_backpropagation(Num_Outputs_FWD,Num_Outputs_BWD);
  for(int p=0;p<length;p++){ propagate_FWD(seq,p); }
  for(int p=length-1;p>=0;p--){ propagate_BWD(seq,p); }
  for(int p=0;p<length;p++){ propagate_MAIN(seq,p);
  if(seq->prediction_error(p)>0){ back_propagate_MAIN(seq,p); } }
  for(int p=length-1;p>=0;p--){ back_propagate_FWD(seq,p); }
  for(int p=0;p<length;p++){ back_propagate_BWD(seq,p); }
}

// Maximization step of the back-propagation method
void Model::maximization(float epsilon)
{
  MAINnet->update_weights(epsilon); MAINnet->reset_gradient();
  FWDnet->update_weights(epsilon);  FWDnet->reset_gradient();
  BWDnet->update_weights(epsilon);  BWDnet->reset_gradient();
}

/*******************************************/
/*                 Methods                 */
/*******************************************/

// Propagate a single position into the forward network
void Model::propagate_FWD(Sequence* seq,int pos)
{
  int nf=Num_Features; int nof=Num_Outputs_FWD;
  for(int i=0;i<nf;i++){ Inputs_FWD[i]=seq->Features[pos][i]; }
  if((pos-1)<0){ for(int i=0;i<nof;i++){ Inputs_FWD[nf+i]=0.0; } }
  else{ for(int i=0;i<nof;i++){ Inputs_FWD[nf+i]=seq->Outputs_FWD[pos-1][i]; } }
  FWDnet->propagate(Inputs_FWD);
  for(int i=0;i<nof;i++){ seq->Outputs_FWD[pos][i]=FWDnet->Outputs[i]; }
}

// Propagate a single position into the backward network
void Model::propagate_BWD(Sequence* seq,int pos)
{
  int nf=Num_Features; int nob=Num_Outputs_BWD;
  for(int i=0;i<nf;i++){ Inputs_BWD[i]=seq->Features[pos][i]; }
  if((pos+1)>=seq->Num_Positions){ for(int i=0;i<nob;i++){ Inputs_BWD[nf+i]=0.0; } }
  else{ for(int i=0;i<nob;i++){ Inputs_BWD[nf+i]=seq->Outputs_BWD[pos+1][i]; } }
  BWDnet->propagate(Inputs_BWD);
  for(int i=0;i<nob;i++){ seq->Outputs_BWD[pos][i]=BWDnet->Outputs[i]; }
}

// Propagate a single position into the main network
void Model::propagate_MAIN(Sequence* seq,int pos)
{
  int nf=Num_Features; int nof=Num_Outputs_FWD; int nob=Num_Outputs_BWD;
  int length=seq->Num_Positions; int index=nf;
  for(int i=0;i<nf;i++){ Inputs_MAIN[i]=seq->Features[pos][i]; }
  for(int p=(pos-Num_AdjPos_FWD);p<=(pos+Num_AdjPos_FWD);p++){
  if((p<0)||(p>=length)){ for(int i=0;i<nof;i++){ Inputs_MAIN[index+i]=0.0; } }
  else{ for(int i=0;i<nof;i++){ Inputs_MAIN[index+i]=seq->Outputs_FWD[p][i]; } } index+=nof; }
  for(int p=(pos-Num_AdjPos_BWD);p<=(pos+Num_AdjPos_BWD);p++){
  if((p<0)||(p>=length)){ for(int i=0;i<nob;i++){ Inputs_MAIN[index+i]=0.0; } }
  else{ for(int i=0;i<nob;i++){ Inputs_MAIN[index+i]=seq->Outputs_BWD[p][i]; } } index+=nob; }
  MAINnet->propagate(Inputs_MAIN);
  for(int i=0;i<Num_Classes;i++){ seq->Outputs_MAIN[pos][i]=MAINnet->Outputs[i]; }
}

// Back-propagate a single position into the forward network
void Model::back_propagate_FWD(Sequence* seq,int pos)
{
  propagate_FWD(seq,pos);
  FWDnet->back_propagate(seq->BackProp_FWD[pos]);
  if((pos-1)>=0){ for(int i=0;i<Num_Outputs_FWD;i++){
  seq->BackProp_FWD[pos-1][i]+=FWDnet->Back_Prop[i]; } }
  FWDnet->update_gradient(Inputs_FWD);
}

// Back-propagate a single position into the backward network
void Model::back_propagate_BWD(Sequence* seq,int pos)
{
  propagate_BWD(seq,pos);
  BWDnet->back_propagate(seq->BackProp_BWD[pos]);
  if((pos+1)<seq->Num_Positions){ for(int i=0;i<Num_Outputs_BWD;i++){
  seq->BackProp_BWD[pos+1][i]+=BWDnet->Back_Prop[i]; } }
  BWDnet->update_gradient(Inputs_BWD);
}

// Back-propagate a single position into the main network
void Model::back_propagate_MAIN(Sequence* seq,int pos)
{
  float* Target=seq->Targets[pos]; int length=seq->Num_Positions;
  int index=Num_Features; int nof=Num_Outputs_FWD; int nob=Num_Outputs_BWD;
  MAINnet->back_propagate(Target); MAINnet->update_gradient(Inputs_MAIN);
  for(int p=(pos-Num_AdjPos_FWD);p<=(pos+Num_AdjPos_FWD);p++)
  { if((p>=0)&&(p<length)){ for(int i=0;i<nof;i++){
  seq->BackProp_FWD[p][i]+=MAINnet->Back_Prop[index+i]; } } index+=nof; }
  for(int p=(pos-Num_AdjPos_BWD);p<=(pos+Num_AdjPos_BWD);p++)
  { if((p>=0)&&(p<length)){ for(int i=0;i<nob;i++){
  seq->BackProp_BWD[p][i]+=MAINnet->Back_Prop[index+i]; } } index+=nob; }
}

// Check the viability of a loaded model
void Model::check_viability()
{
  int v=1; int nF=(2*Num_AdjPos_FWD+1)*Num_Outputs_FWD;
  int nB=(2*Num_AdjPos_BWD+1)*Num_Outputs_BWD;
  if(MAINnet->Num_Inputs!=(Num_Features+nF+nB)){ v=0; }
  if(MAINnet->Num_Hidden!=Num_Hidden){ v=0; }
  if(MAINnet->Num_Outputs!=Num_Classes){ v=0; }
  if(MAINnet->Output_Layer->IsOutputLay!=1){ v=0; }
  if(MAINnet->Hidden_Layer->IsOutputLay!=0){ v=0; }
  if(FWDnet->Num_Inputs!=(Num_Features+Num_Outputs_FWD)){ v=0; }
  if(FWDnet->Num_Hidden!=Num_Hidden_FWD){ v=0; }
  if(FWDnet->Num_Outputs!=Num_Outputs_FWD){ v=0; }
  if(FWDnet->Output_Layer->IsOutputLay!=0){ v=0; }
  if(FWDnet->Hidden_Layer->IsOutputLay!=0){ v=0; }
  if(BWDnet->Num_Inputs!=(Num_Features+Num_Outputs_BWD)){ v=0; }
  if(BWDnet->Num_Hidden!=Num_Hidden_BWD){ v=0; }
  if(BWDnet->Num_Outputs!=Num_Outputs_BWD){ v=0; }
  if(BWDnet->Output_Layer->IsOutputLay!=0){ v=0; }
  if(BWDnet->Hidden_Layer->IsOutputLay!=0){ v=0; }
  if(!v){ cerr << "Inconsistent model, cannot proceed\n"; exit(1); }
}