/***************************** LICENSE START ***********************************

 Copyright 2012 ECMWF and INPE. This software is distributed under the terms
 of the Apache License version 2.0. In applying this license, ECMWF does not
 waive the privileges and immunities granted to it by virtue of its status as
 an Intergovernmental Organization or submit itself to any jurisdiction.

 ***************************** LICENSE END *************************************/

#include "macro.h"
#include "arith.h"
#include "cnetcdf.h"
#include "MvNetCDF.h"
#include <math.h>
#include <limits.h>
#include "MvPath.hpp"

//===========================================================================

CNetCDF::CNetCDF(request *s) : InPool(tnetcdf,s), netCDF_(0),
  current_(1), mode_('r')
{
  r_ = clone_all_requests(s);
}

CNetCDF::~CNetCDF()
{
   const char *t = get_value(r_,"TEMPORARY",0);
   const char *p = get_value(r_,"PATH",0);

   if(t && p)
   {
      if(atoi(t))
         unlink(p);
   }
   free_all_requests(r_);
   unload();
}

const char *CNetCDF::GetFileName()
{
  return get_value(r_,"PATH",0);
}

void CNetCDF::ToRequest(request* &s)
{
  Attach(); // Temp fix: if someone want the request
  s = r_;
}

int CNetCDF::Write(FILE* f)
{
  return CopyFile(get_value(r_,"PATH",0),f);
}

CNetCDF::CNetCDF(const char  *p,int temp,const char mode) :
  InPool(tnetcdf), netCDF_(0),current_(1),mode_(mode)
{

  r_ = empty_request("NETCDF");
  //set_value(r_,"PATH","%s",p);
  set_value(r_,"PATH","%s",FullPathName(p).c_str());
  set_value(r_,"TEMPORARY","%d",temp);
}

void CNetCDF::load()
{
  if ( netCDF_ )
    return;

  netCDF_ = new MvNetCDF(r_, mode_);
}

void CNetCDF::unload()
{
  delete netCDF_;
}

CList* CNetCDF::Variables()
{
  load();

  int count = netCDF_->getNumberOfVariables();
  CList *l = new CList(count);

  for ( int i = 0 ; i < count; i++)
    (*l)[i] = netCDF_->getVariable(i)->name();


  return l;
}

MvRequest CNetCDF::Attributes(bool global)
{
    const size_t NUM_CHARS_IN_ADD_VALUE = 1024;  // Mars function add_value has a restriction
    char shortened_att[NUM_CHARS_IN_ADD_VALUE];                    // on the max number of chars
    MvNcVar *var = global ? GetGlobalVar() : GetVar();

    MvRequest req("ATTRIBUTES");
    // Add any attributes to subrequest.
    int num_attr = var->getNumberOfAttributes();
    if ( num_attr > 0 )
    {
        for ( int j = 0; j < num_attr; j++ )
        {
            MvNcAtt *tmpatt = var->getAttribute(j);

            int numValues = 1;

            if (tmpatt->type() != ncChar )
                numValues = tmpatt->getNumberOfValues();

            for (int k = 0; k <numValues; k++ )
            {
                char *s = tmpatt->as_string(k);
                // do we have to crop the string so that it fits into a request?
                if (strlen(s) > 1024)
                {
                    strncpy(shortened_att, s, NUM_CHARS_IN_ADD_VALUE-1);
                    shortened_att[NUM_CHARS_IN_ADD_VALUE-1] = '\0';
                    s = &(shortened_att[0]);
                }
                req.addValue(tmpatt->name(),s);
            }
        }
    }

    return req;
}


// set the current netCDF variable by name
bool CNetCDF::Current(const char *varName)
{
  load();

  int count = netCDF_->getNumberOfVariables();

  for ( int i = 0 ; i < count; i++)
  {
    const char *thisVarName = netCDF_->getVariable(i)->name();
    if (!strcmp(thisVarName, varName)) // found it?
    {
        Current(i+1);  // 1-based indexing
        return true;
    }
  }
  return false;  // if we got to here then we did not find the variable
}


class CDFCurrentFunction : public Function
{
public:
  CDFCurrentFunction(const char *n) : Function(n,2,tnetcdf,tany)
  {info = "Sets the variable number on which netcdf functions will operate.";};
  virtual Value Execute(int arity,Value *arg);
  virtual int ValidArguments(int arity,Value *arg);
};

int CDFCurrentFunction::ValidArguments(int arity,Value *arg)
{
  if(arity != 2 )  return false;
  if(arg[0].GetType() != tnetcdf) return false;                                // first arg must be a necdf
  if(arg[1].GetType() != tnumber && arg[1].GetType() != tstring) return false; // second arg must be a number or string
  return true;
}

Value CDFCurrentFunction::Execute(int,Value *arg)
{
  CNetCDF *cdf;

  arg[0].GetValue(cdf);

  if (arg[1].GetType() == tnumber)
  {
    int number;

    arg[1].GetValue(number);

    cdf->Current(number);

    return Value(number);
  }
  else // must be string
  {
    const char *varName;

    arg[1].GetValue(varName);

    bool ok = cdf->Current(varName);
    if (!ok)
      return Error("Variable '%s' not found in netCDF.", varName);
    else
      return Value(varName);
  }

}

/////////////////////////////// Function classes. //////////////////////

class CDFVarFunction : public Function
{
public:
  CDFVarFunction(const char *n) : Function(n,1,tnetcdf)
  {info = "Returns a list of the names of the given netcdf file's variables.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFVarFunction::Execute(int,Value *arg)
{
  CNetCDF *cdf;

  arg[0].GetValue(cdf);

  return Value(cdf->Variables());
}

class CDFAttFunction: public Function
{
public:
  CDFAttFunction(const char *n) : Function(n,1,tnetcdf)
  {info = "Returns a definition variable holding the netcdf's metadata.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFAttFunction::Execute(int,Value *arg)
{
  CNetCDF *cdf;

  arg[0].GetValue(cdf);
  return Value(cdf->Attributes(false));  // get the current variable's attributes
}


class CDFGlobalAttFunction: public Function
{
public:
  CDFGlobalAttFunction(const char *n) : Function(n,1,tnetcdf)
  {info = "Returns a definition variable holding the netcdf's global metadata.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFGlobalAttFunction::Execute(int,Value *arg)
{
  CNetCDF *cdf;

  arg[0].GetValue(cdf);
  return Value(cdf->Attributes(true));  // get the global variable's attributes
}



class CDFValuesFunction: public Function
{
  bool listOfDims;

public:
  CDFValuesFunction(const char *n) : Function(n)
  {info = "Returns a list containing all the values for the current netcdf variable.";};
  virtual Value Execute(int arity,Value *arg);
  virtual int ValidArguments(int arity,Value *arg);
};


int CDFValuesFunction::ValidArguments(int arity,Value *arg)
{
  listOfDims = false;

  if( arity < 1)
    return false;

  if( arg[0].GetType() != tnetcdf )
    return false;

  if( arity == 2 )
  {
    listOfDims = true;
    
    if (arg[1].GetType() != tlist)
      return false;
  }

  return true;
}



Value CDFValuesFunction::Execute(int,Value *arg)
{
    CNetCDF *a;

    arg[0].GetValue(a);
    MvNcVar *aVar = a->GetVar();

    if( ! aVar->isValid() )
    {
        return Error("CDFValuesFunction::Execute: invalid variable");
    }

    int  nr_dims = aVar->getNumberOfDimensions();

    if ( nr_dims >  0 )                           //-- arrays of data
    {                                           
        long  num_values = 1;
        long* dim        = aVar->edges();         //-- ~MvNcVar deletes!

        if (aVar->type() == ncChar)               //-- string (char) data - we will return a list of strings
        {
            if (listOfDims)
                return Error("values(netcdf, list) syntax not yet implemented for ncChar variables.");

            for( int d=0; d < nr_dims-1; ++d )    //-- e.g. 5 strings of 10 chars are stored as var(5,10)
                num_values *= dim[d];             //-- count nr of elements

            long string_length = dim[nr_dims-1];  //-- the length of each string
            char str[string_length+1];
            str[string_length] = '\0';            //-- terminating NULL

            CList *list = new CList(num_values);  //-- we need a list to hold strings

            long elem = 0;
            for (long string=0; string < num_values; ++string)  //-- for each string...
            {
                int index = 0;
                for (long j=0; j < string_length; ++j)
                {
                    str[index++] = aVar->as_char(elem++); //-- padding spaces at the end seem to be automagically removed,
                }                                         //-- so we don't need to worry about them!
                (*list)[string] = str;
            }

            return Value(list);
        }

        else         //-- vector (numeric) data
        {
            // say, input list is [0, 1, 1]

            /*
                Two syntaxes for retrieving the values for a variable:
                1) listOfDims=false: simply get all the values for the current variable and
                                     return as a vector
                2) listOfDims=true:  a list specifying the positions along each dimension is
                                     supplied; if one (and only one) of the positions is zero,
                                     then that is the dimension for which we will extract all
                                     the values (and return as a vector); otherwise we will
                                     return just the number at the given position.
                                     e.g. if the position list is given as [0,4,6] for a 3-d
                                     variable, then we will return a vector of all the values
                                     at position (n,4,6) (the user will supply 1-based indexes).
            */

            if (listOfDims)
            {
                vector<long> vPositionList;
		vector<long> vCountList;
                CList *positions;                       // user-supplied list of position indexes
                arg[1].GetValue(positions);

                if (positions->Count() != nr_dims)
                    return Error("values(netcdf, list): size of position list (%d) must be the same as the number of dimensions (%d)",
                                  positions->Count(), nr_dims);

                long indexOfSelectedDimension = -1;

                for (long i = 0;i<positions->Count();i++)    // extract the elements from the list of positions
                {
                    long pos;
		    
                    if ((*positions)[i].GetType() == tstring)  // passing a string?
                    {
                        const char *str;
                        (*positions)[i].GetValue(str);
                        if (!strcmp(str, "all"))    // only accept "all"
			  	pos = 0;                // note that we want all the values for this one
                        else
                            return Error("values(netcdf, list): the only string value allowed in the position list is 'all' ('%s' has been passed).", str);
                    }
                    else if ((*positions)[i].GetType() == tnumber)
                    {
                        (*positions)[i].GetValue(pos);
                        
                        if (pos < 1 || pos > dim[i])
                            return Error("values(netcdf, list): dimension %d should be in the range 1 to %d; %d has been supplied.", i, dim[i],pos);
                    }
                    else
                        return Error("values(netcdf, list): Only numbers and strings are allowed in the input list.");


                    if (pos == 0)                       // selected dimension?
                    {
                        indexOfSelectedDimension = i;
                        pos = 1;                        // we will start at the first element of this dimension
                    }

                    vPositionList.push_back(pos - 1);  // user supplies 1-based indexes, we want 0-based
		    vCountList.push_back((indexOfSelectedDimension != i)? 1 : dim[i]);
                }


                for( int d=0; d < nr_dims; ++d )                    //-- count nr of elements required
		  	num_values *= (indexOfSelectedDimension != d) ? 1 : dim[d];
			
					
                CVector *vec = new CVector( num_values );           // output vector
                vector<double> vd;

                aVar->setCurrent(&(vPositionList[0]));
	        aVar->get(vd,&(vCountList[0]));

                for ( long elem=0; elem < num_values; ++elem )
                    (*vec)[elem] = vd[elem];

                return Value(vec);
            }

            else
            {
                for( int d=0; d < nr_dims; ++d )
                    num_values *= dim[d];               //-- count nr of elements

                CVector *vec = new CVector( num_values );

                for ( long elem=0; elem < num_values; ++elem )
                {
                    double v     = aVar->as_double( elem );
                    (*vec)[elem] = v;
                }
                return Value(vec);
            }

        }
    }

    else
    {                                       //-- scalar data
        double v = aVar->as_double( 0 );
        return Value( v );
    }
}

class CDFValFunction: public Function
{
public:
  CDFValFunction(const char *n) : Function(n,2,tnetcdf,tnumber)
  {info = "Returns the n:th value from the current netcdf variable.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFValFunction::Execute(int,Value *arg)
{
    CNetCDF *a;
    long     elem;

    arg[0].GetValue(a);
    arg[1].GetValue(elem);
    MvNcVar *aVar = a->GetVar();

    if( ! aVar->isValid() )
    {
        return Error("CDFValFunction::Execute: invalid variable!");
    }

    long  num_values = 1;
    int   nr_dims    = aVar->getNumberOfDimensions();
    long* dim        = aVar->edges();     //-- ~MvNcVar deletes!

    if (aVar->type() == ncChar)               //-- string (char) data - we will return a list of strings
    {
        for( int d=0; d < nr_dims-1; ++d )    //-- e.g. 5 strings of 10 chars are stored as var(5,10)
            num_values *= dim[d];             //-- count nr of elements

        if( elem < 1 || elem > num_values )
        {
            return Error("value(netcdf,%d): index out-of-range!", elem);
        }

        long string_length = dim[nr_dims-1];  //-- the length of each string
        char str[string_length+1];
        str[string_length] = '\0';            //-- terminating NULL

        int index = 0;
        elem = string_length * (elem-1);
        for (long j=0; j < string_length; ++j)
        {
            str[index++] = aVar->as_char(elem++); //-- padding spaces at the end seem to be automagically removed,
        }                                         //-- so we don't need to worry about them!

        return Value(str);
    }
    else
    {
        for( int d=0; d < nr_dims; ++d )
            num_values *= dim[d];               //-- count nr of elements

        if( elem < 1 || elem > num_values )
        {
            return Error("value(netcdf,%d): index out-of-range!", elem);
        }

        double d = aVar->as_double( elem-1 );

        return Value( d );
    }
}


class CDFDimFunction: public Function
{
public:
  CDFDimFunction(const char *n) : Function(n,1,tnetcdf)
  {info = "Returns a list of the netcdf's dimensions.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFDimFunction::Execute(int,Value *arg)
{
  CNetCDF *a;

  arg[0].GetValue(a);
  MvNcVar *aVar = a->GetVar();

  long *edges = aVar->edges();
  int nrDims = aVar->getNumberOfDimensions();
  CList *l = new CList(nrDims);

  for( int i = 0; i < nrDims ; i++ )
    (*l)[i] = edges[i];

  return Value(l);

}

class CDFDimNamesFunction: public Function
{
public:
  CDFDimNamesFunction(const char *n) : Function(n,1,tnetcdf)
  {info = "Returns a list of the netcdf's dimension names.";};
  virtual Value Execute(int arity,Value *arg);
};

Value CDFDimNamesFunction::Execute(int,Value *arg)
{
  CNetCDF *a;

  arg[0].GetValue(a);
  MvNcVar *aVar = a->GetVar();

  int nrDims = aVar->getNumberOfDimensions();
  CList *l = new CList(nrDims);

  for( int i = 0; i < nrDims ; i++ )
    (*l)[i] = aVar->getDimension(i)->name();

  return Value(l);

}

// NetCDF  + NetCDF

class CdfCdfBinOp : public Function {

	binproc F_;

public:
	CdfCdfBinOp(const char *n,binproc f) : Function(n,2,tnetcdf,tnetcdf) { F_ = f; }
	virtual Value Execute(int arity,Value *arg);
};

Value CdfCdfBinOp::Execute(int,Value *arg)
{
  CNetCDF *a;
  CNetCDF *b;

  arg[0].GetValue(a);
  arg[1].GetValue(b);

  const char *newName = marstmp();
  copyfile(a->GetFileName(), newName);

  CNetCDF *c = new CNetCDF(newName,1,'u');
  c->Current(a->Current());

  // Get the variables
  MvNcVar *aVar = a->GetVar();
  MvNcVar *bVar = b->GetVar();
  MvNcVar *cVar = c->GetVar();

  vector<double> aValues,bValues;

  aVar->get(aValues,aVar->edges() );
  bVar->get(bValues,aVar->edges() );

  if ( aValues.size() != bValues.size() )
    return Error("Different dimensions on NetCDF variables");

  double *cValues  = new double[aValues.size()];
  for ( unsigned int i = 0; i < aValues.size(); i++ )
    cValues[i] = F_(aValues[i],bValues[i]);

  cVar->put(cValues,cVar->edges());
  delete [] cValues;

  c->Flush();
  return Value(c);

}

///// NetCDF + number
class NumCdfBinOp : public Function
{
  binproc F_;

public:
  NumCdfBinOp(const char *n,binproc f) : Function(n) { F_ = f; }
  virtual Value Execute(int arity,Value *arg);
  virtual int ValidArguments(int arity,Value *arg);
};

Value NumCdfBinOp::Execute(int,Value *arg)
{
  double  d;
  CNetCDF *a;
  if(arg[0].GetType() == tnumber)
    {
      arg[0].GetValue(d);
      arg[1].GetValue(a);
    }
  else
    {
      arg[0].GetValue(a);
      arg[1].GetValue(d);
    }

  const char *newName = marstmp();
  copyfile(a->GetFileName(), newName);
  CNetCDF *b = new CNetCDF(newName,1,'u');
  b->Current(a->Current());

  // Get the variables
  MvNcVar *aVar = a->GetVar();
  MvNcVar *bVar = b->GetVar();

  vector<double> aValues;
  aVar->get(aValues,aVar->edges() );

  double *bValues  = new double[aValues.size()];

  for ( unsigned int i = 0; i < aValues.size(); i++ )
    {
      if ( arg[0].GetType() == tnumber)
	bValues[i] = F_(d,aValues[i]);
      else
	bValues[i] = F_(aValues[i],d);
    }

  bVar->put(bValues,bVar->edges());
  delete [] bValues;

  b->Flush();

  return Value(b);
}

int NumCdfBinOp::ValidArguments(int arity,Value *arg)
{
  if(arity != 2) return false;
  if(arg[0].GetType() == tnumber && arg[1].GetType() == tnetcdf)
    return true;
  if(arg[1].GetType() == tnumber && arg[0].GetType() == tnetcdf)
    return true;
  return false;
}


///// function(NetCDF)
class CdfUniOp : public Function
{
  uniproc F_;

public:
 CdfUniOp (const char *n,uniproc f) : Function(n,1,tnetcdf) { F_ = f; }
  virtual Value Execute(int arity,Value *arg);
};

Value CdfUniOp::Execute(int,Value *arg)
{
  CNetCDF *a;

  arg[0].GetValue(a);

  const char *newName = marstmp();
  copyfile(a->GetFileName(), newName);
  CNetCDF *b = new CNetCDF(newName,1,'u');
  b->Current(a->Current());

  // Get the variables
  MvNcVar *aVar = a->GetVar();
  MvNcVar *bVar = b->GetVar();

  vector<double> aValues;
  aVar->get(aValues,aVar->edges() );

  double *bValues  = new double[aValues.size()];

  for ( unsigned int i = 0; i < aValues.size(); i++ )
  {
    bValues[i] = F_(aValues[i]);
  }

  bVar->put(bValues,bVar->edges());
  delete [] bValues;

  b->Flush();

  return Value(b);
}



static void install(Context *c)
{

  // NetCDF specific functions
  c->AddFunction(new CDFVarFunction("variables"));
  c->AddFunction(new CDFAttFunction("attributes"));
  c->AddFunction(new CDFGlobalAttFunction("global_attributes"));
  c->AddFunction(new CDFCurrentFunction("setcurrent"));
  c->AddFunction(new CDFValuesFunction("values"));
  c->AddFunction(new CDFValFunction("value"));
  c->AddFunction(new CDFDimFunction("dimensions"));
  c->AddFunction(new CDFDimNamesFunction("dimension_names"));

  // Binary op
  int i;
  for(i=0;BinOps[i].symb;i++)
    c->AddFunction(new CdfCdfBinOp(BinOps[i].symb,BinOps[i].proc ));

  for(i=0;BinOps[i].symb;i++)
    c->AddFunction(new NumCdfBinOp(BinOps[i].symb,BinOps[i].proc ));

  // Mult op as Binary op
  for(i=0;MulOps[i].symb;i++)
    c->AddFunction(new CdfCdfBinOp(MulOps[i].symb,MulOps[i].proc ));
  for(i=0;MulOps[i].symb;i++)
    c->AddFunction(new NumCdfBinOp(MulOps[i].symb,MulOps[i].proc ));

	// Unary operations
  for(i=0;UniOps[i].symb;i++)
    c->AddFunction(new CdfUniOp(UniOps[i].symb,UniOps[i].proc ));
}

static Linkage linkage(install);
