package org.datasyslab.geospark.serde;

import com.esotericsoftware.kryo.Kryo;
import com.vividsolutions.jts.geom.GeometryCollection;
import com.vividsolutions.jts.geom.LineString;
import com.vividsolutions.jts.geom.MultiLineString;
import com.vividsolutions.jts.geom.MultiPoint;
import com.vividsolutions.jts.geom.MultiPolygon;
import com.vividsolutions.jts.geom.Point;
import com.vividsolutions.jts.geom.Polygon;
import com.vividsolutions.jts.index.quadtree.Quadtree;
import com.vividsolutions.jts.index.strtree.STRtree;
import org.apache.log4j.Logger;
import org.apache.spark.serializer.KryoRegistrator;
import org.datasyslab.geospark.geometryObjects.Circle;
import org.datasyslab.geospark.geometryObjects.GeometrySerde;

public class GeoSparkKryoRegistrator implements KryoRegistrator {

    final static Logger log = Logger.getLogger(GeoSparkKryoRegistrator.class);

    @Override
    public void registerClasses(Kryo kryo) {
        GeometrySerde serializer = new GeometrySerde();

        log.info("Registering custom serializers for geometry types");

        kryo.register(Point.class, serializer);
        kryo.register(LineString.class, serializer);
        kryo.register(Polygon.class, serializer);
        kryo.register(MultiPoint.class, serializer);
        kryo.register(MultiLineString.class, serializer);
        kryo.register(MultiPolygon.class, serializer);
        kryo.register(GeometryCollection.class, serializer);
        kryo.register(Circle.class, serializer);
        // TODO: Replace the default serializer with default spatial index serializer
        kryo.register(Quadtree.class);
        kryo.register(STRtree.class);
    }
}

@@ -111,307 +111,9 @@
 eom.
-GeometryCollection;%0Aimport com.vividsolutions.jts.geom.LineString;%0Aimport com.vividsolutions.jts.geom.MultiLineString;%0Aimport com.vividsolutions.jts.geom.MultiPoint;%0Aimport com.vividsolutions.jts.geom.MultiPolygon;%0Aimport com.vividsolutions.jts.geom.Point;%0Aimport com.vividsolutions.jts.geom.Polygon
+*
 ;%0Aim
@@ -418,16 +418,82 @@
 rySerde;
+%0Aimport org.datasyslab.geospark.geometryObjects.SpatialIndexSerde;
 %0A%0Apublic
@@ -747,16 +747,95 @@
 Serde();
+%0A        SpatialIndexSerde indexSerializer = new SpatialIndexSerde(serializer);
 %0A%0A      
@@ -1319,32 +1319,83 @@
 s, serializer);%0A
+        kryo.register(Envelope.class, serializer);%0A
         // TODO:
@@ -1496,24 +1496,41 @@
 adtree.class
+, indexSerializer
 );%0A        k
@@ -1555,16 +1555,33 @@
 ee.class
+, indexSerializer
 );%0A    %7D

add register to support envelope and indexes