// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
//
// 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 "paddle/fluid/pir/utils/name_analysis.h"
#include <unordered_map>
#include "paddle/fluid/pir/dialect/kernel/ir/kernel_op.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"

namespace pir {
namespace utils {
namespace name_analysis {

Value GetOutputValueByName(const Program &program, const std::string &name) {
  auto &block = *program.block();
  StrAttribute name_attr = StrAttribute::get(IrContext::Instance(), name);
  Value value;
  for (auto &op : block) {
    if (op.isa<ShadowOutputOp>()) {
      if (op.attribute("output_name") == name_attr) {
        if (value) {
          PADDLE_THROW(common::errors::PreconditionNotMet(
              "More than one shadow output named with %s found.", name));
        }
        value = op.operand_source(0);
      }
    } else if (op.isa<paddle::dialect::FeedOp>() ||
               op.isa<paddle::dialect::DataOp>() ||
               op.isa<paddle::dialect::FetchOp>()) {
      if (op.attribute("name") == name_attr) {
        if (value) {
          PADDLE_THROW(common::errors::PreconditionNotMet(
              "More than one feed/fetch named with %s found.", name));
        }
        value = op.result(0);
      }
    }
  }
  return value;
}

Value GetValueByNameInPhiKernelProgram(const Program &program,
                                       const std::string &name) {
  auto &block = *program.block();
  StrAttribute name_attr = StrAttribute::get(IrContext::Instance(), name);
  Value value;

  for (auto &op : block) {
    if (op.isa<ShadowOutputOp>()) {
      if (op.attribute("output_name") == name_attr) {
        if (value) {
          PADDLE_THROW(common::errors::PreconditionNotMet(
              "More than one shadow output named with %s found.", name));
        }
        value = op.operand_source(0);
      }
    } else if (op.isa<paddle::dialect::PhiKernelOp>()) {
      if (op.attribute("op_name").dyn_cast<StrAttribute>().AsString() ==
          "pd_op.data") {
        if (op.attribute("name") == name_attr) {
          if (value) {
            PADDLE_THROW(common::errors::PreconditionNotMet(
                "More than one feed/fetch named with %s found.", name));
          }
          value = op.result(0);
        }
      }
    }
  }
  return value;
}

Value GetParameterValueByName(const Program &program, const std::string &name) {
  auto &block = *program.block();
  StrAttribute name_attr = StrAttribute::get(IrContext::Instance(), name);
  Value value;
  for (auto &op : block) {
    if (op.isa<ParameterOp>()) {
      if (op.attribute("parameter_name") == name_attr) {
        if (value) {
          PADDLE_THROW(common::errors::PreconditionNotMet(
              "More than one parameter named with %s found.", name));
        }
        value = op.result(0);
      }
    }
  }
  return value;
}

std::unordered_map<std::string, Value> GetAllParameterValues(
    const Program &program) {
  auto &block = *program.block();
  std::unordered_map<std::string, Value> values;
  for (auto &op : block) {
    if (op.isa<ParameterOp>()) {
      values
          [op.attribute("parameter_name").dyn_cast<StrAttribute>().AsString()] =
              op.result(0);
    }
  }
  return values;
}

void SetValueName(Value value, const std::string name) {
  Operation *define_op = value.defining_op();
  if (define_op->isa<ParameterOp>()) {
    define_op->set_attribute("parameter_name",
                             StrAttribute::get(IrContext::Instance(), name));
  } else if (define_op->isa<paddle::dialect::DataOp>()) {
    define_op->set_attribute("name",
                             StrAttribute::get(IrContext::Instance(), name));
  } else if (auto block_arg = value.dyn_cast<BlockArgument>()) {
    PADDLE_THROW(common::errors::InvalidArgument(
        "Can Not set name for BlockArgument! "));
  } else if (value.first_use()) {
    auto nextOp = value.first_use().owner();
    if (nextOp->isa<ShadowOutputOp>()) {
      nextOp->set_attribute("output_name",
                            StrAttribute::get(IrContext::Instance(), name));
    } else {
      PADDLE_THROW(common::errors::InvalidArgument(
          "Currently, we can only set name of Value which is "
          "shadowoutput "));
    }
  } else {
    PADDLE_THROW(common::errors::InvalidArgument(
        "Currently, we can only set name of Value that "
        "is persistable"));
  }
}

std::map<std::string, std::string> RenameValue(Value value,
                                               const std::string &new_name,
                                               Block *block) {
  std::map<std::string, std::string> rename_mapping;
  VLOG(5) << "Starting to rename value to " << new_name;
  // Handle kwarg
  for (auto [name, kwarg] : block->kwargs()) {
    if (kwarg == value) {
      if (name == new_name) {
        break;
      }
      Value new_value;
      if (block->kwargs().count(new_name)) {
        new_value = block->kwargs().at(new_name);
      } else {
        new_value = block->AddKwarg(new_name, value.type());
      }
      value.ReplaceAllUsesWith(new_value);
      block->EraseKwarg(name);
      value = new_value;
      VLOG(5) << "Value is kwarg, rename it from " << name << " to "
              << new_name;
      rename_mapping.insert({name, new_name});
      break;
    }
  }

  // Handle inputs
  auto defining_op = value.defining_op();
  if (defining_op) {
    // Handle DataOp
    if (defining_op->isa<paddle::dialect::DataOp>()) {
      auto name = defining_op->attribute<StrAttribute>("name").AsString();
      if (name != new_name) {
        defining_op->set_attribute(
            "name", StrAttribute::get(IrContext::Instance(), new_name));
        VLOG(5) << "Value is defined by DataOp, rename it from " << name
                << " to " << new_name;
        rename_mapping.insert({name, new_name});
      }
    }

    // Handle ParameterOp
    if (defining_op->isa<ParameterOp>()) {
      auto name =
          defining_op->attribute<StrAttribute>("parameter_name").AsString();
      if (name != new_name) {
        defining_op->set_attribute(
            "parameter_name",
            StrAttribute::get(IrContext::Instance(), new_name));
        VLOG(5) << "Value is defined by ParameterOp, rename it from " << name
                << " to " << new_name;
        rename_mapping.insert({name, new_name});
      }
    }

    // Handle ConstantTensorOp
    if (defining_op->isa<ConstantTensorOp>()) {
      auto name =
          defining_op->attribute<StrAttribute>("tensor_name").AsString();
      if (name != new_name) {
        defining_op->set_attribute(
            "tensor_name", StrAttribute::get(IrContext::Instance(), new_name));
        VLOG(5) << "Value is defined by ConstantTensorOp, rename it from "
                << name << " to " << new_name;
        rename_mapping.insert({name, new_name});
      }
    }
  }

  // Handle outputs
  for (auto iter = value.use_begin(); iter != value.use_end(); ++iter) {
    auto user_op = iter->owner();
    if (user_op->isa<ShadowOutputOp>()) {
      // Handle ShadowOutputOp
      auto name = user_op->attribute<StrAttribute>("output_name").AsString();
      if (name == new_name) {
        continue;
      }
      user_op->set_attribute(
          "output_name", StrAttribute::get(IrContext::Instance(), new_name));
      VLOG(5) << "Value is used by ShadowOutputOp, rename it from " << name
              << " to " << new_name;
      rename_mapping.insert({name, new_name});
    } else if (user_op->isa<SetParameterOp>()) {
      // Handle SetParameterOp
      auto name = user_op->attribute<StrAttribute>("parameter_name").AsString();
      if (name == new_name) {
        continue;
      }
      user_op->set_attribute(
          "parameter_name", StrAttribute::get(IrContext::Instance(), new_name));
      VLOG(5) << "Value is used by SetParameterOp, rename it from " << name
              << " to " << new_name;
      rename_mapping.insert({name, new_name});
    }
  }
  return rename_mapping;
}

std::optional<std::string> GetValueInputName(Value value) {
  std::optional<std::string> name;
  if (auto block_arg = value.dyn_cast<BlockArgument>()) {
    if (block_arg.is_kwarg()) {
      name = block_arg.keyword();
    } else {
      name = "arg_" + std::to_string(block_arg.index());
    }
  } else if (auto param_op = value.defining_op<ParameterOp>()) {
    name = param_op.param_name();
  } else if (auto data_op = value.defining_op<paddle::dialect::DataOp>()) {
    name = data_op.attribute<StrAttribute>("name").AsString();
  } else if (auto constant_op = value.defining_op<ConstantTensorOp>()) {
    name = constant_op.tensor_name();
  }
  return name;
}

std::vector<std::string> GetValueOutputNames(Value value) {
  std::vector<std::string> names;
  for (auto iter = value.use_begin(); iter != value.use_end(); ++iter) {
    if (iter->owner()->isa<ShadowOutputOp>()) {
      names.push_back(
          iter->owner()->attribute<StrAttribute>("output_name").AsString());
    } else if (iter->owner()->isa<SetParameterOp>()) {
      names.push_back(
          iter->owner()->attribute<StrAttribute>("parameter_name").AsString());
    }
  }
  return names;
}

std::vector<std::string> GetValueAllNames(Value value) {
  std::vector<std::string> names;
  std::optional<std::string> input_name = GetValueInputName(value);
  if (input_name.has_value()) {
    names.push_back(input_name.value());
  }

  std::vector<std::string> output_name = GetValueOutputNames(value);
  for (auto &name : output_name) {
    names.push_back(name);
  }

  return names;
}

std::optional<std::string> TryGetValueFirstName(Value value) {
  std::optional<std::string> name;

  auto names = GetValueAllNames(value);
  if (!names.empty()) {
    return names[0];
  }

  return name;
}

std::string GetValueFirstName(Value value) {
  auto name = TryGetValueFirstName(value);

  PADDLE_ENFORCE(name.has_value(),
                 common::errors::InvalidArgument(
                     "Currently, we can only get name of Value from "
                     "DataOp/ParameterOp/BlockArgument/ConstantTensorOp/"
                     "SetParameterOp and ShadowOutputOp."));

  return name.value();
}

std::unordered_map<std::string, Value> GetAllNamedValues(
    const Program &program) {
  std::unordered_map<std::string, Value> named_values;
  std::vector<Value> all_values;
  all_values.insert(all_values.end(),
                    program.block()->args().begin(),
                    program.block()->args().end());
  for (const auto &[k, v] : program.block()->kwargs()) {
    all_values.push_back(v);
  }
  for (auto op : program.block()->ops()) {
    for (auto var : op->results()) {
      all_values.push_back(var);
    }
  }

  for (auto &value : all_values) {
    std::optional<std::string> name = TryGetValueFirstName(value);
    if (!name.has_value()) {
      continue;
    }
    named_values[name.value()] = value;
  }

  return named_values;
}

}  // namespace name_analysis
}  // namespace utils
}  // namespace pir