// Copyright 2025 D. Danopoulos, aie4ml
// SPDX-License-Identifier: Apache-2.0

#pragma once
#include <adf.h>
#include <aie_api/aie.hpp>
#include "parameters.h"

using namespace adf;

template<typename ConfigT>
class dense_base {
public:
  using data_t        = typename ConfigT::data_t;
  using weight_t      = typename ConfigT::weight_t;
  using result_t      = typename ConfigT::result_t;
  using bias_t   = typename ConfigT::bias_t;
  using acc_scalar_t  = typename ConfigT::acc_scalar_t;

  dense_base();
};

/* ── 1. stand-alone (no cascade) ─────────────────────────────────── */
template<typename ConfigT>
class dense_single : public dense_base<ConfigT> {
public:
  using data_t        = typename ConfigT::data_t;
  using weight_t      = typename ConfigT::weight_t;
  using result_t      = typename ConfigT::result_t;
  using acc_scalar_t  = typename dense_base<ConfigT>::acc_scalar_t;
  using bias_t   = typename dense_base<ConfigT>::bias_t;

  void run(input_buffer<data_t>&       ifm,
           const weight_t (&wts)[ConfigT::IN_FEAT_SLICE * ConfigT::OUT_FEAT_SLICE],
           const bias_t (&bias)[ConfigT::OUT_FEAT_SLICE],
           output_buffer<result_t>&      ofm);

  static void registerKernelClass() {
    REGISTER_FUNCTION(dense_single::run);
  }
};

/* ── 2. first tile (cascade out only) ────────────────────────────── */
template<typename ConfigT>
class dense_first : public dense_base<ConfigT> {
public:
  using data_t        = typename ConfigT::data_t;
  using weight_t      = typename ConfigT::weight_t;
  using result_t      = typename ConfigT::result_t;
  using acc_scalar_t  = typename dense_base<ConfigT>::acc_scalar_t;

  void run(input_buffer<data_t>&          ifm,
            const weight_t (&wts)[ConfigT::IN_FEAT_SLICE * ConfigT::OUT_FEAT_SLICE],
            output_cascade<acc_scalar_t>* outCascade);

  static void registerKernelClass() { REGISTER_FUNCTION(dense_first::run); }
};

/* ── 3. middle tile (cascade in + out) ───────────────────────────── */
template<typename ConfigT>
class dense_middle : public dense_base<ConfigT> {
public:
  using data_t        = typename ConfigT::data_t;
  using weight_t      = typename ConfigT::weight_t;
  using result_t      = typename ConfigT::result_t;
  using acc_scalar_t  = typename dense_base<ConfigT>::acc_scalar_t;

  void run(input_buffer<data_t>&          ifm,
            const weight_t (&wts)[ConfigT::IN_FEAT_SLICE * ConfigT::OUT_FEAT_SLICE],
            input_cascade<acc_scalar_t>*  inCascade,
            output_cascade<acc_scalar_t>* outCascade);

  static void registerKernelClass() { REGISTER_FUNCTION(dense_middle::run); }
};

/* ── 4. last tile (cascade in, final output) ─────────────────────── */
template<typename ConfigT>
class dense_last : public dense_base<ConfigT> {
public:
  using data_t        = typename ConfigT::data_t;
  using weight_t      = typename ConfigT::weight_t;
  using result_t      = typename ConfigT::result_t;
  using acc_scalar_t  = typename dense_base<ConfigT>::acc_scalar_t;
  using bias_t   = typename dense_base<ConfigT>::bias_t;

  void run(input_buffer<data_t>&         ifm,
            const weight_t (&wts)[ConfigT::IN_FEAT_SLICE * ConfigT::OUT_FEAT_SLICE],
            input_cascade<acc_scalar_t>* inCascade,
            const bias_t (&bias)[ConfigT::OUT_FEAT_SLICE],
            output_buffer<result_t>&        ofm);

  static void registerKernelClass() { REGISTER_FUNCTION(dense_last::run); }
};
