#include "autoopt/optimizer.hpp"

#include <Eigen/src/Core/Matrix.h>
#include <fftw3.h>

#include <algorithm>
#include <cstddef>
#include <iterator>
#include <print>
#include <stdexcept>
#include <tuple>
#include <variant>

#include "autoopt/ellipse_new.hpp"
#include "autoopt/parabola_new.hpp"

namespace autoopt {

void optimizer::import_file(std::filesystem::path file_path) {
  _data = data{};
  _data->_raw = read_data(file_path);
  _data->lb = *_data->_raw._xs.begin();
  _data->ub = *(_data->_raw._xs.end() - 1);
  plot_input_cb();
  init_bounds_cb();
  update_test_data();
  update_table_cb();
}

void optimizer::set_lower_bound(double lb) {
  if (_data) {
    _data->lb = lb;
  }
  update_test_data();
}

void optimizer::set_upper_bound(double ub) {
  if (_data) {
    _data->ub = ub;
  }
  update_test_data();
}

void optimizer::set_curve_type(int ct) {
  if (ct == _curve.index()) return;
  switch (ct) {
    case CurveType::Parabola:
      _curve = parabola_new{};
      break;
    case CurveType::Ellipse:
      _curve = ellipse_new{};
      break;
  }
  update_table_cb();
}

void optimizer::set_initial(size_t i, double value) {
  _fit_result.reset();
  std::visit(
      [&](auto& curve) -> void {
        if (i >= curve.initial.size()) {
          throw std::runtime_error("index out of bounds");
        }
        curve.initial(i) = value;
      },
      _curve);
  plot_slope_cb();
}

void optimizer::set_delta(size_t i, double value) {
  _fit_result.reset();
  std::visit(
      [&](auto& curve) -> void {
        if (i >= curve.delta.size()) {
          throw std::runtime_error("index out of bounds");
        }
        curve.delta(i) = value;
      },
      _curve);
  plot_slope_cb();
}

void optimizer::flip_curve() {
  std::visit([](auto& curve) { curve.flip = !curve.flip; }, _curve);
  plot_slope_cb();
}

void optimizer::calc_fit() {
  _fit_result =
      std::visit([&](auto& curve) { return curve.fit(*_test_data); }, _curve);
  plot_slope_cb();
  for (auto param : _fit_result->first) {
    std::println("param: {}", param);
  }
  plot_error_cb();
}

void optimizer::set_plot_input_cb(
    std::function<void(Eigen::VectorXd xs, std::vector<Eigen::VectorXd> yss)>
        cb) {
  _plot_input_cb = cb;
}

void optimizer::set_init_bounds_cb(
    std::function<void(double lb, double ub, double step)> cb) {
  _init_bounds_cb = cb;
}

void optimizer::set_plot_slope_cb(
    std::function<void(Eigen::VectorXd xs, Eigen::VectorXd ys,
                       Eigen::VectorXd ys_curve)>
        cb) {
  _plot_slope_cb = cb;
}

void optimizer::set_update_table_cb(
    std::function<
        void(std::vector<std::tuple<std::string, double, double>> entries)>
        cb) {
  _param_table_cb = cb;
}

void optimizer::set_plot_error_cb(
    std::function<void(Eigen::VectorXd xs, Eigen::VectorXd ys)> cb) {
  _plot_error_cb = cb;
}

void optimizer::update_test_data() {
  _fit_result.reset();
  if (!_data) {
    _test_data.reset();
    return;
  }
  auto lb_it = std::lower_bound(_data->_raw._xs.begin(), _data->_raw._xs.end(),
                                _data->lb);
  auto ub_it = std::upper_bound(_data->_raw._xs.begin(), _data->_raw._xs.end(),
                                _data->ub);
  ssize_t n_samples = std::distance(lb_it, ub_it);
  if (n_samples < 1) {
    _test_data.reset();
    plot_input_cb();
    return;
  }

  size_t start_index = std::distance(_data->_raw._xs.begin(), lb_it);

  Eigen::VectorXd xs(n_samples);
  Eigen::VectorXd ys = Eigen::VectorXd::Zero(n_samples);

  std::copy(lb_it, ub_it, xs.begin());
  for (size_t i = 0; i < n_samples; i++) {
    for (auto r : _data->_raw._yss) {
      ys(i) += r(start_index + i);
    }
    ys(i) /= _data->_raw._yss.size();
  }
  double center = (*lb_it + *std::prev(ub_it)) / 2;
  _test_data.emplace(xs, ys, center);
  plot_slope_cb();
}

void optimizer::init_bounds_cb() {
  if (_init_bounds_cb) {
    _init_bounds_cb(_data->_raw._xs(0), _data->_raw._xs.tail(1)(0),
                    _data->_raw._xs(1) - _data->_raw._xs(0));
  }
}

void optimizer::plot_input_cb() {
  if (_plot_input_cb) {
    _plot_input_cb(_data->_raw._xs, _data->_raw._yss);
  }
}

void optimizer::plot_slope_cb() {
  if (_plot_slope_cb) {
    auto ys_curve =
        std::visit([&](auto curve) { return curve.plot(*_test_data); }, _curve);
    _plot_slope_cb(_test_data->_xs, _test_data->_ys, ys_curve);
  }
}

void optimizer::update_table_cb() {
  std::vector<std::tuple<std::string, double, double>> res;
  std::visit(
      [&]<typename T>(T& curve) {
        for (size_t i = 0; i < curve.param_names.size(); i++) {
          res.emplace_back(curve.param_names[i], curve.initial[i],
                           curve.delta[i]);
        }
      },
      _curve);
  _param_table_cb(res);
}

void optimizer::plot_error_cb() {
  if (_plot_error_cb) {
    Eigen::VectorXd ys_curve = std::visit(
        [&](auto curve) { return curve.plot(*_test_data, _fit_result->first); },
        _curve);
    Eigen::VectorXd error = _test_data->_ys - ys_curve;
    _plot_error_cb(_test_data->_xs, error);

    

    fftw_complex* out = fftw_alloc_complex(error.size());
    double* input = &error(0);
    fftw_plan plan = fftw_plan_dft_r2c_1d(error.size(), input, out, FFTW_FORWARD);
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    Eigen::VectorXd out_abs(error.size());
    for (size_t i = 0; i < error.size(); i++) {
      out_abs(i) = out[i][0] * out[i][0] + out[i][1]*out[i][1];
    }
    fftw_free(out);
    std::println("begin fftw");
    for (auto v : out_abs) {
      std::println("  {}", v);
    }
    std::println("end fftw");
  }
}

}  // namespace autoopt