/*
 * Example utility to convert a vertical profile of meteorological
 * data to the ARL format
 */

#include <iostream>
#include <string>
#include <vector>

#include "radixcommand/commandline.hh"
#include "radixcore/stringfunctions.hh"
#include "radixmath/util.hh"

#include "radixio/arldatastream.hh"
#include "radixio/csvfile.hh"

using namespace radix;

/**
 * @brief interpolateValues Interpolate/extrapolate missing values from a data
 * vector Basic capability for now - linear interpolation, extrapolation assumes
 * constant from ends.
 * @param pressureValues Pressure vector for meteorology - requires all values
 * to be present. Used for determining 'interpolation bounds' (there's
 * definitely a better phrase for this)
 * @param valuesToInterpolate
 */
bool interpolateValues(const std::vector<float> &pressureValues,
                       std::vector<float> &valuesToInterpolate)
{
  float missingValue = -9999.f;

  // Ensure pressure and other vector are same length for interpolation
  if (pressureValues.size() != valuesToInterpolate.size())
  {
    std::cout << "Error! Pressure vector is not same size ("
              << pressureValues.size() << ") as vector to interpolate ("
              << valuesToInterpolate.size() << ") - can't use to interpolate"
              << std::endl;
    return false;
  }

  // Ensure we have no missing values in the pressure dataset
  for (float pressure : pressureValues)
  {
    if (pressure == -9999.f)
    {
      std::cout << "Error! Pressure vector contains " << missingValue
                << " values - can't use to interpolate" << std::endl;
      return false;
    }
  }

  // Write out initial values
  std::cout << "Interpolating " << valuesToInterpolate.size()
            << " values; initial:" << std::endl
            << "  ";
  for (float f : valuesToInterpolate)
  {
    std::cout << f << " ";
  }
  std::cout << std::endl;

  // Loop through vector first to find first/last non-missing indices
  size_t firstIndex = valuesToInterpolate.size(), lastIndex = 0;
  bool foundFirst = false;
  for (size_t i = 0; i < valuesToInterpolate.size(); ++i)
  {
    if (valuesToInterpolate[i] != missingValue)
    {
      if (!foundFirst)
      {
        firstIndex = i;
        foundFirst = true;
      }
      lastIndex = i;
    }
  }
  // Fill in values before first and after last indices
  std::cout << "  Found first (" << firstIndex << ") and last (" << lastIndex
            << ") non-missing indices; filling in values before & after..."
            << std::endl;
  for (size_t i = 0; i < firstIndex; ++i)
  {
    valuesToInterpolate[i] = valuesToInterpolate[firstIndex];
  }
  for (size_t i = lastIndex + 1; i < valuesToInterpolate.size(); ++i)
  {
    valuesToInterpolate[i] = valuesToInterpolate[lastIndex];
  }

  // Fill in missing values in central parts of vector
  std::cout << "  Filling in missing data in central part of vector..."
            << std::endl;
  for (size_t i = firstIndex + 1; i < lastIndex; ++i)
  {
    // Search for a missing value
    size_t lastGood = i - 1;
    if (valuesToInterpolate[i] == missingValue)
    {
      // Get the next good value
      while (valuesToInterpolate[i] == missingValue)
      {
        i++;
      }
      size_t nextGood = i;
      // Interpolate between the two good values
      for (size_t j = lastGood + 1; j < nextGood; ++j)
      {
        valuesToInterpolate[j] =
            valuesToInterpolate[lastGood] +
            ((valuesToInterpolate[nextGood] - valuesToInterpolate[lastGood]) *
             ((pressureValues[j] - pressureValues[lastGood]) /
              (pressureValues[nextGood] - pressureValues[lastGood])));
      }
    }
  }

  // Write out final values
  std::cout << "Interpolation complete; final:" << std::endl << "  ";
  for (float f : valuesToInterpolate)
  {
    std::cout << f << " ";
  }
  std::cout << std::endl;

  return true;
}

int main(int argc, char **argv)
{
  std::cout << "************************" << std::endl;
  std::cout << "***** radixsnd2arl *****" << std::endl;
  std::cout << "************************" << std::endl;

  // Set up command line options
  CommandLine commandLine(argc, argv);
  commandLine.addOption("i", "Input csv file containing met data", true);
  commandLine.addOption("clat", "Centre latitude of output ARL file (degrees)",
                        true);
  commandLine.addOption("clon", "Centre longitude of output ARL file (degrees)",
                        true);
  commandLine.addOption("e", "Extent of output ARL file (degrees)", true);
  commandLine.addOption("r", "Resolution of output ARL file (degrees)", true);
  commandLine.addOption("t", "Time of data start (YYYYMMDDHH)", true);
  commandLine.addOption("n", "Number of one hour timesteps to output", false);
  commandLine.addOption("o", "Output ARL file", false);

  // Ensure required options present
  std::vector<std::string> commandErrors;
  if (!commandLine.validate(commandErrors))
  {
    std::cout << "Error in arguments..." << std::endl;
    for (std::string error : commandErrors)
    {
      std::cout << "\t" << error << std::endl;
    }
    std::cout << std::endl;
    commandLine.printParsedLine(std::cout);

    return -1;
  }

  // Get command line options
  std::string inputCsvPath = commandLine.get<std::string>("i");
  std::string outputArlPath =
      commandLine.get<std::string>("o", inputCsvPath + ".bin");
  float extent          = commandLine.get<int>("e");
  float resolution      = commandLine.get<int>("r");
  float centreLat       = commandLine.get<float>("clat");
  float centreLon       = commandLine.get<float>("clon");
  int numberTimesteps   = commandLine.get<int>("n", 1);
  std::string startTime = commandLine.get<std::string>("t");
  // Get the grid size
  int numberGridCells = (int)(extent / resolution);
  if (!numberGridCells % 2 == 1)
  {
    numberGridCells++;
  }
  // Parse the start time
  // TODO: Check string is the right length
  int year  = from_string(startTime.substr(0, 4), 1951);
  int month = from_string(startTime.substr(4, 2), 11);
  int day   = from_string(startTime.substr(6, 2), 19);
  int hour  = from_string(startTime.substr(8, 2), 17);

  std::cout << "Creating ARL-formatted met file with parameters:" << std::endl;
  std::cout << "  " << extent << " by " << extent << " degree grid centred on ("
            << centreLat << "," << centreLon << ") with resolution "
            << resolution << " (" << numberGridCells
            << " cells in each direction)" << std::endl;
  std::cout << "  " << numberTimesteps
            << " timesteps (of 1 hour each) starting from " << year << "-"
            << month << "-" << day << " at " << hour << "00" << std::endl
            << std::endl;

  // If we have correct options, parse input file
  std::cout << "Reading input csv file: " << inputCsvPath << std::endl;

  // Header strings to denote fields
  std::string pressureString = "pres", tempString = "temp", relHumString = "rh",
              dewPtString = "tdew", wDirString = "wdir", wSpdString = "wspd";
  int pressureIndex = -1, tempIndex = -1, relHumIndex = -1, wDirIndex = -1,
      wSpdIndex = -1;

  // Open and read the csv
  CSVFile inputCsv(inputCsvPath);
  std::vector<std::vector<std::string>> inputData;
  bool inputReadSuccess = inputCsv.data(inputData);
  if (!inputReadSuccess)
  {
    std::cout << "Failed to read input csv!" << std::endl;
    std::cout
        << "Please ensure the file exists and has the appropriate permissions."
        << std::endl;
    return -2;
  }
  // Ensure there is data in here
  if (inputData.size() < 5)
  {
    std::cout << "Not enough data in input csv!" << std::endl;
    std::cout
        << "There are only " << inputData.size()
        << " lines in this file - at least 5 are required (4 headers + data)"
        << std::endl;
    return -3;
  }

  // Work out which fields are which data points - should be 2nd line (1)
  std::cout << "Finding data fields..." << std::endl;
  for (int entry = 0; entry < inputData[1].size(); ++entry)
  {
    if (inputData[1][entry] == pressureString)
    {
      std::cout << "  Found pressure at index " << entry << std::endl;
      pressureIndex = entry;
    }
    else if (inputData[1][entry] == tempString)
    {
      std::cout << "  Found temperature at index " << entry << std::endl;
      tempIndex = entry;
    }
    else if (inputData[1][entry] == relHumString)
    {
      std::cout << "  Found relative humidity at index " << entry << std::endl;
      relHumIndex = entry;
    }
    else if (inputData[1][entry] == wDirString)
    {
      std::cout << "  Found wind direction at index " << entry << std::endl;
      wDirIndex = entry;
    }
    else if (inputData[1][entry] == wSpdString)
    {
      std::cout << "  Found wind speed at index " << entry << std::endl;
      wSpdIndex = entry;
    }
  }
  // If we are missing one, we can't continue
  if (pressureIndex == -1 || tempIndex == -1 || relHumIndex == -1 ||
      wDirIndex == -1 || wSpdIndex == -1)
  {
    std::cout << "Missing data fields in input!" << std::endl;
    if (pressureIndex == -1)
    {
      std::cout << "  Missing pressure field (denoted by " << pressureString
                << ")";
    }
    if (tempIndex == -1)
    {
      std::cout << "  Missing temperature field (denoted by " << tempString
                << ")";
    }
    if (relHumIndex == -1)
    {
      std::cout << "  Missing relative humidity field (denoted by "
                << relHumString << ")";
    }
    if (wDirIndex == -1)
    {
      std::cout << "  Missing wind direction field (denoted by " << wDirString
                << ")";
    }
    if (wSpdIndex == -1)
    {
      std::cout << "  Missing wind speed field (denoted by " << wSpdString
                << ")";
    }
    std::cout << "Please ensure these fields are present and rerun.";
    return -4;
  }

  // Read the data
  std::vector<float> inputPressures, inputTemps, inputRelHums, inputWDirs,
      inputWSpds;
  std::cout << "Data fields found - reading data..." << std::endl;
  float missingValue = -9999.f;
  for (int row = 4; row < inputData.size(); ++row)
  {
    if (inputData[row].size() < inputData[1].size())
    {
      std::cout << "  Warning: this row (" << row + 1
                << ") has less entries than the field names row" << std::endl;
      ;
    }

    bool foundPressure = false, foundTemp = false, foundRelHum = false,
         foundWDir = false, foundWSpd = false;
    for (int entry = 0; entry < inputData[row].size(); ++entry)
    {
      float thisValue = from_string(inputData[row][entry], missingValue);

      if (entry == pressureIndex)
      {
        foundPressure = true;
        inputPressures.push_back(thisValue);
      }
      else if (entry == tempIndex)
      {
        foundTemp = true;
        inputTemps.push_back(thisValue);
      }
      else if (entry == relHumIndex)
      {
        foundRelHum = true;
        inputRelHums.push_back(thisValue);
      }
      else if (entry == wSpdIndex)
      {
        foundWSpd = true;
        inputWSpds.push_back(thisValue);
      }
      else if (entry == wDirIndex)
      {
        foundWDir = true;
        inputWDirs.push_back(thisValue);
      }
    }
    // Add missing data if any of the above weren't found
    if (!foundPressure)
    {
      std::cout << "  Warning: couldn't find pressure in row " << row
                << std::endl;
      inputPressures.push_back(missingValue);
    }
    if (!foundTemp)
    {
      std::cout << "  Warning: couldn't find temperature in row " << row
                << std::endl;
      inputTemps.push_back(missingValue);
    }
    if (!foundRelHum)
    {
      std::cout << "  Warning: couldn't find relative humidity in row " << row
                << std::endl;
      inputRelHums.push_back(missingValue);
    }
    if (!foundWSpd)
    {
      std::cout << "  Warning: couldn't find wind speed in row " << row
                << std::endl;
      inputWSpds.push_back(missingValue);
    }
    if (!foundWDir)
    {
      std::cout << "  Warning: couldn't find wind direction in row " << row
                << std::endl;
      inputWDirs.push_back(missingValue);
    }
  }
  std::cout << "Data read complete: " << inputPressures.size()
            << " entries read." << std::endl;

  // Interpolate the data to remove missing values
  std::cout << "Interpolating values to remove missing data..." << std::endl;
  std::cout << "Temperature:" << std::endl;
  interpolateValues(inputPressures, inputTemps);
  std::cout << "Relative humidity:" << std::endl;
  interpolateValues(inputPressures, inputRelHums);
  std::cout << "Wind speed:" << std::endl;
  interpolateValues(inputPressures, inputWSpds);
  std::cout << "Wind direction:" << std::endl;
  interpolateValues(inputPressures, inputWDirs);
  std::cout << "Interpolation complete." << std::endl;

  // Convert the data into ARL format
  std::cout << "Converting data to ARL format..." << std::endl;
  ARLDataStream outputStream(outputArlPath, std::ios::out);
  for (int timestep = 0; timestep < numberTimesteps; ++timestep)
  {
    // Write index header section
    std::cout << "    Writing index headers for timestep " << timestep << "..."
              << std::endl;
    ARLRecordHeader thisRecordHeader;
    thisRecordHeader.year  = year;
    thisRecordHeader.month = month;
    thisRecordHeader.day   = day;
    thisRecordHeader.hour  = hour + timestep;
    thisRecordHeader.ic    = 0;
    thisRecordHeader.il    = 0;
    thisRecordHeader.cgrid = "99";
    thisRecordHeader.kvar  = "INDX";
    thisRecordHeader.nexp  = 0;
    thisRecordHeader.prec  = 0.f;
    thisRecordHeader.var1  = 0.f;
    ARLIndexHeader thisIndexHeader;
    thisIndexHeader.model_id = "NFDB";
    thisIndexHeader.icx      = 0;
    thisIndexHeader.mn       = 0;
    thisIndexHeader.pole_lat = centreLat;
    thisIndexHeader.pole_lon = centreLon;
    thisIndexHeader.ref_lat  = centreLat;
    thisIndexHeader.ref_lon  = centreLon;
    thisIndexHeader.size     = 000.f;
    thisIndexHeader.orient   = 0.f;
    thisIndexHeader.tang_lat = centreLat;
    thisIndexHeader.sync_xp  = (numberGridCells + 1) / 2;
    thisIndexHeader.sync_yp  = (numberGridCells + 1) / 2;
    thisIndexHeader.sync_lat = centreLat;
    thisIndexHeader.sync_lon = centreLon;
    thisIndexHeader.dummy    = 0.f;
    thisIndexHeader.nx       = numberGridCells;
    thisIndexHeader.ny       = numberGridCells;
    thisIndexHeader.nz       = inputPressures.size();
    thisIndexHeader.z_flag   = 2;
    thisIndexHeader.lenh =
        inputPressures.size() * numberGridCells * numberGridCells;
    // Write the headers
    outputStream.write_record_header(thisRecordHeader);
    outputStream.write_index_header(thisRecordHeader, thisIndexHeader);
    std::cout << "    Headers written" << std::endl;

    // Write meteorological variables for each level
    for (int level = 0; level < inputPressures.size(); ++level)
    {
      std::cout << "    Writing meteorological data for timestep " << timestep
                << ", level " << level << "..." << std::endl;

      // Write pressure variables
      {
        thisRecordHeader.kvar = "PRSS";
        thisRecordHeader.var1 = inputPressures[level];
        std::vector<std::vector<float>> thisData(
            numberGridCells,
            std::vector<float>(numberGridCells, inputPressures[level]));
        std::cout << "      pressure...";
        outputStream.write_record_header(thisRecordHeader);
        outputStream.write_record(thisRecordHeader, thisIndexHeader, thisData);
        std::cout << "written" << std::endl;
      }

      // Write temperature levels
      {
        thisRecordHeader.kvar = "TEMP";
        thisRecordHeader.var1 = inputTemps[level];
        std::vector<std::vector<float>> thisData(
            numberGridCells,
            std::vector<float>(numberGridCells, inputTemps[level]));
        std::cout << "      temperature...";
        outputStream.write_record_header(thisRecordHeader);
        outputStream.write_record(thisRecordHeader, thisIndexHeader, thisData);
        std::cout << "written" << std::endl;
      }

      // Write relative humidity levels
      {
        thisRecordHeader.kvar = "RELH";
        thisRecordHeader.var1 = inputRelHums[level];
        std::vector<std::vector<float>> thisData(
            numberGridCells,
            std::vector<float>(numberGridCells, inputRelHums[level]));
        std::cout << "      relative humidity...";
        outputStream.write_record_header(thisRecordHeader);
        outputStream.write_record(thisRecordHeader, thisIndexHeader, thisData);
        std::cout << "written" << std::endl;
      }

      // Write wind data
      // Need to calculate wind u and v components from direction/speed
      float thisWindU =
          (float)(toRadians(cos(90.0 - (inputWDirs[level] - 180.0))) *
                  inputWSpds[level] * 0.51444444444);
      float thisWindV =
          (float)(toRadians(sin(90.0 - (inputWDirs[level] - 180.0))) *
                  inputWSpds[level] * 0.51444444444);
      std::cout << "      Initial wspd = " << inputWSpds[level]
                << " knots, wdir: " << inputWDirs[level] << " degrees"
                << std::endl;
      std::cout << "      Converted wind components: u = " << thisWindU
                << " m/s, v = " << thisWindV << "m/s" << std::endl;
      {
        thisRecordHeader.kvar = "UWND";
        thisRecordHeader.var1 = thisWindU;
        std::vector<std::vector<float>> thisData(
            numberGridCells, std::vector<float>(numberGridCells, thisWindU));
        std::cout << "      wind (u component)...";
        outputStream.write_record_header(thisRecordHeader);
        outputStream.write_record(thisRecordHeader, thisIndexHeader, thisData);
        std::cout << "written" << std::endl;
      }
      {
        thisRecordHeader.kvar = "VWND";
        thisRecordHeader.var1 = thisWindV;
        std::vector<std::vector<float>> thisData(
            numberGridCells, std::vector<float>(numberGridCells, thisWindV));
        std::cout << "      wind (v component)...";
        outputStream.write_record_header(thisRecordHeader);
        outputStream.write_record(thisRecordHeader, thisIndexHeader, thisData);
        std::cout << "written" << std::endl;
      }
    }
    std::cout << " Written timestep " << timestep << std::endl;
  }

  std::cout << "File write complete!" << std::endl;

  return 0;
}