/*
 * 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 "otbLearningApplicationBase.h"
#include "otbWrapperApplicationFactory.h"

#include "otbListSampleGenerator.h"

#include "otbImageToEnvelopeVectorDataFilter.h"
#include "itkPreOrderTreeIterator.h"

// Statistic XML Reader
#include "otbStatisticsXMLFileReader.h"

#include "otbStandardFilterWatcher.h"

// Normalize the samples
#include "otbShiftScaleSampleListFilter.h"

// List sample concatenation
#include "otbConcatenateSampleListFilter.h"

// Balancing ListSample
#include "otbListSampleToBalancedListSampleFilter.h"

#include "itkMersenneTwisterRandomVariateGenerator.h"

// Elevation handler
#include "otbWrapperElevationParametersHandler.h"

namespace otb
{
namespace Wrapper
{

class TrainRegression : public LearningApplicationBase<float, float>
{
public:
  /** Standard class typedefs. */
  typedef TrainRegression Self;
  typedef LearningApplicationBase<float, float> Superclass;
  typedef itk::SmartPointer<Self>       Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;

  /** Standard macro */
  itkNewMacro(Self)

      itkTypeMacro(TrainRegression, otb::Wrapper::LearningApplicationBase)

          typedef Superclass::SampleType SampleType;
  typedef Superclass::ListSampleType       ListSampleType;
  typedef Superclass::TargetSampleType     TargetSampleType;
  typedef Superclass::TargetListSampleType TargetListSampleType;

  typedef Superclass::SampleImageType SampleImageType;
  typedef SampleImageType::PixelType  PixelType;

  // SampleList manipulation
  typedef otb::ListSampleGenerator<SampleImageType, VectorDataType> ListSampleGeneratorType;

  typedef otb::Statistics::ConcatenateSampleListFilter<ListSampleType>       ConcatenateListSampleFilterType;
  typedef otb::Statistics::ConcatenateSampleListFilter<TargetListSampleType> ConcatenateLabelListSampleFilterType;

  // Statistic XML file Reader
  typedef otb::StatisticsXMLFileReader<SampleType> StatisticsReader;

  // Enhance List Sample  typedef otb::Statistics::ListSampleToBalancedListSampleFilter<ListSampleType, LabelListSampleType>      BalancingListSampleFilterType;
  typedef otb::Statistics::ShiftScaleSampleListFilter<ListSampleType, ListSampleType> ShiftScaleFilterType;

  typedef otb::ImageToEnvelopeVectorDataFilter<SampleImageType, VectorDataType> EnvelopeFilterType;

  typedef itk::PreOrderTreeIterator<VectorDataType::DataTreeType> TreeIteratorType;

  typedef itk::Statistics::MersenneTwisterRandomVariateGenerator RandomGeneratorType;

protected:
  TrainRegression()
  {
    this->m_RegressionFlag = true;
  }

private:
  void DoInit() override
  {
    SetName("TrainRegression");
    SetDescription("Train a classifier from multiple images to perform regression.");

    // Documentation
    SetDocLongDescription(
        "This application trains a classifier from multiple input images or a CSV "
        "file, in order to perform regression. Predictors are composed of pixel "
        "values in each band optionally centered and reduced using an XML "
        "statistics file produced by the ComputeImagesStatistics application.\n\n"

        "The output value for each predictor is assumed to be the last band "
        "(or the last column for CSV files). Training and validation predictor "
        "lists are built such that their size is inferior to maximum bounds given "
        "by the user, and the proportion corresponds to the balance parameter. "
        "Several classifier parameters can be set depending on the chosen "
        "classifier. In the validation process, the mean square error is computed "
        "between the ground truth and the estimated model.\n\n"

        "This application is based on LibSVM and on OpenCV Machine Learning "
        "classifiers, and is compatible with OpenCV 2.3.1 and later.");
    SetDocLimitations("None");
    SetDocAuthors("OTB-Team");
    SetDocSeeAlso(
        "OpenCV documentation for machine learning "
        "http://docs.opencv.org/modules/ml/doc/ml.html ");

    AddDocTag(Tags::Deprecated);

    // Group IO
    AddParameter(ParameterType_Group, "io", "Input and output data");
    SetParameterDescription("io", "This group of parameters set the input and output data.");
    AddParameter(ParameterType_InputImageList, "io.il", "Input Image List");
    SetParameterDescription("io.il",
                            "A list of input images. First (n-1) bands should contain the predictor. "
                            "The last band should contain the output value to predict.");
    AddParameter(ParameterType_InputFilename, "io.csv", "Input CSV file");
    SetParameterDescription("io.csv",
                            "Input CSV file containing the predictors, and the output values in last "
                            "column. Only used when no input image is given");
    MandatoryOff("io.csv");

    AddParameter(ParameterType_InputFilename, "io.imstat", "Input XML image statistics file");
    MandatoryOff("io.imstat");
    SetParameterDescription("io.imstat",
                            "Input XML file containing the mean and the standard deviation of the "
                            "input images.");
    AddParameter(ParameterType_OutputFilename, "io.out", "Output regression model");
    SetParameterDescription("io.out", "Output file containing the model estimated (.txt format).");

    AddParameter(ParameterType_Float, "io.mse", "Mean Square Error");
    SetParameterDescription("io.mse", "Mean square error computed with the validation predictors");
    SetParameterRole("io.mse", Role_Output);
    DisableParameter("io.mse");

    // Group Sample list
    AddParameter(ParameterType_Group, "sample", "Training and validation samples parameters");
    SetParameterDescription("sample",
                            "This group of parameters allows one to set training and validation sample "
                            "lists parameters.");

    AddParameter(ParameterType_Int, "sample.mt", "Maximum training predictors");
    // MandatoryOff("mt");
    SetDefaultParameterInt("sample.mt", 1000);
    SetParameterDescription("sample.mt", "Maximum number of training predictors (default = 1000) (no limit = -1).");

    AddParameter(ParameterType_Int, "sample.mv", "Maximum validation predictors");
    // MandatoryOff("mv");
    SetDefaultParameterInt("sample.mv", 1000);
    SetParameterDescription("sample.mv",
                            "Maximum number of validation predictors (default = 1000) "
                            "(no limit = -1).");

    AddParameter(ParameterType_Float, "sample.vtr", "Training and validation sample ratio");
    SetParameterDescription("sample.vtr",
                            "Ratio between training and validation samples (0.0 = all training, "
                            "1.0 = all validation) (default = 0.5).");
    SetParameterFloat("sample.vtr", 0.5);

    Superclass::DoInit();

    AddRANDParameter();

    // Doc example parameter settings
    SetDocExampleParameterValue("io.il", "training_dataset.tif");
    SetDocExampleParameterValue("io.out", "regression_model.txt");
    SetDocExampleParameterValue("io.imstat", "training_statistics.xml");
    SetDocExampleParameterValue("classifier", "libsvm");

    SetOfficialDocLink();
  }

  void DoUpdateParameters() override
  {
    if (HasValue("io.csv") && IsParameterEnabled("io.csv"))
    {
      MandatoryOff("io.il");
    }
    else
    {
      MandatoryOn("io.il");
    }
  }

  void ParseCSVPredictors(std::string path, ListSampleType* outputList)
  {
    std::ifstream ifs;
    ifs.open(path);
    unsigned int       nbCols = 0;
    char               sep    = '\t';
    std::istringstream iss;
    SampleType         elem;
    while (!ifs.eof())
    {
      std::string line;
      std::getline(ifs, line);
      // filter current line
      while (!line.empty() && (line[0] == ' ' || line[0] == '\t'))
      {
        line.erase(line.begin());
      }
      while (!line.empty() && (*(line.end() - 1) == ' ' || *(line.end() - 1) == '\t' || *(line.end() - 1) == '\r'))
      {
        line.erase(line.end() - 1);
      }

      // Avoid commented lines or too short ones
      if (!line.empty() && line[0] != '#')
      {
        std::vector<itksys::String> words = itksys::SystemTools::SplitString(line, sep);
        if (nbCols == 0)
        {
          // detect separator and feature size
          if (words.size() < 2)
          {
            sep   = ' ';
            words = itksys::SystemTools::SplitString(line, sep);
          }
          if (words.size() < 2)
          {
            sep   = ';';
            words = itksys::SystemTools::SplitString(line, sep);
          }
          if (words.size() < 2)
          {
            sep   = ',';
            words = itksys::SystemTools::SplitString(line, sep);
          }
          if (words.size() < 2)
          {
            otbAppLogFATAL(<< "Can't parse CSV file : less than 2 columns or invalid separator (valid separators are tab, space, comma and semi-colon)");
          }
          nbCols = words.size();
          elem.SetSize(nbCols, false);
          outputList->SetMeasurementVectorSize(nbCols);
        }
        else if (words.size() != nbCols)
        {
          otbAppLogWARNING(<< "Skip CSV line, wrong number of columns : got " << words.size() << ", expected " << nbCols);
          continue;
        }
        elem.Fill(0.0);
        for (unsigned int i = 0; i < nbCols; ++i)
        {
          elem[i] = std::stod(words[i]);
        }
        outputList->PushBack(elem);
      }
    }
    ifs.close();
  }

  void DoExecute() override
  {
    GetLogger()->Debug("Entering DoExecute\n");
    // Create training and validation for list samples and label list samples
    ConcatenateListSampleFilterType::Pointer concatenateTrainingSamples   = ConcatenateListSampleFilterType::New();
    ConcatenateListSampleFilterType::Pointer concatenateValidationSamples = ConcatenateListSampleFilterType::New();

    SampleType meanMeasurementVector;
    SampleType stddevMeasurementVector;

    //--------------------------
    // Load measurements from images
    unsigned int nbBands    = 0;
    unsigned int nbFeatures = 0;
    // Iterate over all input images

    FloatVectorImageListType* imageList = GetParameterImageList("io.il");

    // Iterate over all input images
    for (unsigned int imgIndex = 0; imgIndex < imageList->Size(); ++imgIndex)
    {
      FloatVectorImageType::Pointer image = imageList->GetNthElement(imgIndex);
      image->UpdateOutputInformation();

      if (imgIndex == 0)
      {
        nbBands    = image->GetNumberOfComponentsPerPixel();
        nbFeatures = static_cast<unsigned int>(static_cast<int>(nbBands) - 1);
        if (nbBands < 2)
        {
          otbAppLogFATAL(<< "Need at least two bands per image, got " << nbBands);
        }
        else
        {
          if (nbBands != image->GetNumberOfComponentsPerPixel())
          {
            otbAppLogFATAL(<< "Image has a different number of components than "
                              "the first one, expected "
                           << nbBands << ", got " << image->GetNumberOfComponentsPerPixel());
          }
        }
      }

      // Extract image envelope to feed in sampleGenerator
      EnvelopeFilterType::Pointer envelopeFilter = EnvelopeFilterType::New();
      envelopeFilter->SetInput(image);
      envelopeFilter->SetSamplingRate(0);
      if (!image->GetProjectionRef().empty())
      {
        envelopeFilter->SetOutputProjectionRef(image->GetProjectionRef());
      }

      // Setup the DEM Handler
      // otb::Wrapper::ElevationParametersHandler::SetupDEMHandlerFromElevationParameters(this,"elev");

      envelopeFilter->Update();

      VectorDataType::Pointer envelope = envelopeFilter->GetOutput();

      TreeIteratorType itVector(envelope->GetDataTree());
      for (itVector.GoToBegin(); !itVector.IsAtEnd(); ++itVector)
      {
        if (itVector.Get()->IsPolygonFeature())
        {
          itVector.Get()->SetFieldAsInt(std::string("class"), 1);
        }
      }


      // Sample list generator
      ListSampleGeneratorType::Pointer sampleGenerator = ListSampleGeneratorType::New();

      sampleGenerator->SetInput(image);
      sampleGenerator->SetInputVectorData(envelope);

      sampleGenerator->SetClassKey("class");
      sampleGenerator->SetMaxTrainingSize(GetParameterInt("sample.mt"));
      sampleGenerator->SetMaxValidationSize(GetParameterInt("sample.mv"));
      sampleGenerator->SetValidationTrainingProportion(GetParameterFloat("sample.vtr"));
      sampleGenerator->SetBoundByMin(false);
      sampleGenerator->SetPolygonEdgeInclusion(true);

      sampleGenerator->Update();

      // Concatenate training and validation samples from the image
      concatenateTrainingSamples->AddInput(sampleGenerator->GetTrainingListSample());
      concatenateValidationSamples->AddInput(sampleGenerator->GetValidationListSample());
    }

    // if no input image, try CSV
    if (imageList->Size() == 0)
    {
      if (HasValue("io.csv") && IsParameterEnabled("io.csv"))
      {
        ListSampleType::Pointer csvListSample = ListSampleType::New();
        this->ParseCSVPredictors(this->GetParameterString("io.csv"), csvListSample);
        unsigned int totalCSVSize = csvListSample->Size();
        if (totalCSVSize == 0)
        {
          otbAppLogFATAL("No input image and empty CSV file. Missing input data");
        }
        nbBands                                    = csvListSample->GetMeasurementVectorSize();
        nbFeatures                                 = static_cast<unsigned int>(static_cast<int>(nbBands) - 1);
        ListSampleType::Pointer csvTrainListSample = ListSampleType::New();
        ListSampleType::Pointer csvValidListSample = ListSampleType::New();
        csvTrainListSample->SetMeasurementVectorSize(nbBands);
        csvValidListSample->SetMeasurementVectorSize(nbBands);
        double ratio     = this->GetParameterFloat("sample.vtr");
        int    trainSize = static_cast<int>(static_cast<double>(totalCSVSize) * (1.0 - ratio));
        int    validSize = static_cast<int>(static_cast<double>(totalCSVSize) * (ratio));
        if (trainSize > this->GetParameterInt("sample.mt"))
        {
          trainSize = this->GetParameterInt("sample.mt");
        }
        if (validSize > this->GetParameterInt("sample.mv"))
        {
          validSize = this->GetParameterInt("sample.mv");
        }
        double probaTrain = static_cast<double>(trainSize) / static_cast<double>(totalCSVSize);
        double probaValid = static_cast<double>(validSize) / static_cast<double>(totalCSVSize);

        RandomGeneratorType::Pointer randomGenerator = RandomGeneratorType::GetInstance();
        for (unsigned int i = 0; i < totalCSVSize; ++i)
        {
          double random = randomGenerator->GetUniformVariate(0.0, 1.0);
          if (random < probaTrain)
          {
            csvTrainListSample->PushBack(csvListSample->GetMeasurementVector(i));
          }
          else if (random < probaTrain + probaValid)
          {
            csvValidListSample->PushBack(csvListSample->GetMeasurementVector(i));
          }
        }
        concatenateTrainingSamples->AddInput(csvTrainListSample);
        concatenateValidationSamples->AddInput(csvValidListSample);
      }
    }

    // Update
    concatenateTrainingSamples->Update();
    concatenateValidationSamples->Update();

    if (concatenateTrainingSamples->GetOutput()->Size() == 0)
    {
      otbAppLogFATAL("No training samples, cannot perform training.");
    }

    if (concatenateValidationSamples->GetOutput()->Size() == 0)
    {
      otbAppLogWARNING("No validation samples.");
    }

    if (IsParameterEnabled("io.imstat"))
    {
      StatisticsReader::Pointer statisticsReader = StatisticsReader::New();
      statisticsReader->SetFileName(GetParameterString("io.imstat"));
      meanMeasurementVector   = statisticsReader->GetStatisticVectorByName("mean");
      stddevMeasurementVector = statisticsReader->GetStatisticVectorByName("stddev");
      // handle stat file without output normalization
      if (meanMeasurementVector.Size() == nbFeatures)
      {
        meanMeasurementVector.SetSize(nbBands, false);
        meanMeasurementVector[nbFeatures] = 0.0;
        stddevMeasurementVector.SetSize(nbBands, false);
        stddevMeasurementVector[nbFeatures] = 1.0;
      }
    }
    else
    {
      meanMeasurementVector.SetSize(nbBands);
      meanMeasurementVector.Fill(0.);
      stddevMeasurementVector.SetSize(nbBands);
      stddevMeasurementVector.Fill(1.);
    }

    // Shift scale the samples
    ShiftScaleFilterType::Pointer trainingShiftScaleFilter = ShiftScaleFilterType::New();
    trainingShiftScaleFilter->SetInput(concatenateTrainingSamples->GetOutput());
    trainingShiftScaleFilter->SetShifts(meanMeasurementVector);
    trainingShiftScaleFilter->SetScales(stddevMeasurementVector);
    trainingShiftScaleFilter->Update();

    ListSampleType::Pointer rawValidationListSample = ListSampleType::New();

    // Test if the validation test is empty
    if (concatenateValidationSamples->GetOutput()->Size() != 0)
    {
      ShiftScaleFilterType::Pointer validationShiftScaleFilter = ShiftScaleFilterType::New();
      validationShiftScaleFilter->SetInput(concatenateValidationSamples->GetOutput());
      validationShiftScaleFilter->SetShifts(meanMeasurementVector);
      validationShiftScaleFilter->SetScales(stddevMeasurementVector);
      validationShiftScaleFilter->Update();
      rawValidationListSample = validationShiftScaleFilter->GetOutput();
    }

    // Split between predictors and output values
    ListSampleType::Pointer rawlistSample = trainingShiftScaleFilter->GetOutput();
    ListSampleType::Pointer listSample    = ListSampleType::New();
    listSample->SetMeasurementVectorSize(nbFeatures);
    listSample->Resize(rawlistSample->Size());
    TargetListSampleType::Pointer labelListSample = TargetListSampleType::New();
    labelListSample->SetMeasurementVectorSize(1);
    labelListSample->Resize(rawlistSample->Size());

    ListSampleType::Pointer validationListSample = ListSampleType::New();
    validationListSample->SetMeasurementVectorSize(nbFeatures);
    validationListSample->Resize(rawValidationListSample->Size());
    TargetListSampleType::Pointer validationLabeledListSample = TargetListSampleType::New();
    validationLabeledListSample->SetMeasurementVectorSize(1);
    validationLabeledListSample->Resize(rawValidationListSample->Size());

    ListSampleType::MeasurementVectorType       elem;
    TargetListSampleType::MeasurementVectorType outElem;
    for (ListSampleType::InstanceIdentifier i = 0; i < rawlistSample->Size(); ++i)
    {
      elem       = rawlistSample->GetMeasurementVector(i);
      outElem[0] = elem[nbFeatures];
      labelListSample->SetMeasurementVector(i, outElem);
      elem.SetSize(nbFeatures, false);
      listSample->SetMeasurementVector(i, elem);
    }
    for (ListSampleType::InstanceIdentifier i = 0; i < rawValidationListSample->Size(); ++i)
    {
      elem       = rawValidationListSample->GetMeasurementVector(i);
      outElem[0] = elem[nbFeatures];
      validationLabeledListSample->SetMeasurementVector(i, outElem);
      elem.SetSize(nbFeatures, false);
      validationListSample->SetMeasurementVector(i, elem);
    }


    otbAppLogINFO("Number of training samples: " << concatenateTrainingSamples->GetOutput()->Size());
    //--------------------------
    // Split the data set into training/validation set
    ListSampleType::Pointer       trainingListSample        = listSample;
    TargetListSampleType::Pointer trainingLabeledListSample = labelListSample;

    otbAppLogINFO("Size of training set: " << trainingListSample->Size());
    otbAppLogINFO("Size of validation set: " << validationListSample->Size());

    //--------------------------
    // Estimate model
    //--------------------------
    this->Train(trainingListSample, trainingLabeledListSample, GetParameterString("io.out"));

    //--------------------------
    // Performances estimation
    //--------------------------
    ListSampleType::Pointer       performanceListSample;
    TargetListSampleType::Pointer performanceLabeledListSample;

    // Test the input validation set size
    if (validationLabeledListSample->Size() != 0)
    {
      performanceListSample        = validationListSample;
      performanceLabeledListSample = validationLabeledListSample;
    }
    else
    {
      otbAppLogWARNING("The validation set is empty. The performance estimation is done using the input training set in this case.");
      performanceListSample        = trainingListSample;
      performanceLabeledListSample = trainingLabeledListSample;
    }

    TargetListSampleType::Pointer predictedList = this->Classify(performanceListSample, GetParameterString("io.out"));

    otbAppLogINFO("Training performances");
    double                                      mse = 0.0;
    TargetListSampleType::MeasurementVectorType predictedElem;
    for (TargetListSampleType::InstanceIdentifier i = 0; i < performanceListSample->Size(); ++i)
    {
      outElem       = performanceLabeledListSample->GetMeasurementVector(i);
      predictedElem = predictedList->GetMeasurementVector(i);
      mse += (outElem[0] - predictedElem[0]) * (outElem[0] - predictedElem[0]);
    }
    mse /= static_cast<double>(performanceListSample->Size());
    otbAppLogINFO("Mean Square Error = " << mse);
    this->SetParameterFloat("io.mse", mse);
  }
};

} // end namespace Wrapper
} // end namespace otb

OTB_APPLICATION_EXPORT(otb::Wrapper::TrainRegression)
