#pragma once
#include <application.h>
#include <sampling.hpp>
#include <random_sampler_wrapper.hpp>
#include <scenegraph/obj_loader.h>

#include <lights/ambient_light.h>
#include <lights/directional_light.h>
#include <lights/point_light.h>
#include <lights/quad_light.h>
#include <lights/spot_light.h>

using namespace embree;

inline Vec3fa reflect(const Vec3fa& V, const Vec3fa& N) { return 2.0f * dot(V, N) * N - V; }

inline Vec3fa face_forward(const Vec3fa& dir, const Vec3fa& _Ng) {
    const Vec3fa Ng = _Ng;
    return dot(dir, Ng) < 0.0f ? Ng : neg(Ng);
}

inline Light_SampleRes Lights_sample(const Light* self,
                                     const Sample& sp, /*! point to generate the sample for >*/
                                     const Vec2f s) /*! random numbers to generate the sample >*/
{
    LightType ty = self->type;
    switch (ty) {
        case LIGHT_AMBIENT: return AmbientLight_sample(self, sp, s);
        case LIGHT_POINT: return PointLight_sample(self, sp, s);
        case LIGHT_DIRECTIONAL: return DirectionalLight_sample(self, sp, s);
        case LIGHT_SPOT: return SpotLight_sample(self, sp, s);
        case LIGHT_QUAD: return QuadLight_sample(self, sp, s);
        default: {
            Light_SampleRes res;
            res.weight = Vec3fa(0, 0, 0);
            res.dir = Vec3fa(0, 0, 0);
            res.dist = 0;
            res.pdf = inf;
            return res;
        }
    }
}

inline Light_EvalRes Lights_eval(const Light* self,
                                 const Sample& sp,
                                 const Vec3fa& dir) {
    LightType ty = self->type;
    switch (ty) {
        case LIGHT_AMBIENT: return AmbientLight_eval(self, sp, dir);
        case LIGHT_POINT: return PointLight_eval(self, sp, dir);
        case LIGHT_DIRECTIONAL: return DirectionalLight_eval(self, sp, dir);
        case LIGHT_SPOT: return SpotLight_eval(self, sp, dir);
        case LIGHT_QUAD: return QuadLight_eval(self, sp, dir);
        default: {
            Light_EvalRes res;
            res.value = Vec3fa(0, 0, 0);
            res.dist = inf;
            res.pdf = 0.f;
            return res;
        }
    }
}

struct BRDF {
    float Ns; /*< specular exponent */
    float Ni; /*< optical density for the surface (index of refraction) */
    Vec3fa Ka; /*< ambient reflectivity */
    Vec3fa Kd; /*< diffuse reflectivity */
    Vec3fa Ks; /*< specular reflectivity */
    Vec3fa Kt; /*< transmission filter */
    float dummy[30];
};

////////////////////////////////////////////////////////////////////////////////
//                                Lambertian BRDF                             //
////////////////////////////////////////////////////////////////////////////////

struct Lambertian {
    Vec3fa R;
};

inline Vec3fa Lambertian__eval(const Lambertian* This,
                               const Vec3fa& wo, const Sample& dg, const Vec3fa& wi) {
    return This->R * (1.0f / (float) (float(M_PI))) * clamp(dot(wi, dg.Ns));
}

inline Vec3fa Lambertian__sample(const Lambertian* This,
                                 const Vec3fa& wo,
                                 const Sample& dg,
                                 Sample3f& wi,
                                 const Vec2f& s) {
    wi = cosineSampleHemisphere(s.x, s.y, dg.Ns);
    return Lambertian__eval(This, wo, dg, wi.v);
}

inline void Lambertian__Constructor(Lambertian* This, const Vec3fa& R) {
    This->R = R;
}

inline Lambertian make_Lambertian(const Vec3fa& R) {
    Lambertian v;
    Lambertian__Constructor(&v, R);
    return v;
}

////////////////////////////////////////////////////////////////////////////////
//                          Matte Material                                    //
////////////////////////////////////////////////////////////////////////////////

inline void MatteMaterial__preprocess(MatteMaterial* material, BRDF& brdf, const Vec3fa& wo, const Sample& sp) {
}

inline Vec3fa MatteMaterial__eval(MatteMaterial* This, const BRDF& brdf, const Vec3fa& wo, const Sample& sp,
                                  const Vec3fa& wi) {
    Lambertian lambertian = make_Lambertian(Vec3fa((Vec3fa) This->reflectance));
    return Lambertian__eval(&lambertian, wo, sp, wi);
}

inline Vec3fa MatteMaterial__sample(MatteMaterial* This, const BRDF& brdf, const Vec3fa& Lw, const Vec3fa& wo,
                                    const Sample& sp, Sample3f& wi_o, const Vec2f& s) {
    Lambertian lambertian = make_Lambertian(Vec3fa((Vec3fa) This->reflectance));
    return Lambertian__sample(&lambertian, wo, sp, wi_o, s);
}


////////////////////////////////////////////////////////////////////////////////
//                          OBJ Material                                      //
////////////////////////////////////////////////////////////////////////////////

inline void OBJMaterial__preprocess(OBJMaterial* material, BRDF& brdf, const Vec3fa& wo, const Sample& sp) {
    float d = material->d;
    // if (material->map_d) d *= getTextureTexel1f(material->map_d, dg.u, dg.v);
    brdf.Ka = Vec3fa(material->Ka);
    //if (material->map_Ka) { brdf.Ka *= material->map_Ka->get(dg.st); }
    brdf.Kd = d * Vec3fa(material->Kd);
    // if (material->map_Kd) brdf.Kd = brdf.Kd * getTextureTexel3f(material->map_Kd, dg.u, dg.v);
    brdf.Ks = d * Vec3fa(material->Ks);
    //if (material->map_Ks) brdf.Ks *= material->map_Ks->get(dg.st);
    brdf.Ns = material->Ns;
    //if (material->map_Ns) { brdf.Ns *= material->map_Ns.get(dg.st); }
    brdf.Kt = (1.0f - d) * Vec3fa(material->Kt);
    brdf.Ni = material->Ni;
}

inline Vec3fa OBJMaterial__eval(OBJMaterial* material, const BRDF& brdf, const Vec3fa& wo, const Sample& sp,
                                const Vec3fa& wi) {
    Vec3fa R = Vec3fa(0.0f);
    const float Md = max(max(brdf.Kd.x, brdf.Kd.y), brdf.Kd.z);
    const float Ms = max(max(brdf.Ks.x, brdf.Ks.y), brdf.Ks.z);
    const float Mt = max(max(brdf.Kt.x, brdf.Kt.y), brdf.Kt.z);
    if (Md > 0.0f) {
        R = R + (1.0f / float(M_PI)) * clamp(dot(wi, sp.Ns)) * brdf.Kd;
    }
    if (Ms > 0.0f) {
        const Sample3f refl = make_Sample3f(reflect(wo, sp.Ns), 1.0f);
        if (dot(refl.v, wi) > 0.0f) {
            R = R + (brdf.Ns + 2) * float(one_over_two_pi) * powf(max(1e-10f, dot(refl.v, wi)), brdf.Ns) *
                clamp(dot(wi, sp.Ns)) * brdf.Ks;
        }
    }
    if (Mt > 0.0f) {
    }
    return R;
}

inline Vec3fa OBJMaterial__sample(OBJMaterial* material, const BRDF& brdf, const Vec3fa& Lw, const Vec3fa& wo,
                                  const Sample& sp, Sample3f& wi_o, const Vec2f& s) {
    Vec3fa cd = Vec3fa(0.0f);
    Sample3f wid = make_Sample3f(Vec3fa(0.0f), 0.0f);
    if (max(max(brdf.Kd.x, brdf.Kd.y), brdf.Kd.z) > 0.0f) {
        wid = cosineSampleHemisphere(s.x, s.y, sp.Ns);
        cd = float(one_over_pi) * clamp(dot(wid.v, sp.Ns)) * brdf.Kd;
    }

    Vec3fa cs = Vec3fa(0.0f);
    Sample3f wis = make_Sample3f(Vec3fa(0.0f), 0.0f);
    if (max(max(brdf.Ks.x, brdf.Ks.y), brdf.Ks.z) > 0.0f) {
        const Sample3f refl = make_Sample3f(reflect(wo, sp.Ns), 1.0f);
        wis.v = powerCosineSampleHemisphere(brdf.Ns, s);
        wis.pdf = powerCosineSampleHemispherePDF(wis.v, brdf.Ns);
        wis.v = frame(refl.v) * wis.v;
        cs = (brdf.Ns + 2) * float(one_over_two_pi) * powf(max(dot(refl.v, wis.v), 1e-10f), brdf.Ns) *
             clamp(dot(wis.v, sp.Ns)) * brdf.Ks;
    }

    Vec3fa ct = Vec3fa(0.0f);
    Sample3f wit = make_Sample3f(Vec3fa(0.0f), 0.0f);
    if (max(max(brdf.Kt.x, brdf.Kt.y), brdf.Kt.z) > 0.0f) {
        wit = make_Sample3f(neg(wo), 1.0f);
        ct = brdf.Kt;
    }

    const Vec3fa md = Lw * cd / wid.pdf;
    const Vec3fa ms = Lw * cs / wis.pdf;
    const Vec3fa mt = Lw * ct / wit.pdf;

    const float Cd = wid.pdf == 0.0f ? 0.0f : max(max(md.x, md.y), md.z);
    const float Cs = wis.pdf == 0.0f ? 0.0f : max(max(ms.x, ms.y), ms.z);
    const float Ct = wit.pdf == 0.0f ? 0.0f : max(max(mt.x, mt.y), mt.z);
    const float C = Cd + Cs + Ct;

    if (C == 0.0f) {
        wi_o = make_Sample3f(Vec3fa(0, 0, 0), 0);
        return Vec3fa(0, 0, 0);
    }

    const float CPd = Cd / C;
    const float CPs = Cs / C;
    const float CPt = Ct / C;

    if (s.x < CPd) {
        wi_o = make_Sample3f(wid.v, wid.pdf * CPd);
        return cd;
    } else if (s.x < CPd + CPs) {
        wi_o = make_Sample3f(wis.v, wis.pdf * CPs);
        return cs;
    } else {
        wi_o = make_Sample3f(wit.v, wit.pdf * CPt);
        return ct;
    }
}

////////////////////////////////////////////////////////////////////////////////
//                        Metal Material                                      //
////////////////////////////////////////////////////////////////////////////////

// =======================================================
struct FresnelConductor
{
    Vec3fa eta; //!< Real part of refraction index
    Vec3fa k;   //!< Imaginary part of refraction index
};

inline Vec3fa fresnelConductor(const float cosi, const Vec3fa &eta, const Vec3fa &k)
{
    const Vec3fa tmp = eta * eta + k * k;
    const Vec3fa Rpar = (tmp * (cosi * cosi) - 2.0f * eta * cosi + Vec3fa(1.0f)) *
                        rcp(tmp * (cosi * cosi) + 2.0f * eta * cosi + Vec3fa(1.0f));
    const Vec3fa Rper = (tmp - 2.0f * eta * cosi + Vec3fa(cosi * cosi)) *
                        rcp(tmp + 2.0f * eta * cosi + Vec3fa(cosi * cosi));
    return 0.5f * (Rpar + Rper);
}

inline Vec3fa eval(const FresnelConductor &This, const float cosTheta)
{
    return fresnelConductor(cosTheta, This.eta, This.k);
}

inline FresnelConductor make_FresnelConductor(const Vec3fa &eta, const Vec3fa &k)
{
    FresnelConductor m;
    m.eta = eta;
    m.k = k;
    return m;
}

// =======================================================
struct PowerCosineDistribution
{
    float exp;
};

inline float eval(const PowerCosineDistribution &This, const float cosThetaH)
{
    return (This.exp + 2) * (1.0f / (2.0f * (float(M_PI)))) * powf(fabs(cosThetaH), This.exp);
}

inline void sample(const PowerCosineDistribution &This, const Vec3fa &wo, const Vec3fa &N, Sample3f &wi, const Vec2f s)
{
    Vec3fa dir = powerCosineSampleHemisphere(This.exp, s);
    Sample3f wh;
    wh.v = frame(N) * dir;
    wh.pdf = powerCosineSampleHemispherePDF(dir, This.exp);
    Sample3f r = make_Sample3f(reflect(wo, wh.v), 1.0f);
    wi = make_Sample3f(r.v, wh.pdf / (4.0f * fabs(dot(wo, wh.v))));
}

inline PowerCosineDistribution make_PowerCosineDistribution(const float _exp)
{
    PowerCosineDistribution m;
    m.exp = _exp;
    return m;
}

inline void MetalMaterial__preprocess(MetalMaterial *material, BRDF &brdf, const Vec3fa &wo, const Sample &sp)
{
}

inline Vec3fa MetalMaterial__eval(MetalMaterial *material, const BRDF &brdf, const Vec3fa &wo, const Sample &sp, const Vec3fa &wi)
{
    const FresnelConductor fresnel = make_FresnelConductor(Vec3fa(material->eta), Vec3fa(material->k));
    const PowerCosineDistribution distribution = make_PowerCosineDistribution(rcp(material->roughness));

    const float cosThetaO = dot(wo, sp.Ns);
    const float cosThetaI = dot(wi, sp.Ns);
    if (cosThetaI <= 0.0f || cosThetaO <= 0.0f)
        return Vec3fa(0.f);
    const Vec3fa wh = normalize(wi + wo);
    const float cosThetaH = dot(wh, sp.Ns);
    const float cosTheta = dot(wi, wh); // = dot(wo, wh);
    const Vec3fa F = eval(fresnel, cosTheta);
    const float D = eval(distribution, cosThetaH);
    const float G = min(1.0f, min(2.0f * cosThetaH * cosThetaO / cosTheta,
                                  2.0f * cosThetaH * cosThetaI / cosTheta));
    return (Vec3fa(material->reflectance) * F) * D * G * rcp(4.0f * cosThetaO);
}

inline Vec3fa MetalMaterial__sample(MetalMaterial *material, const BRDF &brdf, const Vec3fa &Lw, const Vec3fa &wo, const Sample &sp, Sample3f &wi_o, const Vec2f &s)
{
    const PowerCosineDistribution distribution = make_PowerCosineDistribution(rcp(material->roughness));

    if (dot(wo, sp.Ns) <= 0.0f)
    {
        wi_o = make_Sample3f(Vec3fa(0.0f), 0.0f);
        return Vec3fa(0.f);
    }
    sample(distribution, wo, sp.Ns, wi_o, s);
    if (dot(wi_o.v, sp.Ns) <= 0.0f)
    {
        wi_o = make_Sample3f(Vec3fa(0.0f), 0.0f);
        return Vec3fa(0.f);
    }
    return MetalMaterial__eval(material, brdf, wo, sp, wi_o.v);
}

////////////////////////////////////////////////////////////////////////////////
//                        ReflectiveMetal Material                            //
////////////////////////////////////////////////////////////////////////////////

inline void ReflectiveMetalMaterial__preprocess(ReflectiveMetalMaterial *material, BRDF &brdf, const Vec3fa &wo, const Sample &sp)  
{
}

inline Vec3fa ReflectiveMetalMaterial__eval(ReflectiveMetalMaterial *material, const BRDF &brdf, const Vec3fa &wo, const Sample &sp, const Vec3fa &wi)
{
  return Vec3fa(0.0f);
}

inline Vec3fa ReflectiveMetalMaterial__sample(ReflectiveMetalMaterial *material, const BRDF &brdf, const Vec3fa &Lw, const Vec3fa &wo, const Sample &sp, Sample3f &wi_o, const Vec2f &s)
{
  wi_o = make_Sample3f(reflect(wo,sp.Ns),1.0f);
  return Vec3fa(material->reflectance) * fresnelConductor(dot(wo,sp.Ns),Vec3fa((Vec3fa)material->eta),Vec3fa((Vec3fa)material->k));
}


////////////////////////////////////////////////////////////////////////////////
//                              Material                                      //
////////////////////////////////////////////////////////////////////////////////

inline void Material__preprocess(std::vector<Material *> materials, unsigned int materialID,
                                 BRDF &brdf, const Vec3fa &wo, const Sample &sp)
{
    auto id = materialID;
    {
        if (id < materials.size()) // FIXME: workaround for ISPC bug, location reached with empty execution mask
        {
            Material *material = materials[id];

            switch (material->type)
            {
            case MATERIAL_OBJ:
                OBJMaterial__preprocess((OBJMaterial *)material, brdf, wo, sp);
                break;
            case MATERIAL_MATTE:
                MatteMaterial__preprocess((MatteMaterial *)material, brdf, wo, sp);
                break;
            case MATERIAL_METAL:
                MetalMaterial__preprocess((MetalMaterial *)material, brdf, wo, sp);
                break;
            case MATERIAL_REFLECTIVE_METAL:
                ReflectiveMetalMaterial__preprocess((ReflectiveMetalMaterial *)material, brdf, wo, sp);
                break;
            default:
                break;
            }
        }
    }
}

inline Vec3fa Material__eval(std::vector<Material *> materials, unsigned int materialID,
                             const BRDF &brdf, const Vec3fa &wo, const Sample &sp, const Vec3fa &wi)
{
    Vec3fa c = Vec3fa(0.0f);
    auto id = materialID;
    {
        if (id < materials.size()) // FIXME: workaround for ISPC bug, location reached with empty execution mask
        {
            Material *material = materials[id];
            switch (material->type)
            {
            case MATERIAL_OBJ:
                c = OBJMaterial__eval((OBJMaterial *)material, brdf, wo, sp, wi);
                break;
            case MATERIAL_MATTE:
                c = MatteMaterial__eval((MatteMaterial *)material, brdf, wo, sp, wi);
                break;
            case MATERIAL_METAL:
                c = MetalMaterial__eval((MetalMaterial *)material, brdf, wo, sp, wi);
                break;
            case MATERIAL_REFLECTIVE_METAL:
                c = ReflectiveMetalMaterial__eval((ReflectiveMetalMaterial *)material, brdf, wo, sp, wi);
                break;
            default:
                std::cout << "No Material found" << std::endl;
                c = Vec3fa(0.0f);
            }
        }
    }
    return c;
}

inline Vec3fa Material__sample(std::vector<Material *> materials, unsigned int materialID,
                               const BRDF &brdf, const Vec3fa &Lw, const Vec3fa &wo, const Sample &sp, Sample3f &wi_o,
                               const Vec2f &s)
{
    Vec3fa c = Vec3fa(0.0f);
    auto id = materialID;
    {
        if (id < materials.size()) // FIXME: workaround for ISPC bug, location reached with empty execution mask
        {
            Material *material = materials[id];
            switch (material->type)
            {
            case MATERIAL_OBJ:
                c = OBJMaterial__sample((OBJMaterial *)material, brdf, Lw, wo, sp, wi_o, s);
                break;
            case MATERIAL_MATTE:
                c = MatteMaterial__sample((MatteMaterial *)material, brdf, Lw, wo, sp, wi_o, s);
                break;
            case MATERIAL_METAL:
                c = MetalMaterial__sample((MetalMaterial *)material, brdf, Lw, wo, sp, wi_o, s);
                break;
            case MATERIAL_REFLECTIVE_METAL:
                c = ReflectiveMetalMaterial__sample((ReflectiveMetalMaterial *)material, brdf, Lw, wo, sp, wi_o, s);
                break;
            default:
                wi_o = make_Sample3f(Vec3fa(0.0f), 0.0f);
                c = Vec3fa(0.0f);
                break;
            }
        }
    }
    return c;
}