#include "viz/line_plot.h"
#include "viz/plot_static.h"
#include "implot.h"

namespace {

template <typename T>
void draw_line(const char* title, const T* xs, const T* ys, int count,
               float height, bool auto_x, double x_lo_in, double x_hi_in,
               bool auto_y, double y_lo_in, double y_hi_in) {
    if (count <= 0) return;

    double x_lo, x_hi;
    if (auto_x) {
        T x_min = xs[0], x_max = xs[0];
        for (int i = 1; i < count; i++) {
            if (xs[i] < x_min) x_min = xs[i];
            if (xs[i] > x_max) x_max = xs[i];
        }
        x_lo = static_cast<double>(x_min);
        x_hi = static_cast<double>(x_max);
    } else {
        x_lo = x_lo_in;
        x_hi = x_hi_in;
    }
    if (x_hi - x_lo < 1e-9) x_hi = x_lo + 1.0;

    double y_lo, y_hi;
    if (auto_y) {
        T y_min = ys[0], y_max = ys[0];
        for (int i = 1; i < count; i++) {
            if (ys[i] < y_min) y_min = ys[i];
            if (ys[i] > y_max) y_max = ys[i];
        }
        double dy = static_cast<double>(y_max) - static_cast<double>(y_min);
        if (dy < 1e-9) dy = 1.0;
        y_lo = static_cast<double>(y_min) - dy * 0.05;
        y_hi = static_cast<double>(y_max) + dy * 0.05;
    } else {
        y_lo = y_lo_in;
        y_hi = y_hi_in;
        if (y_hi - y_lo < 1e-9) y_hi = y_lo + 1.0;
    }

    const ImVec2 plot_size(-1.0f, height > 0.0f ? height : 200.0f);

    if (ImPlot::BeginPlot(title, plot_size, plot_static::kPlotFlags)) {
        ImPlot::SetupAxes(nullptr, nullptr, plot_static::kAxisFlags, plot_static::kAxisFlags);
        ImPlot::SetupAxisLimits(ImAxis_X1, x_lo, x_hi, ImPlotCond_Always);
        ImPlot::SetupAxisLimits(ImAxis_Y1, y_lo, y_hi, ImPlotCond_Always);
        ImPlot::PlotLine("##data", xs, ys, count);
        ImPlot::EndPlot();
    }
}

} // namespace

void line_plot(const char* title, const float* xs, const float* ys, int count, float height) {
    draw_line<float>(title, xs, ys, count, height,
                     /*auto_x=*/true, 0.0, 0.0,
                     /*auto_y=*/true, 0.0, 0.0);
}

void line_plot(const char* title, const double* xs, const double* ys, int count, float height) {
    draw_line<double>(title, xs, ys, count, height,
                      /*auto_x=*/true, 0.0, 0.0,
                      /*auto_y=*/true, 0.0, 0.0);
}

void line_plot(const char* title, const float* xs, const float* ys, int count,
               float y_min, float y_max, float height) {
    draw_line<float>(title, xs, ys, count, height,
                     /*auto_x=*/true, 0.0, 0.0,
                     /*auto_y=*/false, (double)y_min, (double)y_max);
}

void line_plot(const char* title, const double* xs, const double* ys, int count,
               double y_min, double y_max, float height) {
    draw_line<double>(title, xs, ys, count, height,
                      /*auto_x=*/true, 0.0, 0.0,
                      /*auto_y=*/false, y_min, y_max);
}

void line_plot(const char* title, const float* xs, const float* ys, int count,
               float x_min, float x_max, float y_min, float y_max, float height) {
    draw_line<float>(title, xs, ys, count, height,
                     /*auto_x=*/false, (double)x_min, (double)x_max,
                     /*auto_y=*/false, (double)y_min, (double)y_max);
}

void line_plot(const char* title, const double* xs, const double* ys, int count,
               double x_min, double x_max, double y_min, double y_max, float height) {
    draw_line<double>(title, xs, ys, count, height,
                      /*auto_x=*/false, x_min, x_max,
                      /*auto_y=*/false, y_min, y_max);
}