#include "ccv.h"

#include "ccv_internal.h"

#include "nnc/ccv_nnc.h"

#include "nnc/ccv_nnc_easy.h"

#include "nnc/ccv_nnc_internal.h"

#ifdef USE_OPENMP

#include <omp.h>

#endif

#ifdef USE_DISPATCH

#include <dispatch/dispatch.h>

#endif

static int _ccv_nnc_categorical_crossentropy_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  assert(input_size == 2);

  const ccv_nnc_tensor_t* a = inputs[0];

  assert(CCV_IS_TENSOR_CONTIGUOUS(a));

  const ccv_nnc_tensor_t* b = inputs[1];

  assert(CCV_IS_TENSOR_CONTIGUOUS(b));

  assert(output_size == 1);

  ccv_nnc_tensor_t* c = outputs[0];

  assert(CCV_IS_TENSOR_CONTIGUOUS(c));

  const int axis_count = ccv_nnc_tensor_nd(a->info.dim);

  const int batch_size = axis_count < 2 ? 
10
 : a->info.dim[0];

  const int count = ccv_nnc_tensor_count(a->info) / batch_size;

  int i;

  if (b->info.datatype == CCV_32F)

  {

    // If has more than 1 axis, then the range is the channel count. Otherwise, if our batch size is 1, then the range is

    // the channel count. Otherwise, the range is 1 (and the only axis is the batch size).

    const int range = ccv_nnc_tensor_nd(b->info.dim) > 1 ? 
ccv_nnc_tensor_get_c(b->info)1
 : 
(8
batch_size == 18
 ? 
b->info.dim[0]0
 : 
18
);

    if (range == 1)

    {

      for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && b->info.dim[i] > 0; 
i++8
)

        { assert(b->info.dim[i] == c->info.dim[i]); }

      const float trim0 = cmd.info.label_smoothing.trim0;

      const float trim1 = cmd.info.label_smoothing.trim1;

      if (trim0 == 0 && 
trim1 == 14
)

      {

        
parallel_for4
(i, batch_size) {

          const int label = (int)(b->data.f32[i] + 0.5);

          assert(label >= 0 && label < count);

          const float p = a->data.f32[i * count + label];

          c->data.f32[i] = -logf(p);

        } parallel_endfor

      } else {

        
parallel_for4
(i, batch_size) {

          const int label = (int)(b->data.f32[i] + 0.5);

          assert(label >= 0 && label < count);

          int j;

          float p = 0;

          float* const ap = a->data.f32 + i * count;

          for (j = 0; j < label; 
j++1.98k
)

            p += -trim0 * logf(ap[j]);

          p += -trim1 * logf(ap[label]);

          for (j = label + 1; j < count; 
j++1.98k
)

            p += -trim0 * logf(ap[j]);

          c->data.f32[i] = p;

        } parallel_endfor

      }

    } else {

      assert(range == count);

      
parallel_for1
(i, batch_size) {

        int j;

        float p = 0;

        float* const bp = b->data.f32 + i * count;

        float* const ap = a->data.f32 + i * count;

        for (j = 0; j < count; 
j++6
)

          p += -bp[j] * logf(ap[j]);

        c->data.f32[i] = p;

      } parallel_endfor

    }

  } else 
if (2
b->info.datatype == CCV_32S2
) {

    for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && b->info.dim[i] > 0; 
i++3
)

      { assert(b->info.dim[i] == c->info.dim[i]); }

    const float trim0 = cmd.info.label_smoothing.trim0;

    const float trim1 = cmd.info.label_smoothing.trim1;

    if (trim0 == 0 && 
trim1 == 11
)

    {

      
parallel_for1
(i, batch_size) {

        const int label = b->data.i32[i];

        assert(label >= 0 && label < count);

        const float p = a->data.f32[i * count + label];

        c->data.f32[i] = -logf(p);

      } parallel_endfor

    } else {

      
parallel_for1
(i, batch_size) {

        const int label = b->data.i32[i];

        assert(label >= 0 && label < count);

        int j;

        float p = 0;

        float* const ap = a->data.f32 + i * count;

        for (j = 0; j < label; 
j++3
)

          p += -trim0 * logf(ap[j]);

        p += -trim1 * logf(ap[label]);

        for (j = label + 1; j < count; 
j++1
)

          p += -trim0 * logf(ap[j]);

        c->data.f32[i] = p;

      } parallel_endfor

    }

  }

  return CCV_NNC_EXEC_SUCCESS;

}

static int _ccv_nnc_categorical_crossentropy_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  assert(input_size >= 3);

  assert(output_size >= 1);

  const ccv_nnc_tensor_t* g = inputs[0];

  assert(!g || !CCV_IS_TENSOR_VIEW(g));

  const ccv_nnc_tensor_t* a = inputs[1];

  assert(CCV_IS_TENSOR_CONTIGUOUS(a));

  const ccv_nnc_tensor_t* b = inputs[2];

  assert(CCV_IS_TENSOR_CONTIGUOUS(b));

  ccv_nnc_tensor_t* h = outputs[0];

  assert(CCV_IS_TENSOR_CONTIGUOUS(h));

  const int axis_count = ccv_nnc_tensor_nd(a->info.dim);

  const int batch_size = axis_count < 2 ? 
10
 : a->info.dim[0];

  const int count = ccv_nnc_tensor_count(a->info) / batch_size;

  int i;

  if (g)

  {

    if (b->info.datatype == CCV_32F)

    {

      // If has more than 1 axis, then the range is the channel count. Otherwise, if our batch size is 1, then the range is

      // the channel count. Otherwise, the range is 1 (and the only axis is the batch size).

      const int range = ccv_nnc_tensor_nd(b->info.dim) > 1 ? 
ccv_nnc_tensor_get_c(b->info)1
 : 
(4
batch_size == 14
 ? 
b->info.dim[0]0
 : 
14
);

      if (range == 1)

      {

        for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && a->info.dim[i] > 0; 
i++8
)

          { assert(a->info.dim[i] == h->info.dim[i]); }

        const float trim0 = cmd.info.label_smoothing.trim0;

        const float trim1 = cmd.info.label_smoothing.trim1;

        if (trim0 == 0 && 
trim1 == 12
)

        {

          
parallel_for2
(i, batch_size) {

            int j;

            const float gp = g->data.f32[i];

            const int label = (int)(b->data.f32[i] + 0.5);

            float* const hp = h->data.f32 + i * count;

            for (j = 0; j < count; 
j++2.00k
)

              hp[j] = 0;

            const float p = a->data.f32[i * count + label];

            hp[label] = -gp / p;

          } parallel_endfor

        } else {

          
parallel_for2
(i, batch_size) {

            int j;

            const float gp = g->data.f32[i];

            const int label = (int)(b->data.f32[i] + 0.5);

            float* const hp = h->data.f32 + i * count;

            float* const ap = a->data.f32 + i * count;

            for (j = 0; j < label; 
j++990
)

              hp[j] = -gp * trim0 / ap[j];

            hp[label] = -gp * trim1 / ap[label];

            for (j = label + 1; j < count; 
j++990
)

              hp[j] = -gp * trim0 / ap[j];

          } parallel_endfor

        }

      } else {

        assert(range == count);

        
parallel_for1
(i, batch_size) {

          int j;

          const float gp = g->data.f32[i];

          float* const hp = h->data.f32 + i * count;

          float* const ap = a->data.f32 + i * count;

          float* const bp = b->data.f32 + i * count;

          for (j = 0; j < count; 
j++6
)

            hp[j] = -gp * bp[j] / ap[j];

        } parallel_endfor

      }

    } else 
if (3
b->info.datatype == CCV_32S3
) {

      for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && a->info.dim[i] > 0; 
i++6
)

        { assert(a->info.dim[i] == h->info.dim[i]); }

      const float trim0 = cmd.info.label_smoothing.trim0;

      const float trim1 = cmd.info.label_smoothing.trim1;

      if (trim0 == 0 && 
trim1 == 12
)

      {

        
parallel_for2
(i, batch_size) {

          int j;

          const float gp = g->data.f32[i];

          const int label = b->data.i32[i];

          float* const hp = h->data.f32 + i * count;

          for (j = 0; j < count; 
j++12
)

            hp[j] = 0;

          const float p = a->data.f32[i * count + label];

          hp[label] = -gp / p;

        } parallel_endfor

      } else {

        
parallel_for1
(i, batch_size) {

          int j;

          const float gp = g->data.f32[i];

          const int label = b->data.i32[i];

          float* const hp = h->data.f32 + i * count;

          float* const ap = a->data.f32 + i * count;

          for (j = 0; j < label; 
j++3
)

            hp[j] = -gp * trim0 / ap[j];

          hp[label] = -gp * trim1 / ap[label];

          for (j = label + 1; j < count; 
j++1
)

            hp[j] = -gp * trim0 / ap[j];

        } parallel_endfor

      }

    }

  } else {

    if (b->info.datatype == CCV_32F)

    {

      // If has more than 1 axis, then the range is the channel count. Otherwise, if our batch size is 1, then the range is

      // the channel count. Otherwise, the range is 1 (and the only axis is the batch size).

      const int range = ccv_nnc_tensor_nd(b->info.dim) > 1 ? ccv_nnc_tensor_get_c(b->info) : (batch_size == 1 ? b->info.dim[0] : 1);

      if (range == 1)

      {

        for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && a->info.dim[i] > 0; i++)

          { assert(a->info.dim[i] == h->info.dim[i]); }

        const float trim0 = cmd.info.label_smoothing.trim0;

        const float trim1 = cmd.info.label_smoothing.trim1;

        if (trim0 == 0 && trim1 == 1)

        {

          parallel_for(i, batch_size) {

            int j;

            const int label = (int)(b->data.f32[i] + 0.5);

            float* const hp = h->data.f32 + i * count;

            for (j = 0; j < count; j++)

              hp[j] = 0;

            const float p = a->data.f32[i * count + label];

            hp[label] = -1. / p;

          } parallel_endfor

        } else {

          parallel_for(i, batch_size) {

            int j;

            const int label = (int)(b->data.f32[i] + 0.5);

            float* const hp = h->data.f32 + i * count;

            float* const ap = a->data.f32 + i * count;

            for (j = 0; j < label; j++)

              hp[j] = -trim0 / ap[j];

            hp[label] = -trim1 / ap[label];

            for (j = label + 1; j < count; j++)

              hp[j] = -trim0 / ap[j];

          } parallel_endfor

        }

      } else {

        assert(range == count);

        parallel_for(i, batch_size) {

          int j;

          float* const hp = h->data.f32 + i * count;

          float* const ap = a->data.f32 + i * count;

          float* const bp = b->data.f32 + i * count;

          for (j = 0; j < count; j++)

            hp[j] = -bp[j] / ap[j];

        } parallel_endfor

      }

    } else if (b->info.datatype == CCV_32S) {

      const float trim0 = cmd.info.label_smoothing.trim0;

      const float trim1 = cmd.info.label_smoothing.trim1;

      if (trim0 == 0 && trim1 == 1)

      {

        for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && a->info.dim[i] > 0; i++)

          { assert(a->info.dim[i] == h->info.dim[i]); }

        parallel_for(i, batch_size) {

          int j;

          const int label = b->data.i32[i];

          float* const hp = h->data.f32 + i * count;

          for (j = 0; j < count; j++)

            hp[j] = 0;

          const float p = a->data.f32[i * count + label];

          hp[label] = -1. / p;

        } parallel_endfor

      } else {

        for (i = 0; i < CCV_NNC_MAX_DIM_ALLOC && a->info.dim[i] > 0; i++)

          { assert(a->info.dim[i] == h->info.dim[i]); }

        parallel_for(i, batch_size) {

          int j;

          const int label = b->data.i32[i];

          float* const hp = h->data.f32 + i * count;

          float* const ap = a->data.f32 + i * count;

          for (j = 0; j < label; j++)

            hp[j] = -trim0 / ap[j];

          hp[label] = -trim1 / ap[label];

          for (j = label + 1; j < count; j++)

            hp[j] = -trim0 / ap[j];

        } parallel_endfor

      }

    }

  }

  return CCV_NNC_EXEC_SUCCESS;

}

REGISTER_COMMAND_BACKEND(CCV_NNC_CATEGORICAL_CROSSENTROPY_FORWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)

{

  registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC | CCV_TENSOR_FORMAT_NCHW;

  registry->tensor_datatypes = CCV_32F | CCV_32S;

  registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;

  registry->algorithms = 1;

  registry->exec = _ccv_nnc_categorical_crossentropy_forw;

}

REGISTER_COMMAND_BACKEND(CCV_NNC_CATEGORICAL_CROSSENTROPY_BACKWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)

{

  registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC | CCV_TENSOR_FORMAT_NCHW;

  registry->tensor_datatypes = CCV_32F | CCV_32S;

  registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;

  registry->algorithms = 1;

  registry->exec = _ccv_nnc_categorical_crossentropy_back;

}