/*
 * Copyright (C) 2005-2020 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 "otbVectorImage.h"
#include "otbImageFileWriter.h"
#include "otbVectorDataFileReader.h"

#include "otbVectorDataToLabelMapWithAttributesFilter.h"
#include "otbSpatialisationFilter.h"
#include "otbImageSimulationMethod.h"
#include "otbAttributesMapLabelObject.h"

#include "otbImageFileReader.h"
#include "itkWeightedCentroidKdTreeGenerator.h"
#include "itkKdTreeBasedKmeansEstimator.h"
#include "otbKMeansImageClassificationFilter.h"

int otbImageSimulationMethodKMeansClassif(int itkNotUsed(argc), char* argv[])
{
  const char*  satRSRFilename   = argv[1];
  unsigned int nbBands          = static_cast<unsigned int>(atoi(argv[2]));
  const char*  rootPath         = argv[3];
  unsigned int radius           = atoi(argv[4]);
  const char*  outfilename      = argv[5];
  const char*  outLabelfilename = argv[6];

  const unsigned int nbClasses            = 4;
  double             ConvergenceThreshold = 0.0001;
  unsigned int       NumberOfIterations   = 1000;

  typedef unsigned short LabelType;
  const unsigned int     Dimension = 2;
  typedef otb::Image<LabelType, Dimension>    LabelImageType;
  typedef otb::VectorImage<double, Dimension> OutputImageType;
  typedef otb::ImageFileWriter<OutputImageType> ImageWriterType;
  typedef otb::ImageFileWriter<LabelImageType>  LabelImageWriterType;
  typedef otb::VectorData<double, Dimension> VectorDataType;
  typedef otb::AttributesMapLabelObject<LabelType, Dimension, std::string> LabelObjectType;
  typedef itk::LabelMap<LabelObjectType>                       LabelMapType;
  typedef otb::SpatialisationFilter<LabelMapType>              SpatialisationFilterType;
  typedef otb::ProspectModel                                   SimulationStep1Type;
  typedef otb::SailModel                                       SimulationStep2Type;
  typedef otb::ProlateInterpolateImageFunction<LabelImageType> FTMType;
  typedef otb::ImageSimulationMethod<VectorDataType, SpatialisationFilterType, SimulationStep1Type, SimulationStep2Type, FTMType, OutputImageType>
      ImageSimulationMethodType;

  typedef OutputImageType::PixelType                                       SampleType;
  typedef itk::Statistics::ListSample<SampleType>                          ListSampleType;
  typedef itk::Statistics::WeightedCentroidKdTreeGenerator<ListSampleType> TreeGeneratorType;
  typedef TreeGeneratorType::KdTreeType                                    TreeType;
  typedef itk::Statistics::KdTreeBasedKmeansEstimator<TreeType>            EstimatorType;

  typedef otb::KMeansImageClassificationFilter<OutputImageType, LabelImageType, nbClasses> ClassificationFilterType;
  typedef itk::ImageRegionIterator<OutputImageType> ImageRegionIteratorType;


  /** Instantiation of pointer objects*/
  ImageWriterType::Pointer           writer               = ImageWriterType::New();
  LabelImageWriterType::Pointer      labelWriter          = LabelImageWriterType::New();
  ImageSimulationMethodType::Pointer imageSimulation      = ImageSimulationMethodType::New();
  SpatialisationFilterType::Pointer  spatialisationFilter = SpatialisationFilterType::New();


  ClassificationFilterType::Pointer classifier = ClassificationFilterType::New();

  SpatialisationFilterType::SizeType objectSize;
  objectSize[0] = 300;
  objectSize[1] = 300;

  SpatialisationFilterType::SizeType nbOjects;
  nbOjects[0] = 2;
  nbOjects[1] = 2;

  std::vector<std::string> pathVector(nbClasses);
  pathVector[0] = "JHU/becknic/rocks/sedimentary/powder/0_75/txt/greywa1f.txt";
  pathVector[1] = "";
  pathVector[2] = "JHU/becknic/water/txt/coarse.txt";
  pathVector[3] = "JHU/becknic/soils/txt/0015c.txt";

  std::vector<std::string> areaVector(nbClasses);
  areaVector[0] = "sedimentaryRock";
  areaVector[1] = "prosail";
  areaVector[2] = "water";
  areaVector[3] = "soils";

  std::vector<LabelType> labels(nbClasses);
  labels[0] = 1;
  labels[1] = 2;
  labels[2] = 3;
  labels[3] = 4;

  spatialisationFilter->SetObjectSize(objectSize);
  spatialisationFilter->SetNumberOfObjects(nbOjects);
  spatialisationFilter->SetPathVector(pathVector);
  spatialisationFilter->SetAreaVector(areaVector);
  spatialisationFilter->SetLabels(labels);


  imageSimulation->SetSpatialisation(spatialisationFilter);
  imageSimulation->SetNumberOfComponentsPerPixel(nbBands);
  imageSimulation->SetSatRSRFilename(satRSRFilename);
  imageSimulation->SetPathRoot(rootPath);
  imageSimulation->SetRadius(radius);
  //    imageSimulation->SetMean();
  //    imageSimulation->SetVariance();
  imageSimulation->UpdateData();

  imageSimulation->GetOutputReflectanceImage()->Update();

  // get all the pixel of the image for KMeans centroid estimation
  OutputImageType::IndexType centroidIndex;
  centroidIndex[0] = objectSize[0] / 2;
  centroidIndex[1] = objectSize[1] / 2;

  EstimatorType::ParametersType initialCentroids(nbBands * nbClasses);
  ImageRegionIteratorType       it(imageSimulation->GetOutputReflectanceImage(), imageSimulation->GetOutputReflectanceImage()->GetLargestPossibleRegion());
  it.GoToBegin();
  ListSampleType::Pointer listSample = ListSampleType::New();
  listSample->SetMeasurementVectorSize(nbBands);
  unsigned int x          = 0;
  unsigned int classIndex = 0;
  while (!it.IsAtEnd())
  {
    if (it.GetIndex() == centroidIndex)
    {
      for (unsigned int j = 0; j < nbBands; ++j)
      {
        initialCentroids[j + classIndex * nbBands] = it.Get()[j];
      }
      classIndex++;

      if (x < (nbOjects[0] - 1))
      {
        centroidIndex[0] += objectSize[0];
        x++;
      }
      else
      {
        x                = 0;
        centroidIndex[0] = objectSize[0] / 2;
        centroidIndex[1] += objectSize[1];
      }
    }
    listSample->PushBack(it.Get());
    ++it;
  }


  TreeGeneratorType::Pointer treeGenerator = TreeGeneratorType::New();
  treeGenerator->SetSample(listSample);
  treeGenerator->SetBucketSize(100 / (10 * nbClasses));
  treeGenerator->Update();

  EstimatorType::Pointer estimator = EstimatorType::New();
  estimator->SetKdTree(treeGenerator->GetOutput());
  estimator->SetParameters(initialCentroids);
  estimator->SetMaximumIteration(NumberOfIterations);
  estimator->SetCentroidPositionChangesThreshold(ConvergenceThreshold);
  estimator->StartOptimization();

  classifier->SetCentroids(estimator->GetParameters());
  classifier->SetInput(imageSimulation->GetOutputReflectanceImage());

  // Write the result to an image file
  writer->SetFileName(outfilename);
  writer->SetInput(imageSimulation->GetOutputReflectanceImage());
  writer->Update();

  labelWriter->SetFileName(outLabelfilename);
  //    labelWriter->SetInput(imageSimulation->GetOutputLabelImage());
  labelWriter->SetInput(classifier->GetOutput());
  labelWriter->Update();


  return EXIT_SUCCESS;
}
