//---------------------------------------------------------------------------//
//  MIT License
//
//  Copyright (c) 2020-2021 Mikhail Komarov <nemo@nil.foundation>
//  Copyright (c) 2020-2021 Nikita Kaskov <nemo@nil.foundation>
//  Copyright (c) 2021-2022 Aleksei Moskvin <alalmoskvin@gmail.com>
//  Copyright (c) 2021 Ilias Khairullin <ilias@nil.foundation>
//
//  Permission is hereby granted, free of charge, to any person obtaining a copy
//  of this software and associated documentation files (the "Software"), to deal
//  in the Software without restriction, including without limitation the rights
//  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
//  copies of the Software, and to permit persons to whom the Software is
//  furnished to do so, subject to the following conditions:
//
//  The above copyright notice and this permission notice shall be included in all
//  copies or substantial portions of the Software.
//
//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
//  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
//  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
//  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
//  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
//  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
//  SOFTWARE.
//---------------------------------------------------------------------------//

#ifndef CRYPTO3_MERKLE_PROOF_HPP
#define CRYPTO3_MERKLE_PROOF_HPP

#include <algorithm>
#include <vector>

#include <boost/variant.hpp>

#include <nil/crypto3/container/merkle/tree.hpp>

namespace nil {
    namespace crypto3 {
        namespace zk {
            namespace components {
                template<typename, typename, std::size_t>
                struct merkle_proof;
            }    // namespace components
        }        // namespace zk
        namespace marshalling {
            namespace types {
                template<typename, typename>
                struct merkle_proof_marshalling;
            }
        }    // namespace marshalling
        namespace containers {
            namespace detail {
                template<typename NodeType, std::size_t Arity = 2>
                struct merkle_proof_impl {
                    typedef NodeType node_type;
                    typedef typename node_type::hash_type hash_type;

                    constexpr static const std::size_t arity = Arity;

                    constexpr static const std::size_t value_bits = node_type::value_bits;
                    typedef typename node_type::value_type value_type;

                    merkle_proof_impl() : _li(0) {};

                    merkle_proof_impl(const merkle_tree<hash_type, arity> &tree, const std::size_t leaf_idx) {
                        _root = tree.root();
                        _path.resize(tree.row_count() - 1);
                        _li = leaf_idx;

                        typename std::vector<layer_type>::iterator v_itr = _path.begin();
                        std::size_t cur_leaf = leaf_idx;
                        std::size_t row_len = tree.leaves();
                        std::size_t row_begin_idx = 0;
                        while (cur_leaf != tree.size() - 1) {    // while it's not _root
                            std::size_t cur_leaf_pos = cur_leaf % arity;
                            std::size_t cur_leaf_arity_pos = (cur_leaf - row_begin_idx) / arity;
                            std::size_t begin_this_arity = cur_leaf - cur_leaf_pos;
                            typename layer_type::iterator a_itr = v_itr->begin();
                            for (size_t i = 0; i < cur_leaf_pos; ++i, ++begin_this_arity, ++a_itr) {
                                *a_itr = path_element_type(tree[begin_this_arity], i);
                            }
                            for (size_t i = cur_leaf_pos + 1; i < arity; ++i, ++begin_this_arity, ++a_itr) {
                                *a_itr = path_element_type(tree[begin_this_arity + 1], i);
                            }
                            v_itr++;
                            cur_leaf = row_len + row_begin_idx + cur_leaf_arity_pos;
                            row_begin_idx += row_len;
                            row_len /= arity;
                        }
                    }

                    template<typename Hashable>
                    bool validate(const Hashable &a) const {
                        value_type d = crypto3::hash<hash_type>(a);
                        for (auto &it : _path) {
                            accumulator_set<hash_type> acc;
                            size_t i = 0;
                            for (; (i < arity - 1) && i == it[i]._position; ++i) {
                                crypto3::hash<hash_type>(it[i]._hash.begin(), it[i]._hash.end(), acc);
                            }
                            crypto3::hash<hash_type>(d.begin(), d.end(), acc);
                            for (; i < arity - 1; ++i) {
                                crypto3::hash<hash_type>(it[i]._hash.begin(), it[i]._hash.end(), acc);
                            }
                            d = accumulators::extract::hash<hash_type>(acc);
                        }
                        return (d == _root);
                    }

                    std::size_t leaf_index() const {
                        return _li;
                    }

                    bool operator==(const merkle_proof_impl &rhs) const {
                        return _li == rhs._li && _root == rhs._root && _path == rhs._path;
                    }
                    bool operator!=(const merkle_proof_impl &rhs) const {
                        return !(rhs == *this);
                    }

                    struct path_element_type {
                        path_element_type(value_type x, size_t pos) : _hash(x), _position(pos) {
                        }
                        path_element_type() {
                        }

                        bool operator==(const path_element_type &rhs) const {
                            return _hash == rhs._hash && _position == rhs._position;
                        }
                        bool operator!=(const path_element_type &rhs) const {
                            return !(rhs == *this);
                        }

                        const value_type &hash() const {
                            return _hash;
                        }

                        std::size_t position() const {
                            return _position;
                        }

                        value_type _hash;
                        std::size_t _position;

                        template<typename, typename>
                        friend class nil::crypto3::marshalling::types::merkle_proof_marshalling;
                    };

                    typedef std::array<path_element_type, Arity - 1> layer_type;
                    typedef std::vector<layer_type> path_type;

                    const value_type &root() const {
                        return _root;
                    }

                    const path_type &path() const {
                        return _path;
                    }

                private:
                    std::size_t _li;
                    value_type _root;
                    path_type _path;

                    template<typename, typename, std::size_t>
                    friend class nil::crypto3::zk::components::merkle_proof;

                    template<typename, typename>
                    friend class nil::crypto3::marshalling::types::merkle_proof_marshalling;
                };
            }    // namespace detail

            template<typename T, std::size_t Arity>
            using merkle_proof =
                typename std::conditional<nil::crypto3::detail::is_hash<T>::value,
                                          detail::merkle_proof_impl<detail::merkle_tree_node<T>, Arity>,
                                          detail::merkle_proof_impl<T, Arity>>::type;
        }    // namespace containers
    }        // namespace crypto3
}    // namespace nil

#endif


// Copyright 2019-2020 Josh Pieper, jjp@pobox.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 "moteus/tool/calibrate.h"

#include <algorithm>
#include <cmath>
#include <deque>

#include <fmt/format.h>

#include <Eigen/Core>

#include <boost/algorithm/string.hpp>

#include "mjlib/base/system_error.h"

namespace moteus {
namespace tool {

namespace {

constexpr double kPi = 3.14159265358979323846;

struct Entry {
  int phase = 0;
  int encoder = 0;
};

struct File {
  std::vector<Entry> phase_up;
  std::vector<Entry> phase_down;
};

File ParseFile(const std::vector<std::string>& lines_in) {
  std::deque<std::string> lines;
  std::copy(lines_in.begin(), lines_in.end(), std::back_inserter(lines));
  if (lines.front() != "CAL start") {
    mjlib::base::system_error::throw_if(
        true, "Does not start with magic line: " + lines.front());
  }

  lines.pop_front();

  File result;

  while (!lines.empty()) {
    const auto line = lines.front();
    lines.pop_front();

    if (line == "CAL done") {
      return result;
    }

    std::vector<std::string> fields;
    boost::split(fields, line, boost::is_any_of(" "));

    if (fields.size() != 3) {
      mjlib::base::system_error::throw_if(
          true, "malformed line: " + line);
    }

    const auto phase = std::stoi(fields[1]);
    const auto encoder = std::stoi(fields[2]);
    const auto entry = Entry{phase, encoder};

    if (fields[0] == "1") {
      result.phase_up.push_back(entry);
    } else if (fields[0] == "2") {
      result.phase_down.push_back(entry);
    } else {
      mjlib::base::system_error::throw_if(
          true, "unknown phase code: " + line);
    }
  }

  mjlib::base::system_error::throw_if(
      true, "Does not end with magic line");
  return {};
}

int WrapInt16(int value) {
  if (value > 32767) {
    return value - 65536;
  } else if (value < -32768) {
    return 65536 + value;
  }
  return value;
}

double WrapNegPiToPi(double value) {
  if (value >= -kPi && value <= kPi) { return value; }
  if (value > 0.0) {
    return std::fmod(value + kPi, 2 * kPi) - kPi;
  } else {
    return std::fmod(value - kPi, 2 * kPi) + kPi;
  }
}

Eigen::VectorXd Unwrap(const Eigen::VectorXd& value) {
  Eigen::VectorXd result(value.size());
  for (int i = 0; i < value.size(); i++) {
    if (i == 0) {
      result(i) = value(i);
    } else {
      result(i) = (result(i - 1) + WrapNegPiToPi(value(i) - result(i - 1)));
    }
  }

  return result;
}

Eigen::VectorXd Linspace(double start, double end, size_t count) {
  Eigen::VectorXd result(count);
  for (size_t i = 0; i < count; i++) {
    result(i) = (
        (i + 1) == count ?
        end :
        i * (end - start) / (count - 1) + start);
  }
  return result;
}

Eigen::VectorXd Range(double start, double end, size_t count) {
  Eigen::VectorXd result(count);
  for (size_t i = 0; i < count; i++) {
    result(i) = (i * (end - start) / count + start);
  }
  return result;
}

Eigen::VectorXd Interpolate(const Eigen::VectorXd& sample_points,
                            const Eigen::VectorXd& x,
                            const Eigen::VectorXd& y) {
  BOOST_ASSERT(x.size() > 1 && x.size() == y.size());

  int xindex = 0;

  Eigen::VectorXd result(sample_points.size());
  for (int i = 0; i < sample_points.size(); i++) {
    const auto point = sample_points[i];
    result(i) = [&]() {
      if (point < x[xindex]) {
        // Whoops?  This can only legimately happen at the very beginning.
        return y[xindex];
      }
      while ((xindex + 2) < x.size() &&
             point >= x[xindex + 1]) {
        xindex++;
      }

      if (point > x[xindex + 1]) {
        // We're past the end.
        return y[xindex + 1];
      }

      // Linearly interpolate.
      const double length = x[xindex + 1] - x[xindex];
      if (length == 0.0) {
        return y[xindex + 1];
      }
      double ratio = (point - x[xindex]) / length;
      return (y[xindex + 1] - y[xindex]) * ratio + y[xindex];
    }();
  }

  return result;
}

Eigen::VectorXd WindowAverage(const Eigen::VectorXd& values, int window_size) {
  Eigen::VectorXd result(values.size());
  for (int i = 0; i < values.size(); i++) {
    int start = std::max<int>(0, i - window_size / 2);
    int end = std::min<int>(values.size(), i + window_size / 2);
    Eigen::VectorXd errs(end - start);
    for (int i = start; i < end; i++) {
      errs[i - start] = WrapNegPiToPi(values[i] - values[start]);
    }
    result[i] = values[start] + errs.mean();
  }

  return result;
}

}

CalibrationResult Calibrate(const std::vector<std::string>& lines) {
  auto file = ParseFile(lines);

  if (file.phase_up.empty() ||
      file.phase_down.empty()) {
    mjlib::base::system_error::throw_if(
        true, "one or more phases was empty");
  }

  double total_delta = 0;
  for (size_t i = 0; i < file.phase_up.size() - 1; i++) {
    total_delta += WrapInt16(file.phase_up[i + 1].encoder -
                             file.phase_up[i].encoder);
  }

  if (std::abs(std::abs(total_delta) - 65536) > 5000) {
    mjlib::base::system_error::throw_if(
        true, "phase_up did not traverse appropriate encoder distance");
  }

  CalibrationResult result;

  // Figure out inversion.
  if (total_delta < 0) {
    result.invert = true;
    for (auto& item : file.phase_up) { item.encoder = 65535 - item.encoder; }
    for (auto& item : file.phase_down) { item.encoder = 65535 - item.encoder; }
    total_delta *= -1;
  }

  // Next, figure out the number of poles.  We compare the total
  // encoder delta to the total phase delta.
  double total_phase = 0;
  for (size_t i = 0; i < file.phase_up.size() - 1; i++) {
    total_phase += WrapInt16(file.phase_up[i + 1].phase -
                             file.phase_up[i].phase);
  }

  const double ratio = total_phase / total_delta;
  const double remainder = std::abs(std::round(ratio) - ratio);

  const double kMaxRemainderError = 0.1;
  if (remainder > kMaxRemainderError) {
    mjlib::base::system_error::throw_if(
        true, fmt::format(
            "encoder not an integral multiple of phase, {} > {}",
            remainder, kMaxRemainderError));
  }

  result.total_phase = total_phase;
  result.total_delta = total_delta;
  result.ratio = ratio;

  result.poles = static_cast<int>(std::round(ratio) * 2);

  // Now we need to figure out the phase offset at select points.  We
  // interpolate and average the phase up and phase down sections.
  std::vector<Entry> phase_up_by_encoder = file.phase_up;
  const auto encoder_sort = [](const auto& lhs, const auto& rhs) {
    return lhs.encoder < rhs.encoder;
  };
  std::sort(phase_up_by_encoder.begin(), phase_up_by_encoder.end(),
            encoder_sort);
  std::vector<Entry> phase_down_by_encoder = file.phase_down;
  std::sort(phase_down_by_encoder.begin(), phase_down_by_encoder.end(),
            encoder_sort);

  double offset = phase_down_by_encoder[0].phase - phase_up_by_encoder[0].phase;
  if (std::abs(offset) > 32767) {
    // We need to shift it so that they start from the same place.
    auto change = static_cast<int>(-65536 * std::round(offset / 65536.0));
    for (auto& item : phase_down_by_encoder) { item.phase += change; }
  }

  Eigen::VectorXd phase_up_encoder(phase_up_by_encoder.size());
  for (size_t i = 0; i < phase_up_by_encoder.size(); i++) {
    phase_up_encoder[i] = phase_up_by_encoder[i].encoder;
  }

  Eigen::VectorXd phase_up_phase(phase_up_by_encoder.size());
  for (size_t i = 0; i < phase_up_by_encoder.size(); i++) {
    phase_up_phase[i] =
        2.0 * kPi / 65536.0 * phase_up_by_encoder[i].phase;
  }
  phase_up_phase = Unwrap(phase_up_phase);

  Eigen::VectorXd phase_down_encoder(phase_down_by_encoder.size());
  for (size_t i = 0; i < phase_down_by_encoder.size(); i++) {
    phase_down_encoder[i] = phase_down_by_encoder[i].encoder;
  }

  Eigen::VectorXd phase_down_phase(phase_down_by_encoder.size());
  for (size_t i = 0; i < phase_down_by_encoder.size(); i++) {
    phase_down_phase[i] =
        2.0 * kPi / 65536.0 * phase_down_by_encoder[i].phase;
  }
  phase_down_phase = Unwrap(phase_down_phase);

  Eigen::VectorXd xpos = Linspace(0, 65535.0, 10000);

  const Eigen::VectorXd pu_interp =
      Interpolate(xpos, phase_up_encoder, phase_up_phase);
  const Eigen::VectorXd pd_interp =
      Interpolate(xpos, phase_down_encoder, phase_down_phase);
  const Eigen::VectorXd avg_interp = 0.5 * (pu_interp + pd_interp);

  const Eigen::VectorXd expected =
      (2.0 * kPi / 65536) * (result.poles / 2) * xpos;

  Eigen::VectorXd err(expected.size());
  for (int i = 0; i < expected.size(); i++) {
    err[i] = WrapNegPiToPi(avg_interp[i] - expected[i]);
  }

  // Make the error seem reasonable, so unwrap if we happen to span
  // the pi boundary.
  if ((err.maxCoeff() - err.minCoeff()) > 1.5 * kPi) {
    for (int i = 0; i < err.size(); i++) {
      err[i] = (err[i] > 0) ? err[i] : (err[i] + 2 * kPi);
    }
  }

  auto avg_window = static_cast<int>(err.size() / result.poles);
  const Eigen::VectorXd avg_err = WindowAverage(err, avg_window);

  const Eigen::VectorXd offset_x = Range(0, 65536, 64);
  const Eigen::VectorXd offsets = Interpolate(offset_x, xpos, avg_err);

  for (int i = 0; i < offsets.size(); i++) {
    result.offset.push_back(offsets[i]);
  }

  return result;
}

}
}