#include "stb_ds.h"
#include <float.h>
#include "HandmadeMath.h"

/* -------------------------------------------------------------------------
   Cubic Spline Basis Matrices
   ------------------------------------------------------------------------- */
static const HMM_Mat4 cubic_hermite_m = {
    2,  -2,  1,  1,
   -3,   3, -2, -1,
    0,   0,  1,  0,
    1,   0,  0,  0
};

static const HMM_Mat4 cubic_hermite_dm = {
    0,   0,  0,  0,
    6,  -6,  3,  3,
   -6,   6, -4, -2,
    0,   0,  1,  0
};

static const HMM_Mat4 cubic_hermite_ddm = {
    0,   0,  0,  0,
    0,   0,  0,  0,
   12, -12,  6,  6,
   -6,   6, -4, -2
};

static const HMM_Mat4 cubic_hermite_dddm = {
    0,   0,  0,  0,
    0,   0,  0,  0,
    0,   0,  0,  0,
   12, -12,  6,  6
};

static const HMM_Mat4 b_spline_m = {
  -1.0f/6,  3.0f/6, -3.0f/6,  1.0f,
   3.0f/6, -6.0f/6,  3.0f/6,  0.0f,
  -3.0f/6,  0.0f,    3.0f/6,  0.0f,
   1.0f/6,  4.0f/6,  1.0f/6,  0.0f
};

static const HMM_Mat4 b_spline_dm = {
   0,     0,     0,    0,
  -3.0f/6, 9.0f/6, -9.0f/6, 3.0f,
   6.0f/6, -12.0f/6, 6.0f/6, 0.0f,
  -3.0f/6,  0.0f,   3.0f/6, 0.0f
};

static const HMM_Mat4 b_spline_ddm = {
   0,    0,    0,    0,
   0,    0,    0,    0,
  -6.0f/6, 18.0f/6, -18.0f/6, 6.0f,
   6.0f/6, -12.0f/6,  6.0f/6, 0.0f
};

static const HMM_Mat4 b_spline_dddm = {
   0,    0,    0,    0,
   0,    0,    0,    0,
   0,    0,    0,    0,
  -6.0f/6, 18.0f/6, -18.0f/6, 6.0f
};

static const HMM_Mat4 bezier_m = {
   -1,  3, -3,  1,
    3, -6,  3,  0,
   -3,  3,  0,  0,
    1,  0,  0,  0
};

static const HMM_Mat4 bezier_dm = {
    0,  0,  0,  0,
   -3,  9, -9,  3,
    6, -12,  6,  0,
   -3,  3,  0,  0
};

static const HMM_Mat4 bezier_ddm = {
    0,  0,   0,  0,
    0,  0,   0,  0,
   -6, 18, -18,  6,
    6, -12,  6,  0
};

static const HMM_Mat4 bezier_dddm = {
    0,  0,  0,  0,
    0,  0,  0,  0,
    0,  0,  0,  0,
   -6, 18, -18,  6
};

/* Catmull–Rom (with tension = 0.5 by default) */
#define CAT_S 0.5f

static const HMM_Mat4 catmull_rom_m = {
  -CAT_S,       2-CAT_S,      CAT_S-2,  CAT_S,
   2*CAT_S,     CAT_S-3,      3-2*CAT_S, -CAT_S,
  -CAT_S,       0,            CAT_S,     0,
   0,           1,            0,         0
};

static const HMM_Mat4 catmull_rom_dm = {
    0,          0,          0,          0,
  -3*CAT_S,   9*CAT_S,    -9*CAT_S,    3*CAT_S,
   4*CAT_S,  -10*CAT_S,    8*CAT_S,   -2*CAT_S,
  -CAT_S,      0,          CAT_S,      0
};

static const HMM_Mat4 catmull_rom_ddm = {
    0,        0,         0,         0,
    0,        0,         0,         0,
  -9*CAT_S, 18*CAT_S, -18*CAT_S,  6*CAT_S,
   4*CAT_S, -10*CAT_S, 8*CAT_S,  -2*CAT_S
};

static const HMM_Mat4 catmull_rom_dddm = {
    0,         0,         0,         0,
    0,         0,         0,         0,
    0,         0,         0,         0,
  -9*CAT_S, 18*CAT_S, -18*CAT_S,  6*CAT_S
};

/* -------------------------------------------------------------------------
   Core “C·T” multiplication:  [ t^3, t^2, t, 1 ] * C
   ------------------------------------------------------------------------- */
HMM_Vec4 spline_CT(HMM_Mat4 *C, float t)
{
    float t2 = t * t;
    float t3 = t2 * t;
    HMM_Vec4 T = { t3, t2, t, 1.0f };
    return HMM_MulM4V4(*C, T);
}

/* Construct the “geometry matrix” G from four 2D points, then multiply by B */
HMM_Mat4 make_C(const HMM_Vec2 *p, const HMM_Mat4 *B)
{
    HMM_Mat4 G = HMM_M4(); // Zeroed out
    // Only fill XY of each column; if you are storing 3D in HMM_Vec4, adapt as needed
    G.Columns[0].XY = p[0];
    G.Columns[1].XY = p[1];
    G.Columns[2].XY = p[2];
    G.Columns[3].XY = p[3];

    return HMM_MulM4(G, *B);
}

/* Evaluate a single-segment cubic spline at parameter d in [0,1].
   p must be 4 control points, m is the cubic basis matrix. */
HMM_Vec2 cubic_spline_d(HMM_Vec2 *p, HMM_Mat4 *m, float d)
{
    HMM_Mat4 C = make_C(p, m);
    HMM_Vec4 v4 = spline_CT(&C, d);
    return v4.XY;
}

/* -------------------------------------------------------------------------
   Convenience single-segment functions for each basis
   (pos / tan / curv / wig all require the appropriate matrix)
   Typically you pass p[0..3] as the 4 relevant control points.
   ------------------------------------------------------------------------- */
HMM_Vec2 cubic_hermite_pos(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&cubic_hermite_m, d); }
HMM_Vec2 cubic_hermite_tan(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&cubic_hermite_dm, d); }
HMM_Vec2 cubic_hermite_curv(HMM_Vec2 *p, float d){ return cubic_spline_d(p, (HMM_Mat4 *)&cubic_hermite_ddm, d); }
HMM_Vec2 cubic_hermite_wig(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&cubic_hermite_dddm, d); }

HMM_Vec2 b_spline_pos(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&b_spline_m, d); }
HMM_Vec2 b_spline_tan(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&b_spline_dm, d); }
HMM_Vec2 b_spline_curv(HMM_Vec2 *p, float d){ return cubic_spline_d(p, (HMM_Mat4 *)&b_spline_ddm, d); }
HMM_Vec2 b_spline_wig(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&b_spline_dddm, d); }

HMM_Vec2 bezier_pos(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&bezier_m, d); }
HMM_Vec2 bezier_tan(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&bezier_dm, d); }
HMM_Vec2 bezier_curv(HMM_Vec2 *p, float d){ return cubic_spline_d(p, (HMM_Mat4 *)&bezier_ddm, d); }
HMM_Vec2 bezier_wig(HMM_Vec2 *p, float d) { return cubic_spline_d(p, (HMM_Mat4 *)&bezier_dddm, d); }

/* -------------------------------------------------------------------------
   Multi-segment sampling (“spline_v2”) for uniform division
   ------------------------------------------------------------------------- */
HMM_Vec2 *spline_v2(HMM_Vec2 *p, HMM_Mat4 *m, int segs)
{
    // For a single 4-point segment, produce 'segs' points along [0..1).
    // If you want the final 1.0 also, you can do <= 1.0 in the loop, etc.
    HMM_Vec2 *ret = NULL;
    if (segs < 2) return NULL;

    HMM_Mat4 C = make_C(p, m);
    float s = 1.0f / (float)segs;
    for (int i = 0; i <= segs; i++)
    {
        float t = s * i;
        arrput(ret, spline_CT(&C, t).XY);
    }
    return ret;
}

/* -------------------------------------------------------------------------
   Adaptive subdivision by min segment length
   (spline2d_min_seg) – for a single 4-point segment
   ------------------------------------------------------------------------- */
HMM_Vec2 *spline2d_min_seg(float u0, float u1, float min_seg, HMM_Mat4 *C, HMM_Vec2 *ret)
{
    HMM_Vec2 a = spline_CT(C, u0).XY;
    HMM_Vec2 b = spline_CT(C, u1).XY;
    if (HMM_DistV2(a, b) > min_seg)
    {
        float umid = 0.5f * (u0 + u1);
        spline2d_min_seg(u0,  umid, min_seg, C, ret);
        spline2d_min_seg(umid, u1,   min_seg, C, ret);
    }
    else
    {
        // We push 'b' so that we don't double-push 'a'
        arrput(ret, b);
    }
    return ret;
}

/* Example: catmull_rom_min_seg -> subdiv for catmull–rom over one segment
   You would decide how to pick the 4 points from a,b,c,d, then run. */
HMM_Vec2 *catmull_rom_min_seg(HMM_Vec2 *a, HMM_Vec2 *b, HMM_Vec2 *c, HMM_Vec2 *d, float min_seg)
{
    HMM_Vec2 *ret = NULL;
    arrsetcap(ret, 1000);

    // Build the matrix for these four points
    HMM_Vec2 p[4] = { *a, *b, *c, *d };
    HMM_Mat4 C = make_C(p, &catmull_rom_m);

    // Always push the starting point (b in your original code was the second ctrl point, etc.)
    // But usually we want the first actual point in the segment:
    arrput(ret, cubic_spline_d(p, (HMM_Mat4*)&catmull_rom_m, 0.0f));

    // Actually subdiv
    spline2d_min_seg(0.0f, 1.0f, min_seg, &C, ret);

    return ret;
}

/* -------------------------------------------------------------------------
   Adaptive subdivision by “max angle” proxy
   (spline2d_min_angle_2) – for single 4-point segment
   ------------------------------------------------------------------------- */
HMM_Vec2 *spline2d_min_angle_2(float u0, float u1, float max_angle, HMM_Mat4 *C, HMM_Vec2 *arr)
{
    // Heuristic approach: sample midpoints, check “chord vs polyline” difference
    float ustep = (u1 - u0) / 4.0f;
    float um0 = u0 + ustep;
    float um1 = u0 + 2.0f * ustep;
    float um2 = u0 + 3.0f * ustep;

    HMM_Vec2 m0 = spline_CT(C, um0).XY;
    HMM_Vec2 m1 = spline_CT(C, um1).XY;
    HMM_Vec2 m2 = spline_CT(C, um2).XY;

    HMM_Vec2 a  = spline_CT(C, u0).XY;
    HMM_Vec2 b  = spline_CT(C, u1).XY;

    // Chord = distance from a to b
    float chord   = HMM_DistV2(a, b);
    // Polyline = a->m0->m1->m2->b
    float cdist   = HMM_DistV2(a,  m0)
                  + HMM_DistV2(m0, m1)
                  + HMM_DistV2(m1, m2)
                  + HMM_DistV2(m2, b);

    // If the difference is bigger than some threshold (max_angle),
    // subdivide. Otherwise, keep it.
    if ((cdist - chord) > max_angle)
    {
        arr = spline2d_min_angle_2(u0,  um1, max_angle, C, arr);
        arr = spline2d_min_angle_2(um1, u1,  max_angle, C, arr);
    }
    else
    {
        // We accept “b” as a new point
        arrput(arr, b);
    }
    return arr;
}

HMM_Vec2 *spline_min_angle(HMM_Vec2 *p, const HMM_Mat4 *B, float min_angle, HMM_Vec2 *arr)
{
    HMM_Mat4 C = make_C(p, B);
    // Subdivide from 0..1
    float u0 = 0.0f, u1 = 1.0f;
    // Usually we want to ensure the start point is in arr:
    HMM_Vec2 startPt = spline_CT(&C, u0).XY;
    if (arrlen(arr) == 0) {
        arrput(arr, startPt);
    }
    // Now subdiv for angle
    arr = spline2d_min_angle_2(u0, u1, min_angle, &C, arr);
    return arr;
}

/* Example: catmull_rom_ma_v2 – uses “min_angle” over multiple segments 
   Each 4 consecutive points is one segment. We do this for all segments. */
HMM_Vec2 *catmull_rom_ma_v2(HMM_Vec2 *cp, float ma)
{
    int n = arrlen(cp);
    if (n < 4) return NULL;

    HMM_Vec2 *ret = NULL;
    // Pre-allocate some capacity
    arrsetcap(ret, (n-3) * 8);

    // For convenience, let's always ensure we push the very first point:
    arrput(ret, cp[0]);

    // For each segment [i, i+1, i+2, i+3], adaptively sample
    // Then move i by 1 each time if you want Catmull–Rom in “overlapped” fashion
    for (int i = 0; i < n - 3; i++)
    {
        // p[i..i+3]
        ret = spline_min_angle(&cp[i], &catmull_rom_m, ma, ret);
    }
    return ret;
}

/* Example: do the same with Bezier in “cubic-bezier” style (control points in groups of 3 “handles”) */
HMM_Vec2 *bezier_cb_ma_v2(HMM_Vec2 *cp, float ma)
{
    int n = arrlen(cp);
    // Typically a Bezier “chain” would use control points in multiples of 3 + 1, etc.
    // E.g. p[0] is start, p[1..3] are control handles for first segment, then p[3..6] for second, etc.
    // Adjust logic to your liking.
    if (n < 4) return NULL;

    HMM_Vec2 *ret = NULL;
    arrsetcap(ret, (n/3) * 8);

    // First point
    arrput(ret, cp[0]);

    // For each cubic Bezier segment: i += 3
    for (int i = 0; i < n - 3; i += 3)
    {
        ret = spline_min_angle(&cp[i], &bezier_m, ma, ret);
    }
    return ret;
}

static HMM_Vec2 catmull_rom_query_internal(HMM_Vec2 *cp, float d, const HMM_Mat4 *M)
{
    int n = arrlen(cp);
    if (n < 4) return HMM_V2(0,0);

    // Number of segments:
    int seg_count = n - 3;
    // Scale d in [0..1] -> which segment?
    float segf = d * seg_count;
    int seg_idx = (int) floorf(segf);
    if (seg_idx >= seg_count) seg_idx = seg_count - 1;
    if (seg_idx < 0) seg_idx = 0;

    // Local parameter in [0..1]
    float u = segf - seg_idx;

    // The control points for that segment are cp[ seg_idx .. seg_idx+3 ]
    return cubic_spline_d(cp + seg_idx, (HMM_Mat4 *)M, u);
}

HMM_Vec2 catmull_rom_pos(HMM_Vec2 *cp, float d)
{
    return catmull_rom_query_internal(cp, d, &catmull_rom_m);
}
HMM_Vec2 catmull_rom_tan(HMM_Vec2 *cp, float d)
{
    return catmull_rom_query_internal(cp, d, &catmull_rom_dm);
}
HMM_Vec2 catmull_rom_curv(HMM_Vec2 *cp, float d)
{
    return catmull_rom_query_internal(cp, d, &catmull_rom_ddm);
}
HMM_Vec2 catmull_rom_wig(HMM_Vec2 *cp, float d)
{
    return catmull_rom_query_internal(cp, d, &catmull_rom_dddm);
}

/* -------------------------------------------------------------------------
   Approximate length of a single 4-point cubic spline segment by
   numeric integration (or sampling) of the velocity magnitude.
   “spline_seglen” below does a quick sampling approach. 
   ------------------------------------------------------------------------- */
float spline_seglen(float t0, float t1, int steps, HMM_Mat4 *Cd, HMM_Mat4 *C)
{
    // Simple uniform sampling of the tangent magnitude
    float total = 0.0f;
    float dt = (t1 - t0) / (float) steps;
    for (int i = 0; i < steps; i++)
    {
        float t = t0 + (i + 0.5f) * dt; // midpoint rule
        // derivative at t
        HMM_Vec2 vel = spline_CT(Cd, t).XY;
        float speed = HMM_LenV2(vel);
        total += speed * dt;
    }
    return total;
}

/* Summation of lengths across all Catmull–Rom segments. */
float catmull_rom_len(HMM_Vec2 *cp)
{
    int stepsPerSegment = 64;
    float len = 0.0f;
    int n = arrlen(cp);
    if (n < 4) return 0.0f;

    for (int i = 0; i < n - 3; i++)
    {
        // Build the position matrix & derivative matrix for this segment
        HMM_Mat4 C  = make_C(&cp[i], &catmull_rom_m);
        HMM_Mat4 Cd = make_C(&cp[i], &catmull_rom_dm);
        // integrate from 0..1
        len += spline_seglen(0.0f, 1.0f, stepsPerSegment, &Cd, &C);
    }
    return len;
}

HMM_Vec2 catmull_rom_closest(HMM_Vec2 *cp, HMM_Vec2 p)
{
    int n = arrlen(cp);
    if (n < 4) return p;

    float bestDist = FLT_MAX;
    HMM_Vec2 bestPt = p;

    int steps = 64; // more steps => more accurate
    for (int seg = 0; seg < n - 3; seg++)
    {
        // Build a single-segment matrix
        HMM_Vec2 segCP[4] = { cp[seg], cp[seg+1], cp[seg+2], cp[seg+3] };
        HMM_Mat4 C  = make_C(segCP, &catmull_rom_m);
        for (int i = 0; i <= steps; i++)
        {
            float t = (float)i / steps;
            HMM_Vec2 pt = spline_CT(&C, t).XY;
            float dist  = HMM_DistV2(p, pt);
            if (dist < bestDist)
            {
                bestDist = dist;
                bestPt   = pt;
            }
        }
    }

    return bestPt;
}