package de.jungblut.clustering;

import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.List;

import javax.imageio.ImageIO;

import org.apache.commons.math3.util.FastMath;

import de.jungblut.datastructure.KDTree;
import de.jungblut.datastructure.KDTree.VectorDistanceTuple;
import de.jungblut.math.DoubleVector;
import de.jungblut.math.dense.DenseDoubleVector;
import de.jungblut.reader.ImageReader;

public final class MeanShiftClustering {

  private static final double SQRT_2_PI = FastMath.sqrt(2 * Math.PI);

  // window size h
  public static List<DoubleVector> cluster(List<DoubleVector> points, double h,
      int maxIterations, boolean verbose) {

    KDTree<Integer> kdTree = new KDTree<>();
    int index = 0;
    for (DoubleVector v : points) {
      kdTree.add(v, index++);
    }

    List<DoubleVector> centers = observeCenters(kdTree, points, h, verbose);
    BitSet convergedSet = new BitSet(centers.size());
    List<DoubleVector> shifts = new ArrayList<>(centers.size());
    for (int i = 0; i < centers.size(); i++) {
      shifts.add(new DenseDoubleVector(points.get(0).getDimension()));
    }
    for (int i = 0; i < maxIterations; i++) {
      int converged = meanShift(kdTree, centers, shifts, convergedSet, h);
      if (verbose) {
        System.out.println("Iteration: " + i
            + " | Number of centers not converged yet: " + converged + "/"
            + centers.size());
      }
      if (converged == 0) {
        break;
      }
    }

    return centers;
  }

  private static int meanShift(KDTree<Integer> kdTree,
      List<DoubleVector> centers, List<DoubleVector> shifts,
      BitSet convergedSet, double h) {
    int remainingConvergence = 0;
    for (int i = 0; i < centers.size(); i++) {
      if (!convergedSet.get(i)) {
        DoubleVector v = centers.get(i);
        List<VectorDistanceTuple<Integer>> neighbours = kdTree
            .getNearestNeighbours(v, h);
        double weightSum = 0d;
        DoubleVector numerator = new DenseDoubleVector(v.getLength());
        for (VectorDistanceTuple<Integer> neighbour : neighbours) {
          if (neighbour.getDistance() < h) {
            double normDistance = neighbour.getDistance() / h;
            weightSum -= gaussianGradient(normDistance);
            numerator = numerator
                .add(neighbour.getVector().multiply(weightSum));
          }
        }
        if (weightSum > 0d) {
          DoubleVector shift = v.divide(numerator);
          DoubleVector oldShift = shifts.get(i);
          // check how the last shift was and if there wasn't too much movement,
          // stay converged
          if (shift.subtract(oldShift).abs().sum() > 1e-5) {
            remainingConvergence++;
            // apply the shift
            centers.set(i, centers.get(i).add(shift));
            shifts.set(i, shift);
          } else {
            convergedSet.set(i);
          }
        }
      }
    }
    return remainingConvergence;
  }

  private static List<DoubleVector> observeCenters(KDTree<Integer> kdTree,
      List<DoubleVector> points, double h, boolean verbose) {
    List<DoubleVector> centers = new ArrayList<>();
    BitSet assignedIndices = new BitSet(kdTree.size());
    // we are doing one pass over the dataset to determine the first centers
    // that are within h range
    for (int i = 0; i < points.size(); i++) {
      if (!assignedIndices.get(i)) {
        DoubleVector v = points.get(i);
        List<VectorDistanceTuple<Integer>> neighbours = kdTree
            .getNearestNeighbours(v, h);
        DoubleVector center = new DenseDoubleVector(v.getLength());
        int added = 0;
        for (VectorDistanceTuple<Integer> neighbour : neighbours) {
          if (!assignedIndices.get(neighbour.getValue())
              && neighbour.getDistance() < h) {
            center = center.add(neighbour.getVector());
            assignedIndices.set(neighbour.getValue());
            added++;
          }
        }
        // so if our sum is positive, we can divide and add the center
        if (added > 1) {
          DoubleVector newCenter = center.divide(added);
          centers.add(newCenter);
          if (verbose && centers.size() % 1000 == 0) {
            System.out.println("#Centers found: " + centers.size());
          }
        }
      }
      assignedIndices.set(i);
    }

    return centers;
  }

  private static double gaussianGradient(double stddev) {
    return -FastMath.exp(-(stddev * stddev) / 2d) / (SQRT_2_PI * stddev);
  }

  public static void main(String[] args) throws Exception {
    DoubleVector[] luv = ImageReader.readImageAsLUV(ImageIO.read(new File(
        "files/img/lenna.png")));
    List<DoubleVector> cluster = MeanShiftClustering.cluster(
        Arrays.asList(luv), 3d, 100, true);
    System.out.println(cluster.size());
  }

}

@@ -30,29 +30,8 @@
 g;%0A%0A
-import java.io.File;%0A
 impo
@@ -58,33 +58,8 @@
 st;%0A
-import java.util.Arrays;%0A
 impo
@@ -107,39 +107,8 @@
 t;%0A%0A
-import javax.imageio.ImageIO;%0A%0A
 impo
@@ -343,47 +343,8 @@
 tor;
-%0Aimport de.jungblut.reader.ImageReader;
 %0A%0Apu
@@ -4435,326 +4435,6 @@
  %7D%0A%0A
-  public static void main(String%5B%5D args) throws Exception %7B%0A    DoubleVector%5B%5D luv = ImageReader.readImageAsLUV(ImageIO.read(new File(%0A        %22files/img/lenna.png%22)));%0A    List%3CDoubleVector%3E cluster = MeanShiftClustering.cluster(%0A        Arrays.asList(luv), 3d, 100, true);%0A    System.out.println(cluster.size());%0A  %7D%0A%0A
 %7D%0A

remove odd main method