/*
 * Copyright (C) 2005-2017 Centre National d'Etudes Spatiales (CNES)
 *
 * This file is part of Orfeo Toolbox
 *
 *     https://www.orfeo-toolbox.org/
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "otbWrapperApplication.h"
#include "otbWrapperApplicationFactory.h"

#include "otbMultiChannelExtractROI.h"
#include "otbExtractROI.h"

#include "otbStreamingStatisticsImageFilter.h"
#include "otbLabelImageToOGRDataSourceFilter.h"
#include "otbOGRFeatureWrapper.h"

#include <time.h>
#include <vcl_algorithm.h>

namespace otb
{
namespace Wrapper
{
class LSMSVectorization : public Application
{
public:
  typedef LSMSVectorization Self;
  typedef Application Superclass;
  typedef itk::SmartPointer<Self> Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;

  typedef FloatVectorImageType              ImageType;
  typedef ImageType::InternalPixelType      ImagePixelType;
  typedef UInt32ImageType                   LabelImageType;
  typedef LabelImageType::InternalPixelType LabelImagePixelType;

  typedef otb::ImageFileReader<LabelImageType> LabelImageReaderType;

  typedef otb::MultiChannelExtractROI <ImagePixelType,ImagePixelType > MultiChannelExtractROIFilterType;
  typedef otb::ExtractROI<LabelImagePixelType,LabelImagePixelType> ExtractROIFilterType;

  typedef otb::StreamingStatisticsImageFilter<LabelImageType> StatisticsImageFilterType;

  typedef itk::ImageRegionConstIterator<LabelImageType> LabelImageIterator;
  typedef itk::ImageRegionConstIterator<ImageType> ImageIterator;

  typedef otb::LabelImageToOGRDataSourceFilter<LabelImageType> LabelImageToOGRDataSourceFilterType;


  itkNewMacro(Self);

  itkTypeMacro(Vectorization, otb::Application);

private:
  void DoInit() override
  {
    SetName("LSMSVectorization");
    SetDescription("This application performs the fourth step of the exact Large-Scale Mean-Shift segmentation workflow [1].");

    SetDocName("Exact Large-Scale Mean-Shift segmentation, step 4");
    SetDocLongDescription("Given a segmentation result (label image), that may come from the"
                          " LSMSSegmentation [2] application (out parameter) or have been"
                          " processed for small regions merging [3] (out parameter), it will"
                          " convert it to a GIS vector file containing one polygon per"
                          " segment. Each polygon contains additional fields: mean and variance of"
                          " each channels from input image (in parameter), segmentation image"
                          " label, number of pixels in the polygon. For large images one can use"
                          " the tilesizex and tilesizey parameters for tile-wise processing, with"
                          " the guarantees of identical results.");
    SetDocLimitations("This application is part of the Large-Scale Mean-Shift segmentation workflow (LSMS) and may not be suited for any other purpose.");
    SetDocAuthors("David Youssefi");

    SetDocSeeAlso( "[1] Michel, J., Youssefi, D., & Grizonnet, M. (2015). Stable"
                   " mean-shift algorithm and its application to the segmentation of"
                   " arbitrarily large remote sensing images. IEEE Transactions on"
                   " Geoscience and Remote Sensing, 53(2), 952-964.\n"
                   "[2] LSMSegmentation\n"
                   "[3] LSMSmallRegionMerging");

    AddDocTag(Tags::Segmentation);
    AddDocTag("LSMS");

    AddParameter(ParameterType_InputImage, "in", "Input Image");
    SetParameterDescription( "in", "The input image, containing initial spectral signatures corresponding to the segmented image (inseg)." );
    AddParameter(ParameterType_InputImage,  "inseg",    "Segmented image");
    SetParameterDescription( "inseg", "Segmented image where each pixel value is the unique integer label of the segment it belongs to.");

    AddParameter(ParameterType_OutputFilename, "out", "Output GIS vector file");
    SetParameterDescription( "out", "The output GIS vector file, representing the vectorized version of the segmented image where the features of the polygons are the radiometric means and variances." );

    AddParameter(ParameterType_Int, "tilesizex", "Size of tiles in pixel (X-axis)");
    SetParameterDescription("tilesizex", "Size of tiles along the X-axis for tile-wise processing.");
    SetDefaultParameterInt("tilesizex", 500);
    SetMinimumParameterIntValue("tilesizex", 1);

    AddParameter(ParameterType_Int, "tilesizey", "Size of tiles in pixel (Y-axis)");
    SetParameterDescription("tilesizey", "Size of tiles along the Y-axis for tile-wise processing.");
    SetDefaultParameterInt("tilesizey", 500);
    SetMinimumParameterIntValue("tilesizey", 1);

    AddRAMParameter();

    // Doc example parameter settings
    SetDocExampleParameterValue("in","maur_rgb.png");
    SetDocExampleParameterValue("inseg","merged.tif");
    SetDocExampleParameterValue("out","vector.shp");
    SetDocExampleParameterValue("tilesizex","256");
    SetDocExampleParameterValue("tilesizey","256");

    SetOfficialDocLink();
  }

  void DoUpdateParameters() override
  {
  }

  void DoExecute() override
  {
    clock_t tic = clock();

    std::string shapefile(GetParameterString("out"));

    unsigned long sizeTilesX = GetParameterInt("tilesizex");
    unsigned long sizeTilesY = GetParameterInt("tilesizey");


    LabelImageType::Pointer labelIn = GetParameterUInt32Image("inseg");
    labelIn->UpdateOutputInformation();

    unsigned long sizeImageX = labelIn->GetLargestPossibleRegion().GetSize()[0];
    unsigned long sizeImageY = labelIn->GetLargestPossibleRegion().GetSize()[1];

    unsigned int nbTilesX = sizeImageX/sizeTilesX + (sizeImageX%sizeTilesX > 0 ? 1 : 0);
    unsigned int nbTilesY = sizeImageY/sizeTilesY + (sizeImageY%sizeTilesY > 0 ? 1 : 0);

    otbAppLogINFO(<<"Number of tiles: "<<nbTilesX<<" x "<<nbTilesY);

    StatisticsImageFilterType::Pointer stats = StatisticsImageFilterType::New();
    stats->SetInput(labelIn);
    stats->GetStreamer()->SetAutomaticAdaptativeStreaming(GetParameterInt("ram"));
    AddProcess(stats->GetStreamer(), "Retrieve region count...");
    stats->Update();
    unsigned int regionCount=stats->GetMaximum();

    ImageType::Pointer imageIn = GetParameterImage("in");
    imageIn->UpdateOutputInformation();

    unsigned long numberOfComponentsPerPixel = imageIn->GetNumberOfComponentsPerPixel();
    std::string projRef = imageIn->GetProjectionRef();

    std::vector<int>nbPixels;
    nbPixels.clear();
    nbPixels.resize(regionCount+1);

    for(LabelImagePixelType curLabel = 1; curLabel <= regionCount; ++curLabel)
      {
      nbPixels[curLabel] = 0;
      }

    ImageType::PixelType defaultValue(numberOfComponentsPerPixel);
    defaultValue.Fill(0);

    std::vector<ImageType::PixelType>sum(regionCount+1,defaultValue);
    std::vector<ImageType::PixelType>sum2(regionCount+1,defaultValue);

    otb::ogr::DataSource::Pointer ogrDS;
    otb::ogr::Layer layer(ITK_NULLPTR, false);

    OGRSpatialReference oSRS(projRef.c_str());
    std::vector<std::string> options;

    ogrDS = otb::ogr::DataSource::New(shapefile, otb::ogr::DataSource::Modes::Overwrite);
    std::string layername = itksys::SystemTools::GetFilenameName(shapefile.c_str());
    std::string extension = itksys::SystemTools::GetFilenameLastExtension(shapefile.c_str());
    layername = layername.substr(0,layername.size()-(extension.size()));
    layer = ogrDS->CreateLayer(layername, &oSRS, wkbMultiPolygon, options);

    OGRFieldDefn labelField("label", OFTInteger);
    layer.CreateField(labelField, true);
    OGRFieldDefn nbPixelsField("nbPixels", OFTInteger);
    layer.CreateField(nbPixelsField, true);

    for(unsigned int comp = 0; comp<numberOfComponentsPerPixel; ++comp){
    std::ostringstream fieldoss;
    fieldoss<<"meanB"<<comp;
    OGRFieldDefn field(fieldoss.str().c_str(), OFTReal);
    layer.CreateField(field, true);
    }

    for(unsigned int comp = 0; comp<numberOfComponentsPerPixel; ++comp){
    std::ostringstream fieldoss;
    fieldoss<<"varB"<<comp;
    OGRFieldDefn field(fieldoss.str().c_str(), OFTReal);
    layer.CreateField(field, true);
    }

    //Vectorization per tile
    otbAppLogINFO(<<"Vectorization ...");
    for(unsigned int row = 0; row < nbTilesY; row++)
      {
      for(unsigned int column = 0; column < nbTilesX; column++)
       {
        unsigned long startX = column*sizeTilesX;
        unsigned long startY = row*sizeTilesY;
        unsigned long sizeX = vcl_min(sizeTilesX,sizeImageX-startX);
          unsigned long sizeY = vcl_min(sizeTilesY,sizeImageY-startY);

        //Tiles extraction of the input image
        MultiChannelExtractROIFilterType::Pointer imageROI = MultiChannelExtractROIFilterType::New();
        imageROI->SetInput(imageIn);
        imageROI->SetStartX(startX);
        imageROI->SetStartY(startY);
        imageROI->SetSizeX(sizeX);
        imageROI->SetSizeY(sizeY);
        imageROI->Update();

        //Tiles extraction of the segmented image
        ExtractROIFilterType::Pointer labelImageROI = ExtractROIFilterType::New();
        labelImageROI->SetInput(labelIn);
        labelImageROI->SetStartX(startX);
        labelImageROI->SetStartY(startY);
        labelImageROI->SetSizeX(sizeX);
        labelImageROI->SetSizeY(sizeY);
        labelImageROI->Update();

        //Sums calculation for the mean and the variance calculation per label
        LabelImageIterator itLabel( labelImageROI->GetOutput(), labelImageROI->GetOutput()->GetLargestPossibleRegion());
        ImageIterator itImage( imageROI->GetOutput(), imageROI->GetOutput()->GetLargestPossibleRegion());
        for (itLabel.GoToBegin(), itImage.GoToBegin(); !itImage.IsAtEnd(); ++itLabel, ++itImage)
          {
          nbPixels[itLabel.Value()]++;
          for(unsigned int comp = 0; comp<numberOfComponentsPerPixel; ++comp)
            {
            sum[itLabel.Value()][comp]+=itImage.Get()[comp];
            sum2[itLabel.Value()][comp]+=itImage.Get()[comp]*itImage.Get()[comp];
            }
          }

        labelImageROI = ExtractROIFilterType::New();
        labelImageROI->SetInput(labelIn);
        labelImageROI->SetStartX(startX);
        labelImageROI->SetStartY(startY);
        labelImageROI->SetSizeX(sizeX+1);
        labelImageROI->SetSizeY(sizeY+1);
        labelImageROI->Update();

        //Raster->Vecteur conversion
        LabelImageToOGRDataSourceFilterType::Pointer labelToOGR = LabelImageToOGRDataSourceFilterType::New();
        labelToOGR->SetInput(labelImageROI->GetOutput());
        labelToOGR->SetInputMask(labelImageROI->GetOutput());
        labelToOGR->SetFieldName("label");
        labelToOGR->Update();

        otb::ogr::DataSource::ConstPointer ogrDSTmp = labelToOGR->GetOutput();
        otb::ogr::Layer layerTmp = ogrDSTmp->GetLayerChecked(0);

        otb::ogr::Layer::const_iterator featIt = layerTmp.begin();
        for(; featIt!=layerTmp.end(); ++featIt)
          {
          otb::ogr::Feature dstFeature(layer.GetLayerDefn());
          dstFeature.SetFrom( *featIt, TRUE );
          layer.CreateFeature( dstFeature );
          }
       }
      }

    //Sorting by increasing label of the features
    std::ostringstream sqloss;
    sqloss.str("");
    sqloss<<"SELECT * FROM \""<<layername<<"\" ORDER BY label";
    otb::ogr::Layer layerTmp=ogrDS->ExecuteSQL(sqloss.str().c_str(), ITK_NULLPTR, ITK_NULLPTR);
    otb::ogr::Feature firstFeature = layerTmp.ogr().GetNextFeature();

    //Geometry fusion
    otbAppLogINFO("Merging polygons across tiles ...");
    while(firstFeature.addr())
      {
      LabelImagePixelType curLabel = firstFeature.ogr().GetFieldAsInteger("label");

      //Creation of a multipolygon where are stored the geometries to be merged
      OGRMultiPolygon geomToMerge;
      geomToMerge.addGeometry(firstFeature.GetGeometry());
      bool merging = true;
      otb::ogr::Feature nextFeature(ITK_NULLPTR);
      bool haveMerged=false;

      while(merging)
        {
        nextFeature = layerTmp.ogr().GetNextFeature();

        if(nextFeature.addr())
          {
          LabelImagePixelType newLabel = nextFeature.ogr().GetFieldAsInteger("label");
          merging=(newLabel==curLabel);

          //Storing of the new geometry if labels are identical
          if(merging)
            {
            geomToMerge.addGeometry(nextFeature.GetGeometry());
            layer.DeleteFeature(nextFeature.GetFID());
            haveMerged=true;
            }
          //If storing made and new label -> polygons fusion
          else if(haveMerged)
            {
            otb::ogr::UniqueGeometryPtr fusionPolygon = otb::ogr::UnionCascaded(geomToMerge);
            firstFeature.SetGeometry(fusionPolygon.get());
            }
          }
        //If end of list : end of loop
        else
          {
          merging=false;
          }
        }

      //Features calculation
      //Number of pixels per label
      firstFeature.ogr().SetField("nbPixels",nbPixels[curLabel]);

      //Radiometric means per label
      for(unsigned int comp = 0; comp<numberOfComponentsPerPixel; ++comp){
      std::ostringstream fieldoss;
      fieldoss<<"meanB"<<comp;
      firstFeature.ogr().SetField(fieldoss.str().c_str(),sum[curLabel][comp]/nbPixels[curLabel]);
      }

      //Variances per label
      for(unsigned int comp = 0; comp<numberOfComponentsPerPixel; ++comp){
      std::ostringstream fieldoss;
      fieldoss<<"varB"<<comp;
      float var = 0;
      if (nbPixels[curLabel]!=1)
        var = (sum2[curLabel][comp]-sum[curLabel][comp]*sum[curLabel][comp]/nbPixels[curLabel])/(nbPixels[curLabel]-1);
      firstFeature.ogr().SetField(fieldoss.str().c_str(),var);
      }

      //Geometries simplification
      otb::ogr::UniqueGeometryPtr geom = otb::ogr::Simplify(*firstFeature.GetGeometry(),0);
      firstFeature.SetGeometryDirectly(otb::ogr::Simplify(*geom,0));

      layer.SetFeature(firstFeature);

      //Next geometry
      firstFeature=nextFeature;
      }

    const OGRErr err = layer.ogr().CommitTransaction();

    if (err != OGRERR_NONE)
    {
    itkExceptionMacro(<< "Unable to commit transaction for OGR layer " << layer.ogr().GetName() << ".");
    }

    if(extension==".shp"){
    sqloss.str("");
    sqloss<<"REPACK "<<layername;
    ogrDS->ogr().ExecuteSQL(sqloss.str().c_str(), ITK_NULLPTR, ITK_NULLPTR);
    }

    ogrDS->SyncToDisk();

    clock_t toc = clock();

    otbAppLogINFO(<<"Elapsed time: "<<(double)(toc - tic) / CLOCKS_PER_SEC<<" seconds");
  }

};
}
}

OTB_APPLICATION_EXPORT(otb::Wrapper::LSMSVectorization)
