// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.

#include "wpi/hardware/led/LEDPattern.hpp"

#include <algorithm>
#include <cmath>
#include <numbers>
#include <unordered_map>
#include <utility>
#include <vector>

#include "wpi/hal/UsageReporting.hpp"
#include "wpi/math/util/MathUtil.hpp"
#include "wpi/util/MathExtras.hpp"
#include "wpi/util/timestamp.hpp"

using namespace wpi;

LEDPattern::LEDPattern(
    std::function<void(wpi::LEDPattern::LEDReader,
                       std::function<void(int, wpi::util::Color)>)>
        impl)
    : m_impl(std::move(impl)) {
  HAL_ReportUsage("LEDPattern", "");
}

void LEDPattern::ApplyTo(
    LEDPattern::LEDReader reader,
    std::function<void(int, wpi::util::Color)> writer) const {
  m_impl(reader, writer);
}

void LEDPattern::ApplyTo(
    std::span<AddressableLED::LEDData> data,
    std::function<void(int, wpi::util::Color)> writer) const {
  ApplyTo(LEDPattern::LEDReader{[=](size_t i) { return data[i]; }, data.size()},
          writer);
}

void LEDPattern::ApplyTo(std::span<AddressableLED::LEDData> data) const {
  ApplyTo(data, [&](int index, wpi::util::Color color) {
    data[index].SetLED(color);
  });
}

LEDPattern LEDPattern::MapIndex(
    std::function<size_t(size_t, size_t)> indexMapper) {
  return LEDPattern{[self = *this, indexMapper](auto data, auto writer) {
    size_t bufLen = data.size();
    self.ApplyTo(
        LEDPattern::LEDReader{
            [=](auto i) { return data[indexMapper(bufLen, i)]; }, bufLen},
        [&](int i, wpi::util::Color color) {
          writer(indexMapper(bufLen, i), color);
        });
  }};
}

LEDPattern LEDPattern::Reversed() {
  return MapIndex([](size_t bufLen, size_t i) { return bufLen - 1 - i; });
}

LEDPattern LEDPattern::OffsetBy(int offset) {
  return MapIndex([offset](size_t bufLen, size_t i) {
    return wpi::math::FloorMod(static_cast<int>(i) + offset,
                               static_cast<int>(bufLen));
  });
}

LEDPattern LEDPattern::ScrollAtRelativeVelocity(wpi::units::hertz_t velocity) {
  // velocity is in terms of LED lengths per second (1_hz = full cycle per
  // second, 0.5_hz = half cycle per second, 2_hz = two cycles per second)
  // Invert and multiply by 1,000,000 to get microseconds
  double periodMicros = 1e6 / velocity.value();

  return MapIndex([=](size_t bufLen, size_t i) {
    auto now = wpi::util::Now();

    // index should move by (bufLen) / (period)
    double t =
        (now % static_cast<int64_t>(std::floor(periodMicros))) / periodMicros;
    int offset = static_cast<int>(std::floor(t * bufLen));

    return wpi::math::FloorMod(static_cast<int>(i) + offset,
                               static_cast<int>(bufLen));
  });
}

LEDPattern LEDPattern::ScrollAtAbsoluteVelocity(
    wpi::units::meters_per_second_t velocity, wpi::units::meter_t ledSpacing) {
  // Velocity is in terms of meters per second
  // Multiply by 1,000,000 to use microseconds instead of seconds
  auto microsPerLed =
      static_cast<int64_t>(std::floor((ledSpacing / velocity).value() * 1e6));

  return MapIndex([=](size_t bufLen, size_t i) {
    auto now = wpi::util::Now();

    // every step in time that's a multiple of microsPerLED will increment
    // the offset by 1
    // cast unsigned int64 `now` to a signed int64 so we can get negative
    // offset values for negative velocities
    auto offset = static_cast<int64_t>(now) / microsPerLed;

    return wpi::math::FloorMod(static_cast<int>(i) + offset,
                               static_cast<int>(bufLen));
  });
}

LEDPattern LEDPattern::Blink(wpi::units::second_t onTime,
                             wpi::units::second_t offTime) {
  auto totalMicros = wpi::units::microsecond_t{onTime + offTime}.to<uint64_t>();
  auto onMicros = wpi::units::microsecond_t{onTime}.to<uint64_t>();

  return LEDPattern{[=, self = *this](auto data, auto writer) {
    if (wpi::util::Now() % totalMicros < onMicros) {
      self.ApplyTo(data, writer);
    } else {
      LEDPattern::Off().ApplyTo(data, writer);
    }
  }};
}

LEDPattern LEDPattern::Blink(wpi::units::second_t onTime) {
  return LEDPattern::Blink(onTime, onTime);
}

LEDPattern LEDPattern::SynchronizedBlink(std::function<bool()> signal) {
  return LEDPattern{[=, self = *this](auto data, auto writer) {
    if (signal()) {
      self.ApplyTo(data, writer);
    } else {
      LEDPattern::Off().ApplyTo(data, writer);
    }
  }};
}

LEDPattern LEDPattern::Breathe(wpi::units::second_t period) {
  auto periodMicros = wpi::units::microsecond_t{period};

  return LEDPattern{[periodMicros, self = *this](auto data, auto writer) {
    self.ApplyTo(data, [&writer, periodMicros](int i, wpi::util::Color color) {
      double t = (wpi::util::Now() % periodMicros.to<uint64_t>()) /
                 periodMicros.to<double>();
      double phase = t * 2 * std::numbers::pi;

      // Apply the cosine function and shift its output from [-1, 1] to [0, 1]
      // Use cosine so the period starts at 100% brightness
      double dim = (std::cos(phase) + 1) / 2.0;

      writer(i, wpi::util::Color{color.red * dim, color.green * dim,
                                 color.blue * dim});
    });
  }};
}

LEDPattern LEDPattern::OverlayOn(const LEDPattern& base) {
  return LEDPattern{[self = *this, base](auto data, auto writer) {
    // write the base pattern down first...
    base.ApplyTo(data, writer);

    // ... then, overwrite with illuminated LEDs from the overlay
    self.ApplyTo(data, [&](int i, wpi::util::Color color) {
      if (color.red > 0 || color.green > 0 || color.blue > 0) {
        writer(i, color);
      }
    });
  }};
}

LEDPattern LEDPattern::Blend(const LEDPattern& other) {
  return LEDPattern{[self = *this, other](auto data, auto writer) {
    // Apply the current pattern down as normal...
    self.ApplyTo(data, writer);

    other.ApplyTo(data, [&](int i, wpi::util::Color color) {
      // ... then read the result and average it with the output from the other
      // pattern
      writer(i, wpi::util::Color{(data[i].r / 255.0 + color.red) / 2,
                                 (data[i].g / 255.0 + color.green) / 2,
                                 (data[i].b / 255.0 + color.blue) / 2});
    });
  }};
}

LEDPattern LEDPattern::Mask(const LEDPattern& mask) {
  return LEDPattern{[self = *this, mask](auto data, auto writer) {
    // Apply the current pattern down as normal...
    self.ApplyTo(data, writer);

    mask.ApplyTo(data, [&](int i, wpi::util::Color color) {
      auto currentColor = data[i];
      // ... then perform a bitwise AND operation on each channel to apply the
      // mask
      writer(i, wpi::util::Color{
                    currentColor.r & static_cast<uint8_t>(255 * color.red),
                    currentColor.g & static_cast<uint8_t>(255 * color.green),
                    currentColor.b & static_cast<uint8_t>(255 * color.blue)});
    });
  }};
}

LEDPattern LEDPattern::AtBrightness(double relativeBrightness) {
  return LEDPattern{[relativeBrightness, self = *this](auto data, auto writer) {
    self.ApplyTo(data, [&](int i, wpi::util::Color color) {
      writer(i, wpi::util::Color{color.red * relativeBrightness,
                                 color.green * relativeBrightness,
                                 color.blue * relativeBrightness});
    });
  }};
}

// Static constants and functions

LEDPattern LEDPattern::Off() {
  return LEDPattern::Solid(wpi::util::Color::BLACK);
}

LEDPattern LEDPattern::Solid(const wpi::util::Color color) {
  return LEDPattern{[=](auto data, auto writer) {
    auto bufLen = data.size();
    for (size_t i = 0; i < bufLen; i++) {
      writer(i, color);
    }
  }};
}

LEDPattern LEDPattern::ProgressMaskLayer(
    std::function<double()> progressFunction) {
  return LEDPattern{[=](auto data, auto writer) {
    double progress = std::clamp(progressFunction(), 0.0, 1.0);
    auto bufLen = data.size();
    size_t max = bufLen * progress;

    for (size_t led = 0; led < max; led++) {
      writer(led, wpi::util::Color::WHITE);
    }
    for (size_t led = max; led < bufLen; led++) {
      writer(led, wpi::util::Color::BLACK);
    }
  }};
}

LEDPattern LEDPattern::Steps(
    std::span<const std::pair<double, wpi::util::Color>> steps) {
  if (steps.size() == 0) {
    // no colors specified
    return LEDPattern::Off();
  }
  if (steps.size() == 1 && steps[0].first == 0) {
    // only one color specified, just show a static color
    return LEDPattern::Solid(steps[0].second);
  }

  return LEDPattern{[steps = std::vector(steps.begin(), steps.end())](
                        auto data, auto writer) {
    auto bufLen = data.size();

    // precompute relevant positions for this buffer so we don't need to do a
    // check on every single LED index
    std::unordered_map<int, wpi::util::Color> stopPositions;

    for (auto step : steps) {
      stopPositions[std::floor(step.first * bufLen)] = step.second;
    }
    auto currentColor = wpi::util::Color::BLACK;
    for (size_t led = 0; led < bufLen; led++) {
      if (stopPositions.contains(led)) {
        currentColor = stopPositions[led];
      }
      writer(led, currentColor);
    }
  }};
}

LEDPattern LEDPattern::Steps(
    std::initializer_list<std::pair<double, wpi::util::Color>> steps) {
  return Steps(std::span{steps.begin(), steps.end()});
}

LEDPattern LEDPattern::Gradient(GradientType type,
                                std::span<const wpi::util::Color> colors) {
  if (colors.size() == 0) {
    // no colors specified
    return LEDPattern::Off();
  }
  if (colors.size() == 1) {
    // only one color specified, just show a static color
    return LEDPattern::Solid(colors[0]);
  }

  return LEDPattern{[type, colors = std::vector(colors.begin(), colors.end())](
                        auto data, auto writer) {
    size_t numSegments = colors.size();
    auto bufLen = data.size();
    int ledsPerSegment = 0;
    switch (type) {
      case GradientType::CONTINUOUS:
        ledsPerSegment = bufLen / numSegments;
        break;
      case GradientType::DISCONTINUOUS:
        ledsPerSegment = (bufLen - 1) / (numSegments - 1);
        break;
    }

    for (size_t led = 0; led < bufLen; led++) {
      int colorIndex = (led / ledsPerSegment) % numSegments;
      int nextColorIndex = (colorIndex + 1) % numSegments;
      double t = std::fmod(led / static_cast<double>(ledsPerSegment), 1.0);

      auto color = colors[colorIndex];
      auto nextColor = colors[nextColorIndex];

      wpi::util::Color gradientColor{
          wpi::util::Lerp(color.red, nextColor.red, t),
          wpi::util::Lerp(color.green, nextColor.green, t),
          wpi::util::Lerp(color.blue, nextColor.blue, t)};
      writer(led, gradientColor);
    }
  }};
}

LEDPattern LEDPattern::Gradient(
    GradientType type, std::initializer_list<wpi::util::Color> colors) {
  return Gradient(type, std::span{colors.begin(), colors.end()});
}

LEDPattern LEDPattern::Rainbow(int saturation, int value) {
  return LEDPattern{[=](auto data, auto writer) {
    auto bufLen = data.size();
    for (size_t led = 0; led < bufLen; led++) {
      int hue = ((led * 180) / bufLen) % 180;
      writer(led, wpi::util::Color::FromHSV(hue, saturation, value));
    }
  }};
}
