#include <stdint.h>
#include <immintrin.h>
#include <math.h>
#include "randombytes.h"
#include "params.h"
#include "ntt.h"
#include "poly.h"

void poly_uniform(poly * restrict r,const int32_t *pdata) {
  int i;
  int32_t t;

  for(i=0;i<N;i++) {
    do {
      randombytes((unsigned char *)&t,4);
      t &= 0x1FFFFFFF;
    } while(t >= pdata[P]);
    r->coeffs[i] = t;
  }
}

void poly_binary(poly *r) {
  int i,t;
  const uint32_t L = 0xA815;
  unsigned char buf[N/2];

  randombytes(buf,N/2);

  for(i=0;i<N/2;i++) {
    t = buf[i];
    r->coeffs[2*i+0]  = (L >> (t & 0xF)) & 0x3;
    r->coeffs[2*i+0] -= 1;

    r->coeffs[2*i+1]  = (L >> (t >> 4)) & 0x3;
    r->coeffs[2*i+1] -= 1;
  }
}

double poly_sprod(const poly *a,const poly *b) {
  int i;
  int64_t t=0;

  for(i=0;i<N;i++)
    t += (int64_t)a->coeffs[i]*b->coeffs[i];

  return t;
}

double poly_norm(const poly *a) {
  int i;
  double d=0;

  for(i=0;i<N;i++)
    d += (int64_t)a->coeffs[i]*a->coeffs[i];

  return sqrt(d);
}

void poly_reduce(poly * restrict r,const int32_t *pdata) {
  int i;
  __m256i pmask,f0,f1,g0,g1;
  const __m256i v = _mm256_set1_epi32(pdata[V]);
  const __m256i s = _mm256_set1_epi32(1U << 26);
  const __m256i p = _mm256_set1_epi32(pdata[P]);
  __m128i t = _mm_cvtsi32_si128(pdata[MASK]);
  t = _mm_cvtepu8_epi32(t);
  pmask = _mm256_castsi128_si256(t);

  for(i=0;i<N;i+=8) {
    f0 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&r->coeffs[i]));
    f1 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&r->coeffs[i+4]));
    g0 = _mm256_mul_epi32(f0,v);
    g1 = _mm256_mul_epi32(f1,v);
    g0 = _mm256_srli_epi64(g0,32);
    g1 = _mm256_srli_epi64(g1,32);
    g0 = _mm256_add_epi32(g0,s);
    g1 = _mm256_add_epi32(g1,s);
    g0 = _mm256_srai_epi32(g0,27);
    g1 = _mm256_srai_epi32(g1,27);
    g0 = _mm256_mul_epi32(g0,p);
    g1 = _mm256_mul_epi32(g1,p);
    f0 = _mm256_sub_epi32(f0,g0);
    f1 = _mm256_sub_epi32(f1,g1);
    f0 = _mm256_permutevar8x32_epi32(f0,pmask);
    f1 = _mm256_permutevar8x32_epi32(f1,pmask);
    _mm_store_si128((__m128i *)&r->coeffs[i],_mm256_castsi256_si128(f0));
    _mm_store_si128((__m128i *)&r->coeffs[i+4],_mm256_castsi256_si128(f1));
  }
}

void poly_csubp(poly * restrict r,const int32_t *pdata) {
  int i;
  __m256i f,t;
  const __m256i p = _mm256_set1_epi32(pdata[P]);

  for(i=0;i<N;i+=8) {
    f = _mm256_load_si256((__m256i *)&r->coeffs[i]);
    f = _mm256_sub_epi32(f,p);
    t = _mm256_srai_epi32(f,31);
    t = _mm256_and_si256(t,p);
    f = _mm256_add_epi32(f,t);
    _mm256_store_si256((__m256i *)&r->coeffs[i],f);
  }
}

void poly_caddp(poly * restrict r,const int32_t *pdata) {
  int i;
  __m256i f,t;
  const __m256i p = _mm256_set1_epi32(pdata[P]);

  for(i=0;i<N;i+=8) {
    f = _mm256_load_si256((__m256i *)&r->coeffs[i]);
    t = _mm256_srai_epi32(f,31);
    t = _mm256_and_si256(t,p);
    f = _mm256_add_epi32(f,t);
    _mm256_store_si256((__m256i *)&r->coeffs[i],f);
  }
}

void poly_neg(poly *r,const poly *a) {
  int i;
  __m256i f;
  const __m256i zero = _mm256_set1_epi32(0);

  for(i=0;i<N;i+=8) {
    f = _mm256_load_si256((__m256i *)&a->coeffs[i]);
    f = _mm256_sub_epi32(zero,f);
    _mm256_store_si256((__m256i *)&r->coeffs[i],f);
  }
}

void poly_add(poly *r,const poly *a,const poly *b) {
  int i;
  __m256i f,g;

  for(i=0;i<N;i+=8) {
    f = _mm256_load_si256((__m256i *)&a->coeffs[i]);
    g = _mm256_load_si256((__m256i *)&b->coeffs[i]);
    f = _mm256_add_epi32(f,g);
    _mm256_store_si256((__m256i *)&r->coeffs[i],f);
  }
}

void poly_sub(poly *r,const poly *a,const poly *b) {
  int i;
  __m256i f,g;

  for(i=0;i<N;i+=8) {
    f = _mm256_load_si256((__m256i *)&a->coeffs[i]);
    g = _mm256_load_si256((__m256i *)&b->coeffs[i]);
    f = _mm256_sub_epi32(f,g);
    _mm256_store_si256((__m256i *)&r->coeffs[i],f);
  }
}

void poly_ntt(poly * restrict a,const int32_t *pdata) {
  unsigned int start, i;

  for(i=0;i<N/8;i+=4)
    ntt_levels0t2_avx(a->coeffs+i,a->coeffs+i,pdata,pdata+ZETAS0T2+1);
  for(start=0;start<N;start+=512)
    for(i=start;i<start+64;i+=4)
      ntt_levels3t5_avx(a->coeffs+i,a->coeffs+i,pdata,
                        pdata+ZETAS3T5+7*start/512);
  for(start=0;start<N;start+=64)
    for(i=start;i<start+8;i+=4)
      ntt_levels6t8_avx(a->coeffs+i,a->coeffs+i,pdata,
                        pdata+ZETAS6T8+7*start/64);
  for(start=0;start<N;start+=32)
    ntt_levels9t11_avx(a->coeffs+start,a->coeffs+start,
                       pdata,pdata+ZETAS9T11+7*start/8);
}

void poly_invntt(poly * restrict a,const int32_t *pdata) {
  unsigned int start, i;

  for(start=0;start<N;start+=32)
    invntt_levels0t2_avx(a->coeffs+start,a->coeffs+start,
                         pdata,pdata+ZETAS9T11+3556-7*start/8);
  for(start=0;start<N;start+=64)
    for(i=start;i<start+8;i+=4)
      invntt_levels3t5_avx(a->coeffs+i,a->coeffs+i,
                           pdata,pdata+ZETAS6T8+441-7*start/64);
  for(start=0;start<N;start+=512)
    for(i=start;i<start+64;i+=4)
      invntt_levels6t8_avx(a->coeffs+i,a->coeffs+i,pdata,
                           pdata+ZETAS3T5+49-7*start/512);
  for(i=0;i<N/8;i+=4)
    invntt_levels9t11_avx(a->coeffs+i,a->coeffs+i,pdata,pdata+ZETAS0T2+1);
}

void poly_pointwise(poly *r,
                    const poly *a,
                    const poly *b,
                    const int32_t *pdata)
{
  int i;
  __m256i pmask,f0,f1,g0,g1;
  const __m256i p = _mm256_set1_epi32(pdata[P]);
  const __m256i pinv = _mm256_set1_epi32(pdata[PINV]);
  __m128i t = _mm_cvtsi32_si128(pdata[MASK]);
  t = _mm_cvtepu8_epi32(t);
  pmask = _mm256_castsi128_si256(t);

  for(i=0;i<N;i+=8) {
    f0 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i]));
    f1 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i+4]));
    g0 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&b->coeffs[i]));
    g1 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&b->coeffs[i+4]));
    f0 = _mm256_mul_epi32(f0,g0);
    f1 = _mm256_mul_epi32(f1,g1);
    g0 = _mm256_mul_epi32(f0,pinv);
    g1 = _mm256_mul_epi32(f1,pinv);
    g0 = _mm256_mul_epi32(g0,p);
    g1 = _mm256_mul_epi32(g1,p);
    f0 = _mm256_sub_epi32(f0,g0);
    f1 = _mm256_sub_epi32(f1,g1);
    f0 = _mm256_srli_epi64(f0,32);
    f1 = _mm256_srli_epi64(f1,32);
    f0 = _mm256_permutevar8x32_epi32(f0,pmask);
    f1 = _mm256_permutevar8x32_epi32(f1,pmask);
    _mm_store_si128((__m128i *)&r->coeffs[i],_mm256_castsi256_si128(f0));
    _mm_store_si128((__m128i *)&r->coeffs[i+4],_mm256_castsi256_si128(f1));
  }
}

void poly_scale(poly *r,const poly *a,const int32_t *pdata,int32_t s) {
  int i;
  __m256i pmask,f0,f1,g0,g1;
  const __m256i l = _mm256_set1_epi32(s);
  const __m256i p = _mm256_set1_epi32(pdata[P]);
  const __m256i pinv = _mm256_set1_epi32(pdata[PINV]);
  __m128i t = _mm_cvtsi32_si128(pdata[MASK]);
  t = _mm_cvtepu8_epi32(t);
  pmask = _mm256_castsi128_si256(t);

  for(i=0;i<N;i+=8) {
    f0 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i]));
    f1 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i+4]));
    f0 = _mm256_mul_epi32(f0,l);
    f1 = _mm256_mul_epi32(f1,l);
    g0 = _mm256_mul_epi32(f0,pinv);
    g1 = _mm256_mul_epi32(f1,pinv);
    g0 = _mm256_mul_epi32(g0,p);
    g1 = _mm256_mul_epi32(g1,p);
    f0 = _mm256_sub_epi32(f0,g0);
    f1 = _mm256_sub_epi32(f1,g1);
    f0 = _mm256_srli_epi64(f0,32);
    f1 = _mm256_srli_epi64(f1,32);
    f0 = _mm256_permutevar8x32_epi32(f0,pmask);
    f1 = _mm256_permutevar8x32_epi32(f1,pmask);
    _mm_store_si128((__m128i *)&r->coeffs[i],_mm256_castsi256_si128(f0));
    _mm_store_si128((__m128i *)&r->coeffs[i+4],_mm256_castsi256_si128(f1));
  }
}

void poly_quot(poly * restrict r,const poly *a,const int32_t *pdata) {
  int i;
  __m256i pmask,f0,f1;
  const __m256i v = _mm256_set1_epi32(pdata[V]);
  __m128i t = _mm_cvtsi32_si128(pdata[MASK]);
  t = _mm_cvtepu8_epi32(t);
  pmask = _mm256_castsi128_si256(t);

  for(i=0;i<N;i+=8) {
    f0 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i]));
    f1 = _mm256_cvtepu32_epi64(_mm_load_si128((__m128i *)&a->coeffs[i+4]));
    f0 = _mm256_mul_epu32(f0,v);
    f1 = _mm256_mul_epu32(f1,v);
    f0 = _mm256_srli_epi64(f0,55);
    f1 = _mm256_srli_epi64(f1,55);
    f0 = _mm256_permutevar8x32_epi32(f0,pmask);
    f1 = _mm256_permutevar8x32_epi32(f1,pmask);
    _mm_store_si128((__m128i *)&r->coeffs[i],_mm256_castsi256_si128(f0));
    _mm_store_si128((__m128i *)&r->coeffs[i+4],_mm256_castsi256_si128(f1));
  }
}

void poly_sigmam1(poly *r,const poly *a) {
  int i;
  __m256i f,g;
  const __m256i pmask = _mm256_set_epi32(0,1,2,3,4,5,6,7);
  const __m256i zero = _mm256_set1_epi32(0);

  r->coeffs[0] = a->coeffs[0];
  for(i=1;i<N/2;i+=8) {
    f = _mm256_loadu_si256((__m256i *)&a->coeffs[i]);
    g = _mm256_load_si256((__m256i *)&a->coeffs[N-7-i]);
    f = _mm256_permutevar8x32_epi32(f,pmask);
    g = _mm256_permutevar8x32_epi32(g,pmask);
    f = _mm256_sub_epi32(zero,f);
    g = _mm256_sub_epi32(zero,g);
    _mm256_storeu_si256((__m256i *)&r->coeffs[i],g);
    _mm256_store_si256((__m256i *)&r->coeffs[N-7-i],f);
  }
}

static void poly_sigma(poly *r,const poly *a,unsigned int k) {
  int i,j;

  for(i=j=0;i<N;i+=1,j+=k) {
    r->coeffs[j%N] = (j&N) ? -a->coeffs[i]: a->coeffs[i];
  }
}

void poly_sigma5(poly *r,const poly *a) {
  poly t;

  t = *a;
  poly_sigma(r,&t,5);
}

void poly_sigma5inv(poly *r,const poly *a) {
  poly t;

  t = *a;
  poly_sigma(r,&t,3277);
}
