/*
Copyright 2015 Ian Jauslin

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
As of version 1.0, the mean of a monomial is computed directly
*/

#include "mean.h"

#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include "definitions.cpp"
#include "tools.h"
#include "polynomial.h"
#include "rational.h"
#include "array.h"
#include "fields.h"
#include "number.h"

// mean of a monomial
int mean(Int_Array monomial, Polynomial* out, Fields_Table fields, Polynomial_Matrix propagator){
  int sign=1;
  // +/- internal fields
  Int_Array internal_plus;
  Int_Array internal_minus;

  // init out
  *out=polynomial_one();

  // sort first
  monomial_sort(monomial, 0, monomial.length-1, fields, &sign);
  polynomial_multiply_Qscalar(*out, quot(sign,1));
  // get internals
  // (*out).monomials is the first element of out but it only has 1 element
  // first, free (*out).monomials[0]
  free_Int_Array((*out).monomials[0]);
  get_internals(monomial, &internal_plus, &internal_minus, (*out).monomials, fields);

  if(internal_plus.length>0 && internal_minus.length>0){
    mean_internal(internal_plus, internal_minus, out, propagator, fields);
  }

  free_Int_Array(internal_plus);
  free_Int_Array(internal_minus);
  return(0);
}

// compute the mean of a monomial of internal fields (with split + and -)
int mean_internal(Int_Array internal_plus, Int_Array internal_minus, Polynomial* out, Polynomial_Matrix propagator, Fields_Table fields){
  if(internal_plus.length!=internal_minus.length){
    fprintf(stderr,"error: monomial contains unmatched +/- fields\n");
    exit(-1);
  }
  int n=internal_minus.length;
  // pairing as an array of positions
  int* pairing=calloc(n,sizeof(int));
  // specifies which indices are available for pairing
  int* mask=calloc(n,sizeof(int));
  // signature of the permutation
  int pairing_sign;
  // sign from mixing - and + together
  int mixing_sign;
  Polynomial num;
  Polynomial num_summed=polynomial_zero();
  // propagator matrix indices
  int* indices_minus=calloc(n,sizeof(int));
  int* indices_plus=calloc(n,sizeof(int));
  int i;
  // whether the next pairing exists
  int exists_next=0;

  // indices
  for(i=0;i<n;i++){
    indices_plus[i]=intlist_find_err(propagator.indices, propagator.length, internal_plus.values[i]);
    indices_minus[i]=intlist_find_err(propagator.indices, propagator.length, -internal_minus.values[i]);
  }

  // init pairing and mask
  exists_next=init_pairing(pairing, mask, n, propagator, indices_plus, indices_minus)+1;

  // initial sign
  pairing_sign=permutation_signature(pairing,n);

  // mixing sign (from ordering psi+psi-): (-1)^{n(n+1)/2}
  if((n*(n+1))/2 %2 ==0){
    mixing_sign=1;
  }
  else{
    mixing_sign=-1;
  }

  // loop over pairings
  // loop ends when the first pairing leaves the array
  while(exists_next==1){
    num=polynomial_one();
    // propagator product for the current pairing (only simplify after all pairings)
    for(i=0;i<n;i++){
      polynomial_prod_chain_nosimplify(propagator.matrix[indices_plus[i]][indices_minus[pairing[i]]],&num, fields);
    }
    polynomial_multiply_Qscalar(num,quot(mixing_sign*pairing_sign,1));
    polynomial_concat_noinit_inplace(num,&num_summed);

    exists_next=next_pairing(pairing, mask, n, propagator, indices_plus, indices_minus)+1;
    pairing_sign=permutation_signature(pairing,n);
  }

  // only simplify in mean_symbols
  polynomial_prod_chain_nosimplify(num_summed,out,fields);
  free_Polynomial(num_summed);
  free(pairing);
  free(mask);
  free(indices_plus);
  free(indices_minus);
  return(0);
}

// first pairing with a non-vanishing propagator
int init_pairing(int* pairing, int* mask, int n, Polynomial_Matrix propagator, int* indices_plus, int* indices_minus){
  // index we want to increment
  int move=0;
  int i;
  for(i=0;i<n;i++){
    pairing[i]=-1;
    mask[i]=0;
  }

  // loop until move is out of range
  while(move>=0 && move<n){
    // move
    pairing[move]=next_wick(move, pairing[move], mask, n, propagator, indices_plus, indices_minus);
    // if the next term does not exist, then move previous index
    if(pairing[move]==-1){
      move--;
    }
    // else move next index
    else{
      move++;
    }
  }

  // if move==-1, then there is no first term, return -1
  if(move==-1){
    return(-1);
  }
  // if the first term exists
  return(0);
}

// next pairing with a non-vanishing propagator
int next_pairing(int* pairing, int* mask, int n, Polynomial_Matrix propagator, int* indices_plus, int* indices_minus){
  // index we want to increment
  int move=n-1;

  // last index
  mask[pairing[n-1]]=0;

  // loop until move is out of range
  while(move>=0 && move<n){
    // move
    pairing[move]=next_wick(move, pairing[move], mask, n, propagator, indices_plus, indices_minus);
    // if the next term does not exist, then move previous index
    if(pairing[move]==-1){
      move--;
    }
    // else move next index
    else{
      move++;
    }
  }

  // if move==-1, then there is no next term, return -1
  if(move==-1){
    return(-1);
  }
  // if the next term exists
  return(0);
}

// next term in the Wick expansion
int next_wick(int index, int start, int* mask, int n, Polynomial_Matrix propagator, int* indices_plus, int* indices_minus){
  int i;
  // unset mask
  if(start>=0 && start<n){
    mask[start]=0;
  }
  // find next position
  for(i=start+1;i<n;i++){
    // if the propagator doesn't vanish
    if(mask[i]==0 && polynomial_is_zero(propagator.matrix[indices_plus[index]][indices_minus[i]])==0){
      mask[i]=1;
      return(i);
    }
  }
  // no next term
  return(-1);
}


/* Older function: propagator as number
// compute the mean of a monomial of internal fields (with split + and -)
// compute all contractions
int mean_internal_slow(Int_Array internal_plus, Int_Array internal_minus, Number* outnum, Polynomial_Matrix propagator){
  if(internal_plus.length!=internal_minus.length){
    fprintf(stderr,"error: monomial contains unmatched +/- fields\n");
    exit(-1);
  }
  int n=internal_minus.length;
  // pairing as an array of positions
  int* pairing=calloc(n,sizeof(int));
  // specifies which indices are available for pairing
  int* mask=calloc(n,sizeof(int));
  // signature of the permutation
  int pairing_sign;
  // sign from mixing - and + together
  int mixing_sign;
  Number num;
  Number num_summed=number_zero();
  // propagator matrix indices
  int index1, index2;
  int i,j,k,l;

  // init pairing and mask
  for(i=0;i<n-1;i++){
    pairing[i]=i;
    mask[i]=1;
  }
  pairing[n-1]=n-1;
  pairing_sign=1;

  // mixing sign: (-1)^{n(n+1)/2}
  if((n*(n+1))/2 %2 ==0){
    mixing_sign=1;
  }
  else{
    mixing_sign=-1;
  }

  // loop over pairings
  // loop ends when the first pairing leaves the array
  while(pairing[0]<n){
    num=number_one();
    // propagator product for the current pairing
    for(i=0;i<n;i++){
      // indices within the propagator matrix
      index1=intlist_find_err(propagator.indices, propagator.length, internal_plus.values[i]);
      index2=intlist_find_err(propagator.indices, propagator.length, -internal_minus.values[pairing[i]]);

      number_prod_chain(propagator.matrix[index1][index2],&num);
    }
    number_Qprod_chain(quot(mixing_sign*pairing_sign,1),&num);
    number_add_chain(num,&num_summed);
    free_Number(num);

    // next pairing
    // last element of the pairing that we can move
    for(i=n-1;i>=0;i--){
      // move i-th
      mask[pairing[i]]=0;
      // search for next possible position
      for(j=pairing[i]+1;j<n;j++){
	if(mask[j]==0){
	  break;
	}
      }

      // if the next position exists
      if(j<n){
	// sign correction: change sign by (-1)^{1+(n-i)(n-i-1)/2}
	// actually (-1)^{1+(n-1-i)(n-1-i-1)/2
	if(((n-i-1)*(n-i-2))/2 % 2==0){
	  pairing_sign*=-1;
	}

	pairing[i]=j;
	mask[j]=1;
	// put the remaining pairings at their left-most possible values
	if(i<n-1){
	  k=i+1;
	  for(l=0;l<n;l++){
	    if(mask[l]==0){
	      mask[l]=1;
	      pairing[k]=l;
	      k++;
	      // if exhausted all indices
	      if(k>=n){
		break;
	      }
	    }
	  }
	}
	// if the next position was found, then don't try to move the previous pairings
	break;
      }
      // if no next position is found, store the pairing outside the array (so the algorithm stops when the first pairing is outside the array)
      else{
	pairing[i]=n;
      }
    }
  }

  number_prod_chain(num_summed,outnum);
  free_Number(num_summed);
  free(pairing);
  free(mask);
  return(0);
}
*/


// get lists of internal fields from a monomial (separate + and -)
// requires the monomial to be sorted (for the sign to be correct)
int get_internals(Int_Array monomial, Int_Array* internal_plus, Int_Array* internal_minus, Int_Array* others, Fields_Table fields){
  int i;
  init_Int_Array(internal_plus, monomial.length);
  init_Int_Array(internal_minus, monomial.length);
  init_Int_Array(others, monomial.length);
  for (i=0;i<monomial.length;i++){
    if(int_array_find(abs(monomial.values[i]),fields.internal)>=0){
      // split +/- fields
      if(monomial.values[i]>0){
	int_array_append(monomial.values[i],internal_plus);
      }
      else{
	int_array_append(monomial.values[i],internal_minus);
      }
    }
    else{
      int_array_append(monomial.values[i], others);
    }
  }
  return(0);
}


// compute the mean of a monomial containing symbolic expressions
// keep track of which means were already computed
int mean_symbols(Int_Array monomial, Polynomial* output, Fields_Table fields, Polynomial_Matrix propagator, Groups groups, Identities* computed){
  Int_Array symbol_list;
  int i;
  int power;
  int* current_term;
  Polynomial mean_num;
  Int_Array tmp_monomial;
  Number tmp_num;
  Int_Array base_monomial;
  int sign;
  // whether or not the next term exists
  int exists_next=0;
  // simplify polynomial periodically
  int simplify_freq=1;
  Polynomial mean_poly;

  init_Polynomial(output, POLY_SIZE);

  // check whether the mean was already computed
  for(i=0;i<(*computed).length;i++){
    if(int_array_cmp((*computed).lhs[i], monomial)==0){
      // write polynomial
      polynomial_concat((*computed).rhs[i], output);
      return(0);
    }
  }

  init_Int_Array(&symbol_list, monomial.length);
  init_Int_Array(&base_monomial, monomial.length);

  // generate symbols list
  for(i=0;i<monomial.length;i++){
    if(field_type(monomial.values[i], fields)==FIELD_SYMBOL){
      int_array_append(intlist_find_err(fields.symbols.indices, fields.symbols.length, monomial.values[i]), &symbol_list);
    }
    else{
      int_array_append(monomial.values[i], &base_monomial);
    }
  }
  power=symbol_list.length;

  // trivial case
  if(power==0){
    mean(monomial, &mean_num, fields, propagator);
    polynomial_concat_noinit(mean_num, output);
    
    free_Int_Array(symbol_list);
    free_Int_Array(base_monomial);
    return(0);
  }
  else{
    // initialize current term to a position that has no repetitions
    current_term=calloc(power,sizeof(int));
    exists_next=init_prod(current_term, symbol_list, fields, power, base_monomial)+1;
  }

  // loop over terms; the loop stops when all the pointers are at the end of the first symbol
  while(exists_next==1){
    // construct monomial
    int_array_cpy(base_monomial, &tmp_monomial);
    tmp_num=number_one();
    for(i=0;i<power;i++){
      int_array_concat(fields.symbols.expr[symbol_list.values[i]].monomials[current_term[i]], &tmp_monomial);
      number_prod_chain(fields.symbols.expr[symbol_list.values[i]].nums[current_term[i]], &tmp_num);
    }
    // check whether the monomial vanishes
    if(check_monomial_match(tmp_monomial, fields)==1){
      // sort monomial
      sign=1;
      monomial_sort(tmp_monomial, 0, tmp_monomial.length-1, fields, &sign);
      number_Qprod_chain(quot(sign,1), &tmp_num);

      // mean
      mean_groups(tmp_monomial, &mean_poly, fields, propagator, groups, computed);

      // write to output
      polynomial_multiply_scalar(mean_poly,tmp_num);
      polynomial_concat_noinit_inplace(mean_poly, output);
    }

    free_Number(tmp_num);
    free_Int_Array(tmp_monomial);

    // next term
    exists_next=next_prod(current_term, symbol_list, fields, power, base_monomial)+1;

    
    // simplfiy every 25 steps (improves both memory usage and performance)
    if(simplify_freq %25 ==0){
      polynomial_simplify(output, fields);
      simplify_freq=0;
    }
    simplify_freq++;
  }

  // simplify
  polynomial_simplify(output, fields);

  // write computed
  identities_append(monomial, *output, computed);

  // free memory
  free(current_term);
  free_Int_Array(symbol_list);
  free_Int_Array(base_monomial);
  return(0);
}

// first term in product with no repetitions
int init_prod(int* current_term, Int_Array symbol_list, Fields_Table fields, int power, Int_Array base_monomial){
  // index we want to increment
  int move=0;
  // tmp monomial
  Int_Array monomial;
  int i;

  init_Int_Array(&monomial, base_monomial.length+5*power);
  int_array_cpy_noinit(base_monomial, &monomial);
  // init current term
  for(i=0;i<power;i++){
    current_term[i]=-1;
  }

  // loop until move is out of range
  while(move>=0 && move<power){
    // move
    current_term[move]=next_term_norepeat(current_term[move], fields.symbols.expr[symbol_list.values[move]], &monomial, fields);
    // if the next term does not exist, then move previous index
    if(current_term[move]==-1){
      move--;
    }
    // else move next index
    else{
      move++;
    }
  }

  free_Int_Array(monomial);
  // if move==-1, then there is no first term, return -1
  if(move==-1){
    return(-1);
  }
  // if the next term exists
  return(0);
}

// next term in product with no repetitions
int next_prod(int* current_term, Int_Array symbol_list, Fields_Table fields, int power, Int_Array base_monomial){
  // index we want to increment
  int move=power-1;
  // tmp monomial
  Int_Array monomial;
  int i;

  // init monomial
  init_Int_Array(&monomial, base_monomial.length+5*power);
  int_array_cpy_noinit(base_monomial, &monomial);
  for(i=0;i<=move;i++){
    int_array_concat(fields.symbols.expr[symbol_list.values[i]].monomials[current_term[i]],&monomial);
  }

  // loop until move is out of range
  while(move>=0 && move<power){
    // move
    current_term[move]=next_term_norepeat(current_term[move], fields.symbols.expr[symbol_list.values[move]], &monomial, fields);
    // if the next term does not exist, then move previous index
    if(current_term[move]==-1){
      move--;
    }
    // else move next index
    else{
      move++;
    }
  }

  free_Int_Array(monomial);
  // if move==-1, then there is no next term, return -1
  if(move==-1){
    return(-1);
  }
  // if the next term exists
  return(0);
}

// find the next term in a polynomial that can be multiplied to a given monomial and add it to the monomial
int next_term_norepeat(int start, Polynomial polynomial, Int_Array* monomial, Fields_Table fields){
  int i;
  // remove last term from monomial
  if(start>=0 && start<polynomial.length){
    (*monomial).length-=polynomial.monomials[start].length;
  }
  // find next position
  for(i=start+1;i<polynomial.length;i++){
    // if no repetitions
    if(check_monomial_addterm(*monomial,polynomial.monomials[i],fields)==1){
      // append to monomial
      int_array_concat(polynomial.monomials[i], monomial);
      return(i);
    }
  }
  // no next term
  return(-1);
}


// signature of a permutation
int permutation_signature(int* permutation, int n){
  int* tmp_array=calloc(n,sizeof(int));
  int i;
  int ret=1;
  for(i=0;i<n;i++){
    tmp_array[i]=permutation[i];
  }
  sort_fermions(tmp_array, 0, n-1, &ret);
  free(tmp_array);
  return(ret);
}

// sort a list of anti-commuting variables
int sort_fermions(int* array, int begin, int end, int* sign){
  int i;
  int tmp;
  int index;
  // the pivot: middle of the monomial
  int pivot=(begin+end)/2;

  // if the monomial is non trivial
  if(begin<end){
    // send pivot to the end
    if(pivot!=end){
      tmp=array[end];
      array[end]=array[pivot];
      array[pivot]=tmp;
      *sign*=-1;
    }

    // loop over the others
    for(i=begin, index=begin;i<end;i++){
      // compare with pivot
      if(array[i]<array[end]){
	// if smaller, exchange with reference index
	if(i!=index){
	  tmp=array[i];
	  array[i]=array[index];
	  array[index]=tmp;
	  *sign*=-1;
	}
	// move reference index
	index++;
      }
    }
    // put pivot (which we had sent to the end) in the right place
    if(end!=index){
      tmp=array[end];
      array[end]=array[index];
      array[index]=tmp;
      *sign*=-1;
    }

    // recurse
    sort_fermions(array, begin, index-1, sign);
    sort_fermions(array, index+1, end, sign);
  }
  return(0);
}


// mean while factoring groups out
int mean_groups(Int_Array monomial, Polynomial* output, Fields_Table fields, Polynomial_Matrix propagator, Groups groups, Identities* computed){
  Polynomial num_mean;
  Int_Array tmp_monomial;
  int i;
  int group=-2;
  int next_group=-2;
  Polynomial tmp_poly;
  int sign;

  init_Polynomial(output, MONOMIAL_SIZE);

  // check whether there are symbols
  // requires the symbols to be at the end of the monomial
  if(monomial.length==0 || field_type(monomial.values[monomial.length-1], fields)!=FIELD_SYMBOL){
    // mean
    mean(monomial, &num_mean, fields, propagator);
    // add to output
    polynomial_concat_noinit(num_mean,output);
  }
  else{
    // sort into groups
    if(groups.length>0){
      sign=1;
      monomial_sort_groups(monomial, 0, monomial.length-1, fields, groups, &sign);
    }
    // construct groups and take mean
    init_Int_Array(&tmp_monomial, MONOMIAL_SIZE);
    for(i=0;i<=monomial.length;i++){
      // new group
      if(i<monomial.length){
	next_group=find_group(monomial.values[i], groups);
      }
      // if group changes, take mean
      if((i>0 && next_group!=group) || i==monomial.length){
	mean_symbols(tmp_monomial, &tmp_poly, fields, propagator, groups, computed);
	// if zero
	if(polynomial_is_zero(tmp_poly)==1){
	  // set output to 0
	  free_Polynomial(*output);
	  init_Polynomial(output, 1);
	  free_Polynomial(tmp_poly);
	  break;
	}
	// add to output
	if(polynomial_is_zero(*output)==1){
	  polynomial_concat(tmp_poly, output);
	}
	else{
	  polynomial_prod_chain_nosimplify(tmp_poly, output, fields);
	}
	free_Polynomial(tmp_poly);
	
	// reset tmp_monomial
	free_Int_Array(tmp_monomial);
	init_Int_Array(&tmp_monomial, MONOMIAL_SIZE);
      }

      // add to monomial
      if(i<monomial.length){
	int_array_append(monomial.values[i], &tmp_monomial);
      }
      group=next_group;
    }

    // sign correction
    if(sign==-1){
      polynomial_multiply_Qscalar(*output,quot(sign,1));
    }

    free_Int_Array(tmp_monomial);

  }

  return(0);
}



// mean of a polynomial

// argument struct for multithreaded mean
struct mean_args{
  Polynomial* polynomial;
  Fields_Table fields;
  Polynomial_Matrix propagator;
  Groups groups;
};

// multithreaded
int polynomial_mean_multithread(Polynomial* polynomial, Fields_Table fields, Polynomial_Matrix propagator, Groups groups, int threads){
  int i;
  Polynomial thread_polys[threads];
  pthread_t thread_ids[threads];
  struct mean_args args[threads];
  int len=(*polynomial).length;


  // alloc
  for(i=0;i<threads;i++){
    init_Polynomial(thread_polys+i,(*polynomial).length/threads+1);
    // arguments
    args[i].fields=fields;
    args[i].propagator=propagator;
    args[i].groups=groups;
  }

  // split polynomial
  for(i=0;i<len;i++){
    polynomial_append((*polynomial).monomials[i], (*polynomial).factors[i], (*polynomial).nums[i], thread_polys+(i % threads));
  }

  // start threads
  for(i=0;i<threads;i++){
    args[i].polynomial=thread_polys+i;
    pthread_create(thread_ids+i, NULL, polynomial_mean_thread, (void*)(args+i));
  }

  free_Polynomial(*polynomial);
  init_Polynomial(polynomial, len);

  // wait for completion and join
  for(i=0;i<threads;i++){
    pthread_join(thread_ids[i], NULL);
    polynomial_concat_noinit(thread_polys[i], polynomial);
  }

  polynomial_simplify(polynomial, fields);

  return(0);
}

// mean for one of the threads
void* polynomial_mean_thread(void* mean_args){
  struct mean_args *args=mean_args;
  polynomial_mean((*args).polynomial,(*args).fields,(*args).propagator,(*args).groups);
  return(NULL);
}

// single threaded version
int polynomial_mean(Polynomial* polynomial, Fields_Table fields, Polynomial_Matrix propagator, Groups groups){
  int i,j;
  Polynomial output;
  Polynomial tmp_poly;
  // a list of already computed means
  Identities computed;

  init_Polynomial(&output, (*polynomial).length);
  init_Identities(&computed, EQUATION_SIZE);

  remove_unmatched_plusminus(polynomial, fields);

  // mean of each monomial
  for(i=0;i<(*polynomial).length;i++){
    fprintf(stderr,"computing %d of %d means\n",i,(*polynomial).length-1);
    mean_groups((*polynomial).monomials[i], &tmp_poly, fields, propagator, groups, &computed);

    // write factors
    for(j=0;j<tmp_poly.length;j++){
      int_array_concat((*polynomial).factors[i], tmp_poly.factors+j);
      number_prod_chain((*polynomial).nums[i], tmp_poly.nums+j);
    }

    // add to output
    polynomial_concat_noinit(tmp_poly, &output);
    // simplify (simplify here in order to keep memory usage low)
    polynomial_simplify(&output, fields);
  }
  free_Identities(computed);

  free_Polynomial(*polynomial);
  *polynomial=output;

  return(0);
}