#pragma once

#include <array>
#include <cstdint>
#include <stdexcept>
#include <vector>

namespace xcom {

using PlanRow = std::array<std::uint32_t, 9>;

struct Plan {
  PlanRow sender{};
  PlanRow receiver{};
};

struct FanOutPlan {
  PlanRow sender{};
  std::vector<PlanRow> receivers;
};

struct Transfer {
  unsigned source;
  unsigned destination;
};

using TransferRounds = std::vector<std::vector<Transfer>>;

enum class FanOutStrategy {
  ProperMulticast,
  SingleSendCandidate,
  Serialized,
  RelayTree,
};

struct FanOutCost {
  unsigned launches;
  std::uint64_t estimatedCycles;
  std::uint64_t wordsSent;
};

struct FanOutDispatch {
  FanOutStrategy strategy;
  FanOutCost cost;
  TransferRounds rounds;
};

inline unsigned logicalToPhysical(unsigned logical) {
  const unsigned pair = logical / 2;
  const unsigned lane = logical & 1u;
  const unsigned block = pair / 23;
  unsigned row = pair % 23;
  if (block & 1u) {
    row = 22 - row;
  }
  const unsigned col = (block / 2) * 4 + (block & 1u);
  return row * 64 + col + lane * 2;
}

namespace detail {

inline int foldedColumn(unsigned physical) {
  const int column = static_cast<int>((physical >> 2) & 15u);
  return column > 7 ? column ^ 15 : column;
}

// Signed route displacement used by the SDK's XCOM timing helpers.
inline int routeDisplacement(unsigned source, unsigned destination) {
  const int sourceRaw = static_cast<int>((source >> 2) & 15u);
  const int destinationRaw = static_cast<int>((destination >> 2) & 15u);
  const int sourceColumn = sourceRaw > 7 ? sourceRaw ^ 15 : sourceRaw;
  const int destinationColumn =
      destinationRaw > 7 ? destinationRaw ^ 15 : destinationRaw;
  const int sourceMux = sourceColumn + ((sourceRaw >> 3) ^ (source & 1u));
  const int base = (destinationColumn - sourceMux) * 6;
  const unsigned destinationLane = destination & 3u;
  const unsigned destinationHalf = destinationRaw >> 3;
  if (destinationLane > 1) {
    return base + (destinationHalf == (destination & 1u) ? 2 : 4);
  }
  return base + (destinationHalf == destinationLane ? 1 : 5);
}

inline unsigned direction(unsigned source, unsigned destination) {
  return routeDisplacement(source, destination) < 1 ? 2u : 1u;
}

inline int timeToMux(unsigned source, unsigned destination) {
  const int sourceRaw = static_cast<int>((source >> 2) & 15u);
  const int destinationRaw = static_cast<int>((destination >> 2) & 15u);
  const int sourceLow = static_cast<int>((source >> 2) & 7u);
  const int displacement = routeDisplacement(source, destination);

  const int sourceEdge =
      sourceRaw > 7 ? (sourceRaw * 4) ^ 60 : sourceRaw * 4;
  const int destinationEdge =
      destinationRaw > 7 ? (destinationRaw * 4) ^ 60 : destinationRaw * 4;
  const int local = ((source >> 2) & 8u) | ((source >> 3) & 3u);
  const int crossing = local - destinationRaw + ((sourceLow >> 1) ^ 3);
  const bool sameRegion = ((source ^ destination) & 0x20u) == 0;
  const int turn = sameRegion ? sourceLow + 1 : 16 - sourceLow;
  const int groupDelta =
      (static_cast<int>((source >> 6) & 31u) -
       static_cast<int>((destination >> 6) & 31u)) *
      2;
  return crossing + sourceEdge + turn - destinationEdge + groupDelta +
         (displacement < 0 ? -displacement : displacement) - 34;
}

inline std::uint32_t delay(unsigned cycles) {
  return 0x40a00000u | (cycles & 0x7ffffu);
}

inline std::uint32_t delayPic(unsigned a, unsigned b = 0, unsigned c = 0) {
  return 0x60000000u | ((a << 19) & 0x03f80000u) |
         ((b << 18) & 0x00040000u) | (c & 0x3ffffu);
}

inline std::uint32_t delayXPic(unsigned a, unsigned b = 0, unsigned c = 0) {
  return 0x64000000u | ((a << 14) & 0x03ffc000u) |
         ((b << 13) & 0x00002000u) | (c & 0x1fffu);
}

inline std::uint32_t send(unsigned countMinusOne, unsigned direction,
                          unsigned baseWord = 0) {
  return 0x78000000u | ((countMinusOne << 21) & 0x07e00000u) |
         ((baseWord << 3) & 0x001ffff8u) | (direction & 7u);
}

inline std::uint32_t sendOff(unsigned countMinusOne, unsigned direction,
                             unsigned baseWord = 0) {
  return 0x70000000u | ((countMinusOne << 21) & 0x07e00000u) |
         (((countMinusOne >> 6) << 14) & 0x000fc000u) |
         ((baseWord << 3) & 0x00003ff8u) | (direction & 7u);
}

inline std::uint32_t sync(unsigned zone) {
  return 0x41800000u | (zone & 0xfu);
}

} // namespace detail

inline Plan assemble(unsigned senderLogical, unsigned receiverLogical,
                     unsigned count) {
  if (senderLogical >= 1472 || receiverLogical >= 1472) {
    throw std::invalid_argument("logical tile must be in [0, 1471]");
  }
  if (senderLogical == receiverLogical) {
    throw std::invalid_argument("sender and receiver must differ");
  }
  // The receive count occupies a 12-bit delayxpic field after a 53-word prefix.
  if (count == 0 || count > 4148) {
    throw std::invalid_argument("exchange count must be in [1, 4148]");
  }

  const unsigned sender = logicalToPhysical(senderLogical);
  const unsigned receiver = logicalToPhysical(receiverLogical);
  const unsigned direction = detail::direction(sender, receiver);
  const int muxTime = detail::timeToMux(sender, receiver);
  const unsigned receiverPhase = 2 * (receiver >> 6);
  // A delay immediate N advances to the event scheduled N + 1 cycles later.
  const int senderDelay = 111 - muxTime;
  if (senderDelay < -1 || senderDelay > 0x7ffff) {
    throw std::runtime_error("computed sender delay is not encodable");
  }

  Plan result;
  result.sender[0] = detail::sync(3);
  unsigned senderEnd = 1;
  if (senderDelay >= 0) {
    result.sender[senderEnd++] =
        detail::delay(static_cast<unsigned>(senderDelay));
  }
  const unsigned firstPacket = count < 64 ? count : 64;
  result.sender[senderEnd++] = detail::send(firstPacket - 1, direction);
  if (count > 64) {
    result.sender[senderEnd++] = detail::sendOff(count - 65, direction);
  }
  const int trailingDelay = 4 - senderDelay - static_cast<int>(count);
  if (trailingDelay >= 0) {
    result.sender[senderEnd++] =
        detail::delay(static_cast<unsigned>(trailingDelay));
  }
  result.sender[senderEnd] = 0x43a00000u; // br $m10

  result.receiver[0] = 1; // patch code word 1 with the physical source ID
  result.receiver[1] = 0x41800003u;
  result.receiver[2] = detail::delayXPic(112);
  if (count <= 51) {
    result.receiver[3] = detail::delayXPic(count - 1, 0, 0x640);
    result.receiver[4] = detail::delayPic(51 - count + receiverPhase);
    result.receiver[5] = detail::delay(count + 4);
    result.receiver[6] = 0x43a00000u;
  } else if (count == 52) {
    // The SDK emits an alignment-sensitive dual-issue bundle here. Splitting
    // the same 52 receive cycles avoids depending on the plan buffer address.
    result.receiver[3] = detail::delayPic(50 + receiverPhase);
    result.receiver[4] = detail::delayXPic(0, 0, 0x640);
    result.receiver[5] = detail::delay(56);
    result.receiver[6] = 0x43a00000u;
  } else {
    result.receiver[3] = detail::delayPic(51 + receiverPhase);
    result.receiver[4] = detail::delayXPic(count - 53, 0, 0x640);
    result.receiver[5] = detail::delay(56);
    result.receiver[6] = 0x43a00000u;
  }
  return result;
}

// Produces a candidate single-send plan. Equal timing and direction are
// necessary, but receivers must also occupy compatible fabric branches.
inline FanOutPlan assembleSingleSendFanOutCandidate(
    unsigned senderLogical, const std::vector<unsigned> &receiverLogical,
    unsigned count) {
  if (receiverLogical.empty()) {
    throw std::invalid_argument("fan-out requires at least one receiver");
  }
  FanOutPlan result;
  for (unsigned receiver : receiverLogical) {
    const Plan pointToPoint = assemble(senderLogical, receiver, count);
    if (result.receivers.empty()) {
      result.sender = pointToPoint.sender;
    } else if (result.sender != pointToPoint.sender) {
      throw std::invalid_argument(
          "single-send fan-out requires one timing/direction class");
    }
    result.receivers.push_back(pointToPoint.receiver);
  }
  return result;
}

inline FanOutPlan assembleMulticast(
    unsigned senderLogical, const std::vector<unsigned> &receiverLogical,
    unsigned count, unsigned scheduleOffset = 0) {
  if (receiverLogical.empty()) {
    throw std::invalid_argument("multicast requires at least one receiver");
  }
  if (senderLogical >= 1472 || count == 0 || count > 4148) {
    throw std::invalid_argument("invalid multicast source or count");
  }

  std::array<bool, 1472> used{};
  used[senderLogical] = true;
  const unsigned sourcePhysical = logicalToPhysical(senderLogical);
  std::vector<int> muxTimes;
  muxTimes.reserve(receiverLogical.size());
  int minimumMuxTime = 0x7fffffff;
  for (unsigned receiver : receiverLogical) {
    if (receiver >= 1472 || used[receiver]) {
      throw std::invalid_argument("multicast tiles must be distinct and valid");
    }
    used[receiver] = true;
    const int muxTime =
        detail::timeToMux(sourcePhysical, logicalToPhysical(receiver));
    muxTimes.push_back(muxTime);
    if (muxTime < minimumMuxTime) minimumMuxTime = muxTime;
  }

  const int naturalStart = minimumMuxTime < 0 ? -minimumMuxTime : 0;
  if (scheduleOffset > 4095u - static_cast<unsigned>(naturalStart)) {
    throw std::invalid_argument("multicast schedule offset is not encodable");
  }
  const int startCycle = naturalStart + static_cast<int>(scheduleOffset);
  const int senderDelay = startCycle - 1;
  FanOutPlan result;
  unsigned senderWord = 0;
  result.sender[senderWord++] = 0x41800003u;
  if (senderDelay >= 0)
    result.sender[senderWord++] =
        detail::delay(static_cast<unsigned>(senderDelay));
  const unsigned firstPacket = count < 64 ? count : 64;
  result.sender[senderWord++] = detail::send(firstPacket - 1, 3);
  if (count > 64)
    result.sender[senderWord++] = detail::sendOff(count - 65, 3);
  const int trailingDelay =
      4 - senderDelay - static_cast<int>(count);
  if (trailingDelay >= 0)
    result.sender[senderWord++] =
        detail::delay(static_cast<unsigned>(trailingDelay));
  result.sender[senderWord] = 0x43a00000u;

  for (unsigned i = 0; i < receiverLogical.size(); ++i) {
    const unsigned receiverPhysical = logicalToPhysical(receiverLogical[i]);
    const int receiveCycle = startCycle + muxTimes[i];
    if (receiveCycle < 0 || receiveCycle > 4095) {
      throw std::invalid_argument("multicast receive cycle is not encodable");
    }
    const unsigned receiverPhase = 2 * (receiverPhysical >> 6);
    PlanRow row{};
    row[0] = 0x41800003u;
    row[1] = detail::delayXPic(static_cast<unsigned>(receiveCycle), 0,
                               sourcePhysical);
    if (count <= 51) {
      row[2] = detail::delayXPic(count - 1, 0, 0x640);
      row[3] = detail::delayPic(51 - count + receiverPhase) | 0x00014000u;
      row[4] = detail::delay(count + 4);
      row[5] = 0x43a00000u;
    } else if (count == 52) {
      row[2] = detail::delayPic(50 + receiverPhase) | 0x00014000u;
      row[3] = detail::delayXPic(0, 0, 0x640);
      row[4] = detail::delay(56);
      row[5] = 0x43a00000u;
    } else {
      row[2] = detail::delayPic(51 + receiverPhase) | 0x00014000u;
      row[3] = detail::delayXPic(count - 53, 0, 0x640);
      row[4] = detail::delay(56);
      row[5] = 0x43a00000u;
    }
    result.receivers.push_back(row);
  }
  return result;
}

inline FanOutPlan assembleSerializedFanOut(
    unsigned senderLogical, const std::vector<unsigned> &receiverLogical,
    unsigned count) {
  if (receiverLogical.empty()) {
    throw std::invalid_argument("fan-out requires at least one receiver");
  }
  const unsigned instructionsPerSend = count > 64 ? 2 : 1;
  if (2 + receiverLogical.size() * instructionsPerSend + 1 > 9) {
    throw std::invalid_argument("serialized fan-out does not fit a plan row");
  }

  std::vector<Plan> pointToPoint;
  pointToPoint.reserve(receiverLogical.size());
  for (unsigned receiver : receiverLogical) {
    pointToPoint.push_back(assemble(senderLogical, receiver, count));
  }

  const unsigned sourcePhysical = logicalToPhysical(senderLogical);
  const unsigned firstReceiverPhysical = logicalToPhysical(receiverLogical[0]);
  const int firstSendCycle =
      112 - detail::timeToMux(sourcePhysical, firstReceiverPhysical);
  if (firstSendCycle < 1) {
    throw std::runtime_error("first serialized send precedes the plan entry");
  }

  FanOutPlan result;
  unsigned senderWord = 0;
  result.sender[senderWord++] = 0x41800003u;
  result.sender[senderWord++] =
      detail::delay(static_cast<unsigned>(firstSendCycle - 1));
  const unsigned firstPacket = count < 64 ? count : 64;
  for (unsigned i = 0; i < receiverLogical.size(); ++i) {
    const unsigned receiverPhysical = logicalToPhysical(receiverLogical[i]);
    const unsigned direction =
        detail::direction(sourcePhysical, receiverPhysical);
    result.sender[senderWord++] = detail::send(firstPacket - 1, direction);
    if (count > 64) {
      result.sender[senderWord++] = detail::sendOff(count - 65, direction);
    }

    const int receiveCycle =
        firstSendCycle + static_cast<int>(i * count) +
        detail::timeToMux(sourcePhysical, receiverPhysical);
    if (receiveCycle < 0 || receiveCycle > 4095) {
      throw std::invalid_argument("serialized receive cycle is not encodable");
    }
    PlanRow receiverPlan = pointToPoint[i].receiver;
    receiverPlan[2] =
        detail::delayXPic(static_cast<unsigned>(receiveCycle));
    result.receivers.push_back(receiverPlan);
  }
  result.sender[senderWord] = 0x43a00000u;
  return result;
}

inline TransferRounds buildRelayTree(
    unsigned senderLogical, const std::vector<unsigned> &receiverLogical) {
  if (senderLogical >= 1472) {
    throw std::invalid_argument("sender tile must be in [0, 1471]");
  }
  std::array<bool, 1472> used{};
  used[senderLogical] = true;
  for (unsigned receiver : receiverLogical) {
    if (receiver >= 1472 || used[receiver]) {
      throw std::invalid_argument("fan-out tiles must be distinct and valid");
    }
    used[receiver] = true;
  }
  std::vector<unsigned> informed = {senderLogical};
  std::size_t next = 0;
  TransferRounds rounds;
  while (next < receiverLogical.size()) {
    std::vector<Transfer> round;
    const std::size_t availableSenders = informed.size();
    for (std::size_t i = 0;
         i < availableSenders && next < receiverLogical.size(); ++i) {
      round.push_back({informed[i], receiverLogical[next]});
      informed.push_back(receiverLogical[next++]);
    }
    rounds.push_back(std::move(round));
  }
  return rounds;
}

inline FanOutCost estimateRelayTree(const TransferRounds &rounds,
                                    unsigned count) {
  // A point-to-point row starts receiving at cycle 112. The physical-row
  // phase can add at most 44 cycles on IPU21.
  constexpr unsigned launchOverhead = 156;
  std::uint64_t transfers = 0;
  for (const auto &round : rounds) transfers += round.size();
  return {static_cast<unsigned>(rounds.size()),
          rounds.size() * static_cast<std::uint64_t>(launchOverhead + count),
          transfers * count};
}

inline FanOutDispatch dispatchFanOut(
    unsigned senderLogical, const std::vector<unsigned> &receiverLogical,
    unsigned count, bool = false) {
  if (receiverLogical.empty()) {
    throw std::invalid_argument("fan-out requires at least one receiver");
  }
  assembleMulticast(senderLogical, receiverLogical, count);
  TransferRounds directRounds(1);
  for (unsigned receiver : receiverLogical)
    directRounds[0].push_back({senderLogical, receiver});
  return {FanOutStrategy::ProperMulticast,
          {1, 156u + count, static_cast<std::uint64_t>(count)},
          std::move(directRounds)};
}

} // namespace xcom