#include <math.h>
#include <iostream>
#include <sstream>
#include "em_sqdist.h"

EMSqdist::EMSqdist() :
  EMCircular(),
  clusters()
{

}

EMSqdist::EMSqdist(int cluster_count_) :
  EMCircular(cluster_count_),
  clusters()
{
  double s = 0;
  for (unsigned int i = 0; i < cluster_count; ++i) {
    clusters.push_back(Cluster());
    s += clusters[i].weight;
    double r = double(rand())/RAND_MAX;
    double f = 2*M_PI*double(rand())/RAND_MAX;
    clusters[i].xx = r*cos(f);
    clusters[i].yy = r*sin(f);
  }

  for (unsigned int i = 0; i < cluster_count; ++i) {
    clusters[i].weight /= s;
  }
}

void EMSqdist::expectation() {
  std::cerr << "Performing expectation ..." << std::endl;

  for (unsigned int i = 0; i < timestamps.size(); ++i) {
    double s = 0;
    for (unsigned int j = 0; j < clusters.size(); ++j) {
      double a = clusters[j].density_at(timestamps[i].phase) * clusters[j].weight;
      timestamps[i].alpha[j] = a;
      s += a;
    }

    for (unsigned int j = 0; j < clusters.size(); ++j) {
      timestamps[i].alpha[j] /= s;
    }
  }
}

double EMSqdist::maximisation() {
  double shift = 0;
  std::cerr << "Performing maximisation ..." << std::endl;
  for (unsigned int i = 0; i < clusters.size(); ++i) {
    double last_xx = clusters[i].xx;
    double last_yy = clusters[i].yy;
    double last_weight = clusters[i].weight;

    double s = 0;
    for (unsigned int j = 0; j < timestamps.size(); ++j) {
      s += timestamps[j].alpha[i];
    }
    clusters[i].weight = s / timestamps.size();

    /*double mean_re = 0;
    double mean_im = 0;
    for (unsigned int j = 0; j < timestamps.size(); ++j) {
      mean_re += cos(timestamps[j].phase) * timestamps[j].alpha[i];
      mean_im += sin(timestamps[j].phase) * timestamps[j].alpha[i];
      //std::cout << timestamps[j] << std::endl;
    }
    clusters[i].xx = mean_re / s;
    clusters[i].yy = mean_im / s;
    double norm = sqrt(clusters[i].xx*clusters[i].xx + clusters[i].yy*clusters[i].yy);
    if (norm > 0.99) {
      clusters[i].xx *= 0.99/norm;
      clusters[i].yy *= 0.99/norm;
    }

    clusters[i].xx *= 0.99;
    clusters[i].yy *= 0.99;*/
    /*size_t iter = 0;
    int status;

    const gsl_multimin_fdfminimizer_type *T;
    gsl_multimin_fdfminimizer *solver;

    gsl_vector *x;
    gsl_multimin_function_fdf my_func;

    MinParams params(&timestamps, i);
    my_func.n = 2;
    my_func.f = my_f;
    my_func.df = my_df;
    my_func.fdf = my_fdf;
    my_func.params = &params;

    x = gsl_vector_alloc(2);
    double xs = last_xx;
    double ys = last_yy;
    inv_tanh_ext(xs, ys);
    gsl_vector_set(x, 0, xs);
    gsl_vector_set(x, 1, ys);

    T = gsl_multimin_fdfminimizer_conjugate_fr;
    solver = gsl_multimin_fdfminimizer_alloc(T, 2);

    gsl_multimin_fdfminimizer_set(solver, &my_func, x, 0.3, 1e-7);

    do {
      iter++;
      status = gsl_multimin_fdfminimizer_iterate(solver);
      if (status) break;
      status = gsl_multimin_test_gradient (solver->gradient, 1e-3);
    }
    while (status == GSL_CONTINUE && iter < 100);

    clusters[i].xx = gsl_vector_get (solver->x, 0);
    clusters[i].yy = gsl_vector_get (solver->x, 1);
    EMSqdist::tanh_ext(clusters[i].xx, clusters[i].yy);

    gsl_multimin_fdfminimizer_free (solver);
    gsl_vector_free (x);*/

    mle(i, s, clusters[i].xx, clusters[i].yy);
    double norm = sqrt(clusters[i].xx*clusters[i].xx + clusters[i].yy*clusters[i].yy);
    if (norm > 0.999) {
      clusters[i].xx *= 0.99/norm;
      clusters[i].yy *= 0.99/norm;
    }

    double delta_xx = last_xx - clusters[i].xx;
    double delta_yy = last_yy - clusters[i].yy;
    double delta_weight = last_weight - clusters[i].weight;
    shift += delta_xx*delta_xx + delta_yy*delta_yy + delta_weight*delta_weight;
  }
  return sqrt(shift);
}

void EMSqdist::train() {
  double shift = 555;
  int i = 0;
  double dl;
  do {
    double l = get_loglikelihood();
    std::ofstream density("results/density" + std::to_string(i));
    for (double phase = -M_PI; phase < M_PI; phase += 0.001) {
      double d = get_density_at(phase);
      density << phase << " " << d << std::endl;
    }
    density.close();
    likelihood << get_loglikelihood() << std::endl;

    expectation();
    shift = maximisation();
    ++i;
    dl = get_loglikelihood() - l;
    std::cout << "dl = " << dl << std::endl;
  } while (dl > 1);
}

void EMSqdist::add_time(double phase) {
  timestamps.push_back(Timestamp(phase, cluster_count));
}

EMSqdist::Cluster::Cluster() :
  xx(),
  yy(),
  weight(float(rand()) / RAND_MAX)
{
  double r = double(rand())/RAND_MAX;
  double f = 2*M_PI*double(rand())/RAND_MAX;
  xx = r*cos(f);
  yy = r*sin(f);
}

EMSqdist::Cluster::Cluster(double xx_, double yy_, double weight_) :
  xx(xx_),
  yy(yy_),
  weight()
{

}

void EMSqdist::Cluster::print() {
  std::cout << "(" << xx << ", " << yy << ", " << weight << ")";
}

double EMSqdist::Cluster::density_at(double phase) const {
  double dx = xx - cos(phase);
  double dy = yy - sin(phase);
  return (1 - xx*xx - yy*yy) / ((dx*dx + dy*dy) * 2*M_PI);
}

void EMSqdist::print() {
  std::cout << "[";
  double sum_w = 0;
  for (unsigned int i = 0; i < clusters.size(); ++i) {
    if (i) {
      std::cout << ", ";
    }
    clusters[i].print();
    sum_w += clusters[i].weight;
  }
  std::cout << "] (w = " << sum_w << ")" << std::endl;
}

double EMSqdist::get_density_at(double phase) const {
  double result = 0;
  for (unsigned int i = 0; i < clusters.size(); ++i) {
		result += clusters[i].weight * clusters[i].density_at(phase);
  }
  return result;
}

void EMSqdist::tanh_ext(double& x, double& y) {
  if (x==0 && y==0) {
    return;
  }

  double r = sqrt(x*x + y*y);
  double coef = tanh(r) / r;
  x *= coef;
  y *= coef;
}

void EMSqdist::inv_tanh_ext(double& x, double& y) {
  if (x==0 && y==0) {
    return;
  }

  double r = sqrt(x*x + y*y);
  double coef = atanh(r) / r;
  x *= coef;
  y *= coef;
}

double EMSqdist::my_f(const gsl_vector *v, void *params) {
  std::vector<Timestamp>* timestamps = ((MinParams*)params)->timestamps;
  int k = ((MinParams*)params)->k;

  double x = gsl_vector_get(v, 0);
  double y = gsl_vector_get(v, 1);
  EMSqdist::tanh_ext(x, y);

  double result = 0;
  double xc, ys;
  for (unsigned int i = 0; i < timestamps->size(); ++i) {
    xc = cos((*timestamps)[i].phase) - x;
    ys = sin((*timestamps)[i].phase) - y;
    result += (*timestamps)[i].alpha[k] * (log(1-x*x-y*y) - log(xc*xc + ys*ys));
  }

  //std::cout << x << " " << y << " " << result << std::endl;
  return -result;
}

/* The gradient of f, df = (df/dx, df/dy). */
void EMSqdist::my_df(const gsl_vector *v, void *params, gsl_vector *df) {
  std::vector<Timestamp>* timestamps = ((MinParams*)params)->timestamps;
  int k = ((MinParams*)params)->k;

  double x = gsl_vector_get(v, 0);
  double y = gsl_vector_get(v, 1);
  EMSqdist::tanh_ext(x, y);

  double dx = 0;
  double dy = 0;
  double xc, ys;
  for (unsigned int i = 0; i < timestamps->size(); ++i) {
    xc = cos((*timestamps)[i].phase) - x;
    ys = sin((*timestamps)[i].phase) - y;
    dx += (*timestamps)[i].alpha[k] * (x/(1 - x*x - y*y) - xc/(xc*xc + ys*ys));
    dy += (*timestamps)[i].alpha[k] * (y/(1 - x*x - y*y) - ys/(xc*xc + ys*ys));
  }

  gsl_vector_set(df, 0, 2*dx);
  gsl_vector_set(df, 1, 2*dy);
}

/* Compute both f and df together. */
void EMSqdist::my_fdf(const gsl_vector *v, void *params, double *f, gsl_vector *df) {
  std::vector<Timestamp>* timestamps = ((MinParams*)params)->timestamps;
  int k = ((MinParams*)params)->k;

  double x = gsl_vector_get(v, 0);
  double y = gsl_vector_get(v, 1);
  EMSqdist::tanh_ext(x, y);

  double dx = 0;
  double dy = 0;
  double result = 0;
  double xc, ys;
  for (unsigned int i = 0; i < timestamps->size(); ++i) {
    xc = cos((*timestamps)[i].phase) - x;
    ys = sin((*timestamps)[i].phase) - y;
    dx += (*timestamps)[i].alpha[k] * (x/(1 - x*x - y*y) - xc/(xc*xc + ys*ys));
    dy += (*timestamps)[i].alpha[k] * (y/(1 - x*x - y*y) - ys/(xc*xc + ys*ys));
    result += (*timestamps)[i].alpha[k] * (log(1-x*x-y*y) - log(xc*xc + ys*ys));
  }

  gsl_vector_set(df, 0, 2*dx);
  gsl_vector_set(df, 1, 2*dy);
  *f = -result;
}

EMSqdist::MinParams::MinParams(std::vector<Timestamp>* timestamps_, int k_) :
  timestamps(timestamps_),
  k(k_)
{

}

void EMSqdist::U(double zr, double zi, double fr, double fi, double w, double& vr, double& vi) {
  double a = zr - fr;
  double b = zi - fi;
  double c = 1 - zr*fr - zi*fi;
  double d = zr*fi - zi*fr;
  double e = w/(c*c + d*d);
  vr += e*(a*c + b*d);
  vi += e*(b*c - a*d);
}

void EMSqdist::mle(int c, double s, double& xhat, double& yhat) {
  xhat = 0;
  yhat = 0;
  int i = 0;
  while (true) {
    double vr = 0, vi = 0;
    double xhat_ = 0, yhat_ = 0;
    for (unsigned int j = 0; j < timestamps.size(); ++j) {
      U(cos(timestamps[j].phase), sin(timestamps[j].phase), xhat, yhat, timestamps[j].alpha[c], vr, vi);
    }
    U(vr/s, vi/s, -xhat, -yhat, 1, xhat_, yhat_);
    double dx = xhat - xhat_;
    double dy = yhat - yhat_;
    xhat = xhat_;
    yhat = yhat_;
    double d = sqrt(dx*dx + dy*dy);
    if (d < 1e-7) {
      return;
    }
    ++i;
  }
}
