#include <windows.h>
#include "avisynth.h"

static const AVS_Linkage* AVS_linkage = 0;

class Footprint {

    int index;
    BYTE* map;
    DWORD* xy;
    int width;
    int height;
    const BYTE* altp;
    BYTE* dstp;
    int alt_pitch;
    int dst_pitch;

public:
    Footprint(int w, int h, const BYTE* a, BYTE* d, int ap, int dp, IScriptEnvironment* env)
        : width(w), height(h), altp(a), dstp(d), alt_pitch(ap), dst_pitch(dp), index(-1) {
        map = (BYTE*)calloc(width * height, 1);
        xy = (DWORD*)malloc(width * height * sizeof(DWORD));
        if (!map || !xy) {
            env->ThrowError("failed to allocate Footprint.");
        }
    }

    ~Footprint() {
        free(map);
        free(xy);
        map = 0;
        xy = 0;
    }

    int get_width() {
        return width;
    }
    int get_height() {
        return height;
    }
    void push(int x, int y) {
        map[x + y * width] = 0xFF;
        dstp[x + y * dst_pitch] = altp[x + y * alt_pitch];
        xy[++index] = ((WORD)x << 16) | (WORD)y;
    }

    void pop(int* x, int* y) {
        *x = xy[index] >> 16;
        *y = xy[index--] & 0x0000FFFF;
    }

    bool is_empty() {
        return index == -1;
    }

    int passed(int x, int y) {
        return (int)map[x + y * width];
    }
};

class Hysteresis : public GenericVideoFilter {
    PClip child2;
    bool chroma;
    int num_planes;

public:
    Hysteresis(PClip c, PClip c2, bool chroma, IScriptEnvironment* env);
    ~Hysteresis() {}
    PVideoFrame __stdcall GetFrame(int n, IScriptEnvironment* env);
};

Hysteresis::Hysteresis(PClip c, PClip c2, bool ch, IScriptEnvironment* env)
    : GenericVideoFilter(c), child2(c2), chroma(ch)
{
    if (!vi.IsPlanar()) {
        env->ThrowError("not planar format");
    }
    const VideoInfo& vi2 = child2->GetVideoInfo();
    if (vi.IsSameColorspace(vi2) == false) {
        env->ThrowError("not the same csp");
    }
    if (vi.width != vi2.width || vi.height != vi2.height) {
        env->ThrowError("not the same resolution");
    }
    num_planes = vi.IsY8() ? 1 : 3;
}

static void proc(Footprint* fp, int x, int y, const BYTE* altp, int alt_pitch)
{
    fp->push(x, y);

    int posx, posy;
    while (!fp->is_empty()) {
        fp->pop(&posx, &posy);

        int x_min = posx > 0 ? posx - 1 : 0;
        int x_max = posx < fp->get_width() - 1 ? posx + 1 : posx;
        int y_min = posy > 0 ? posy - 1 : 0;
        int y_max = posy < fp->get_height() - 1 ? posy + 1 : posy;

        for (int yy = y_min; yy <= y_max; yy++) {
            for (int xx = x_min; xx <= x_max; xx++) {
                if (altp[xx + yy * alt_pitch] && !fp->passed(xx, yy)) {
                    fp->push(xx, yy);
                }
            }
        }
    }
}

PVideoFrame __stdcall Hysteresis::GetFrame(int n, IScriptEnvironment* env)
{
    const int planes[] = {PLANAR_Y, PLANAR_U, PLANAR_V};

    PVideoFrame base = child->GetFrame(n, env);
    PVideoFrame alt = child2->GetFrame(n, env);
    PVideoFrame dst = env->NewVideoFrame(vi);

    for (int i = 0; i < num_planes; i++) {
        if (i > 0 && !chroma) {
            return dst;
        }
        int p = planes[i];
        int width = base->GetRowSize(p);
        int height = base->GetHeight(p);
        int base_pitch = base->GetPitch(p);
        int alt_pitch = alt->GetPitch(p);
        int dst_pitch = dst->GetPitch(p);
        const BYTE* basep = base->GetReadPtr(p);
        const BYTE* altp = alt->GetReadPtr(p);
        BYTE* dstp = dst->GetWritePtr(p);

        Footprint *fp = new Footprint(width, height, altp, dstp, alt_pitch, dst_pitch, env);
        memset(dstp, 0, dst_pitch * height);

        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                if (basep[x + y * base_pitch] && altp[x + y * alt_pitch] &&
                    !fp->passed(x, y)) {
                    proc(fp, x, y, altp, alt_pitch);
                }
            }
        }
        delete fp;
    }
    return dst;
}

AVSValue __cdecl
create_hysteresis(AVSValue args, void* user_data, IScriptEnvironment* env)
{
    return new Hysteresis(args[0].AsClip(), args[1].AsClip(),
                           args[2].AsBool(false), env);
}

extern "C" __declspec(dllexport) const char* __stdcall
AvisynthPluginInit3(IScriptEnvironment* env, const AVS_Linkage* const vectors)
{
    AVS_linkage = vectors;
    env->AddFunction("Hysteresis", "cc[chroma]b", create_hysteresis, 0);
    return "hysteresis mask filter";
}