#include "tlbmc/thermal/controller/pid_controller.h"

#include <cmath>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "g3/macros.h"
#include <nlohmann/json_fwd.hpp>
#include "thermal_config.pb.h"
#include "thermal_controller_specification.pb.h"
#include "tlbmc/thermal/controller/algorithm/pid.h"
#include "tlbmc/thermal/controller/optimizer/pid/pid_optimizer.h"
#include "tlbmc/thermal/zone_manager.h"

namespace milotic_tlbmc {
namespace thermal {

absl::StatusOr<std::unique_ptr<PidController>> PidController::Create(
    PidControllerParameters params) {
  if (params.input_sensors.empty()) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "No input sensor is assigned to PID controller `%s`!", params.id));
  }

  auto pid_controller = absl::WrapUnique(
      params.pid_optimizer_info.has_value()
          ? new PidController(params.id, params.zone_manager,
                              std::move(params.pid_loop), params.setpoint,
                              std::move(params.input_sensors),
                              std::move(params.pid_optimizer_info.value()))
          : new PidController(params.id, params.zone_manager,
                              std::move(params.pid_loop), params.setpoint,
                              std::move(params.input_sensors)));

  ECCLESIA_RETURN_IF_ERROR(
      pid_controller->SetInputProcessType(params.input_process_type));

  return pid_controller;
}

void PidController::ProcessThermalLoop() {
  double input = ProcessInput();

  if (pid_optimizer_info_.optimizer != nullptr &&
      !zone_owner_->GetFailsafeMode()) {
    if (!pid_optimizer_info_.states.pid_optimizer_metadata.initialized) {
      pid_optimizer_info_.optimizer->Initialize(
          pid_optimizer_info_.states.pid_optimizer_init_params);
      pid_optimizer_info_.states.pid_optimizer_metadata.initialized = true;
    }
    if (!pid_optimizer_info_.optimizer->FinishedTuning()) {
      zone_owner_->TryLogThermalDebugMessage(absl::StrFormat(
          "PID controller %s is tuning PID coefficients.", GetId()));
      pid_optimizer_info_.states.pid_optimizer_tuning_params.input = input;
      double output = pid_optimizer_info_.optimizer->TunePid(
          pid_optimizer_info_.states.pid_optimizer_tuning_params);
      // For further optimization, we may use actual system time instead of the
      // PID loop config to improve the tuning accuracy.
      // However, PID is already an overly simplified model, and such a
      // precision loss may be negligible in practice.
      pid_optimizer_info_.states.pid_optimizer_tuning_params.current_time_ms +=
          static_cast<uint32_t>(pid_loop_->GetSampleTime() * 1000);
      pid_loop_->SetCoeffProportional(
          pid_optimizer_info_.optimizer->GetCoeffP());
      pid_loop_->SetCoeffIntegral(pid_optimizer_info_.optimizer->GetCoeffI());
      pid_loop_->SetCoeffDerivative(pid_optimizer_info_.optimizer->GetCoeffD());

      ProcessOutput(output);
      return;
    }
  }

  double setpoint = GetSetpoint();

  ProcessOutput(GetPidLoop()->ExecutePidLoop(
      input, setpoint, zone_owner_->GetDebugPidEnabled()));
}

double PidController::ProcessInput() {
  const std::vector<PidControllerSensorInfo>& input_sensors = GetInputSensors();
  // If no input sensor is assigned to this PID, always use `setpoint_` as the
  // input.
  if (input_sensors.empty()) {
    zone_owner_->TryLogThermalDebugMessage(
        absl::StrFormat("PID controller %s is assigned with no input sensor. "
                        "Use setpoint %lf as the input.\n",
                        GetId(), GetSetpoint()));
    return GetSetpoint();
  }
  std::string leader_sensor = input_sensors[0].sensor_key();
  double input_value = InitializeInputBasedOnPidThermalType();
  bool is_input_acceptable = false;

  for (const PidControllerSensorInfo& sensor : input_sensors) {
    double sensor_value =
        zone_owner_->GetSensorReadings(sensor.sensor_key()).scaled_value;

    // The sensor reading must not be NaN or infinite.
    // Note that it is fine to get negative values.
    if (!std::isfinite(sensor_value)) {
      continue;
    }

    // Perform temperature to margin conversion if needed.
    double org_sensor_value = sensor_value;
    sensor_value = ProcessSensorReadingForMarginPid(sensor_value, sensor);
    if (org_sensor_value != sensor_value) {
      zone_owner_->TryLogThermalDebugMessage(absl::StrFormat(
          "Converting temperature to margin: (temperature=%lf, Tjmax=%lf, "
          "margin=%lf)\n",
          org_sensor_value, sensor.max_junction_temperature(), sensor_value));
    }

    double org_input_value = input_value;
    input_value =
        IntegrateSensorReadingBasedOnPidThermalType(input_value, sensor_value);

    if (org_input_value != input_value) {
      leader_sensor = sensor.sensor_key();
    }

    is_input_acceptable = true;
  }

  if (is_input_acceptable) {
    zone_owner_->TryLogThermalDebugMessage(absl::StrFormat(
        "PID controller %s chose the input value %lf from sensor %s\n", GetId(),
        input_value, leader_sensor));
    return input_value;
  }

  // If no valid input available, use the setpoint.
  // This provides safe behavior until the input becomes acceptable.
  zone_owner_->TryLogThermalDebugMessage(absl::StrFormat(
      "PID controller %s chose the setpoint %lf as the input value\n", GetId(),
      GetSetpoint()));
  return GetSetpoint();
}

void PidController::ProcessOutput(double output) {
  zone_owner_->AddSetpoint(GetId(), output);
  zone_owner_->TryLogThermalDebugMessage(absl::StrFormat(
      "PID controller %s outputs setpoint: %lf\n", GetId(), output));
  if (zone_owner_->GetDebugEnabled()) {
    zone_owner_->AppendSampledData(absl::StrCat(output, ","));
  }
}

void PidController::SetCoefficients(
    const PidControllerCoefficients& coefficients) {
  if (coefficients.coeff_proportional.has_value()) {
    pid_loop_->SetCoeffProportional(*coefficients.coeff_proportional);
  }
  if (coefficients.coeff_integral.has_value()) {
    pid_loop_->SetCoeffIntegral(*coefficients.coeff_integral);
  }
  if (coefficients.coeff_derivative.has_value()) {
    pid_loop_->SetCoeffDerivative(*coefficients.coeff_derivative);
  }
  if (coefficients.feed_forward_offset.has_value()) {
    pid_loop_->SetFeedForwardOffset(*coefficients.feed_forward_offset);
  }
  if (coefficients.feed_forward_gain.has_value()) {
    pid_loop_->SetFeedForwardGain(*coefficients.feed_forward_gain);
  }
  if (coefficients.setpoint.has_value()) {
    setpoint_ = kDefaultPidSetpointValue;
    if (coefficients.setpoint > 0) {
      setpoint_ = *coefficients.setpoint;
    }
  }
}

void PidController::PreparePidOptimizer() {
  // This function will only be called when `pid_optimizer_info_` is present.
  // Therefore, no need for nullability check.
  pid_optimizer_info_.states.pid_optimizer_init_params.coeff_p =
      pid_loop_->GetCoeffProportional();
  pid_optimizer_info_.states.pid_optimizer_init_params.coeff_i =
      pid_loop_->GetCoeffIntegral();
  pid_optimizer_info_.states.pid_optimizer_init_params.coeff_d =
      pid_loop_->GetCoeffDerivative();

  pid_optimizer_info_.states.pid_optimizer_metadata =
      optimizer::PidOptimizerMetadata{
          .initialized = false,
      };
}

nlohmann::json PidController::ToJson() const {
  nlohmann::json json;
  json["DCoefficient"] = pid_loop_->GetCoeffDerivative();
  json["FFGainCoefficient"] = pid_loop_->GetFeedForwardGain();
  json["FFOffCoefficient"] = pid_loop_->GetFeedForwardOffset();
  json["ICoefficient"] = pid_loop_->GetCoeffIntegral();
  json["ILimitMax"] = pid_loop_->GetIntegralLimitMax();
  json["ILimitMin"] = pid_loop_->GetIntegralLimitMin();
  json["OutLimitMax"] = pid_loop_->GetOutputLimitMax();
  json["OutLimitMin"] = pid_loop_->GetOutputLimitMin();
  json["PCoefficient"] = pid_loop_->GetCoeffProportional();
  json["SetPoint"] = GetSetpoint();
  json["SlewNeg"] = pid_loop_->GetSlewNeg();
  json["SlewPos"] = pid_loop_->GetSlewPos();
  json["Inputs"] = nlohmann::json::array();
  for (const auto& input : input_sensors_) {
    json["Inputs"].push_back(input.sensor_key());
  }
  return json;
}

}  // namespace thermal
}  // namespace milotic_tlbmc