/*************************************************************************************
                           The MIT License

   BWA-MEM2  (Sequence alignment using Burrows-Wheeler Transform),
   Copyright (C) 2019  Intel Corporation, Heng Li.

   Permission is hereby granted, free of charge, to any person obtaining
   a copy of this software and associated documentation files (the
   "Software"), to deal in the Software without restriction, including
   without limitation the rights to use, copy, modify, merge, publish,
   distribute, sublicense, and/or sell copies of the Software, and to
   permit persons to whom the Software is furnished to do so, subject to
   the following conditions:

   The above copyright notice and this permission notice shall be
   included in all copies or substantial portions of the Software.

   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
   MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
   BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
   ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
   CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
   SOFTWARE.

Authors: Vasimuddin Md <vasimuddin.md@intel.com>; Sanchit Misra <sanchit.misra@intel.com>;
*****************************************************************************************/

#ifndef _KSWV_H_
#define  _KSWV_H_

#include <stdlib.h>
#include <stdint.h>
#include <assert.h>
#include "macro.h"

#if !MAINY
#include "ksw.h"
#include "bandedSWA.h"
#else
    #if defined(__ARM_NEON) || defined(__aarch64__) || defined(APPLE_SILICON)
        #include "simd_compat.h"
    #else
        #include <immintrin.h>
    #endif
#endif

/* LIKELY/UNLIKELY come from macro.h (included above). */

#define MAX_SEQ_LEN_REF_SAM 2048
#define MAX_SEQ_LEN_QER_SAM 512

#if MAINY
#define KSW_XBYTE  0x10000
#define KSW_XSTOP  0x20000
#define KSW_XSUBO  0x40000
#define KSW_XSTART 0x80000

#ifdef USE_MALLOC_WRAPPERS
#  include "malloc_wrap.h"
#endif


#define MAX_SEQ_LEN_EXT 256

/* SIMD width definitions based on architecture */
#if defined(__ARM_NEON) || defined(__aarch64__) || defined(APPLE_SILICON)
    /* ARM/Apple Silicon - 128-bit NEON vectors */
    #ifndef SIMD_WIDTH8
    #define SIMD_WIDTH8 16   /* 128-bit / 8-bit = 16 elements */
    #endif
    #ifndef SIMD_WIDTH16
    #define SIMD_WIDTH16 8   /* 128-bit / 16-bit = 8 elements */
    #endif
#elif __AVX512BW__
    #define SIMD_WIDTH8 64
    #define SIMD_WIDTH16 32
#elif __AVX2__
    /* AVX2 - 256-bit vectors. Stub kernel for now (phase 1); the real
     * vector kernel lands on the c6i iteration in phase 2. */
    #define SIMD_WIDTH8 32
    #define SIMD_WIDTH16 16
#endif

#define max(x, y) ((x)>(y)?(x):(y))
#define min(x, y) ((x)>(y)?(y):(x))

#define MAX_NUM_PAIRS 1000000
#define MAX_NUM_PAIRS_ALLOC 20000

#define DEFAULT_AMBIG -1

typedef struct dnaSeqPair
{
	int32_t idr, idq, id;
	int32_t len1, len2;
	int32_t h0;
	int seqid, regid;
	int score; // best score
	int te, qe; // target end and query end
	int score2, te2; // second best score and ending position on the target
	int tb, qb; // target start and query start
}SeqPair;

typedef struct {
	int qlen, slen;
	uint8_t shift, mdiff, max, size;
	__m128i *qp, *H0, *H1, *E, *Hmax;
} kswq_t;

typedef struct {
	int score; // best score
	int te, qe; // target end and query end
	int score2, te2; // second best score and ending position on the target
	int tb, qb; // target start and query start
} kswr_t;


const kswr_t g_defr = { 0, -1, -1, -1, -1, -1, -1 };

#define __max_16(ret, xx) do { \
		(xx) = _mm_max_epu8((xx), _mm_srli_si128((xx), 8)); \
		(xx) = _mm_max_epu8((xx), _mm_srli_si128((xx), 4)); \
		(xx) = _mm_max_epu8((xx), _mm_srli_si128((xx), 2)); \
		(xx) = _mm_max_epu8((xx), _mm_srli_si128((xx), 1)); \
    	(ret) = _mm_extract_epi16((xx), 0) & 0x00ff; \
	} while (0)

#define DP  6
#define DP1 7
#define DP2 8
#define DP3 9

#endif

#include "kernel_dispatch.h"  /* must come before kswv class so per-tier compiles see the renamed symbol */

#include <memory>

/* Abstract interface for the vectorized KSW batch kernel.
 * Concrete implementations are per-tier mangled subclasses (kswv compiled
 * with KERNEL_VARIANT=_sse41/_sse42/_avx/_avx2/_avx512bw on x86, unsuffixed
 * on arm64). Construct via make_kswv().
 *
 * getScores8/getScores16 are declared unconditionally here even though the
 * concrete class only defines them on builds with vector support. Every
 * tier we ship satisfies this; if a future build disables vector kernels,
 * these decls would need conditional guards to match the cpp side. */
class Ikswv {
public:
    virtual ~Ikswv() = default;

    virtual void getScores8(SeqPair *pairArray,
                            uint8_t *seqBufRef,
                            uint8_t *seqBufQer,
                            kswr_t* aln,
                            int32_t numPairs,
                            uint16_t numThreads,
                            int phase) = 0;

    virtual void getScores16(SeqPair *pairArray,
                             uint8_t *seqBufRef,
                             uint8_t *seqBufQer,
                             kswr_t* aln,
                             int32_t numPairs,
                             uint16_t numThreads,
                             int phase) = 0;

};

/* Factory: returns a per-tier concrete kswv. Construction args mirror the
 * kswv ctor exactly. */
std::unique_ptr<Ikswv> make_kswv(
    int o_del, int e_del, int o_ins, int e_ins,
    int8_t w_match, int8_t w_mismatch,
    int numThreads, int32_t maxRefLen, int32_t maxQerLen);

/* Per-tier factory function forward declarations.
 * Defined in kswv.<tier>.o (each kernel TU compile).
 * Called from simd_dispatch.cpp to construct the right per-tier concrete
 * class without exposing the class layout to the dispatcher TU (which
 * previously caused an ODR size mismatch / heap corruption). */
#if defined(__x86_64__) || defined(__i386__)
extern "C" Ikswv *make_kswv_kernel_sse41(int, int, int, int, int8_t, int8_t,
                                         int, int32_t, int32_t);
extern "C" Ikswv *make_kswv_kernel_sse42(int, int, int, int, int8_t, int8_t,
                                         int, int32_t, int32_t);
extern "C" Ikswv *make_kswv_kernel_avx(int, int, int, int, int8_t, int8_t,
                                       int, int32_t, int32_t);
extern "C" Ikswv *make_kswv_kernel_avx2(int, int, int, int, int8_t, int8_t,
                                        int, int32_t, int32_t);
extern "C" Ikswv *make_kswv_kernel_avx512bw(int, int, int, int, int8_t, int8_t,
                                            int, int32_t, int32_t);
#endif


class kswv final : public Ikswv {
public:

	kswv(const int o_del, const int e_del, const int o_ins,
		 const int e_ins, const int8_t w_match, const int8_t w_mismatch,
		 int numThreads, int32_t maxRefLen, int32_t maxQerLen);

	~kswv() override;

	// kswv owns heap buffers (rowMax8/16, F/H/E vectors, etc.) freed in the
	// destructor. Allowing copy/move would alias those allocations and make
	// double-free / use-after-free trivial to introduce. Mirrors the same
	// guard on BandedPairWiseSW.
	kswv(const kswv&)            = delete;
	kswv& operator=(const kswv&) = delete;
	kswv(kswv&&)                 = delete;
	kswv& operator=(kswv&&)      = delete;

	void getScores8(SeqPair *pairArray,
					uint8_t *seqBufRef,
					uint8_t *seqBufQer,
					kswr_t* aln,
					int32_t numPairs,
					uint16_t numThreads,
					int phase) override;

	void getScores16(SeqPair *pairArray,
					 uint8_t *seqBufRef,
					 uint8_t *seqBufQer,
					 kswr_t* aln,
					 int32_t numPairs,
					 uint16_t numThreads,
					 int phase) override;

	void kswvScalarWrapper(SeqPair *seqPairArray,
                           uint8_t *seqBufRef,
                           uint8_t *seqBufQer,
                           kswr_t* aln,
                           int numPairs,
                           int nthreads,
                           bool sw, int tid);

	kswq_t* ksw_qinit(int size, int qlen, uint8_t *query, int m, const int8_t *mat);
	
private:
#if defined(__ARM_NEON) || defined(__aarch64__) || defined(APPLE_SILICON)
	/* ARM/NEON implementations (native NEON for hot paths) */
	void kswvBatchWrapper8(SeqPair *pairArray,
						   uint8_t *seqBufRef,
						   uint8_t *seqBufQer,
						   kswr_t* aln,
						   int32_t numPairs,
						   uint16_t numThreads,
						   int phase);

	int kswv_neon_u8(uint8_t seq1SoA[],
				     uint8_t seq2SoA[],
				     int16_t nrow,
				     int16_t ncol,
				     SeqPair *p,
				     kswr_t *aln,
				     int po_ind,
				     uint16_t tid,
				     int32_t numPairs,
				     int phase);

	void kswvBatchWrapper16(SeqPair *pairArray,
							uint8_t *seqBufRef,
							uint8_t *seqBufQer,
							kswr_t* aln,
							int32_t numPairs,
							uint16_t numThreads,
							int phase);

	int kswv_neon_16(int16_t seq1SoA[],
                     int16_t seq2SoA[],
                     int16_t nrow,
                     int16_t ncol,
                     SeqPair *p,
                     kswr_t* aln,
                     int po_ind,
                     uint16_t tid,
                     int32_t numPairs,
                     int phase);

#elif ((!__AVX512BW__) & (__AVX2__))
	/* AVX2 (256-bit, 32-lane u8) batched mate-rescue SW kernel.
	 * Modeled on the corrected NEON kernel (not AVX-512, which has a
	 * pre-existing coord/score2 bug class that the NEON port uncovered
	 * and fixed; see PR 18). All four NEON bug fixes are pre-applied:
	 *  (1) te tracked in two half-width int16 vectors (_lo/_hi), since
	 *      32 u8 lanes need 32 int16 slots but a __m256i int16 holds 16.
	 *  (2) per-lane freeze mask once a pair hits KSW_XSTOP.
	 *  (3) score2 scan with per-lane len1/low/high/qe exclusion.
	 *  (4) minsc filter on rowMax in the score2 scan. */
	void kswvBatchWrapper8_avx2(SeqPair *pairArray,
								uint8_t *seqBufRef,
								uint8_t *seqBufQer,
								kswr_t* aln,
								int32_t numPairs,
								uint16_t numThreads,
								int phase);

	int kswv256_u8(uint8_t seq1SoA[],
				   uint8_t seq2SoA[],
				   int16_t nrow,
				   int16_t ncol,
				   SeqPair *p,
				   kswr_t *aln,
				   int po_ind,
				   uint16_t tid,
				   int32_t numPairs,
				   int phase);

#elif __AVX512BW__
	void kswvBatchWrapper8(SeqPair *pairArray,
						   uint8_t *seqBufRef,
						   uint8_t *seqBufQer,
						   kswr_t* aln,
						   int32_t numPairs,
						   uint16_t numThreads,
						   int phase);

	int kswv512_u8(uint8_t seq1SoA[],
				   uint8_t seq2SoA[],
				   int16_t nrow,
				   int16_t ncol,
				   SeqPair *p,
				   kswr_t *aln,
				   int po_ind,
				   uint16_t tid,
				   int32_t numPairs,
				   int phase);

	void kswvBatchWrapper16(SeqPair *pairArray,
							uint8_t *seqBufRef,
							uint8_t *seqBufQer,
							kswr_t* aln,
							int32_t numPairs,
							uint16_t numThreads,
							int phase);

	int kswv512_16(int16_t seq1SoA[],
                   int16_t seq2SoA[],
                   int16_t nrow,
                   int16_t ncol,
                   SeqPair *p,
                   kswr_t* aln,
                   int po_ind,
                   uint16_t tid,
                   int32_t numPairs,
                   int phase);
#endif
	
	kswr_t kswvScalar_u8(kswq_t *q, int tlen, const uint8_t *target,
						int _o_del, int _e_del, int _o_ins, int _e_ins,
						int xtra);  // the first gap costs -(_o+_e)
	
	kswr_t kswvScalar_i16(kswq_t *q, int tlen, const uint8_t *target,
						  int _o_del, int _e_del, int _o_ins, int _e_ins,
						  int xtra); // the first gap costs -(_o+_e)
	
	void bwa_fill_scmat(int8_t mat[25]);
		
	int m;
	int o_del, o_ins, e_del, e_ins;
	// const int8_t *mat;

	int8_t w_match;
	int8_t w_mismatch;
	int8_t w_open;
	int8_t w_extend;
	int8_t w_ambig;
	uint8_t *F8;
	uint8_t *H8_0, *H8_max, *H8_1;
	uint8_t *rowMax8;
	
	int16_t *F16;
	int16_t *H16_0, *H16_max, *H16_1;
	int16_t *rowMax16;
	int32_t maxRefLen, maxQerLen;
	
	int g_qmax;
	int64_t sort1Ticks;
	int64_t setupTicks;
	int64_t swTicks;
	int64_t sort2Ticks;
};

#endif
