//#####################################################################
// Copyright 2021, Craig Schroeder.
// This file is part of PhysBAM whose distribution is governed by the license contained in the accompanying file PHYSBAM_COPYRIGHT.txt.
//#####################################################################
#include "QT_CAMERA.h"
#include "QT_FRAGMENT_SHADER.h"
#include "QT_GEOMETRY_SHADER.h"
#include "QT_MATH.h"
#include "QT_RENDER_UTILS.h"
#include "QT_SHADER.h"
#include "QT_SHADER_FLAGS.h"
#include "QT_VERTEX_SHADER.h"
#include <QOpenGLShaderProgram>
#include <QVector3D>
#include <string>
namespace PhysBAM{
//#####################################################################
// Constructor
//#####################################################################
SHADER::
SHADER(QObject* obj,SHADER_FLAGS flags)
    :flags(flags)
{
    ARRAY<std::pair<const char*,const char*>> al,ul;
    if(flags.col==1) ul.Append({"vec3","col"});
    if(flags.col==2) al.Append({"vec3","col"});
    if(flags.bg_col==1) ul.Append({"vec3","bgcol"});
    if(flags.bg_col==2) al.Append({"vec3","bgcol"});
    if(flags.special==SHADER_TYPE::segment || flags.special==SHADER_TYPE::arrow) al.Append({"vec3","dpos"});
    if(flags.special==SHADER_TYPE::sphere || flags.special==SHADER_TYPE::regular_poly) al.Append({"float","radius"});
    if((flags.special==SHADER_TYPE::box || flags.special==SHADER_TYPE::diamond) && flags.dim==2)
        ul.Append({"vec2","box_edges"});
    if((flags.special==SHADER_TYPE::box || flags.special==SHADER_TYPE::diamond) && flags.dim==3)
        ul.Append({"vec3","box_edges"});
    if(flags.special==SHADER_TYPE::regular_poly)
    {
        ul.Append({"int","num_sides"});
        ul.Append({"int","initial_rotate"});
    }
    if(flags.special==SHADER_TYPE::matrix && flags.dim==2)
    {
        al.Append({"vec2","mat0"});
        al.Append({"vec2","mat1"});
    }
    if(flags.special==SHADER_TYPE::matrix && flags.dim==3)
    {
        al.Append({"vec3","mat0"});
        al.Append({"vec3","mat1"});
        al.Append({"vec3","mat2"});
    }
    if(flags.special!=SHADER_TYPE::regular) ul.Append({"float","scale"});

    auto vs=Generate_Vertex_Shader(true,al,nullptr);
    auto gs=Generate_Geometry_Shader(flags);
    auto fs=Generate_Fragment_Shader(flags);
    Init(obj,flags,true,al,ul,vs.c_str(),gs.c_str(),fs.c_str());
}
//#####################################################################
// Function Init
//#####################################################################
void SHADER::
Init(QObject* obj,SHADER_FLAGS flags_,bool use_pos,
    const ARRAY<std::pair<const char*,const char*>>& attr_list,
    const ARRAY<std::pair<const char*,const char*>>& uniform_list,
    const char* vs,const char* gs,const char* fs)
{
    flags=flags_;
    // LOG::printf("vs:\n%P\n",vs);
    // LOG::printf("gs:\n%P\n",gs);
    // LOG::printf("fs:\n%P\n",fs);

    program = new QOpenGLShaderProgram(obj);
    program->addShaderFromSourceCode(QOpenGLShader::Vertex,vs);
    program->addShaderFromSourceCode(QOpenGLShader::Geometry,gs);
    program->addShaderFromSourceCode(QOpenGLShader::Fragment,fs);
    program->link();

    auto do_attr=[this](const char* s)
        {attrs.Set(s,program->attributeLocation(LOG::sprintf("vs_%P",s).c_str()));};
    auto do_unif=[this](const char* s)
        {uniforms.Set(s,program->uniformLocation(s));};
    for(auto s:attr_list) do_attr(s.second);
    for(auto s:uniform_list) do_unif(s.second);

    do_unif("world_to_gl");
    do_unif("obj_to_world");
    do_unif("obj_to_world_scale");
    do_unif("camera_pos");
    if(use_pos)
    {
        do_attr("pos");
        do_unif("slice_plane_0");
        do_unif("slice_plane_1");
    }
    if(flags.select) do_unif("first_id");
    else if(!flags.flat_col)
    {
        do_unif("ambient_brightness");
        do_unif("shininess");
        do_unif("specular_color");
        do_unif("light_pos");
    }
}
//#####################################################################
// Destructor
//#####################################################################
SHADER::
~SHADER()
{
//    delete program;
}
//#####################################################################
// Constructor
//#####################################################################
template<class T> SHADER::ATTR::
ATTR(SH& sh,const char* name,ARRAY_VIEW<T> X,int part)
    :sh(sh)
{
    typedef std::remove_const_t<T> U;typedef SCALAR_TYPE<U> S;
    if(!name || !sh.shader.attrs.Get(name,attr)) return;
    auto type=sizeof(S)==sizeof(double)?GL_DOUBLE:GL_FLOAT;
    if constexpr(IS_SCALAR<U>::value)
    {
        sh.gl.glVertexAttribPointer(attr,1,type,GL_FALSE,0,X.Get_Array_Pointer());
    }
    else if constexpr(IS_VECTOR<U>::value)
    {
        sh.gl.glVertexAttribPointer(attr,T::m,type,GL_FALSE,0,X.Get_Array_Pointer());
    }
    else
    {
        static_assert(IS_MATRIX<U>::value);
        assert((unsigned)part<T::n);
        sh.gl.glVertexAttribPointer(attr,T::m,type,GL_FALSE,T::m*T::n*sizeof(S),(S*)X.Get_Array_Pointer()+T::m*part);
    }
    sh.gl.glEnableVertexAttribArray(attr);
}
//#####################################################################
// Constructor
//#####################################################################
SHADER::ATTR::
ATTR(SH& sh,const char* name,const RU_COLOR& color)
    :sh(sh)
{
    if(!name || !color.valid) return;
    if(color.pf.m)
    {
        if(!sh.shader.attrs.Get(name,attr)) return;
        sh.gl.glVertexAttribPointer(attr,3,GL_FLOAT,GL_FALSE,0,color.pf.Get_Array_Pointer());
        sh.gl.glEnableVertexAttribArray(attr);
    }
    else if(color.pd.m)
    {
        if(!sh.shader.attrs.Get(name,attr)) return;
        sh.gl.glVertexAttribPointer(attr,3,GL_DOUBLE,GL_FALSE,0,color.pd.Get_Array_Pointer());
        sh.gl.glEnableVertexAttribArray(attr);
    }
    else sh.Set_Uniform(name,to_qt(color.color));
    // Don't store attr in the uniform case; no attribute to disable.
}
//#####################################################################
// Destructor
//#####################################################################
SHADER::ATTR::
~ATTR()
{
    if(attr>=0) sh.gl.glDisableVertexAttribArray(attr);
}
//#####################################################################
// Constructor
//#####################################################################
SHADER::SH::
SH(const SHADER& shader,const QT_RENDER_UTILS& ru,SHADER_FLAGS flags)
    :shader(shader),gl(ru.gl)
{
    shader.program->bind();

    // make segments and points easier to select.
    float depth_offset=ru.depth_offset;
    if(flags.select && flags.out_dim==1) depth_offset-=.001f;
    if(flags.select && flags.out_dim==0) depth_offset-=.0015f;
    QMatrix4x4 offset;
    offset.setToIdentity();
    offset.translate({0,0,depth_offset});
    Set_Uniform("world_to_gl",offset*to_qt(ru.world_to_gl));

    Set_Uniform("obj_to_world",to_qt(ru.obj_to_world));
    Set_Uniform("obj_to_world_scale",ru.obj_to_world_scale);
    Set_Uniform("camera_pos",to_qt(ru.camera.Camera_Frame_To_World({})));
    if(flags.select) Set_Uniform("first_id",ru.pick.Next());
    else if(!flags.flat_col)
    {
        Set_Uniform("ambient_brightness",ru.ambient_brightness);
        Set_Uniform("shininess",ru.shininess);
        Set_Uniform("specular_color",to_qt(ru.specular_color));
        QVector3D world_light_pos[2];
        for(int i=0;i<2;i++) world_light_pos[i]=to_qt(ru.world_light_pos[i]);
        Set_Uniform_Array("light_pos",world_light_pos,2);
    }
    Set_Uniform("slice_plane_0",to_qt(ru.clip_plane0));
    Set_Uniform("slice_plane_1",to_qt(ru.clip_plane1));
}
//#####################################################################
// Destructor
//#####################################################################
SHADER::SH::
~SH()
{
    shader.program->release();
}
//#####################################################################
// Function Set_Uniform
//#####################################################################
template<class T> void SHADER::SH::
Set_Uniform(const char* name,const T& v)
{
    if(auto*p=shader.uniforms.Get_Pointer(name))
        if(*p!=-1)
            shader.program->setUniformValue(*p,v);
}
//#####################################################################
// Function Set_Uniform_Array
//#####################################################################
template<class T> void SHADER::SH::
Set_Uniform_Array(const char* name,const T* v,int length)
{
    if(auto*p=shader.uniforms.Get_Pointer(name))
        if(*p!=-1)
            shader.program->setUniformValueArray(*p,v,length);
}
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<double,2>,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<float,2>,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<double,3>,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<float,3>,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<double,2> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<double,3> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<float,2> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<VECTOR<float,3> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<double const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<float const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<MATRIX<double,2,2> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<MATRIX<double,3,3> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<MATRIX<float,2,2> const,int>,int);
template SHADER::ATTR::ATTR(SHADER::SH&,const char*,ARRAY_VIEW<MATRIX<float,3,3> const,int>,int);
template void SHADER::SH::Set_Uniform<QGenericMatrix<2,2,float> >(const char*,QGenericMatrix<2,2,float> const&);
template void SHADER::SH::Set_Uniform<QGenericMatrix<3,3,float> >(const char*,QGenericMatrix<3,3,float> const&);
template void SHADER::SH::Set_Uniform<int>(const char*,int const&);
template void SHADER::SH::Set_Uniform<float>(const char*,float const&);
}
