/*
   Copyright (C) 2005-2007,2010 Benjamin Redelings

This file is part of BAli-Phy.

BAli-Phy 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, or (at your option) any later
version.

BAli-Phy 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 BAli-Phy; see the file COPYING.  If not see
<http://www.gnu.org/licenses/>.  */

/**
 * @file dp-array.C
 *
 * @brief This file contains routines for 1-dimensional dynamic programming.
 */

#include <cmath>
#include "dp-array.H"
#include "math/pow2.H"
#include "probability/choose.H"
#include "util.H"

using std::vector;

// We can ignore scale(i) here, because it factors out.
log_double_t DParray::path_P(const vector<int>& g_path) const 
{
  const int I = size()-1;
  int i=I;
  int l=g_path.size()-1;
  int state2 = g_path[l];

  vector<double> transition(n_dp_states());

  log_double_t Pr=1.0;
  while (l>0) 
  {
    for(int state1=0;state1<n_dp_states();state1++)
      transition[state1] = (*this)(i,state1)*GQ(state1,state2);

    int state1 = g_path[l-1];
    double p = choose_P(state1,transition);
    assert(p > 0.0);

    if (di(state1)) i--;

    l--;
    state2 = state1;
    Pr *= p;
  }

  // include probability of choosing 'Start' vs ---+ !
  for(int state1=0;state1<n_dp_states();state1++)
    transition[state1] = (*this)(0,state1) * GQ(state1,state2);

  // Get the probability that the previous state was 'Start'
  double p=0.0;
  for(int state1=0;state1<n_dp_states();state1++)  
    if (not silent(state1))
      p += choose_P(state1,transition);

  Pr *= p;

  assert(Pr > 0.0);
  return Pr;
}

void DParray::forward() {
  for(int i=0;i<size();i++)
    forward(i);
}


vector<int> DParray::sample_path() const {
  vector<int> path;
  
  const int I = size()-1;
  int i = I;

  int state2 = endstate();

  vector<double> transition(n_dp_states());

  while(i >= 0) 
  {
    path.push_back(state2);
    for(int state1=0;state1<n_dp_states();state1++)
      transition[state1] = (*this)(i,state1)*GQ(state1,state2);

    int state1 = -1;
    try {
      state1 = choose_scratch(transition);
    }
    catch (choose_exception<double>& c)
    {
      std::cerr<<"I = "<<I<<"\n";
      std::cerr<<"i = "<<i<<"\n";
      for(int state1=0;state1<n_dp_states();state1++)
	std::cerr<<"transition["<<state1<<"] = "<<transition[state1]<<std::endl;

      c.prepend(__PRETTY_FUNCTION__);
      throw c;
    }

    if (di(state1)) i--;

    state2 = state1;
  }
  assert(i+di(state2)==0);

  std::reverse(path.begin(),path.end());

#ifndef NDEBUG_DP
  check_sampling_probability(path);
#endif

  return path;
}

log_double_t DParray::Pr_sum_all_paths() const {
  const int I = size()-1;

  double total = 0.0;
  for(int state1=0;state1<n_dp_states();state1++)
    total += (*this)(I,state1) * GQ(state1,endstate());

  return pow(log_double_t(2.0),scale(I)) * total;
}

/// Here we set the SEQUENCE length 'l'.  This gives rise to a
/// DP array of length 'l+1'.
void DParray::set_length(int l)
{
  length = l+1;
  state_array::resize(length, n_dp_states());

  for(int i=0;i<n_dp_states();i++)
    (*this)(0,i) = start_P[i];

#if 0
  // Set the start probabilities.
  for(int i=0;i<n_dp_states();i++)
  {
    int s = dp_order(i);
    if (s < start_P.size())
      (*this)(0,i) = start_P[s];
    else
      (*this)(0,i) = 0;
  }
#endif

  Pr_total = 0;
}

DParray::DParray(int l, const HMM& H)
  :DPengine(H)
{
  set_length(l);
}

// We can ignore scale(i) here, because it factors out.
log_double_t DParrayConstrained::path_P(const vector<int>& g_path) const 
{
  const int I = size()-1;
  int i=I;
  int l=g_path.size()-1;
  int state2 = g_path[l];

  vector<double> transition(n_dp_states());

  log_double_t Pr=1.0;
  while (l>0) 
  {
    transition.resize(states(i).size());
    for(int s1=0;s1<transition.size();s1++) {
      int state1 = states(i)[s1];
      transition[s1] = (*this)(i,state1)*GQ(state1,state2);
    }

    int state1 = g_path[l-1];
    auto s1 = find_index(states(i),state1);
    double p = choose_P(*s1,transition);
    assert(p > 0.0);

    if (di(state1)) i--;

    l--;
    state2 = state1;
    Pr *= p;
    assert(Pr > 0.0);
  }

  // In column 0, all states are allowed:
  // - silent states can't contradict anything
  // - non-silent states reprent the start state.
  transition.resize(states(0).size());

  // include probability of choosing 'Start' vs ---+ !
  for(int state1=0;state1<n_dp_states();state1++)
    transition[state1] = (*this)(0,state1) * GQ(state1,state2);

  // Get the probability that the previous state was 'Start'
  double p=0.0;
  for(int state1=0;state1<n_dp_states();state1++)  
    if (not silent(state1))
      p += choose_P(state1,transition);

  Pr *= p;

  assert(Pr > 0.0);
  return Pr;
}

vector<int> DParrayConstrained::sample_path() const {
  vector<int> path;
  
  const int I = size()-1;
  int i = I;

  int state2 = endstate();

  vector<double> transition(n_dp_states());

  while(i >= 0) {
    path.push_back(state2);
    transition.resize(states(i).size());
    for(int s1=0;s1<transition.size();s1++) {
      int state1 = states(i)[s1];
      transition[s1] = (*this)(i,state1)*GQ(state1,state2);
    }

    int s1 = -1;
    try {
      s1 = choose_scratch(transition);
    }
    catch (choose_exception<double>& c)
    {
      std::cerr<<"I = "<<I<<"\n";
      std::cerr<<"i = "<<i<<"\n";
      for(int state1=0;state1<n_dp_states();state1++)
	std::cerr<<"transition["<<state1<<"] = "<<transition[state1]<<std::endl;

      c.prepend(__PRETTY_FUNCTION__);
      throw c;
    }
    int state1 = states(i)[s1];

    if (di(state1)) i--;

    state2 = state1;
  }
  assert(i+di(state2)==0);

  std::reverse(path.begin(),path.end());

#ifndef NDEBUG_DP
  check_sampling_probability(path);
#endif

  return path;
}

inline void DParrayConstrained::forward(int i2) 
{
  assert(i2<size());

  // determine initial scale for this cell
  if (i2==0)
    scale(i2) = 0;
  else
    scale(i2) = scale(i2-1);

  double maximum = 0;

  for(int s2=0;s2<states(i2).size();s2++) 
  {
    int S2 = states(i2)[s2];

    int i1 = i2;
    if (di(S2)) // we need to change this to (state_emit[S2]&~hidden).any()
      i1--;

    //----- don't go off the boundary -----//
    if (i1<0) continue;

    //---- compute arrival probability ----//
    unsigned MAX = states(i1).size();
    if (not di(S2)) MAX = s2;

    double temp = 0;
    for(int s1=0;s1<MAX;s1++) {
      int S1 = states(i1)[s1];

      temp += (*this)(i1,S1) * GQ(S1,S2);
    }

    // record maximum 
    if (temp > maximum) maximum = temp;

    // store the result
    (*this)(i2,S2) = temp;
  }

  //------- if exponent is too low, rescale ------//
  if (maximum > 0 and maximum < fp_scale::lo_cutoff) {
    int logs = -(int)log2(maximum);
    double scale_ = pow2(logs);
    for(int s2=0;s2<states(i2).size();s2++)  
    {
      int S2 = states(i2)[s2];
      (*this)(i2,S2) *= scale_;
    }
    scale(i2) -= logs;
  }
}

void DParrayConstrained::forward() {
  for(int i=0;i<size();i++)
    forward(i);
}


int DParrayConstrained::order_of_computation() const {
  unsigned total=0;
  for(int c=0;c<allowed_states.size()-1;c++)
    total += allowed_states[c].size() * allowed_states[c+1].size();

  return total;
}


void DParrayConstrained::prune() {

  std::abort();

  unsigned order1 = order_of_computation();

  // For column
  for(int c = 1;c<allowed_states.size();c++) {

    // and for each allowed state in that column
    for(int s2=allowed_states[c].size()-1;s2 >= 0;s2--) {

      // check to see whether any states in the previous column are connected to it
      bool used = false;
      for(int s1 = 0;s1<allowed_states[c-1].size() and not used;s1++) {
	if (connected(allowed_states[c-1][s1],allowed_states[c][s2]))
	    used = true;
      }

      // if nothing is connected, then remove it
      if (not used)
	allowed_states[c].erase(allowed_states[c].begin()+s2);
    }
  }

  unsigned order2 = order_of_computation();
  std::cerr<<" order1 = "<<order1<<"    order2 = "<<order2<<"  fraction = "<<double(order2)/double(order1)<<std::endl;
}

log_double_t DParrayConstrained::Pr_sum_all_paths() const 
{
  const int I = size()-1;

  double total = 0.0;
  for(int s1=0;s1<states(I).size();s1++) {
    int state1 = states(I)[s1];
    total += (*this)(I,state1) * GQ(state1,endstate());
  }

  return pow(log_double_t(2.0),scale(I)) * total;
}

void DParrayConstrained::set_length(int l)
{
  DParray::set_length(l);
  allowed_states.resize(size());
}

DParrayConstrained::DParrayConstrained(int l, const HMM& H)
  :DParray(l,H)
{ 
  allowed_states.resize(size());
}
