/*
 * Copyright (C) 2005-2022 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 "otbSamplingRateCalculatorList.h"
#include "otbStatisticsXMLFileReader.h"

namespace otb
{
namespace Wrapper
{

class MultiImageSamplingRate : public Application
{
public:
  /** Standard class typedefs. */
  typedef MultiImageSamplingRate        Self;
  typedef Application                   Superclass;
  typedef itk::SmartPointer<Self>       Pointer;
  typedef itk::SmartPointer<const Self> ConstPointer;

  /** Standard macro */
  itkNewMacro(Self);

  itkTypeMacro(MultiImageSamplingRate, otb::Application);

  /** typedef */
  typedef otb::SamplingRateCalculatorList RateCalculatorListType;

  typedef RateCalculatorListType::ClassCountMapType ClassCountMapType;
  typedef RateCalculatorListType::MapRateType       MapRateType;

  typedef itk::VariableLengthVector<float>              MeasurementType;
  typedef otb::StatisticsXMLFileReader<MeasurementType> XMLReaderType;

private:
  MultiImageSamplingRate()
  {
    m_CalculatorList = RateCalculatorListType::New();
  }

  void DoInit() override
  {
    SetName("MultiImageSamplingRate");
    SetDescription("Compute sampling rate for an input set of images.");

    // Documentation
    SetDocLongDescription(
        "The application computes sampling rates for a set of"
        " input images. Before calling this application, each pair of image and "
        "training vectors has to be analysed with the application "
        "PolygonClassStatistics. The statistics file is then used to compute the "
        "sampling rates for each class in each image. Several types of sampling "
        " are implemented. Each one is a combination of a mono-image strategy "
        "and a multi-image mode. The mono-image strategies are:\n\n"

        "* smallest (default): select the same number of sample in each "
        "class so that the smallest one is fully sampled.\n"
        "* constant: select the same number of samples N in each class "
        "(with N below or equal to the size of the smallest class).\n"
        "* byclass: set the required number for each class manually, with an "
        "input CSV file (first column is class name, second one is the required "
        "samples number).\n\n"
        "The multi-image modes (mim) are proportional, equal and custom. The custom "
        "mode lets the user choose the distribution of samples among the "
        "images. The different behaviours are described below. Ti(c) and Ni(c) "
        " refers respectively to the total number and needed number of samples in "
        "image i for class c. Let's call L the total number of images.\n\n"
        "* strategy = all\n\n"
        "    - Same behaviour for all modes: take all samples\n\n"
        "* strategy = constant:"
        " let's call M the global number of samples required per class."
        " For each image i and each class c:\n\n"
        "    - if mim = proportional, then Ni( c ) = M * Ti( c ) / sum_k( Tk(c) )\n\n"
        "    - if mim = equal       , then Ni( c ) = M / L\n\n"
        "    - if mim = custom      , then Ni( c ) = Mi where Mi is the custom requested number of samples for image i\n\n"
        "* strategy = byClass :"
        " let's call M(c) the global number of samples for class c)."
        " For each image i and each class c:\n\n"
        "    - if mim = proportional, then Ni( c ) = M(c) * Ti( c ) / sum_k( Tk(c) )\n\n"
        "    - if mim = equal       , then Ni( c ) = M(c) / L\n\n"
        "    - if mim = custom      , then Ni( c ) = Mi(c) where Mi(c) is the custom requested number of samples for image i and class c\n\n"
        "* strategy = percent :"
        " For each image i and each class c:\n\n"
        "    - if mim = proportional, then Ni( c ) = p * Ti( c ) where p is the global percentage of samples\n\n"
        "    - if mim = equal       , then Ni( c ) = p * sum_k(Tk(c)]/L where p is the global percentage of samples\n\n"
        "    - if mim = custom      , then Ni( c ) = p(i) * Ti(c) where p(i) is the percentage of samples for image i. c\n\n"
        "* strategy = total :"
        " For each image i and each class c:\n\n"
        "    - if mim = proportional, then Ni( c ) = total * (sum_k(Ti(k))/sum_kl(Tl(k))) * (Ti(c)/sum_k(Ti(k))) where total is the total number of samples "
        "specified.\n\n"
        "    - if mim = equal       , then Ni( c ) = (total / L) * (Ti(c)/sum_k(Ti(k))) where total is the total number of samples specified.\n\n"
        "    - if mim = custom      , then Ni( c ) = total(i) * (Ti(c)/sum_k(Ti(k))) where total(i) is the total number of samples specified for image i. \n\n"
        "* strategy = smallest class\n\n"
        "    - if mim = proportional, then the smallest class size (computed globally) is used for the strategy constant+proportional.\n\n"
        "    - if mim = equal       , then the smallest class size (computed globally) is used for the strategy constant+equal.\n\n"
        "    - if mim = custom      , then the smallest class is computed and used for each image separately.");
    SetDocLimitations("None");
    SetDocAuthors("OTB-Team");
    SetDocSeeAlso(" ");

    AddDocTag(Tags::Learning);

    AddParameter(ParameterType_InputFilenameList, "il", "Input statistics");
    SetParameterDescription("il", "List of statistics files for each input image.");

    AddParameter(ParameterType_OutputFilename, "out", "Output sampling rates");
    SetParameterDescription("out",
                            "Output filename storing sampling rates (CSV "
                            "format with class name, required samples, total samples, and rate). "
                            "The given filename will be used with a suffix to indicate the "
                            "corresponding input index (for instance: rates.csv will give rates_1.csv"
                            ", rates_2.csv, ...).");

    AddParameter(ParameterType_Choice, "strategy", "Sampling strategy");

    AddChoice("strategy.byclass", "Set samples count for each class");
    SetParameterDescription("strategy.byclass", "Set samples count for each class");

    AddParameter(ParameterType_InputFilenameList, "strategy.byclass.in", "Number of samples by class");
    SetParameterDescription("strategy.byclass.in",
                            "Number of samples by class "
                            "(CSV format with class name in 1st column and required samples in the 2nd)."
                            "In the case of the custom multi-image mode, several inputs may be given for each image.");

    AddChoice("strategy.constant", "Set the same samples counts for all classes");
    SetParameterDescription("strategy.constant", "Set the same samples counts for all classes");

    AddParameter(ParameterType_String, "strategy.constant.nb", "Number of samples for all classes");
    SetParameterDescription("strategy.constant.nb",
                            "Number of samples for all classes."
                            "In the case of the custom multi-image mode, several values can be given for each image.");

    AddChoice("strategy.smallest", "Set same number of samples for all classes, with the smallest class fully sampled");
    SetParameterDescription("strategy.smallest", "Set same number of samples for all classes, with the smallest class fully sampled");

    AddChoice("strategy.percent", "Use a percentage of the samples available for each class");
    SetParameterDescription("strategy.percent", "Use a percentage of the samples available for each class");

    AddParameter(ParameterType_String, "strategy.percent.p", "The percentage(s) to use");
    SetParameterDescription("strategy.percent.p",
                            "The percentage(s) to use "
                            "In the case of the custom multi-image mode, several values can be given for each image.");

    AddChoice("strategy.total", "Set the total number of samples to generate, and use class proportions.");
    SetParameterDescription("strategy.total", "Set the total number of samples to generate, and use class proportions.");

    AddParameter(ParameterType_String, "strategy.total.v", "The number of samples to generate");
    SetParameterDescription("strategy.total.v",
                            "The number of samples to generate"
                            "In the case of the custom multi-image mode, several values can be given for each image.");

    AddChoice("strategy.all", "Use all samples");
    SetParameterDescription("strategy.all", "Use all samples");

    // Default strategy : smallest
    SetParameterString("strategy", "smallest");

    AddParameter(ParameterType_Choice, "mim", "Multi-Image Mode");

    AddChoice("mim.proportional", "Proportional");
    SetParameterDescription("mim.proportional", "Split the required number of samples proportionally");

    AddChoice("mim.equal", "equal");
    SetParameterDescription("mim.equal", "Equal split of the required number of samples");

    AddChoice("mim.custom", "Custom");
    SetParameterDescription("mim.custom", "Split the number of samples based on the user's choice.");

    // Doc example parameter settings
    SetDocExampleParameterValue("il", "stats_1.xml stats_2.xml");
    SetDocExampleParameterValue("out", "rates.csv");
    SetDocExampleParameterValue("strategy", "smallest");
    SetDocExampleParameterValue("mim", "proportional");

    SetOfficialDocLink();
  }

  void DoUpdateParameters() override
  {
  }

  void DoExecute() override
  {
    // Clear state
    m_CalculatorList->Clear();
    std::vector<std::string> inputs     = this->GetParameterStringList("il");
    unsigned int             nbInputs   = inputs.size();
    XMLReaderType::Pointer   statReader = XMLReaderType::New();
    for (unsigned int i = 0; i < nbInputs; ++i)
    {
      m_CalculatorList->PushBack(otb::SamplingRateCalculator::New());
      statReader->SetFileName(inputs[i]);
      ClassCountMapType classCount = statReader->GetStatisticMapByName<ClassCountMapType>("samplesPerClass");
      m_CalculatorList->SetNthClassCount(i, classCount);
    }

    // Cautions : direct mapping between the enum PartitionType and the choice order
    RateCalculatorListType::PartitionType partitionMode = static_cast<RateCalculatorListType::PartitionType>(this->GetParameterInt("mim"));

    unsigned int minParamSize = 1;
    if (partitionMode == RateCalculatorListType::CUSTOM)
    {
      // Check we have enough inputs for the custom mode
      minParamSize = nbInputs;
    }

    switch (this->GetParameterInt("strategy"))
    {
    // byclass
    case 0:
    {
      std::vector<std::string>       requiredFiles = this->GetParameterStringList("strategy.byclass.in");
      std::vector<ClassCountMapType> requiredCounts;
      if (requiredFiles.size() < minParamSize)
      {
        otbAppLogFATAL("Missing arguments in strategy.byclass.in to process sampling rates");
      }
      otbAppLogINFO("Sampling strategy : set number of samples for each class");
      for (unsigned int i = 0; i < minParamSize; i++)
      {
        requiredCounts.push_back(otb::SamplingRateCalculator::ReadRequiredSamples(requiredFiles[i]));
      }
      m_CalculatorList->SetNbOfSamplesByClass(requiredCounts, partitionMode);
    }
    break;
    // constant
    case 1:
    {
      std::vector<itksys::String> parts = itksys::SystemTools::SplitString(this->GetParameterString("strategy.constant.nb"), ' ');
      std::vector<unsigned long>  countList;
      for (unsigned int i = 0; i < parts.size(); i++)
      {
        if (!parts[i].empty())
        {
          std::string::size_type pos1 = parts[i].find_first_not_of(" \t");
          std::string::size_type pos2 = parts[i].find_last_not_of(" \t");
          std::string            value(parts[i].substr(pos1, pos2 - pos1 + 1));
          countList.push_back(boost::lexical_cast<unsigned long>(parts[i]));
        }
      }
      if (countList.size() < minParamSize)
      {
        otbAppLogFATAL("Missing arguments in strategy.constant.nb to process sampling rates");
      }
      otbAppLogINFO("Sampling strategy : set a constant number of samples for all classes");
      m_CalculatorList->SetNbOfSamplesAllClasses(countList, partitionMode);
    }
    break;
    // smallest class
    case 2:
    {
      otbAppLogINFO("Sampling strategy : fit the number of samples based on the smallest class");
      m_CalculatorList->SetMinimumNbOfSamplesByClass(partitionMode);
    }
    break;
    // percent
    case 3:
    {
      std::vector<itksys::String> parts = itksys::SystemTools::SplitString(this->GetParameterString("strategy.percent.p"), ' ');
      std::vector<double>         percentList;
      for (unsigned int i = 0; i < parts.size(); i++)
      {
        if (!parts[i].empty())
        {
          std::string::size_type pos1 = parts[i].find_first_not_of(" \t");
          std::string::size_type pos2 = parts[i].find_last_not_of(" \t");
          std::string            value(parts[i].substr(pos1, pos2 - pos1 + 1));
          percentList.push_back(boost::lexical_cast<double>(parts[i]));

          if (percentList.back() < 0 || percentList.back() > 1)
          {
            otbAppLogFATAL("Percent parameter should be in range [0,1]");
          }
        }
      }
      if (percentList.size() < minParamSize)
      {
        otbAppLogFATAL("Missing arguments in strategy.percent.p to process sampling rates");
      }
      otbAppLogINFO("Sampling strategy : set a percentage of samples to be used.");
      m_CalculatorList->SetPercentageOfSamples(percentList, partitionMode);
    }
    break;

    // total
    case 4:
    {
      std::vector<itksys::String> parts = itksys::SystemTools::SplitString(this->GetParameterString("strategy.total.v"), ' ');
      std::vector<unsigned long>  totalList;
      for (unsigned int i = 0; i < parts.size(); i++)
      {
        if (!parts[i].empty())
        {
          std::string::size_type pos1 = parts[i].find_first_not_of(" \t");
          std::string::size_type pos2 = parts[i].find_last_not_of(" \t");
          std::string            value(parts[i].substr(pos1, pos2 - pos1 + 1));
          totalList.push_back(boost::lexical_cast<unsigned long>(parts[i]));
        }
      }
      if (totalList.size() < minParamSize)
      {
        otbAppLogFATAL("Missing arguments in strategy.total.v to process sampling rates");
      }
      otbAppLogINFO("Sampling strategy : set a constant number of samples for all classes");
      m_CalculatorList->SetTotalNumberOfSamples(totalList, partitionMode);
    }
    break;
    // all samples
    case 5:
    {
      otbAppLogINFO("Sampling strategy : take all samples");
      m_CalculatorList->SetAllSamples(partitionMode);
    }
    break;
    default:
      otbAppLogFATAL("Strategy mode unknown :" << this->GetParameterString("strategy"));
      break;
    }

    std::ostringstream oss;
    std::string        outputPath(this->GetParameterString("out"));
    std::string        outputBase    = outputPath.substr(0, outputPath.find_last_of('.'));
    std::string        outputExt     = outputPath.substr(outputPath.find_last_of('.'), std::string::npos);
    unsigned int       overflowCount = 0;
    bool               noSamples     = true;
    for (unsigned int i = 0; i < nbInputs; i++)
    {
      // Print results
      oss.str(std::string(""));
      oss << " className  requiredSamples  totalSamples  rate\n";
      MapRateType rates = m_CalculatorList->GetRatesByClass(i);
      if (!rates.empty())
      {
        noSamples = false;
      }
      MapRateType::const_iterator itRates = rates.begin();
      for (; itRates != rates.end(); ++itRates)
      {
        otb::SamplingRateCalculator::TripletType tpt = itRates->second;
        oss << itRates->first << "\t" << tpt.Required << "\t" << tpt.Tot << "\t" << tpt.Rate;
        if (tpt.Required > tpt.Tot)
        {
          overflowCount++;
          oss << "\t[OVERFLOW]";
        }
        oss << std::endl;
      }
      otbAppLogINFO("Sampling rates for image " << i + 1 << " : " << oss.str());
      // Output results to disk
      oss.str(std::string(""));
      oss << outputBase << "_" << i + 1 << outputExt;
      m_CalculatorList->GetNthElement(i)->Write(oss.str());
    }
    if (noSamples)
    {
      otbAppLogFATAL("No samples found in the inputs!");
    }
    if (overflowCount)
    {
      std::string plural(overflowCount > 1 ? "s" : "");
      otbAppLogWARNING(<< overflowCount << " case" << plural << " of overflow detected! (requested number of samples higher than total available samples)");
    }
  }

  RateCalculatorListType::Pointer m_CalculatorList;
};

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

OTB_APPLICATION_EXPORT(otb::Wrapper::MultiImageSamplingRate)