/*
Copyright 2015-2022 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.
*/

#include "number.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <stdarg.h>
// define MPFR_USE_VA_LIST to enable the use of mpfr_inits and mpfr_clears
#define MPFR_USE_VA_LIST
#include <mpfr.h>
#include "istring.h"
#include "definitions.cpp"
#include "tools.h"
#include "rational.h"
#include "array.h"
#include "parse_file.h"

// init
int init_Number(Number* number, int memory){
  (*number).scalars=calloc(memory,sizeof(Q));
  (*number).base=calloc(memory,sizeof(int));
  (*number).memory=memory;
  (*number).length=0;
  return(0);
}
int free_Number(Number number){
  free(number.scalars);
  free(number.base);
  return(0);
}

// copy
int number_cpy(Number input, Number* output){
  init_Number(output,input.length);
  number_cpy_noinit(input,output);
  return(0);
}
int number_cpy_noinit(Number input, Number* output){
  int i;
  if((*output).memory<input.length){
    fprintf(stderr,"error: trying to copy a number of length %d to another with memory %d\n",input.length, (*output).memory);
    exit(-1);
  }
  for(i=0;i<input.length;i++){
    (*output).scalars[i]=input.scalars[i];
    (*output).base[i]=input.base[i];
  }
  (*output).length=input.length;
  return(0);
}

// resize memory
int number_resize(Number* number, int newsize){
  Number new_number;
  init_Number(&new_number, newsize);
  number_cpy_noinit(*number,&new_number);
  free_Number(*number);
  *number=new_number;
  return(0);
}

// add a value
int number_append(Q scalar, int base, Number* output){
  if((*output).length>=(*output).memory){
    number_resize(output,2*(*output).memory);
  }
  (*output).scalars[(*output).length]=scalar;
  (*output).base[(*output).length]=base;
  (*output).length++;
  // not optimal
  number_sort(*output,0,(*output).length-1);
  return(0);
}

// concatenate
int number_concat(Number input, Number* output){
  int i;
  int offset=(*output).length;
  if((*output).length+input.length>(*output).memory){
    // make it longer than needed by (*output).length (for speed)
    number_resize(output,2*(*output).length+input.length);
  }
  for(i=0;i<input.length;i++){
    (*output).scalars[offset+i]=input.scalars[i];
    (*output).base[offset+i]=input.base[i];
  }
  (*output).length=offset+input.length;
  return(0);
}

// special numbers
Number number_zero(){
  Number ret;
  // don't allocate 0 memory since zero's will usually be used to add to
  init_Number(&ret,NUMBER_SIZE);
  return(ret);
}
Number number_one(){
  Number ret=number_zero();
  number_add_elem(quot(1,1),1,&ret);
  return(ret);
}

// find a base element
int number_find_base(int base, Number number){
  return(intlist_find(number.base, number.length, base));
}

// sort
int number_sort(Number number, int begin, int end){
  int i;
  int index;
  // the pivot: middle of the array
  int pivot=(begin+end)/2;
  // if the array is non trivial
  if(begin<end){
    // send pivot to the end
    number_exchange_terms(pivot, end, number);

    // loop over the others
    for(i=begin, index=begin;i<end;i++){
      // compare with pivot
      if(number.base[i]<number.base[end]){
	// if smaller, exchange with reference index
	number_exchange_terms(i, index, number);
	// move reference index
	index++;
      }
    }
    // put pivot (which we had sent to the end) in the right place
    number_exchange_terms(index, end, number);

    // recurse
    number_sort(number, begin, index-1);
    number_sort(number, index+1, end);
  }
  return(0);
}

// exchange terms (for sorting)
int number_exchange_terms(int index1, int index2, Number number){
  Q qtmp;
  int tmp;

  qtmp=number.scalars[index1];
  number.scalars[index1]=number.scalars[index2];
  number.scalars[index2]=qtmp;

  tmp=number.base[index1];
  number.base[index1]=number.base[index2];
  number.base[index2]=tmp;

  return(0);
}

// checks whether two numbers are equal
// requires both numbers to be sorted
int number_compare(Number x1, Number x2){
  int i;
  if(x1.length!=x2.length){
    return(0);
  }
  for(i=0;i<x1.length;i++){
    if(x1.base[i]!=x2.base[i] || Q_cmp(x1.scalars[i],x2.scalars[i])!=0){
      return(0);
    }
  }
  return(1);
}


// add (write result to second element)
int number_add_chain(Number input, Number* inout){
  int i;
  for(i=0;i<input.length;i++){
    number_add_elem(input.scalars[i], input.base[i], inout);
  }
  return(0);
}
// add a single element
int number_add_elem(Q scalar, int base, Number* inout){
  int index;
  index=number_find_base(base,*inout);
  if(index>=0){
    (*inout).scalars[index]=Q_add((*inout).scalars[index], scalar);
  }
  else{
    number_append(scalar, base, inout);
  }
  return(0);
}
// create a new number
int number_add(Number x1, Number x2, Number* out){
  number_cpy(x1,out);
  number_add_chain(x2,out);
  return(0);
}
// return the number
Number number_add_ret(Number x1, Number x2){
  Number out;
  number_add(x1,x2,&out);
  return(out);
}

// multiply
int number_prod(Number x1, Number x2, Number* out){
  int i,j;
  int div;
  Q new_scalar;
  int new_base;
  init_Number(out, x1.length);
  for(i=0;i<x1.length;i++){
    for(j=0;j<x2.length;j++){
      new_scalar=Q_prod(x1.scalars[i], x2.scalars[j]);
      // simplify the base
      div=gcd(x1.base[i], x2.base[j]);
      new_base=(x1.base[i]/div)*(x2.base[j]/div);
      new_scalar.numerator*=div;

      number_add_elem(new_scalar, new_base, out);
    }
  }
  return(0);
}
// write to second number
int number_prod_chain(Number input, Number* inout){
  Number tmp;
  number_prod(input,*inout,&tmp);
  free_Number(*inout);
  *inout=tmp;
  return(0);
}
// return result
Number number_prod_ret(Number x1, Number x2){
  Number ret;
  number_prod(x1,x2,&ret);
  return(ret);
}

// multiply by a rational
int number_Qprod_chain(Q q, Number* inout){
  int i;
  for(i=0;i<(*inout).length;i++){
    (*inout).scalars[i]=Q_prod(q,(*inout).scalars[i]);
  }
  return(0);
}
// write to output
int number_Qprod(Q q, Number x, Number* inout){
  number_cpy(x,inout);
  number_Qprod_chain(q,inout);
  return(0);
}
// return result
Number number_Qprod_ret(Q q, Number x){
  Number ret;
  number_Qprod(q,x,&ret);
  return(ret);
}

// quotient of two numbers
// recursively simplify numerator/denominator until denominator only has one term
// the output is set to 'numerator'
// both numerator and denominator may be changed by this function
// this algorithm is not optimal in cases where denominator has several terms, in particular if their bases are large
// it is optimal if denominator only has one term
int number_quot_inplace(Number* numerator, Number* denominator){
  Number tmp;
  int i,factor;

  switch((*denominator).length){
  // error
  case 0:
    fprintf(stderr,"error: attempting to invert 0\n");
    exit(-1);
    break;

  // trivial case
  case 1:
    // trivial base
    if((*denominator).base[0]==1){
      number_Qprod_chain(Q_inverse((*denominator).scalars[0]), numerator);
    }
    // non-trivial base
    else{
      // set tmp=1/denominator
      tmp=number_one();
      tmp.base[0]=(*denominator).base[0];
      tmp.scalars[0].denominator=(*denominator).scalars[0].numerator*abs((*denominator).base[0]);
      if((*denominator).base[0]>0){
	tmp.scalars[0].numerator=(*denominator).scalars[0].denominator;
      }
      else if((*denominator).base[0]<0){
	tmp.scalars[0].numerator=-(*denominator).scalars[0].denominator;
      }
      else{
	fprintf(stderr,"error: attempting to invert 0\n");
	exit(-1);
      }
      number_prod_chain(tmp, numerator);
    }
    break;

  default:
    // find non-trivial basis
    for(i=0;(*denominator).base[i]==1;i++){}

    // smallest prime factor of the base
    if((*denominator).base[i]<0){
      factor=-1;
    }
    else if((*denominator).base[i]==2){
      factor=2;
    }
    else{
      // COMMENT: this is not optimal, but, provided the basis is not too large, this should not be a problem
      for(factor=3;is_factor(factor, (*denominator).base[i])==0;factor=factor+2){}
    }

    // tmp is set to ((terms that k does not divide)-(terms that k divides))
    // tmp will be multiplied to numerator and denominator
    init_Number(&tmp,(*denominator).length);

    // find all terms whose base is a multiple of k
    for(i=0;i<(*denominator).length;i++){
      if(is_factor(factor, (*denominator).base[i])){
	// add to tmp with a - sign
	number_add_elem(quot(-(*denominator).scalars[i].numerator,(*denominator).scalars[i].denominator), (*denominator).base[i], &tmp);
      }
      else{
	// add to tmp
	number_add_elem((*denominator).scalars[i], (*denominator).base[i], &tmp);
      }
    }

    number_prod_chain(tmp, numerator);
    number_prod_chain(tmp, denominator);
    free_Number(tmp);

    // recurse
    number_quot_inplace(numerator, denominator);
  }

  return(0);
}

// not inplace
int number_quot(Number x1, Number x2, Number* output){
  Number numerator, denominator;
  number_cpy(x1, &numerator);
  number_cpy(x2, &denominator);
  number_quot_inplace(&numerator, &denominator);
  *output=numerator;
  free_Number(denominator);
  return(0);
}
int number_quot_chain(Number* inout, Number x2){
  Number tmp;
  number_quot(*inout,x2,&tmp);
  free_Number(*inout);
  *inout=tmp;
  return(0);
}
Number number_quot_ret(Number x1, Number x2){
  Number ret;
  number_quot(x1, x2, &ret);
  return(ret);
}


// remove 0's
int number_simplify(Number in, Number* out){
  int i;
  init_Number(out, in.length);
  for(i=0;i<in.length;i++){
    if(in.scalars[i].numerator!=0){
      number_append(in.scalars[i],in.base[i], out);
    }
  }
  return(0);
}


// check whether a number is 0
int number_is_zero(Number x){
  int i;
  for(i=0;i<x.length;i++){
    if(x.scalars[i].numerator!=0 && x.base[i]!=0){
      return(0);
    }
  }
  return(1);
}


// approximate numerical expression
long double number_double_val(Number x){
  int i;
  long double ret=0.;
  long double b;
  for(i=0;i<x.length;i++){
    if(x.scalars[i].numerator!=0){
      b=1.0*x.base[i];
      ret+=Q_double_value(x.scalars[i])*sqrt(b);
    }
  }
  return(ret);
}
// approximate numerical expression (as mpfr float)
int number_mpfr_val(mpfr_t out, Number x){
  int i;
  // auxiliary variables (do not initialize A)
  mpfr_t A,b,c;
  mpfr_inits(b,c, (mpfr_ptr)NULL);

  mpfr_init(out);
  mpfr_set_zero(out,1);

  for(i=0;i<x.length;i++){
    if(x.scalars[i].numerator!=0){
      mpfr_sqrt_ui(b, x.base[i], MPFR_RNDN);
      Q_mpfr_value(A, x.scalars[i]);
      mpfr_mul(c, A, b, MPFR_RNDN);
      mpfr_add(b, out, c, MPFR_RNDN);
      mpfr_set(out, b, MPFR_RNDN);
    }
  }

  mpfr_clears(A,b,c, (mpfr_ptr)NULL);

  return(0);
}


// print to string
int number_sprint(Number number, Char_Array* out){
  int i;
  Number simp;
  number_simplify(number, &simp);
  for(i=0;i<simp.length;i++){
    if(i>0){
      char_array_snprintf(out," + ");
    }
    if(simp.length>1 || (simp.length==1 && simp.base[0]!=1)){
      char_array_append('(',out);
    }
    Q_sprint(simp.scalars[i], out);
    if(simp.length>1 || (simp.length==1 && simp.base[0]!=1)){
      char_array_append(')',out);
    }
    
    if(simp.base[i]!=1){
      char_array_snprintf(out,"s{%d}",simp.base[i]);
    }
  }

  free_Number(simp);
  return(0);
}

// print to stdout
int number_print(Number number){
  Char_Array buffer;
  init_Char_Array(&buffer,5*number.length);
  number_sprint(number, &buffer);
  printf("%s",char_array_to_str_noinit(&buffer));
  return(0);
}

#define PP_NULL_MODE 0
#define PP_NUM_MODE 1
#define PP_SQRT_MODE 2
// read from a string
int str_to_Number(char* str, Number* number){
  char* ptr;
  int mode;
  char* buffer=calloc(str_len(str)+1,sizeof(char));
  char* buffer_ptr=buffer;
  Q num;
  int base;
  int ret;
  int j;
  char* aux_str;
  int aux_free=0;

  init_Number(number, NUMBER_SIZE);

  // check whether the string is blank (return 0 in that case)
  for(ptr=str;*ptr!='\0';ptr++){
    if(*ptr!=' ' && *ptr!='\n'){
      break;
    }
  }
  // blank string
  if(*ptr=='\0'){
    number_add_elem(quot(0,1), 1, number);
    free(buffer);
    return(0);
  }

  // init num and base
  num=quot(1,1);
  base=1;

  // check whether the str only contains a rational number, and add parentheses
  // keep rtack of the length of str
  for(j=0,ptr=str;*ptr!='\0';j++,ptr++){
    if((*ptr<'0' || *ptr>'9') && *ptr!='-' && *ptr!='/'){
      break;
    }
  }
  // only rational
  if(*ptr=='\0'){
    aux_str=calloc(j+3,sizeof(char));
    aux_str[0]='(';
    for(j=0,ptr=str;*ptr!='\0';ptr++,j++){
      aux_str[j+1]=*ptr;
    }
    aux_str[j+1]=')';
    aux_str[j+2]='\0';
    aux_free=1;
  }
  else{
    aux_str=str;
  }

  mode=PP_NULL_MODE;
  for(ptr=aux_str;*ptr!='\0';ptr++){
    switch(*ptr){
    // read number
    case '(':
      if(mode==PP_NULL_MODE){
	// init base
	base=1;
	mode=PP_NUM_MODE;
      }
      else{
	fprintf(stderr,"syntax error: misplaced '(' in number '%s'\n",str);
	exit(-1);
      }
      break;
    case ')':
      if(mode==PP_NUM_MODE){
	str_to_Q(buffer,&num);
	buffer_ptr=buffer;
	*buffer_ptr='\0';
	mode=PP_NULL_MODE;
      }
      else{
	fprintf(stderr,"syntax error: mismatched ')' in number '%s'\n",str);
	exit(-1);
      }
      break;

    // read sqrt
    case '{':
      // init num
      if(num.numerator==0){
	num=quot(1,1);
      }
      if(mode==PP_NULL_MODE){
	mode=PP_SQRT_MODE;
      }
      else{
	fprintf(stderr,"syntax error: misplaced '{' in number '%s'\n",str);
	exit(-1);
      }
      break;
    case '}':
      if(mode==PP_SQRT_MODE){
	ret=read_int(buffer,&base);
	if(ret<0){
	  fprintf(stderr,"syntax error: number base should be an integer, got '%s' in '%s'",buffer,str);
	  exit(-1);
	}
	buffer_ptr=buffer;
	*buffer_ptr='\0';
	mode=PP_NULL_MODE;
      }
      else{
	fprintf(stderr,"syntax error: mismatched '}' in number '%s'\n",str);
	exit(-1);
      }
      break;

    // write num
    case '+':
      if(mode==PP_NULL_MODE){
	number_add_elem(num, base, number);
	// re-init num and base
	num=quot(0,1);
	base=1;
	}
      else{
	fprintf(stderr,"syntax error: misplaced '+' in number '%s'\n",str);
	exit(-1);
      }
      break;

    // ignore 's', ' ' and '\n'
    case 's':break;
    case ' ':break;
    case '\n':break;
    default:
      if(mode!=PP_NULL_MODE){
	buffer_ptr=str_addchar(buffer_ptr,*ptr);
      }
      else{
	fprintf(stderr,"syntax error: unrecognized character '%c' in number '%s'\n",*ptr,str);
	exit(-1);
      }
      break;
    }
  }

  if(mode==PP_NULL_MODE){
    number_add_elem(num, base, number);
  }
  else{
    fprintf(stderr,"syntax error: mismatched '(' in number '%s'\n",str);
    exit(-1);
  }


  if(aux_free==1){
    free(aux_str);
  }
  free(buffer);
  return(0);
}

// with Char_Array input
int char_array_to_Number(Char_Array cstr_num,Number* output){
  char* buffer;
  char_array_to_str(cstr_num,&buffer);
  str_to_Number(buffer, output);
  free(buffer);
  return(0);
}


// -------------------- Number_Matrix ---------------------

// init
int init_Number_Matrix(Number_Matrix* matrix, int length){
  int i,j;
  (*matrix).matrix=calloc(length,sizeof(Number*));
  (*matrix).indices=calloc(length,sizeof(int));
  for(i=0;i<length;i++){
    (*matrix).matrix[i]=calloc(length,sizeof(Number));
    for(j=0;j<length;j++){
      (*matrix).matrix[i][j]=number_zero();
    }
  }
  (*matrix).length=length;
  return(0);
}
int free_Number_Matrix(Number_Matrix matrix){
  int i,j;
  for(i=0;i<matrix.length;i++){
    for(j=0;j<matrix.length;j++){
      free_Number(matrix.matrix[i][j]);
    }
    free(matrix.matrix[i]);
  }
  free(matrix.matrix);
  free(matrix.indices);
  return(0);
}

// Pauli matrices
int Pauli_matrix(int i, Number_Matrix* output){
  init_Number_Matrix(output,2);
  switch(i){
  case 1:
    number_add_elem(quot(1,1),1,(*output).matrix[0]+1);
    number_add_elem(quot(1,1),1,(*output).matrix[1]+0);
    break;
  case 2:
    number_add_elem(quot(-1,1),-1,(*output).matrix[0]+1);
    number_add_elem(quot(1,1),-1,(*output).matrix[1]+0);
    break;
  case 3:
    number_add_elem(quot(1,1),1,(*output).matrix[0]+0);
    number_add_elem(quot(-1,1),1,(*output).matrix[1]+1);
    break;
  default:
    fprintf(stderr,"error: requested %d-th pauli matrix\n",i);
    exit(-1);
  }
  return(0);
}