#include <math.h>
#include <iostream>
#include <sstream>
#include <fstream>
#include "em_gaussian.h"

EMGaussian::EMGaussian(int cluster_count_) :
  cluster_count(cluster_count_),
  clusters(),
  measurements()
{
  double s = 0;
  for (int i = 0; i < cluster_count; ++i) {
    clusters.push_back(Cluster());
    s += clusters[i].weight;
  }

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

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

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

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

double EMGaussian::maximisation() {
  double shift = 0;
  std::cerr << "Performing maximisation ..." << std::endl;

  for (unsigned int i = 0; i < clusters.size(); ++i) {
    double last_mu = clusters[i].mu;
    double last_sigma = clusters[i].sigma;
    double last_weight = clusters[i].weight;

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

    double mean = 0;
    for (unsigned int j = 0; j < measurements.size(); ++j) {
      mean += measurements[j].value * measurements[j].alpha[i];
    }
    mean /= s;

    double deviation = 0;
    for (unsigned int j = 0; j < measurements.size(); ++j) {
      double foo = measurements[j].value - mean;
      deviation += foo*foo * measurements[j].alpha[i];
    }
    deviation /= s;

    clusters[i].mu = mean;
    clusters[i].sigma = sqrt(deviation);
    if (clusters[i].sigma < 0.1) clusters[i].sigma = 0.1;

    double delta_mu = last_mu - clusters[i].mu;
    double delta_sigma = last_sigma - clusters[i].sigma;
    double delta_weight = last_weight - clusters[i].weight;
    shift += delta_mu*delta_mu + delta_sigma*delta_sigma + delta_weight*delta_weight;
  }
  return sqrt(shift);
}

void EMGaussian::train() {
  double shift = 555;
  std::ofstream likelihood("results/likelihood");
  std::cout << "Clustering values ..." << std::endl;
  do {
    expectation();
    shift = maximisation();
    std::cout << "shift = " << shift << std::endl;
    likelihood << get_loglikelihood() << std::endl;
  } while (shift > 1E-3);
  likelihood.close();
}

void EMGaussian::add_time(double value) {
  measurements.push_back(Measurement(value, cluster_count));
}

EMGaussian::Measurement::Measurement(double value_, int cluster_count) :
  value(value_),
  alpha()
{
  double s = 0;
  double r;

  for (int i = cluster_count; i; --i) {
    r = float(rand()) / RAND_MAX;
    s += r;
    alpha.push_back(r);
  }

  for (int i = cluster_count - 1; i >= 0; --i) {
    alpha[i] /= s;
  }
}

EMGaussian::Cluster::Cluster() :
  mu(2*float(rand()) / RAND_MAX - 1),
  sigma(float(rand()) / RAND_MAX),
  weight(float(rand()) / RAND_MAX)
{

}

EMGaussian::Cluster::Cluster(double mu_, double sigma_, double weight_) :
  mu(mu_),
  sigma(sigma_),
  weight()
{

}

void EMGaussian::Cluster::print() {
  std::cout << "(" << mu << ", " << sigma << ", " << weight << ")";
}

double EMGaussian::Cluster::density_at(double value) const {
  double t = value - mu;
  return exp(-t*t / (2*sigma*sigma)) / (sigma * sqrt(2*M_PI));
}

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

int EMGaussian::get_cluster_count() const {
  return cluster_count;
}

double EMGaussian::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;
}

double EMGaussian::get_loglikelihood() const {
  double result = 0;
  for (unsigned int i = 0; i < measurements.size(); ++i) {
    result += log(get_density_at(measurements[i].value));
  }
  return result;
}
