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

EMTesedist::EMTesedist() :
  EMCircular(),
  clusters()
{

}

EMTesedist::EMTesedist(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 EMTesedist::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);
      timestamps[i].alpha[j] = a;
      s += a;
    }

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

double EMTesedist::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;
    }
    mean_re /= s;
    mean_im /= s;
    double m12 = mean_re*mean_re + mean_im*mean_im;
    double m1 = sqrt(m12);
    double x = (1 - sqrt(1-m12)) / m1;
    if (x > 0.99) {
      x = 0.99;
    }
    if (x != 0) {
      clusters[i].xx = mean_re * x / m1;
      clusters[i].yy = mean_im * x / m1;
    }

    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 EMTesedist::train() {
  double shift = 555;
  do {
    expectation();
    shift = maximisation();
    std::cout << "shift = " << shift << std::endl;
    likelihood << get_loglikelihood() << std::endl;
  } while (shift > 1E-3);
}

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

EMTesedist::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);
}

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

}

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

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

void EMTesedist::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 EMTesedist::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;
}
