// Copyright © 2023 Apple Inc.

#pragma once

#include <metal_stdlib>

#include "mlx/backend/metal/kernels/bf16.h"

using namespace metal;

template <typename T>
struct complex_t;

template <typename T>
static constexpr constant bool is_complex_v = false;

template <typename T>
static constexpr constant bool is_complex_v<complex_t<T>> = true;

// Metal accepts explicit bfloat casts that is_convertible_v reports as false.
template <typename From, typename To>
static constexpr constant bool is_lane_convertible_v =
    is_convertible_v<From, To> ||
    (is_same_v<To, bfloat16_t> && is_convertible_v<From, float>) ||
    (is_same_v<From, bfloat16_t> && is_convertible_v<float, To>);

template <typename T>
struct complex_t {
  using value_type = T;

  T real;
  T imag;

  // Constructors
  constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {};
  constexpr complex_t() thread : real(0), imag(0) {};
  constexpr complex_t() threadgroup : real(0), imag(0) {};

  // Conversions from scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) thread : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) threadgroup : real(static_cast<T>(x)),
                                         imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) device : real(static_cast<T>(x)),
                                    imag(static_cast<T>(0)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(U x) constant : real(static_cast<T>(x)),
                                      imag(static_cast<T>(0)) {}

  // Conversions between complex types
  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) thread : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) threadgroup
      : real(static_cast<T>(x.real)),
        imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) device : real(static_cast<T>(x.real)),
                                               imag(static_cast<T>(x.imag)) {}

  template <
      typename U,
      typename = typename enable_if<
          !is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
  constexpr complex_t(complex_t<U> x) constant : real(static_cast<T>(x.real)),
                                                 imag(static_cast<T>(x.imag)) {}

  // Conversions to and from two-lane vectors (the FFT lane representation)
  constexpr complex_t(vec<T, 2> v) thread : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) threadgroup : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) device : real(v.x), imag(v.y) {};
  constexpr complex_t(vec<T, 2> v) constant : real(v.x), imag(v.y) {};

  constexpr operator vec<T, 2>() const thread {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const threadgroup {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const device {
    return vec<T, 2>(real, imag);
  }

  constexpr operator vec<T, 2>() const constant {
    return vec<T, 2>(real, imag);
  }

  // Conversions to scalar types
  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const thread {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const threadgroup {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const device {
    return static_cast<U>(real);
  }

  template <
      typename U,
      typename = typename enable_if<
          !is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
  constexpr operator U() const constant {
    return static_cast<U>(real);
  }
};

using complex32_t = complex_t<half>;
using complex64_t = complex_t<float>;

static_assert(sizeof(complex32_t) == 2 * sizeof(half));
static_assert(sizeof(complex64_t) == 2 * sizeof(float));
static_assert(sizeof(complex_t<bfloat16_t>) == 2 * sizeof(bfloat16_t));

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> x) {
  return {-x.real, -x.imag};
}

template <typename T>
constexpr bool operator>=(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag);
}

template <typename T>
constexpr bool operator>(complex_t<T> a, complex_t<T> b) {
  return (a.real > b.real) || (a.real == b.real && a.imag > b.imag);
}

template <typename T>
constexpr bool operator<=(complex_t<T> a, complex_t<T> b) {
  return operator>=(b, a);
}

template <typename T>
constexpr bool operator<(complex_t<T> a, complex_t<T> b) {
  return operator>(b, a);
}

template <typename T>
constexpr bool operator==(complex_t<T> a, complex_t<T> b) {
  return a.real == b.real && a.imag == b.imag;
}

template <typename T>
constexpr complex_t<T> operator+(complex_t<T> a, complex_t<T> b) {
  return {a.real + b.real, a.imag + b.imag};
}

template <typename T>
constexpr thread complex_t<T>& operator+=(
    thread complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr threadgroup complex_t<T>& operator+=(
    threadgroup complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <typename T>
constexpr device complex_t<T>& operator+=(
    device complex_t<T>& a,
    complex_t<T> b) {
  a.real += b.real;
  a.imag += b.imag;
  return a;
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(U a, complex_t<T> b) {
  return {static_cast<T>(a) + b.real, b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(complex_t<T> a, U b) {
  return {a.real + static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator-(complex_t<T> a, complex_t<T> b) {
  return {a.real - b.real, a.imag - b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(U a, complex_t<T> b) {
  return {static_cast<T>(a) - b.real, -b.imag};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(complex_t<T> a, U b) {
  return {a.real - static_cast<T>(b), a.imag};
}

template <typename T>
constexpr complex_t<T> operator*(complex_t<T> a, complex_t<T> b) {
  return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
}

template <typename T>
constexpr complex_t<T> operator/(complex_t<T> a, complex_t<T> b) {
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = a.real * b.real + a.imag * b.imag;
  auto y = a.imag * b.real - a.real * b.imag;
  return {x / denom, y / denom};
}

template <
    typename T,
    typename U,
    enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator/(U a, complex_t<T> b) {
  auto scalar = static_cast<T>(a);
  auto denom = b.real * b.real + b.imag * b.imag;
  auto x = scalar * b.real;
  auto y = -scalar * b.imag;
  return {x / denom, y / denom};
}

template <typename T>
constexpr complex_t<T> operator%(complex_t<T> a, complex_t<T> b) {
  auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
  auto imag = a.imag - (b.imag * static_cast<int64_t>(a.imag / b.imag));
  if (real != 0 && (real < 0 != b.real < 0)) {
    real += b.real;
  }
  if (imag != 0 && (imag < 0 != b.imag < 0)) {
    imag += b.imag;
  }
  return {real, imag};
}

static_assert(
    (complex_t<half>{1.0h, 2.0h} * complex_t<half>{3.0h, 4.0h}).real == -5.0h);
static_assert(
    (complex_t<bfloat16_t>{bfloat16_t(1.0f), bfloat16_t(2.0f)} *
     complex_t<bfloat16_t>{bfloat16_t(3.0f), bfloat16_t(4.0f)})
        .real == bfloat16_t(-5.0f));
