///////////////////////////////////////////////////////////////////////
// File:        static_shape.h
// Description: Defines the size of the 4-d tensor input/output from a network.
// Author:      Ray Smith
// Created:     Fri Oct 14 09:07:31 PST 2016
//
// (C) Copyright 2016, Google Inc.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
// http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
///////////////////////////////////////////////////////////////////////

#ifndef TESSERACT_LSTM_STATIC_SHAPE_H_
#define TESSERACT_LSTM_STATIC_SHAPE_H_

#include "serialis.h" // for TFile
#include "tprintf.h"  // for tprintf

namespace tesseract {

// Enum describing the loss function to apply during training and/or the
// decoding method to apply at runtime.
enum LossType {
  LT_NONE,     // Undefined.
  LT_CTC,      // Softmax with standard CTC for training/decoding.
  LT_SOFTMAX,  // Outputs sum to 1 in fixed positions.
  LT_LOGISTIC, // Logistic outputs with independent values.
};

// Simple class to hold the tensor shape that is known at network build time
// and the LossType of the loss function.
class StaticShape {
public:
  StaticShape() : batch_(0), height_(0), width_(0), depth_(0), loss_type_(LT_NONE) {}
  int batch() const {
    return batch_;
  }
  void set_batch(int value) {
    batch_ = value;
  }
  int height() const {
    return height_;
  }
  void set_height(int value) {
    height_ = value;
  }
  int width() const {
    return width_;
  }
  void set_width(int value) {
    width_ = value;
  }
  int depth() const {
    return depth_;
  }
  void set_depth(int value) {
    depth_ = value;
  }
  LossType loss_type() const {
    return loss_type_;
  }
  void set_loss_type(LossType value) {
    loss_type_ = value;
  }
  void SetShape(int batch, int height, int width, int depth) {
    batch_ = batch;
    height_ = height;
    width_ = width;
    depth_ = depth;
  }

  void Print() const {
    tprintf("Batch=%d, Height=%d, Width=%d, Depth=%d, loss=%d\n", batch_, height_, width_, depth_,
            loss_type_);
  }

  bool DeSerialize(TFile *fp) {
    int32_t tmp = LT_NONE;
    bool result = fp->DeSerialize(&batch_) && fp->DeSerialize(&height_) &&
                  fp->DeSerialize(&width_) && fp->DeSerialize(&depth_) && fp->DeSerialize(&tmp);
    loss_type_ = static_cast<LossType>(tmp);
    return result;
  }

  bool Serialize(TFile *fp) const {
    int32_t tmp = loss_type_;
    return fp->Serialize(&batch_) && fp->Serialize(&height_) && fp->Serialize(&width_) &&
           fp->Serialize(&depth_) && fp->Serialize(&tmp);
  }

private:
  // Size of the 4-D tensor input/output to a network. A value of zero is
  // allowed for all except depth_ and means to be determined at runtime, and
  // regarded as variable.
  // Number of elements in a batch, or number of frames in a video stream.
  int32_t batch_;
  // Height of the image.
  int32_t height_;
  // Width of the image.
  int32_t width_;
  // Depth of the image. (Number of "nodes").
  int32_t depth_;
  // How to train/interpret the output.
  LossType loss_type_;
};

} // namespace tesseract

#endif // TESSERACT_LSTM_STATIC_SHAPE_H_