#include "TestCheck.hpp"
#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include <chrono>
#include <functional>
#include <map>
#include <set>
#include <iostream>

using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
using E = PimExecutionEvent;
static bool sawJointMove = false, sawSuccessiveMoves = false;

// Test-only workloads use stage-boundary host transfers and local direct links.
// The same queue
// engine and complete-state evidence rank reference and searched schedules.
static WholeScheduleCost queueCost(const ComputeGraph &graph,
    const MergeScheduleResult &schedule, const PimInstructionTiming &target) {
  PimExecutionProgram program(target.processorCount);
  std::vector<std::vector<size_t>> tasks(target.processorCount);
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    tasks[schedule.computeToCpuMap.lookup(graph.nodes[task].instance)].push_back(task);
  for (auto &core : tasks)
    llvm::sort(core, [&](size_t a, size_t b) {
      return schedule.computeToCpuSlotMap.lookup(graph.nodes[a].instance)
        < schedule.computeToCpuSlotMap.lookup(graph.nodes[b].instance);
    });
  struct Transfer { size_t source, target, reader, writer, reg; bool host; };
  std::vector<Transfer> transfers;
  std::vector<std::set<size_t>> readers(target.processorCount), writers(target.processorCount);
  // Branched synthetic workloads use host routing to avoid inventing a second
  // blocking-transfer scheduler in this test-only objective.
  bool branching = llvm::any_of(graph.successors, [](const auto &edges) { return edges.size() > 1; });
  for (size_t task = 0; task < graph.nodes.size(); ++task)
    for (const auto &[predecessor, dependency] : graph.predecessors[task]) {
      size_t writer = schedule.computeToCpuMap.lookup(graph.nodes[predecessor].instance);
      size_t reader = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
      if (writer == reader) continue;
      bool host = branching || schedule.processorStages[writer] != schedule.processorStages[reader];
      transfers.push_back({predecessor, task, reader, writer, 1 + transfers.size(), host});
      if (host) { readers[writer].insert(reader); writers[reader].insert(writer); }
    }
  for (size_t core = 0; core < program.size(); ++core) {
    auto &events = program[core];
    for (size_t writer : writers[core]) events.emplace_back(E::Signal, 0, writer, 0, 1);
    if (schedule.processorStages[core] == 0) {
      size_t peers = 0;
      for (size_t peer = 0; peer < program.size(); ++peer)
        if (peer != core && schedule.processorStages[peer] == 0) {
          events.emplace_back(E::Signal, 0, peer, 100, 1); ++peers;
        }
      if (peers) events.emplace_back(E::Wait, 0, 0, 100, peers);
    }
    bool reuseWait = false;
    for (size_t task : tasks[core]) {
      for (const auto &transfer : transfers)
        if (transfer.target == task) {
          if (transfer.host) {
            events.emplace_back(E::Wait, 0, 0, transfer.reg, 1);
            events.emplace_back(E::Load, 64);
          } else events.emplace_back(E::Receive, 64, transfer.writer);
        }
      E work(E::Work, graph.nodes[task].cost); work.workClass = E::Vector;
      events.push_back(work);
      for (const auto &transfer : transfers)
        if (transfer.source == task) {
          if (!transfer.host) {
            events.emplace_back(E::Send, 64, transfer.reader); continue;
          }
          if (!reuseWait) {
            events.emplace_back(E::Wait, 0, 0, 0, readers[core].size()); reuseWait = true;
          }
          events.emplace_back(E::Store, 64);
          events.emplace_back(E::Signal, 0, transfer.reader, transfer.reg, 1);
        }
    }
  }
  PimPeriodicOptions options;
  options.verifyPeriodicState = true; options.searchGuidance = true;
  auto estimate = evaluatePimExecution(program, target, options);
  auto cost = std::move(estimate.cost);
  cost.loweringSupport = StaticLoweringSupport::Supported;
  cost.selectionMargin = 0.021;
  TEST_CHECK(cost.eligible());
  return cost;
}

static PimInstructionTiming hardware(size_t profile) {
  PimInstructionTiming target;
  target.processorCount = 5;
  target.residentWeightCapacity = 2;
  target.estimatePipelineDependencySpan = [](llvm::ArrayRef<Cost> spans) { return spans.front(); };
  target.peerPermissionThenSourceRead = true;
  target.signalIssueNs = profile + 1;
  target.controlMessageBytes = 1;
  target.sharedMemoryWidthBytes = 64;
  target.sharedMemoryReadLatencyNs = 17 + profile * 9;
  target.sharedMemoryWriteLatencyNs = 23 + profile * 7;
  target.synchronizationRegisterCount = 128;
  target.sharedMemoryRequestLatencyNs.assign(5, 2);
  target.sharedMemoryResponseLatencyNs.assign(5, 3);
  target.interProcessorLatencyNs.resize(25);
  for (size_t a = 0; a < 5; ++a)
    for (size_t b = 0; b < 5; ++b) {
      const size_t irregular[] = {0, 1, 5, 6, 20};
      size_t x = profile ? irregular[a] : a, y = profile ? irregular[b] : b;
      target.interProcessorLatencyNs[a * 5 + b] = x > y ? x - y : y - x;
    }
  return target;
}

static void runReference(size_t profile, size_t fixture) {
  auto target = hardware(profile);
  ComputeGraph graph;
  const std::vector<std::vector<Cost>> work{{900, 70, 800, 90, 700},
    {30, 700, 110, 600, 90}, {1000, 50, 40, 900, 60},
    {90, 80, 100, 110, 120}, {100, 800, 100, 800, 100},
    {900, 800, 40, 30, 20}};
  size_t count = work[fixture].size();
  graph.predecessors.resize(count); graph.successors.resize(count);
  for (size_t task = 0; task < count; ++task) {
    ResidentWeight weight; weight.opaqueLane = fixture == 4 ? task % 3 : task % 2;
    graph.nodes.push_back({{nullptr, static_cast<uint32_t>(task), 1}, work[fixture][task],
      {weight}, task});
    auto edge = [&](size_t source) {
      graph.predecessors[task].push_back({source, {64, 1}});
      graph.successors[source].push_back({task, {64, 1}});
      graph.publicationTransfers.push_back({source, task, {}, 0, {}, 64, 1});
    };
    if (fixture == 1) {
      if (task == 1 || task == 2) edge(0);
      if (task == 3) { edge(1); edge(2); }
      if (task == 4) edge(3);
    } else if (task) edge(task - 1);
    if (fixture == 2 && task == 4) edge(1);
  }
  if (fixture == 3) {
    target.sharedMemoryReadLatencyNs = 400;
    target.sharedMemoryWriteLatencyNs = 600;
  }
  std::string error;
  MergeScheduleResult initial;
  TEST_CHECK(mlir::succeeded(runPipelineScheduler(graph, initial, 3, target, error)));
  auto evaluate = [&](const MergeScheduleResult &schedule) { return queueCost(graph, schedule, target); };
  auto begin = std::chrono::steady_clock::now();
  auto search = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
  double searchSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count();
  TEST_CHECK(mlir::succeeded(search));
  auto repeat = optimizePipelineSchedule(graph, initial, 3, target, evaluate, error);
  TEST_CHECK(mlir::succeeded(repeat) && search->schedules.size() == repeat->schedules.size());
  for (size_t i = 0; i < search->schedules.size(); ++i) {
    TEST_CHECK(search->schedules[i].computeToCpuMap == repeat->schedules[i].computeToCpuMap);
    TEST_CHECK(search->schedules[i].computeToCpuSlotMap == repeat->schedules[i].computeToCpuSlotMap);
    TEST_CHECK(search->schedules[i].processorStages == repeat->schedules[i].processorStages);
    TEST_CHECK(search->costs.empty() || search->costs[i].primary == repeat->costs[i].primary);
    TEST_CHECK(search->costs.empty() || search->costs[i].alternative == repeat->costs[i].alternative);
  }
  double reference = std::numeric_limits<double>::infinity();
  size_t referenceEvaluations = 0;
  std::vector<size_t> stages(count), budgets(3, 1);
  std::function<void(size_t)> partitions = [&](size_t task) {
    if (task != count) {
      for (size_t stage = 0; stage <= (task ? 2 : 0); ++stage) {
        stages[task] = stage; partitions(task + 1);
      }
      return;
    }
    std::vector<bool> used(3);
    for (size_t stage : stages) used[stage] = true;
    if (used[2] && !used[1]) return;
    for (size_t core0 = 1; core0 + 1 < target.processorCount; ++core0)
      for (size_t core1 = 1; core0 + core1 < target.processorCount; ++core1) {
        budgets = {core0, core1, target.processorCount - core0 - core1};
        MergeScheduleResult schedule;
        if (mlir::failed(packPipelineSchedule(graph, stages, budgets, target, schedule, error))) continue;
        auto cost = evaluate(schedule); ++referenceEvaluations;
        reference = std::min(reference, std::max(*cost.primary, *cost.alternative));
      }
  };
  auto referenceBegin = std::chrono::steady_clock::now();
  partitions(0);
  double referenceSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - referenceBegin).count();
  std::string reason;
  size_t winner = search->costs.empty() ? 0 : selectPipelineCandidate(search->costs, reason);
  auto selected = evaluate(search->schedules[winner]);
  double cost = std::max(*selected.primary, *selected.alternative);
  double bestEvaluated = cost;
  for (const auto &estimate : search->costs)
    bestEvaluated = std::min(bestEvaluated, std::max(*estimate.primary, *estimate.alternative));
  TEST_CHECK(std::isfinite(reference) && cost + 1e-6 >= reference);
  TEST_CHECK(bestEvaluated + 1e-6 >= reference);
  TEST_CHECK(cost <= reference * (1 + kPipelineSelectionMargin) + 1e-6);
  TEST_CHECK(search->costs.size() <= kPipelineQueueEvaluationBudget
    && search->packingAttempts <= 4 * kPipelineQueueEvaluationBudget
    && search->consideredAssignments <= 4 * kPipelineQueueEvaluationBudget);
  size_t depth = 0;
  for (size_t i = winner; i; i = search->parents[i]) {
    ++depth;
    const auto &current = search->schedules[i], &parent = search->schedules[search->parents[i]];
    bool cutChanged = false;
    for (const auto &node : graph.nodes)
      cutChanged |= current.processorStages[current.computeToCpuMap.lookup(node.instance)]
        != parent.processorStages[parent.computeToCpuMap.lookup(node.instance)];
    std::vector<size_t> currentBudgets(3), parentBudgets(3);
    for (size_t core = 0; core < target.processorCount; ++core) {
      ++currentBudgets[current.processorStages[core]];
      ++parentBudgets[parent.processorStages[core]];
    }
    sawJointMove |= cutChanged && currentBudgets != parentBudgets;
  }
  sawSuccessiveMoves |= depth >= 2;
  std::cout << "profile=" << profile << " fixture=" << fixture
    << " reference_ns=" << reference << " selected_ns=" << cost
    << " gap_percent=" << 100 * (cost / reference - 1)
    << " best_evaluated_ns=" << bestEvaluated
    << " best_evaluated_gap_percent=" << 100 * (bestEvaluated / reference - 1)
    << " reference_evaluations=" << referenceEvaluations
    << " search_evaluations=" << search->costs.size() << " refinement_depth=" << depth
    << " search_seconds=" << searchSeconds << " reference_seconds=" << referenceSeconds
    << " total_seconds=" << std::chrono::duration<double>(std::chrono::steady_clock::now() - begin).count()
    << '\n';
  auto unavailable = [&](const MergeScheduleResult &) { WholeScheduleCost value; value.reason = "test unavailable timing"; return value; };
  auto unknown = optimizePipelineSchedule(graph, initial, 3, target, unavailable, error);
  TEST_CHECK(search->costs.empty() || (mlir::failed(unknown) && error.find("test unavailable timing") != std::string::npos));
  auto unproved = [&](const MergeScheduleResult &schedule) {
    auto cost = evaluate(schedule);
    cost.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
    cost.reason = "test unproved recurrence";
    return cost;
  };
  auto invalid = optimizePipelineSchedule(graph, initial, 3, target, unproved, error);
  TEST_CHECK(search->costs.empty() || (mlir::failed(invalid) && error.find("test unproved recurrence") != std::string::npos));
  if (!profile && !fixture) {
    auto badGuidance = [&](const MergeScheduleResult &schedule) {
      auto cost = evaluate(schedule); cost.processorTaskWorkNs.resize(target.processorCount + 1);
      return cost;
    };
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, badGuidance, error)));
    TEST_CHECK(error.find("physical processor count") != std::string::npos);
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, 0)));
    TEST_CHECK(mlir::failed(optimizePipelineSchedule(graph, initial, 3, target, evaluate, error, kPipelineQueueEvaluationBudget + 1)));
  }
}

int main() {
  for (size_t profile = 0; profile < 2; ++profile)
    for (size_t fixture = 0; fixture < 6; ++fixture) runReference(profile, fixture);
  TEST_CHECK(sawJointMove && sawSuccessiveMoves);
}