/*    This file is a part of Calyp project
 *    Copyright (C) 2014-2019  by Joao Carreira   (jfmcarreira@gmail.com)
 *                                Luis Lucas      (luisfrlucas@gmail.com)
 *
 *    This program is free software; you can redistribute it and/or modify
 *    it under the terms of the GNU General Public License as published by
 *    the Free Software Foundation; either version 2 of the License, or
 *    (at your option) any later version.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    GNU General Public License for more details.
 *
 *   You should have received a copy of the GNU General Public License along
 *   with this program; if not, write to the Free Software Foundation, Inc.,
 *   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

/**
 * \file     CalypPythonModuleIf.cpp
 * \brief    Calyp modules interface for python
 */

#include "CalypPythonModuleIf.h"

bool CalypPythonModuleIf::m_sbPythonInitialized = false;

// OpenCV
#include <opencv2/core/mat.hpp>
#include <opencv2/opencv.hpp>

// Python
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>

static PyObject* opencv_error = 0;

static int failmsg( const char* fmt, ... )
{
  char str[1000];

  va_list ap;
  va_start( ap, fmt );
  vsnprintf( str, sizeof( str ), fmt, ap );
  va_end( ap );

  PyErr_SetString( PyExc_TypeError, str );
  return 0;
}

struct ArgInfo
{
  const char* name;
  bool outputarg;
  // more fields may be added if necessary

  ArgInfo( const char* name_, bool outputarg_ )
      : name( name_ )
      , outputarg( outputarg_ ) {}

  // to match with older pyopencv_to function signature
  operator const char*() const { return name; }
};

class PyAllowThreads
{
public:
  PyAllowThreads()
      : _state( PyEval_SaveThread() ) {}
  ~PyAllowThreads()
  {
    PyEval_RestoreThread( _state );
  }

private:
  PyThreadState* _state;
};

class PyEnsureGIL
{
public:
  PyEnsureGIL()
      : _state( PyGILState_Ensure() ) {}
  ~PyEnsureGIL()
  {
    PyGILState_Release( _state );
  }

private:
  PyGILState_STATE _state;
};

#define ERRWRAP2( expr )                       \
  try                                          \
  {                                            \
    PyAllowThreads allowThreads;               \
    expr;                                      \
  }                                            \
  catch( const cv::Exception& e )              \
  {                                            \
    PyErr_SetString( opencv_error, e.what() ); \
    return 0;                                  \
  }

using namespace cv;

class NumpyAllocator : public MatAllocator
{
public:
  NumpyAllocator() { stdAllocator = Mat::getStdAllocator(); }
  ~NumpyAllocator() {}

  UMatData* allocate( PyObject* o, int dims, const int* sizes, int type, size_t* step ) const
  {
    UMatData* u = new UMatData( this );
    u->data = u->origdata = (uchar*)PyArray_DATA( (PyArrayObject*)o );
    npy_intp* _strides = PyArray_STRIDES( (PyArrayObject*)o );
    for( int i = 0; i < dims - 1; i++ )
      step[i] = (size_t)_strides[i];
    step[dims - 1] = CV_ELEM_SIZE( type );
    u->size = sizes[0] * step[0];
    u->userdata = o;
    return u;
  }

  UMatData* allocate( int dims0, const int* sizes, int type, void* data, size_t* step, int flags, UMatUsageFlags usageFlags ) const
  {
    if( data != 0 )
    {
      // issue #6969: CV_Error(Error::StsAssert, "The data should normally be NULL!");
      // probably this is safe to do in such extreme case
      return stdAllocator->allocate( dims0, sizes, type, data, step, flags, usageFlags );
    }
    PyEnsureGIL gil;

    int depth = CV_MAT_DEPTH( type );
    int cn = CV_MAT_CN( type );
    const int f = (int)( sizeof( size_t ) / 8 );
    int typenum = depth == CV_8U ? NPY_UBYTE : depth == CV_8S ? NPY_BYTE : depth == CV_16U ? NPY_USHORT : depth == CV_16S ? NPY_SHORT : depth == CV_32S ? NPY_INT : depth == CV_32F ? NPY_FLOAT : depth == CV_64F ? NPY_DOUBLE : f * NPY_ULONGLONG + ( f ^ 1 ) * NPY_UINT;
    int i, dims = dims0;
    cv::AutoBuffer<npy_intp> _sizes( dims + 1 );
    for( i = 0; i < dims; i++ )
      _sizes[i] = sizes[i];
    if( cn > 1 )
      _sizes[dims++] = cn;
    PyObject* o = PyArray_SimpleNew( dims, _sizes, typenum );
    if( !o )
      CV_Error_( Error::StsError, ( "The numpy array of typenum=%d, ndims=%d can not be created", typenum, dims ) );
    return allocate( o, dims0, sizes, type, step );
  }

  bool allocate( UMatData* u, int accessFlags, UMatUsageFlags usageFlags ) const
  {
    return stdAllocator->allocate( u, accessFlags, usageFlags );
  }

  void deallocate( UMatData* u ) const
  {
    if( !u )
      return;
    PyEnsureGIL gil;
    CV_Assert( u->urefcount >= 0 );
    CV_Assert( u->refcount >= 0 );
    if( u->refcount == 0 )
    {
      PyObject* o = (PyObject*)u->userdata;
      Py_XDECREF( o );
      delete u;
    }
  }

  const MatAllocator* stdAllocator;
};

NumpyAllocator g_numpyAllocator;

template <typename T>
static bool pyopencv_to( PyObject* obj, T& p, const char* name = "<unknown>" );

template <typename T>
static PyObject* pyopencv_from( const T& src );

enum
{
  ARG_NONE = 0,
  ARG_MAT = 1,
  ARG_SCALAR = 2
};

// special case, when the converter needs full ArgInfo structure
static bool pyopencv_to( PyObject* o, Mat& m, const ArgInfo info )
{
  bool allowND = true;
  if( !o || o == Py_None )
  {
    if( !m.data )
      m.allocator = &g_numpyAllocator;
    return true;
  }

  //    if( PyInt_Check(o) )
  //    {
  //        double v[] = {static_cast<double>(PyInt_AsLong((PyObject*)o)), 0., 0., 0.};
  //        m = Mat(4, 1, CV_64F, v).clone();
  //        return true;
  //    }
  if( PyFloat_Check( o ) )
  {
    double v[] = {PyFloat_AsDouble( (PyObject*)o ), 0., 0., 0.};
    m = Mat( 4, 1, CV_64F, v ).clone();
    return true;
  }
  if( PyTuple_Check( o ) )
  {
    int i, sz = (int)PyTuple_Size( (PyObject*)o );
    m = Mat( sz, 1, CV_64F );
    for( i = 0; i < sz; i++ )
    {
      PyObject* oi = PyTuple_GET_ITEM( o, i );
      //            if( PyInt_Check(oi) )
      //                m.at<double>(i) = (double)PyInt_AsLong(oi);
      //            else
      if( PyFloat_Check( oi ) )
        m.at<double>( i ) = (double)PyFloat_AsDouble( oi );
      else
      {
        failmsg( "%s is not a numerical tuple", info.name );
        m.release();
        return false;
      }
    }
    return true;
  }

  if( !PyArray_Check( o ) )
  {
    failmsg( "%s is not a numpy array, neither a scalar", info.name );
    return false;
  }

  PyArrayObject* oarr = (PyArrayObject*)o;

  bool needcopy = false, needcast = false;
  int typenum = PyArray_TYPE( oarr ), new_typenum = typenum;
  int type = typenum == NPY_UBYTE ? CV_8U :
                                    typenum == NPY_BYTE ? CV_8S :
                                                          typenum == NPY_USHORT ? CV_16U :
                                                                                  typenum == NPY_SHORT ? CV_16S :
                                                                                                         typenum == NPY_INT ? CV_32S :
                                                                                                                              typenum == NPY_INT32 ? CV_32S :
                                                                                                                                                     typenum == NPY_FLOAT ? CV_32F :
                                                                                                                                                                            typenum == NPY_DOUBLE ? CV_64F : -1;

  if( type < 0 )
  {
    if( typenum == NPY_INT64 || typenum == NPY_UINT64 || typenum == NPY_LONG )
    {
      needcopy = needcast = true;
      new_typenum = NPY_INT;
      type = CV_32S;
    }
    else
    {
      failmsg( "%s data type = %d is not supported", info.name, typenum );
      return false;
    }
  }

#ifndef CV_MAX_DIM
  const int CV_MAX_DIM = 32;
#endif

  int ndims = PyArray_NDIM( oarr );
  if( ndims >= CV_MAX_DIM )
  {
    failmsg( "%s dimensionality (=%d) is too high", info.name, ndims );
    return false;
  }

  int size[CV_MAX_DIM + 1];
  size_t step[CV_MAX_DIM + 1];
  size_t elemsize = CV_ELEM_SIZE1( type );
  const npy_intp* _sizes = PyArray_DIMS( oarr );
  const npy_intp* _strides = PyArray_STRIDES( oarr );
  bool ismultichannel = ndims == 3 && _sizes[2] <= CV_CN_MAX;

  for( int i = ndims - 1; i >= 0 && !needcopy; i-- )
  {
    // these checks handle cases of
    //  a) multi-dimensional (ndims > 2) arrays, as well as simpler 1- and 2-dimensional cases
    //  b) transposed arrays, where _strides[] elements go in non-descending order
    //  c) flipped arrays, where some of _strides[] elements are negative
    // the _sizes[i] > 1 is needed to avoid spurious copies when NPY_RELAXED_STRIDES is set
    if( ( i == ndims - 1 && _sizes[i] > 1 && (size_t)_strides[i] != elemsize ) ||
        ( i < ndims - 1 && _sizes[i] > 1 && _strides[i] < _strides[i + 1] ) )
      needcopy = true;
  }

  if( ismultichannel && _strides[1] != (npy_intp)elemsize * _sizes[2] )
    needcopy = true;

  if( needcopy )
  {
    if( info.outputarg )
    {
      failmsg( "Layout of the output array %s is incompatible with cv::Mat (step[ndims-1] != elemsize or step[1] != elemsize*nchannels)", info.name );
      return false;
    }

    if( needcast )
    {
      o = PyArray_Cast( oarr, new_typenum );
      oarr = (PyArrayObject*)o;
    }
    else
    {
      oarr = PyArray_GETCONTIGUOUS( oarr );
      o = (PyObject*)oarr;
    }

    _strides = PyArray_STRIDES( oarr );
  }

  // Normalize strides in case NPY_RELAXED_STRIDES is set
  size_t default_step = elemsize;
  for( int i = ndims - 1; i >= 0; --i )
  {
    size[i] = (int)_sizes[i];
    if( size[i] > 1 )
    {
      step[i] = (size_t)_strides[i];
      default_step = step[i] * size[i];
    }
    else
    {
      step[i] = default_step;
      default_step *= size[i];
    }
  }

  // handle degenerate case
  if( ndims == 0 )
  {
    size[ndims] = 1;
    step[ndims] = elemsize;
    ndims++;
  }

  if( ismultichannel )
  {
    ndims--;
    type |= CV_MAKETYPE( 0, size[2] );
  }

  if( ndims > 2 && !allowND )
  {
    failmsg( "%s has more than 2 dimensions", info.name );
    return false;
  }

  m = Mat( ndims, size, type, PyArray_DATA( oarr ), step );
  m.u = g_numpyAllocator.allocate( o, ndims, size, type, step );
  m.addref();

  if( !needcopy )
  {
    Py_INCREF( o );
  }
  m.allocator = &g_numpyAllocator;

  return true;
}

template <>
bool pyopencv_to( PyObject* o, Mat& m, const char* name )
{
  return pyopencv_to( o, m, ArgInfo( name, 0 ) );
}

template <typename _Tp, int m, int n>
bool pyopencv_to( PyObject* o, Matx<_Tp, m, n>& mx, const ArgInfo info )
{
  Mat tmp;
  if( !pyopencv_to( o, tmp, info ) )
  {
    return false;
  }

  tmp.copyTo( mx );
  return true;
}

template <typename _Tp, int m, int n>
bool pyopencv_to( PyObject* o, Matx<_Tp, m, n>& mx, const char* name )
{
  return pyopencv_to( o, mx, ArgInfo( name, 0 ) );
}

template <typename T>
bool pyopencv_to( PyObject* o, Ptr<T>& p, const char* name )
{
  if( !o || o == Py_None )
    return true;
  p = makePtr<T>();
  return pyopencv_to( o, *p, name );
}

template <>
PyObject* pyopencv_from( const Mat& m )
{
  if( !m.data )
    Py_RETURN_NONE;
  Mat temp, *p = (Mat*)&m;
  if( !p->u || p->allocator != &g_numpyAllocator )
  {
    temp.allocator = &g_numpyAllocator;
    ERRWRAP2( m.copyTo( temp ) );
    p = &temp;
  }
  PyObject* o = (PyObject*)p->u->userdata;
  Py_INCREF( o );
  return o;
}

bool CalypPythonModuleIf::create( std::vector<CalypFrame*> apcFrameList )
{
  _BASIC_MODULE_API_2_CHECK_

  // initialize Python embedding
  if( !m_sbPythonInitialized )
  {
    Py_Initialize();

    // Initialize numpy library
    import_array1( -1 );
    m_sbPythonInitialized = true;
  }

#if PY_MAJOR_VERSION >= 3
  PyObject* pName = PyUnicode_DecodeFSDefault( m_pchPythonFunctionName );
#else
  PyObject* pName = PyString_FromString( m_pchPythonFunctionName );
#endif
  m_pyModule = PyImport_Import( pName );
  m_pyFunction = PyObject_GetAttrString( m_pyModule, "ClpModule" );
  Py_DECREF( pName );
  if( !m_pyFunction || !PyCallable_Check( m_pyFunction ) )
  {
    return false;
  }
  unsigned pelFormat = m_bConvertToGray ? CalypFrame::findPixelFormat( "GRAY" ) : apcFrameList[0]->getPelFormat();
  m_pcOutputFrame = new CalypFrame( apcFrameList[0]->getWidth(), apcFrameList[0]->getHeight(), pelFormat, apcFrameList[0]->getBitsPel() );
  return true;
}

CalypFrame* CalypPythonModuleIf::process( std::vector<CalypFrame*> apcFrameList )
{
  m_pcOutputFrame->reset();
  if( m_pyFunction && PyCallable_Check( m_pyFunction ) )
  {
    Mat cvMatOutput;
    Mat cvMatInput[m_uiNumberOfFrames];
    PyObject* pyInputArg = PyTuple_New( apcFrameList.size() );
    for( unsigned i = 0; i < apcFrameList.size(); i++ )
    {
      apcFrameList[0]->toMat( cvMatInput[i], m_bConvertToGray, false );

      PyObject* pyMat = pyopencv_from( cvMatInput[i] );
      /* pValue reference stolen here: */
      PyTuple_SetItem( pyInputArg, i, pyMat );
    }
    // execute the function
    PyObject* pyResult = PyEval_CallObject( m_pyFunction, pyInputArg );
    if( !pyResult )
      return m_pcOutputFrame;

    pyopencv_to( pyResult, cvMatOutput );

    // decrement the object references
    Py_XDECREF( pyResult );
    Py_XDECREF( pyInputArg );

    m_pcOutputFrame->fromMat( cvMatOutput, CLP_LUMA );
  }
  return m_pcOutputFrame;
}

void CalypPythonModuleIf::destroy()
{
  Py_XDECREF( m_pyFunction );
  //#if PY_MAJOR_VERSION >= 3
  //  Py_FinalizeEx();
  //#else
  //  Py_Finalize();
  //#endif
  if( m_pcOutputFrame )
    delete m_pcOutputFrame;
  m_pcOutputFrame = NULL;
}