#include <ai.h>
#include <string>
#include <algorithm>
#include <cmath>
using namespace std;

AI_SHADER_NODE_EXPORT_METHODS(apCrossBlend)
enum apCrossBlendParams
{
    p_surface,
    p_enable,
    p_ray_max_length,
    p_ray_nsamples,
    p_ray_min_dist,
    p_ray_max_dist,
    p_ray_max_depth,
    p_trace_sets,
    p_verbose,
};

struct ShaderData
{
    AtSampler *sampler;
    bool enable;
    float ray_max_length;
    int ray_nsamples;
    float ray_min_dist;
    float ray_max_dist;
    int ray_max_depth;
    AtString trace_sets;
    bool verbose;
    const int seed = 3158863998;
    const int ndim = 2;
};

struct SgCache
{
    SgCache(AtShaderGlobals *sg)
        : sg(sg),
          N(sg->N),
          Ng(sg->Ng),
          Ns(sg->Ns)
    {
    }
    void restore()
    {
        sg->N = N;
        sg->Ng = Ng;
        sg->Ns = Ns;
    }
    AtShaderGlobals *sg;
    AtVector N;
    AtVector Ng;
    AtVector Ns;
};

inline AtVector uniformSampleHemisphere(float u1, float u2)
{
    float y = u1;
    float r = sqrtf(max(0.f, 1.f - y * y));
    float phi = 2 * AI_PI * u2;
    float x = r * cosf(phi);
    float z = r * sinf(phi);
    return AtVector(x, y, z);
}

inline void concentricSampleDisk(float u1, float u2, float &dx, float &dy)
{
    float r, theta;
    // Map uniform random numbers to $[-1,1]^2$
    float sx = 2 * u1 - 1;
    float sy = 2 * u2 - 1;

    // Map square to $(r,\theta)$

    // Handle degeneracy at the origin
    if (sx == 0.0 && sy == 0.0)
    {
        dx = 0.0;
        dy = 0.0;
        return;
    }
    if (sx >= -sy)
    {
        if (sx > sy)
        {
            // Handle first region of disk
            r = sx;
            if (sy > 0.0)
                theta = sy / r;
            else
                theta = 8.0f + sy / r;
        }
        else
        {
            // Handle second region of disk
            r = sy;
            theta = 2.0f - sx / r;
        }
    }
    else
    {
        if (sx <= sy)
        {
            // Handle third region of disk
            r = -sx;
            theta = 4.0f - sy / r;
        }
        else
        {
            // Handle fourth region of disk
            r = -sy;
            theta = 6.0f + sx / r;
        }
    }
    theta *= float(AI_PI) / 4.f;
    dx = r * cosf(theta);
    dy = r * sinf(theta);
}

node_parameters
{
    AiParameterClosure("surface");
    AiParameterBool("enable", true);
    AiParameterFlt("ray_max_length", 1.0);
    AiParameterInt("ray_nsamples", 64);
    AiParameterFlt("ray_min_dist", 0.001);
    AiParameterFlt("ray_max_dist", 1000);
    AiParameterInt("ray_max_depth", 10);
    AiParameterStr("trace_sets", "");
    AiParameterBool("verbose", true);
}

node_initialize
{
    ShaderData *data = (ShaderData *)AiMalloc(sizeof(ShaderData));
    const int seed = 3158863998;
    const int nsamples = AiNodeGetInt(node, "ray_nsamples");
    const int ndim = 2;
    data->sampler = AiSampler(seed, nsamples, ndim);
    AiNodeSetLocalData(node, data);
}

node_finish
{
    ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
    AiSamplerDestroy(data->sampler);
    AiFree(data);
}

node_update
{
    ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
    data->enable = AiNodeGetBool(node, "enable");
    data->ray_max_length = AiNodeGetFlt(node, "ray_max_length");
    data->ray_nsamples = AiNodeGetInt(node, "ray_nsamples");
    data->ray_min_dist = AiNodeGetFlt(node, "ray_min_dist");
    data->ray_max_dist = AiNodeGetFlt(node, "ray_max_dist");
    data->ray_max_depth = AiNodeGetInt(node, "ray_max_depth");
    data->trace_sets = AiNodeGetStr(node, "trace_sets");
    data->verbose = AiNodeGetBool(node, "verbose");
    if (data->sampler)
        AiSamplerDestroy(data->sampler);
    data->sampler = AiSampler(data->seed, data->ray_nsamples, data->ndim);
}

shader_evaluate
{
    ShaderData *data = (ShaderData *)AiNodeGetLocalData(node);
    AtClosureList surface = AiShaderEvalParamClosure(p_surface);
    if (!data->enable)
    {
        sg->out.CLOSURE() = surface;
        return;
    }

    if (sg->Rt & AI_RAY_SHADOW)
        return;

    AtVector n = AiV3Normalize(sg->N);
    AtVector tangent = sg->dPdu;
    AtVector bitangent = sg->dPdv;

    SgCache sg_cache(sg);
    string shading_object_name = AiNodeGetName(sg->Op);
    AtRay ray;
    AtVector rp = sg->P;
    // AtVector n_sum = {0.f, 0.f, 0.f};

    float samples[2];
    AtSamplerIterator *iterator = AiSamplerIterator(data->sampler, sg);
    AtVector sample_N, sample_Nf, sample_Ns;
    AtClosureList sample_shader;
    float ray_length = AI_BIG;

    AiShaderGlobalsSetTraceSet(sg, data->trace_sets, false);

    int sample_count = AiSamplerGetSampleCount(iterator);
    int nsmp = 0;
    while (AiSamplerGetSample(iterator, samples))
    {
        // float u, v;
        // concentricSampleDisk(samples[0], samples[1], u, v);
        // AtVector sample_dir = u * tangent + v * bitangent;
        // sample_dir = AiV3Normalize(sample_dir);
        AtVector sample_dir = AiV3Normalize(uniformSampleHemisphere(samples[0], samples[1]));
        sample_dir += n;
        sample_dir *= -1;

        ray = AiMakeRay(AI_RAY_ALL_DIFFUSE, rp, &sample_dir, data->ray_max_dist, sg);
        float sample_length;

        AtShaderGlobals *sample_point = AiShaderGlobals();

        if (!AiTraceProbe(ray, sample_point))
        {
            AiShaderGlobalsDestroy(sample_point);
            nsmp++;
            continue;
        }

        string sample_object_name = AiNodeGetName(sample_point->Op);
        if (sample_object_name == shading_object_name)
        {
            AiShaderGlobalsDestroy(sample_point);
            nsmp++;
            continue;
        }

        sample_length = AiV3Length(sample_point->P - sg->P);
        if (data->verbose)
        {
            AiMsgDebug("## id  : %d", nsmp);
            AiMsgDebug("## rp  : %f, %f, %f", rp.x, rp.y, rp.z);
            AiMsgDebug("## dir : %f, %f, %f", sample_dir.x, sample_dir.y, sample_dir.z);
            AiMsgDebug("## smp : %f, %f, %f", sample_point->P.x, sample_point->P.y, sample_point->P.z);
            AiMsgDebug("## len : %f", sample_length);
        }
        if (sample_length < ray_length)
        {
            ray_length = min(ray_length, sample_length);

            sample_shader = sample_point->out.CLOSURE();
            sample_N = sample_point->N;
            sample_Nf = sample_point->Nf;
            sample_Ns = sample_point->Ns;
            nsmp++;
        }
    }

    AiShaderGlobalsUnsetTraceSet(sg);

    if (ray_length > data->ray_max_length)
    {
        sg->out.CLOSURE() = AiShaderEvalParamClosure(p_surface);
        return;
    }

    float rv = ray_length / data->ray_max_length;

    // Interploate normal
    sg->N = ((1.0 - rv) * sample_N) + (rv * sg->N);
    sg->Nf = ((1.0 - rv) * sample_Nf) + (rv * sg->Nf);
    sg->Ns = ((1.0 - rv) * sample_Ns) + (rv * sg->Ns);

    // Interpolate shading
    // surface.add(AiClosureTransparent(sg, AtRGB(rv)));
    // sample_shader.add(AiClosureTransparent(sg, AI_RGB_WHITE - AtRGB(rv)));

    AtClosureList result;
    if (rv < 1)
    {
        if (data->verbose)
            AiMsgInfo("rv : %f", rv);
        // result.add(sample_shader);
        //  sample_shader.add(AiClosureTransparent(sg, AI_RGB_WHITE - AtRGB(0.5)));
        sg->out.CLOSURE() = AiClosureEmission(sg, AI_RGB_RED);
        // return;
    }
    else
    {
        sg->out.CLOSURE() = AiClosureEmission(sg, AI_RGB_GREEN);

        // sg->out.CLOSURE() = surface;
    }
    // result.add(surface);
    // result.add(sample_shader);
    // sg->out.CLOSURE() = result;

    sg_cache.restore();
}

node_loader
{
    if (i > 0)
        return false;

    node->methods = apCrossBlend;
    node->output_type = AI_TYPE_CLOSURE;
    node->name = "apCrossBlend";
    node->node_type = AI_NODE_SHADER;
    strcpy(node->version, AI_VERSION);
    return true;
}
