#include "init.h"
#include "navier-stokes.h"
#include "io.h"
#include <stdlib.h>
#include <math.h>

// random initial condition
int init_random (
  _Complex double* u0,
  double init_en,
  int K1,
  int K2,
  double L,
  int seed,
  bool irreversible
){
  int kx,ky;
  double rescale;
  double x,y;

  srand(seed);

  // random init (set half, then the other half are the conjugates)
  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      if (kx!=0 || ky>0){
	x=-0.5+((double) rand())/RAND_MAX;
	y=-0.5+((double) rand())/RAND_MAX;
	u0[klookup_sym(kx,ky,K2)]=x+y*I;
      }
    }
  }

  if (irreversible) {
    // fix the energy
    rescale=compute_energy(u0, K1, K2);
  } else {
    // fix the enstrophy
    rescale=compute_enstrophy(u0, K1, K2, L);
  }
  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u0[klookup_sym(kx,ky,K2)]=u0[klookup_sym(kx,ky,K2)]*sqrt(init_en/rescale);
    }
  }

  return 0;
}

// Gaussian initial condition
int init_gaussian (
  _Complex double* u0,
  double init_en,
  int K1,
  int K2,
  double L,
  bool irreversible
){
  int kx,ky;
  double rescale;

  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u0[klookup_sym(kx,ky,K2)]=(kx*kx+ky*ky)*exp(-(kx*kx+ky*ky));
    }
  }

  if (irreversible) {
    // fix the energy
    rescale=compute_energy(u0, K1, K2);
  } else {
    // fix the enstrophy
    rescale=compute_enstrophy(u0, K1, K2, L);
  }
  for(kx=0;kx<=K1;kx++){
    for(ky=-K2;ky<=K2;ky++){
      u0[klookup_sym(kx,ky,K2)]=u0[klookup_sym(kx,ky,K2)]*sqrt(init_en/rescale);
    }
  }

  return 0;
}

// Initialize from file
int init_file (
  _Complex double* u0,
  int K1,
  int K2,
  FILE* initfile
){
  read_u(u0, K1, K2, initfile);
  return 0;
}