/*============================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center (DKFZ)
All rights reserved.

Use of this source code is governed by a 3-clause BSD license that can be
found in the LICENSE file.

============================================================================*/

/**
 * \file mitkModelTestFixture.h
 * \brief Provides a reusable test fixture base class for validating MITK pharmacokinetic / model fit models
 *        against JSON-based reference data.
 */

#ifndef mitkModelTestFixture_h
#define mitkModelTestFixture_h


#include <mitkTestingMacros.h>
#include <mitkTestingConfig.h>
#include <mitkTestFixture.h>
#include <itksys/SystemTools.hxx>
#include <mitkModelBase.h>
#include <mitkModelFitException.h>

#include <nlohmann/json.hpp>
#include <iostream>
#include <fstream>
#include <mitkVector.h>

using json = nlohmann::json;

namespace mitk
{
  /**
   * \brief Test fixture base class for model fit unit tests.
   *
   * Provides static helper methods that load model reference data from JSON files
   * and compare a model's profile, signal output, and derived parameters against
   * the reference values. Subclass this fixture, set up your model in \c setUp(),
   * and use the provided comparison methods in your test cases.
   *
   * The expected JSON structure contains:
   * - A "profile" object with numberOfParameters, parameterNames, parameterScales,
   *   parameterUnits, static parameter metadata, functionString, classID, etc.
   * - A "modelValues" array with entries containing timeGrid, modelParameterValues,
   *   staticParameterValues, signal, and derivedParameterValues.
   *
   * \sa mitk::ModelBase, mitk::TestFixture
   */
  class mitkModelTestFixture : public mitk::TestFixture
  {
  public:
    /**
     * \brief Parse a JSON file from the test data directory.
     *
     * \param path Relative path to the JSON file within the MITK test data directory.
     * \return The parsed JSON object. Returns an empty object if the file cannot be opened or parsed.
     */
    static json ParseJSONFile(const std::string path)
    {
      std::string pathToFile = GetTestDataFilePath(path);
      std::ifstream file(pathToFile);
      json json_obj;
      if (!file.is_open())
      {
        MITK_ERROR << "Could not open \"" << pathToFile << "\"!";
      }
      try
      {
        json_obj = nlohmann::json::parse(file);
      }
      catch (const json::exception& e)
      {
        MITK_ERROR << "Could not parse JSON file!";
      }
      return json_obj;
    }

    /**
     * \brief Extract model parameter values from a JSON object into a ParametersType array.
     *
     * Reads the "modelParameterValues" array from the JSON object and returns it
     * as a ModelBase::ParametersType.
     *
     * \param modelValues_json_obj JSON object containing a "modelParameterValues" array.
     * \return The parsed parameter values.
     */
    static ModelBase::ParametersType ParseTestParameters(const json modelValues_json_obj)
    {
      ModelBase::ParametersType testparameters;
      testparameters.SetSize(modelValues_json_obj["modelParameterValues"].size());
      for (unsigned long i = 0; i < modelValues_json_obj["modelParameterValues"].size(); ++i)
      {
        testparameters[i] = modelValues_json_obj["modelParameterValues"][i];
      }
      return testparameters;
    }

    /**
     * \brief Parse static parameter names and values from JSON into a StaticParameterMapType.
     *
     * \param profile_json_obj      JSON object containing "staticParameterNames".
     * \param modelValues_json_obj  JSON object containing "staticParameterValues".
     * \return A map of static parameter names to their value vectors.
     */
    static ModelBase::StaticParameterMapType ParseStaticParameters(const json profile_json_obj, const json modelValues_json_obj)
    {
      ModelBase::StaticParameterMapType staticParameterMap;
      for (unsigned long i = 0; i < modelValues_json_obj["staticParameterValues"].size(); ++i)
      {
        ModelBase::StaticParameterValuesType staticParameterValues;
        ModelBase::ParameterNameType staticParameterName;
        staticParameterName = profile_json_obj["staticParameterNames"][i];
        for (unsigned int j = 0; j < modelValues_json_obj["staticParameterValues"][i].size(); ++j)
        {
          staticParameterValues.push_back(modelValues_json_obj["staticParameterValues"][i][j]);
        }
        staticParameterMap.insert(std::make_pair(staticParameterName, staticParameterValues));
      }
      return staticParameterMap;
    }

    /**
     * \brief Convenience method that parses static parameters from JSON and applies them to a model.
     *
     * \param testmodel             The model instance to configure.
     * \param profile_json_obj      JSON object containing "staticParameterNames".
     * \param modelValues_json_obj  JSON object containing "staticParameterValues".
     */
    static void SetStaticParametersForTest(mitk::ModelBase::Pointer testmodel, const json profile_json_obj, const json modelValues_json_obj)
    {
      mitk::ModelBase::StaticParameterMapType staticParameterMap;
      staticParameterMap = ParseStaticParameters(profile_json_obj, modelValues_json_obj);
      testmodel->SetStaticParameters(staticParameterMap);
    }

    /**
     * \brief Assert that a model's profile (parameter names, scales, units, etc.) matches a JSON reference.
     *
     * Checks numberOfParameters, parameterNames, parameterScales, parameterUnits,
     * derived parameters, static parameters, functionString, classID, modelDisplayName,
     * modelType, axis names, and axis units.
     *
     * \param testmodel        The model instance to validate.
     * \param profile_json_obj JSON object containing the expected profile values.
     */
    static void CompareModelAndReferenceProfile(const mitk::ModelBase::Pointer testmodel, const json profile_json_obj)
    {
      CPPUNIT_ASSERT_MESSAGE("Checking number of parameters in model.", testmodel->GetNumberOfParameters() == profile_json_obj["numberOfParameters"].get<unsigned int>());
      for (unsigned long i = 0; i < profile_json_obj["numberOfParameters"]; i++)
      {
        CPPUNIT_ASSERT_MESSAGE("Checking parameter names.", testmodel->GetParameterNames()[i] == profile_json_obj["parameterNames"][i].get<std::string>());
        CPPUNIT_ASSERT_MESSAGE("Checking parameter scales.", testmodel->GetParameterScales()[testmodel->GetParameterNames()[i]] == profile_json_obj["parameterScales"][i].get<double>());
        CPPUNIT_ASSERT_MESSAGE("Checking parameter units.", testmodel->GetParameterUnits()[testmodel->GetParameterNames()[i]] == profile_json_obj["parameterUnits"][i].get<std::string>());
      }
      CPPUNIT_ASSERT_MESSAGE("Checking number of derived parameters in model.", testmodel->GetNumberOfDerivedParameters() == profile_json_obj["numberOfDerivedParameters"].get<unsigned int>());
      for (unsigned long i = 0; i < profile_json_obj["numberOfDerivedParameters"]; i++)
      {
        CPPUNIT_ASSERT_MESSAGE("Checking derived parameter names.", testmodel->GetDerivedParameterNames()[i] == profile_json_obj["derivedParameterNames"][i].get<std::string>());
        CPPUNIT_ASSERT_MESSAGE("Checking derived parameter scales.", testmodel->GetDerivedParameterScales()[testmodel->GetDerivedParameterNames()[i]] == profile_json_obj["derivedParameterScales"][i].get<double>());
        CPPUNIT_ASSERT_MESSAGE("Checking derived parameter units.", testmodel->GetDerivedParameterUnits()[testmodel->GetDerivedParameterNames()[i]] == profile_json_obj["derivedParameterUnits"][i].get<std::string>());
      }
      CPPUNIT_ASSERT_MESSAGE("Checking number of static parameters in model.", testmodel->GetNumberOfStaticParameters() == profile_json_obj["numberOfStaticParameters"].get<unsigned int>());
      for (unsigned long i = 0; i < profile_json_obj["numberOfStaticParameters"]; i++)
      {
        CPPUNIT_ASSERT_MESSAGE("Checking static parameter names.", testmodel->GetStaticParameterNames()[i] == profile_json_obj["staticParameterNames"][i].get<std::string>());
        CPPUNIT_ASSERT_MESSAGE("Checking static parameter units.", testmodel->GetStaticParameterUnits()[testmodel->GetStaticParameterNames()[i]] == profile_json_obj["staticParameterUnits"][i].get<std::string>());
      }
      CPPUNIT_ASSERT_MESSAGE("Checking function string.", testmodel->GetFunctionString() == profile_json_obj["functionString"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking class ID.", testmodel->GetClassID() == profile_json_obj["classID"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking model display name.", testmodel->GetModelDisplayName() == profile_json_obj["modelDisplayName"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking model type.", testmodel->GetModelType() == profile_json_obj["modelType"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking x name.", testmodel->GetXName() == profile_json_obj["xName"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking x axis name.", testmodel->GetXAxisName() == profile_json_obj["xAxisName"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking x axis unit.", testmodel->GetXAxisUnit() == profile_json_obj["xAxisUnit"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking y axis name.", testmodel->GetYAxisName() == profile_json_obj["yAxisName"].get<std::string>());
      CPPUNIT_ASSERT_MESSAGE("Checking y axis unit.", testmodel->GetYAxisUnit() == profile_json_obj["yAxisUnit"].get<std::string>());
    }

    /**
     * \brief Assert that a model's computed signal matches a JSON reference for all parameter sets.
     *
     * For each entry in the "modelValues" array, sets the static parameters and time grid,
     * computes the signal via \c GetSignal(), and compares each sample against the
     * reference "signal" array with a tolerance of 1e-6.
     *
     * \param testmodel             The model instance to validate.
     * \param modelValues_json_obj  JSON object containing the "modelValues" array.
     * \param profile_json_obj      JSON object containing static parameter names.
     */
    static void CompareModelAndReferenceSignal(mitk::ModelBase::Pointer testmodel, const json modelValues_json_obj, const json profile_json_obj)
    {

      for (unsigned int j = 0; j < modelValues_json_obj["modelValues"].size(); j++)
      {
        json modelValues_json_obj_current = modelValues_json_obj["modelValues"][j];

        SetStaticParametersForTest(testmodel, profile_json_obj, modelValues_json_obj_current);

        // Set time grid
        mitk::ModelBase::TimeGridType timeGrid;
        timeGrid.SetSize(modelValues_json_obj_current["timeGrid"].size());
        for (unsigned long i = 0; i < modelValues_json_obj_current["timeGrid"].size(); ++i)
        {
          timeGrid[i] = modelValues_json_obj_current["timeGrid"][i];
        }
        testmodel->SetTimeGrid(timeGrid);

        // Parse test parameters
        mitk::ModelBase::ParametersType testparameters;
        testparameters = ParseTestParameters(modelValues_json_obj_current);

        //Generate signal
        mitk::ModelBase::ModelResultType signal;
        signal = testmodel->GetSignal(testparameters);


        //ComputeModelfunction is called within GetSignal(), therefore no explicit testing of ComputeModelFunction()
        // The reference signal has been calculated in a matlab script.
        std::stringstream ss;
        ss << "Checking signal for model parameter set " << j << ".";
        std::string message = ss.str();
        for (unsigned long i = 0; i < signal.size(); i++)
        {
          CPPUNIT_ASSERT_MESSAGE(message, mitk::Equal(signal[i], modelValues_json_obj_current["signal"][i], 1e-6, true) == true);
        }
      }
    }

    /**
     * \brief Assert that a model's derived parameters match a JSON reference for all parameter sets.
     *
     * For each entry in the "modelValues" array, sets the time grid, computes the
     * derived parameters via \c GetDerivedParameters(), and compares each value against
     * the reference "derivedParameterValues" array with a tolerance of 1e-6.
     *
     * \param testmodel             The model instance to validate.
     * \param modelValues_json_obj  JSON object containing the "modelValues" array.
     */
    static void CompareModelAndReferenceDerivedParameters(const mitk::ModelBase::Pointer testmodel, json modelValues_json_obj)
    {
      for (unsigned int j = 0; j < modelValues_json_obj["modelValues"].size(); j++)
      {
        json modelValues_json_obj_current = modelValues_json_obj["modelValues"][j];

        // Set time grid
        mitk::ModelBase::TimeGridType timeGrid;
        timeGrid.SetSize(modelValues_json_obj_current["timeGrid"].size());
        for (unsigned long i = 0; i < modelValues_json_obj_current["timeGrid"].size(); ++i)
        {
          timeGrid[i] = modelValues_json_obj_current["timeGrid"][i];
        }
        testmodel->SetTimeGrid(timeGrid);
        mitk::ModelBase::ParametersType testparameters;
        testparameters = mitk::mitkModelTestFixture::ParseTestParameters(modelValues_json_obj_current);
        mitk::ModelBase::DerivedParameterMapType derivedParameterMap;
        derivedParameterMap = testmodel->GetDerivedParameters(testparameters);
        std::stringstream ss;
        ss << "Checking derived parameters for model parameter set " << j << ".";
        std::string message = ss.str();
        for (unsigned long i = 0; i < modelValues_json_obj_current["derivedParameterValues"].size(); i++)
        {
          CPPUNIT_ASSERT_MESSAGE(message, mitk::Equal(derivedParameterMap[testmodel->GetDerivedParameterNames()[i]], modelValues_json_obj_current["derivedParameterValues"][i], 1e-6, true) == true);
        }
      }
    }
  };
}
#endif // mitkModelTestFixture_h
