#include <odinpara/reco.h>

#include <odindata/image.h>
#include <odindata/data.h>
#include <odindata/gridding.h>
#include <odindata/fitting.h>

/**
  * \page gencoil Generate a coil file
  * \verbinclude cmdline-utils/gencoil.usage
  */


void usage() {
  STD_cout << STD_endl;
  STD_cout << "gencoil: Generates a virtual coil suitable for sequence simulation in ODIN" << STD_endl;
  STD_cout << configInfo() << STD_endl;
  STD_cout << STD_endl;
  STD_cout << "Usage and options:" << STD_endl;
  STD_cout << "  Generate a radially inhomogenous coil:" << STD_endl;
  STD_cout << "    gencoil -rad -n <inplane-size> -fov <FOV> -R <radial-inhomogeneity[%]> -o <Coil-file>" << STD_endl;
  STD_cout << "  Generate an array coil by rotating pre-existing image file:" << STD_endl;
  STD_cout << "    gencoil -rot -n <inplane-size> -fov <FOV> -nc <numof-rotated-subcoils> -i <input-image> -sl <selected-slice> -o <Coil-file>" << STD_endl;
  STD_cout << "  Generate an array coil consisting of simple loops:" << STD_endl;
  STD_cout << "    gencoil -arr -n <inplane-size> -nc <numof-rotated-loops>  -fov <FOV> -R <array-radius> -o <Coil-file>" << STD_endl;
  STD_cout << "  Generate coil with a B1 gradient:" << STD_endl;
  STD_cout << "    gencoil -grad -n <inplane-size> -fov <FOV> -g <B1-gradient-endpoint> -o <Coil-file>" << STD_endl;
  STD_cout << "Other options:" << STD_endl;
  STD_cout << "\t" << LogBase::get_usage() << STD_endl;
  STD_cout << "\t" << helpUsage() << STD_endl;
}

int main(int argc, char* argv[]) {
  LogBase::set_log_levels(argc,argv);
  Log<Para> odinlog("gencoil","main");

  if(hasHelpOption(argc,argv)) {usage(); return 0;}

  Range all=Range::all();

  char optval[ODIN_MAXCHAR];

  STD_string coil_fname;
  if(getCommandlineOption(argc,argv,"-o",optval,ODIN_MAXCHAR)) coil_fname=optval; else {usage();exit(0);}

  CoilSensitivity sens("Coil Sensitivity");

  bool valid_mode=false;


  /////////////////////////////////////////////////////////////////////////////

  if(isCommandlineOption(argc,argv,"-rad")) {

    int n;
    float R,fov;
    if(getCommandlineOption(argc,argv,"-n",optval,ODIN_MAXCHAR)) n=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-R",optval,ODIN_MAXCHAR)) R=0.01*atof(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-fov",optval,ODIN_MAXCHAR)) fov=atof(optval); else {usage();exit(0);}

    carray sens_map(1,1,n,n);

    int nx_center=n/2;
    int ny_center=n/2;

    double a=-4.0*R;
    double c=1.0+R;

    STD_complex b1factor;

    ndim nn=sens_map.get_extent();
    ndim indexvec;

    for(unsigned int i=0; i<nn.total(); i++) {
      indexvec=nn.index2extent(i);

      int ix=indexvec[3];
      int iy=indexvec[2];

      double rx=secureDivision(fabs(double(ix-nx_center)),nx_center);
      double ry=secureDivision(fabs(double(iy-ny_center)),ny_center);
      double radius=sqrt(rx*rx+ry*ry);

      double re=a*radius*radius+c;
      if(re<0.0) re=0.0;

      b1factor=STD_complex(re,0.0);
      sens_map(indexvec)=b1factor;
    }

    sens.set_sensitivity_map(sens_map,fov,fov,fov);

    valid_mode=true;
  }

  /////////////////////////////////////////////////////////////////////////////

  if(isCommandlineOption(argc,argv,"-rot")) {

    int n,nc,sl;
    float fov;
    STD_string input_fname;
    if(getCommandlineOption(argc,argv,"-n",optval,ODIN_MAXCHAR)) n=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-nc",optval,ODIN_MAXCHAR)) nc=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-fov",optval,ODIN_MAXCHAR)) fov=atof(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-i",optval,ODIN_MAXCHAR)) input_fname=optval; else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-sl",optval,ODIN_MAXCHAR)) sl=atoi(optval); else {usage();exit(0);}

    Image img;
    ODINLOG(odinlog,infoLog) << "Loading image file ... " << STD_endl;
    if(img.load(input_fname)<0) {
      ODINLOG(odinlog,errorLog) << "unable to load ODIN image file " << input_fname << STD_endl;
      return -1;
    }

    TinyMatrix<float,2,2> rotation;
    TinyVector<float,2> offset;
    offset=0.0;

    carray sens_map(nc,1,n,n);

    Data<float,4> imgdata(img.get_magnitude());

    TinyVector<int,2> newshape(n,n);

    ODINLOG(odinlog,infoLog) << "Congridding map ... " << STD_endl;
    Data<float,2> onecoil=imgdata(0,sl,all,all);
    onecoil.congrid(newshape);
    float maxmap=max(onecoil);
    float sum=0.0;
    float nvals=0.0;
    for(unsigned int i=0; i<onecoil.numElements(); i++) {
      float val=onecoil(onecoil.create_index(i));
      if(val>0.1*maxmap) {
        sum+=val;
        nvals+=1.0;
      }
    }
    float mean=sum/nvals;
    onecoil/=(2.0*mean); // factor to account for overlapping adjacent coils


    ODINLOG(odinlog,infoLog) << "Polynomial fit ... " << STD_endl;
    onecoil=polyniomial_fit(onecoil,onecoil,2,6.0);

    Data<float,2> onecoil_transform(newshape);


    for(int ic=0; ic<nc; ic++) {
      ODINLOG(odinlog,infoLog) << "Rotating coil " << ic << " ... " << STD_endl;
      float phi=2.0*PII*float(ic)/float(nc);
      rotation(0,0)=cos(phi);  rotation(0,1)=sin(phi);
      rotation(1,0)=-sin(phi); rotation(1,1)=cos(phi);

      CoordTransformation<float,2> rotate(onecoil_transform.shape(), rotation, offset, 2.0);
      onecoil_transform=rotate(onecoil);

      for(int iy=0; iy<n; iy++) {
        for(int ix=0; ix<n; ix++) {
          sens_map(ic,0,iy,ix)=STD_complex(onecoil_transform(iy,ix));
        }
      }
    }

    sens.set_sensitivity_map(sens_map,fov,fov,fov);

    valid_mode=true;
  }


  /////////////////////////////////////////////////////////////////////////////


  if(isCommandlineOption(argc,argv,"-arr")) {

    int n,nc;
    float fov,R;
    if(getCommandlineOption(argc,argv,"-n",optval,ODIN_MAXCHAR)) n=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-nc",optval,ODIN_MAXCHAR)) nc=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-fov",optval,ODIN_MAXCHAR)) fov=atof(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-R",optval,ODIN_MAXCHAR)) R=atof(optval); else {usage();exit(0);}


    carray sens_map(nc,1,n,n);

    int nx_center=n/2;
    int ny_center=n/2;

    ndim nn=sens_map.get_extent();
    ndim indexvec;

    float overlap=2.0;
    float loopR=overlap*PII*R/float(nc); //estimate sub-loop size from overall radius and num of sub-loops


    Array<float,2> rotmat[nc];
    Array<float,1> offset[nc];

    ODINLOG(odinlog,infoLog) << "Preparing rotmat/offset ... " << STD_endl;
    for(int icoil=0; icoil<nc; icoil++) {
      float ang=2.0*PII*float(icoil)/float(nc);

      rotmat[icoil].resize(3,3);
      rotmat[icoil]=0.0;
      rotmat[icoil](0,0)=cos(ang); rotmat[icoil](0,1)=-sin(ang);
      rotmat[icoil](1,0)=sin(ang); rotmat[icoil](1,1)=cos(ang);
      rotmat[icoil](2,2)=1.0;

      offset[icoil].resize(3);
      offset[icoil](0)=R*cos(ang);
      offset[icoil](1)=R*sin(ang);
      offset[icoil](2)=0.0;
    }


    ODINLOG(odinlog,infoLog) << "Calculating sens_map ..." << STD_endl;
    for(unsigned int i=0; i<nn.total(); i++) {
      indexvec=nn.index2extent(i);

      int icoil=indexvec[0];
      int iy=indexvec[2];
      int ix=indexvec[3];

      Array<float,1> x(3);
      x(0)=fov*secureDivision(double(ix-nx_center),nx_center);
      x(1)=fov*secureDivision(double(iy-ny_center),ny_center);
      x(2)=0.0;


      int nloopsegments=100;
      float dphi=2.0*PII/nloopsegments;
      Array<float,1> dl(3);
      Array<float,1> r(3);
      Array<float,1> r0(3);
      Array<float,1> dB(3);
      STD_complex B1(0.0);
      for(int iseg=0; iseg<nloopsegments; iseg++) {
        float phi=2.0*PII*float(iseg)/float(nloopsegments);
        dl(0)=0.0;
        dl(1)= loopR*cos(phi)*dphi;
        dl(2)=-loopR*sin(phi)*dphi;

        dl=matrix_product(rotmat[icoil], dl);

        r0(0)=0.0;
        r0(1)=loopR*sin(phi);
        r0(2)=loopR*cos(phi);

        r0=matrix_product(rotmat[icoil], r0) + offset[icoil];

        r=x-r0;

        dB=vector_product(dl,r);

        float r2=sum(r*r);
        if(r2>0.0) {
          B1+=STD_complex(dB(0),dB(1))/r2;
        }

      }
      sens_map(indexvec)=B1/float(2.0*PII); // will be normalized to ~ 1 at center of loop
    }

    sens.set_sensitivity_map(sens_map,fov,fov,fov);



    valid_mode=true;
  }


  /////////////////////////////////////////////////////////////////////////////

  if(isCommandlineOption(argc,argv,"-grad")) {

    int n;
    float fov,grad;
    if(getCommandlineOption(argc,argv,"-n",optval,ODIN_MAXCHAR)) n=atoi(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-g",optval,ODIN_MAXCHAR)) grad=atof(optval); else {usage();exit(0);}
    if(getCommandlineOption(argc,argv,"-fov",optval,ODIN_MAXCHAR)) fov=atof(optval); else {usage();exit(0);}

    carray sens_map(1,1,n,n);

    int nx_center=n/2;

    ndim nn=sens_map.get_extent();
    ndim indexvec;

    for(unsigned int i=0; i<nn.total(); i++) {
      indexvec=nn.index2extent(i);
      int ix=indexvec[3];
      sens_map(indexvec)=STD_complex(1.0+grad*float(ix-nx_center)/float(nx_center),0.0);
    }

    sens.set_sensitivity_map(sens_map,fov,fov,fov);

    valid_mode=true;
  }



  if(!valid_mode) {usage();exit(0);}

  ODINLOG(odinlog,infoLog) << "Writing coil file ... " << STD_endl;
  sens.write(coil_fname);

  return 0;
}