/*
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.
*/

#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"

// 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);
}

// inverse
int number_inverse_inplace(Number* inout){
  int i;
  for(i=0;i<(*inout).length;i++){
    if((*inout).base[i]>0){
      (*inout).scalars[i]=Q_inverse((*inout).scalars[i]);
      (*inout).scalars[i].denominator*=(*inout).base[i];
    }
    else if((*inout).base[i]<0){
      (*inout).scalars[i]=Q_inverse((*inout).scalars[i]);
      (*inout).scalars[i].denominator*=-(*inout).base[i];
      (*inout).scalars[i].numerator*=-1;
    }
    else{
      fprintf(stderr,"error: attempting to invert 0\n");
      exit(-1);
    }
  }
  return(0);
}
// write to output
int number_inverse(Number input, Number* output){
  number_cpy(input,output);
  number_inverse_inplace(output);
  return(0);
}
// return result
Number number_inverse_ret(Number x){
  Number ret;
  number_inverse(x,&ret);
  return(ret);
}

// quotient
int number_quot(Number x1, Number x2, Number* output){
  Number inv;
  number_inverse(x2, &inv);
  number_prod(x1, inv, output);
  free_Number(inv);
  return(0);
}
int number_quot_chain(Number x1, Number* inout){
  number_inverse_inplace(inout);
  number_prod_chain(x1, inout);
  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",buffer.str);
  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;
  // whether there are parentheses in the string
  int exist_parenthesis=0;

  init_Number(number, NUMBER_SIZE);

  // init num and base
  // init to 0 so that if str is empty, then the number is set to 0
  num=quot(0,1);
  base=1;

  mode=PP_NULL_MODE;
  for(ptr=str;*ptr!='\0';ptr++){
    switch(*ptr){
    // read number
    case '(':
      if(mode==PP_NULL_MODE){
	// init base
	base=1;
	mode=PP_NUM_MODE;
	exist_parenthesis=1;
      }
      break;
    case ')':
      if(mode==PP_NUM_MODE){
	str_to_Q(buffer,&num);
	buffer_ptr=buffer;
	*buffer_ptr='\0';
	mode=PP_NULL_MODE;
      }
      break;

    // read sqrt
    case '{':
      // init num
      if(num.numerator==0){
	num=quot(1,1);
      }
      if(mode==PP_NULL_MODE){
	mode=PP_SQRT_MODE;
      }
      // if there is a square root, then do not read a fraction (see end of loop)
      exist_parenthesis=1;
      break;
    case '}':
      if(mode==PP_SQRT_MODE){
	sscanf(buffer,"%d",&base);
	buffer_ptr=buffer;
	*buffer_ptr='\0';
	mode=PP_NULL_MODE;
      }
      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;
	}
      break;

    default:
      if(mode!=PP_NULL_MODE){
	buffer_ptr=str_addchar(buffer_ptr,*ptr);
      }
      break;
    }
  }

  // last step
  if(mode==PP_NULL_MODE){
    if(exist_parenthesis==0){
      str_to_Q(str, &num);
    }
    number_add_elem(num, base, number);
  }

  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);
}