/*
 * Copyright 2017 Jan Dorazil <deu439@gmail.com>
 *
 * 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.
 *
 */

#include "readdata.hpp"

#include <tinyxml.h>

const char *data_entry = "data";
const char *unit_entry = "dataunit";

// Must keep the same order as in DataUnit
const char *entries[] = {
  "template",
  "template-roi",
  "input",
  "ground-truth",
  "output"
};

#define NENT  5

ReadData::ReadData(const std::string &xmlfile)
{
    TiXmlDocument doc(xmlfile);
    if (!doc.LoadFile()) return;

    TiXmlHandle hDoc(&doc);
    TiXmlHandle hUnit(0);
    TiXmlElement *pElem0, *pElem1;
  
    // Iterate through data_entry
    pElem0 = hDoc.FirstChild( data_entry ).FirstChild().Element();
    for( pElem0; pElem0; pElem0=pElem0->NextSiblingElement())
    {
        // Check if this is unit_entry
        std::string sKey(pElem0->Value());
        if( sKey.compare( unit_entry ) )
            continue;

        // Iterate through unit_entry and create record
        RecordMap record;
        hUnit = TiXmlHandle(pElem0);
        
        pElem1 = hUnit.FirstChild().Element();
        for( pElem1; pElem1; pElem1=pElem1->NextSiblingElement() )
        {
            const char *pKey = pElem1->Value();
            const char *pText = pElem1->GetText();
            if( pKey && pText ) {
                record[pKey] = pText;
            }
        }
        
        // Add record to the record vector
        mRecords.push_back( record );
    }
    
    // Store path
    size_t end = xmlfile.find_last_of("/");
    if( end != std::string::npos )
      prefix = xmlfile.substr(0, end + 1);
    else
      prefix = std::string("");
}

DataUnit ReadData::getUnit(int idx)
{
  DataUnit unit;

  for( int i = 0; i < NENT; i++ ){
    RecordMap &record = mRecords[idx];
    RecordMap::iterator iter = record.find(entries[i]);
    if( iter != record.end() ){
      // Build file path
      std::string path = prefix + iter->second;
      
      // Read image - read masks in grayscale, dont read output image
      switch( i ){
        case GTRU:
        case TROI:
          unit.mMat[i] = imread(path, IMREAD_GRAYSCALE);
          break;
        case TEMP:
        case INPU:
          unit.mMat[i] = imread(path, IMREAD_COLOR);
          break;
        default:
          break;
      }
      
      // Check whether we have succeeded
      if( unit.mMat[i].empty() && i != OUTP )
        throw std::invalid_argument(path);
      
      // Store path
      unit.mPath[i] = path.c_str();
    }else{
      unit.mPath[i] = std::string();
    }
  }
  
  return unit;
}


ReadData::iterator ReadData::begin()
{
  return ReadData::iterator( *this,  0 );
}

ReadData::iterator ReadData::end()
{
  return ReadData::iterator( *this, mRecords.size() );
}

DataIterator::DataIterator(ReadData& reader, int idx)
  : mReader(reader), mIdx(idx)
{
}
