/*
  This file is for public-key generation
*/

#include "pk_gen.h"

#include "benes.h"
#include "controlbits.h"
#include "fft.h"
#include "params.h"
#include "util.h"

#include <immintrin.h>
#include <stdint.h>

#define min(a, b) (((a) < (b)) ? (a) : (b))

static void de_bitslicing(uint64_t *out, vec128 in[][GFBITS]) {
    int i, j, r;
    uint64_t u = 0;

    for (i = 0; i < (1 << GFBITS); i++) {
        out[i] = 0 ;
    }

    for (i = 0; i < 64; i++) {
        for (j = GFBITS - 1; j >= 0; j--) {
            u = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(in[i][j], 0);
            for (r = 0; r < 64; r++) {
                out[i * 128 + 0 * 64 + r] <<= 1;
                out[i * 128 + 0 * 64 + r] |= (u >> r) & 1;
            }
            u = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(in[i][j], 1);
            for (r = 0; r < 64; r++) {
                out[i * 128 + 1 * 64 + r] <<= 1;
                out[i * 128 + 1 * 64 + r] |= (u >> r) & 1;
            }
        }
    }
}

static void to_bitslicing_2x(vec128 out0[][GFBITS], vec128 out1[][GFBITS], const uint64_t *in) {
    int i, j, k, r;
    uint64_t u[2] = {0};

    for (i = 0; i < 64; i++) {
        for (j = GFBITS - 1; j >= 0; j--) {
            for (k = 0; k < 2; k++) {
                for (r = 63; r >= 0; r--) {
                    u[k] <<= 1;
                    u[k] |= (in[i * 128 + k * 64 + r] >> (j + GFBITS)) & 1;
                }
            }

            out1[i][j] = PQCLEAN_MCELIECE6960119F_SSE_vec128_set2x(u[0], u[1]);
        }

        for (j = GFBITS - 1; j >= 0; j--) {
            for (k = 0; k < 2; k++) {
                for (r = 63; r >= 0; r--) {
                    u[k] <<= 1;
                    u[k] |= (in[i * 128 + k * 64 + r] >> j) & 1;
                }
            }

            out0[i][GFBITS - 1 - j] = PQCLEAN_MCELIECE6960119F_SSE_vec128_set2x(u[0], u[1]);
        }
    }
}

static void transpose_64x64(uint64_t *out, const uint64_t *in) {
    int i, j, s, d;

    uint64_t x, y;
    uint64_t masks[6][2] = {
        {0x5555555555555555, 0xAAAAAAAAAAAAAAAA},
        {0x3333333333333333, 0xCCCCCCCCCCCCCCCC},
        {0x0F0F0F0F0F0F0F0F, 0xF0F0F0F0F0F0F0F0},
        {0x00FF00FF00FF00FF, 0xFF00FF00FF00FF00},
        {0x0000FFFF0000FFFF, 0xFFFF0000FFFF0000},
        {0x00000000FFFFFFFF, 0xFFFFFFFF00000000}
    };

    for (i = 0; i < 64; i++) {
        out[i] = in[i];
    }

    for (d = 5; d >= 0; d--) {
        s = 1 << d;

        for (i = 0; i < 64; i += s * 2) {
            for (j = i; j < i + s; j++) {
                x = (out[j] & masks[d][0]) | ((out[j + s] & masks[d][0]) << s);
                y = ((out[j] & masks[d][1]) >> s) | (out[j + s] & masks[d][1]);

                out[j + 0] = x;
                out[j + s] = y;
            }
        }
    }
}

/* return number of trailing zeros of the non-zero input in */
static inline int ctz(uint64_t in) {
    return (int)_tzcnt_u64(in);
}

static inline uint64_t same_mask(uint16_t x, uint16_t y) {
    uint64_t mask;

    mask = x ^ y;
    mask -= 1;
    mask >>= 63;
    mask = -mask;

    return mask;
}

static int mov_columns(uint64_t mat[][ ((SYS_N + 127) / 128) * 2 ], uint32_t *perm) {
    int i, j, k, s, block_idx, row, tail;
    uint64_t buf[64], ctz_list[32], t, d, mask;

    row = GFBITS * SYS_T - 32;
    block_idx = row / 64;
    tail = row % 64;

    // extract the 32x64 matrix

    for (i = 0; i < 32; i++) {
        buf[i] = (mat[ row + i ][ block_idx + 0 ] >> tail) |
                 (mat[ row + i ][ block_idx + 1 ] << (64 - tail));
    }

    // compute the column indices of pivots by Gaussian elimination.
    // the indices are stored in ctz_list

    for (i = 0; i < 32; i++) {
        t = buf[i];
        for (j = i + 1; j < 32; j++) {
            t |= buf[j];
        }

        if (t == 0) {
            return -1;    // return if buf is not full rank
        }

        ctz_list[i] = s = ctz(t);

        for (j = i + 1; j < 32; j++) {
            mask = (buf[i] >> s) & 1;
            mask -= 1;
            buf[i] ^= buf[j] & mask;
        }
        for (j =   0; j <  i; j++) {
            mask = (buf[j] >> s) & 1;
            mask = -mask;
            buf[j] ^= buf[i] & mask;
        }
        for (j = i + 1; j < 32; j++) {
            mask = (buf[j] >> s) & 1;
            mask = -mask;
            buf[j] ^= buf[i] & mask;
        }
    }

    // updating permutation

    for (j = 0;   j < 32; j++) {
        for (k = j + 1; k < 64; k++) {
            d = perm[ row + j ] ^ perm[ row + k ];
            d &= same_mask(k, ctz_list[j]);
            perm[ row + j ] ^= d;
            perm[ row + k ] ^= d;
        }
    }

    // moving columns of mat according to the column indices of pivots

    for (i = 0; i < GFBITS * SYS_T; i += 64) {

        for (j = 0; j < min(64, GFBITS * SYS_T - i); j++) {
            buf[j] = (mat[ i + j ][ block_idx + 0 ] >> tail) |
                     (mat[ i + j ][ block_idx + 1 ] << (64 - tail));
        }

        transpose_64x64(buf, buf);

        for (j = 0; j < 32; j++) {
            for (k = j + 1; k < 64; k++) {
                d = buf[ j ] ^ buf[ k ];
                d &= same_mask(k, ctz_list[j]);
                buf[ j ] ^= d;
                buf[ k ] ^= d;
            }
        }

        transpose_64x64(buf, buf);

        for (j = 0; j < min(64, GFBITS * SYS_T - i); j++) {
            mat[ i + j ][ block_idx + 0 ] = (mat[ i + j ][ block_idx + 0 ] << (64 - tail) >> (64 - tail)) | (buf[j] << tail);
            mat[ i + j ][ block_idx + 1 ] = (mat[ i + j ][ block_idx + 1 ] >> tail << tail) | (buf[j] >> (64 - tail));
        }
    }

    return 0;
}

#define NBLOCKS1_H ((SYS_N + 63) / 64)
#define NBLOCKS2_H ((SYS_N + 127) / 128)
#define NBLOCKS1_I ((GFBITS * SYS_T + 63) / 64)
int PQCLEAN_MCELIECE6960119F_SSE_pk_gen(unsigned char *pk, uint32_t *perm, const unsigned char *sk) {
    const int block_idx = NBLOCKS1_I - 1;
    int tail = (GFBITS * SYS_T) % 64;

    int i, j, k;
    int row, c;

    uint64_t mat[ GFBITS * SYS_T ][ NBLOCKS2_H * 2 ];

    uint64_t mask;

    vec128 irr_int[ GFBITS ];

    vec128 consts[64][ GFBITS ];
    vec128 eval[ 64 ][ GFBITS ];
    vec128 prod[ 64 ][ GFBITS ];
    vec128 tmp[ GFBITS ];

    uint64_t list[1 << GFBITS];
    uint64_t one_row[ NBLOCKS2_H * 2 ];

    // compute the inverses

    PQCLEAN_MCELIECE6960119F_SSE_irr_load(irr_int, sk);

    PQCLEAN_MCELIECE6960119F_SSE_fft(eval, irr_int);

    PQCLEAN_MCELIECE6960119F_SSE_vec128_copy(prod[0], eval[0]);

    for (i = 1; i < 64; i++) {
        PQCLEAN_MCELIECE6960119F_SSE_vec128_mul(prod[i], prod[i - 1], eval[i]);
    }

    PQCLEAN_MCELIECE6960119F_SSE_vec128_inv(tmp, prod[63]);

    for (i = 62; i >= 0; i--) {
        PQCLEAN_MCELIECE6960119F_SSE_vec128_mul(prod[i + 1], prod[i], tmp);
        PQCLEAN_MCELIECE6960119F_SSE_vec128_mul(tmp, tmp, eval[i + 1]);
    }

    PQCLEAN_MCELIECE6960119F_SSE_vec128_copy(prod[0], tmp);

    // fill matrix

    de_bitslicing(list, prod);

    for (i = 0; i < (1 << GFBITS); i++) {
        list[i] <<= GFBITS;
        list[i] |= i;
        list[i] |= ((uint64_t) perm[i]) << 31;
    }

    PQCLEAN_MCELIECE6960119F_SSE_sort_63b(1 << GFBITS, list);

    to_bitslicing_2x(consts, prod, list);

    for (i = 0; i < (1 << GFBITS); i++) {
        perm[i] = list[i] & GFMASK;
    }

    for (j = 0; j < NBLOCKS2_H; j++) {
        for (k = 0; k < GFBITS; k++) {
            mat[ k ][ 2 * j + 0 ] = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(prod[ j ][ k ], 0);
            mat[ k ][ 2 * j + 1 ] = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(prod[ j ][ k ], 1);
        }
    }

    for (i = 1; i < SYS_T; i++) {
        for (j = 0; j < NBLOCKS2_H; j++) {
            PQCLEAN_MCELIECE6960119F_SSE_vec128_mul(prod[j], prod[j], consts[j]);

            for (k = 0; k < GFBITS; k++) {
                mat[ i * GFBITS + k ][ 2 * j + 0 ] = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(prod[ j ][ k ], 0);
                mat[ i * GFBITS + k ][ 2 * j + 1 ] = PQCLEAN_MCELIECE6960119F_SSE_vec128_extract(prod[ j ][ k ], 1);
            }
        }
    }

    // gaussian elimination

    for (row = 0; row < PK_NROWS; row++) {
        i = row >> 6;
        j = row & 63;

        if (row == GFBITS * SYS_T - 32) {
            if (mov_columns(mat, perm)) {
                return -1;
            }
        }

        for (k = row + 1; k < PK_NROWS; k++) {
            mask = mat[ row ][ i ] >> j;
            mask &= 1;
            mask -= 1;

            for (c = 0; c < NBLOCKS1_H; c++) {
                mat[ row ][ c ] ^= mat[ k ][ c ] & mask;
            }
        }

        if ( ((mat[ row ][ i ] >> j) & 1) == 0 ) { // return if not systematic
            return -1;
        }

        for (k = 0; k < row; k++) {
            mask = mat[ k ][ i ] >> j;
            mask &= 1;
            mask = -mask;

            for (c = 0; c < NBLOCKS1_H; c++) {
                mat[ k ][ c ] ^= mat[ row ][ c ] & mask;
            }
        }

        for (k = row + 1; k < PK_NROWS; k++) {
            mask = mat[ k ][ i ] >> j;
            mask &= 1;
            mask = -mask;

            for (c = 0; c < NBLOCKS1_H; c++) {
                mat[ k ][ c ] ^= mat[ row ][ c ] & mask;
            }
        }
    }

    for (row = 0; row < PK_NROWS; row++) {
        for (k = block_idx; k < NBLOCKS1_H; k++) {
            one_row[k] = mat[ row ][k];
        }

        for (k = block_idx; k < NBLOCKS1_H - 1; k++) {
            one_row[k] = (one_row[k] >> tail) | (one_row[k + 1] << (64 - tail));
            PQCLEAN_MCELIECE6960119F_SSE_store8(pk, one_row[k]);
            pk += 8;
        }

        one_row[k] >>= tail;
        PQCLEAN_MCELIECE6960119F_SSE_store_i(pk, one_row[k], PK_ROW_BYTES % 8);

        pk[ (PK_ROW_BYTES % 8) - 1 ] &= (1 << (PK_NCOLS % 8)) - 1; // removing redundant bits

        pk += PK_ROW_BYTES % 8;
    }

    //

    return 0;
}