import torch
import numpy as np
from PIL import Image
import sys
import OpenGL.GL as gl
import glfw
import ctypes


VERTEX_SHADER = """
#version 330 core

layout (location = 0) in vec3 aPos;
layout (location = 1) in vec2 aTexCoord;

out vec2 TexCoord;

void main()
{
    gl_Position = vec4(aPos, 1.0);
    TexCoord = aTexCoord;
}
"""

FRAGMENT_SHADER = """
#version 330 core

out vec4 FragColor;

in vec2 TexCoord;

uniform sampler2D iChannel0;
uniform sampler2D iChannel1;
uniform vec3 iResolution;
uniform float iTime;
uniform float iAngleNum;
uniform float iSampNum;
uniform float iLineWidth;
uniform float iVignette;

#define Res0 textureSize(iChannel0, 0)
#define Res1 textureSize(iChannel1, 0)
#define Res  iResolution.xy

#define randSamp iChannel1
#define colorSamp iChannel0

vec4 getRand(vec2 pos)
{
    return textureLod(iChannel1, pos / Res1 / iResolution.y * 1080., 0.0);
}

vec4 getCol(vec2 pos)
{
    vec2 uv = ((pos - Res.xy * .5) / Res.y * Res0.y) / Res0.xy + .5;
    vec4 c1 = texture(iChannel0, uv);
    vec4 e = smoothstep(vec4(-0.05), vec4(-0.0), vec4(uv, vec2(1) - uv));
    c1 = mix(vec4(1, 1, 1, 0), c1, e.x * e.y * e.z * e.w);
    float d = clamp(dot(c1.xyz, vec3(-.5, 1., -.5)), 0.0, 1.0);
    vec4 c2 = vec4(.7);
    return min(mix(c1, c2, 1.8 * d), .7);
}

vec4 getColHT(vec2 pos)
{
    return smoothstep(.95, 1.05, getCol(pos) * .8 + .2 + getRand(pos * .7));
}

float getVal(vec2 pos)
{
    vec4 c = getCol(pos);
    return pow(dot(c.xyz, vec3(.333)), 1.) * 1.;
}

vec2 getGrad(vec2 pos, float eps)
{
    vec2 d = vec2(eps, 0);
    return vec2(
        getVal(pos + d.xy) - getVal(pos - d.xy),
        getVal(pos + d.yx) - getVal(pos - d.yx)
    ) / eps / 2.;
}

#define PI2 6.28318530717959

void main()
{
    vec2 pos = TexCoord * iResolution.xy + 4.0 * sin(iTime * 1. * vec2(1, 1.7)) * iResolution.y / 400.;
    vec3 col = vec3(0);
    vec3 col2 = vec3(0);
    float sum = 0.;
    for (int i = 0; i < int(iAngleNum); i++)
    {
        float ang = PI2 / iAngleNum * (float(i) + .8);
        vec2 v = vec2(cos(ang), sin(ang));
        for (int j = 0; j < int(iSampNum); j++)
        {
            vec2 dpos = v.yx * vec2(1, -1) * float(j) * iLineWidth * iResolution.y / 400.;
            vec2 dpos2 = v.xy * float(j * j) / iSampNum * .5 * iLineWidth * iResolution.y / 400.;
            vec2 g;
            float fact;
            float fact2;

            for (float s = -1.; s <= 1.; s += 2.)
            {
                vec2 pos2 = pos + s * dpos + dpos2;
                vec2 pos3 = pos + (s * dpos + dpos2).yx * vec2(1, -1) * 2.;
                g = getGrad(pos2, .4);
                fact = dot(g, v) - .5 * abs(dot(g, v.yx * vec2(1, -1)));
                fact2 = dot(normalize(g + vec2(.0001)), v.yx * vec2(1, -1));

                fact = clamp(fact, 0., .05);
                fact2 = abs(fact2);

                fact *= 1. - float(j) / iSampNum;
                col += fact;
                col2 += fact2 * getColHT(pos3).xyz;
                sum += fact2;
            }
        }
    }
    col /= iSampNum * iAngleNum * .75 / sqrt(iResolution.y);
    col2 /= sum;
    col.x *= (.6 + .8 * getRand(pos * .7).x);
    col.x = 1. - col.x;
    col.x *= col.x * col.x;

    vec2 s = sin(pos.xy * .1 / sqrt(iResolution.y / 400.));
    vec3 karo = vec3(1);
    karo -= .5 * vec3(.25, .1, .1) * dot(exp(-s * s * 80.), vec2(1));
    float r = length(pos - iResolution.xy * .5) / iResolution.x;
    float vign = 1. - r * r * r * iVignette;
    FragColor = vec4(vec3(col.x * col2 * karo * vign), 1);
}
"""

class FL_PaperDrawn:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            },
            "optional": {
                "angle_num": ("FLOAT", {"default": 3.0, "min": 1.0, "max": 10.0, "step": 1.0}),
                "samp_num": ("FLOAT", {"default": 2.2, "min": 1.0, "max": 10.0, "step": 0.1}),
                "line_width": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}),
                "vignette": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.1}),
                "fps": ("INT", {"default": 30, "min": 1, "max": 120, "step": 1}),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_shader"
    CATEGORY = "🏵️Fill Nodes"

    def apply_shader(self, image, angle_num, samp_num, line_width, vignette, fps):
        result = []
        total_images = len(image)
        frame_time = 1.0 / fps

        for i, img in enumerate(image, start=1):
            img = self.t2p(img)
            result_img = self.process_image(img, angle_num, samp_num, line_width, vignette, i * frame_time)
            result_img = self.p2t(result_img)
            result.append(result_img)

            # Update the print log
            progress = i / total_images * 100
            sys.stdout.write(f"\rProcessing images: {progress:.2f}%")
            sys.stdout.flush()

        # Print a new line after the progress log
        print()

        return (torch.cat(result, dim=0),)

    def process_image(self, image, angle_num, samp_num, line_width, vignette, time):
        # Convert the PIL image to a numpy array
        img_array = np.array(image).astype(np.float32) / 255.0

        # Create a white image for iChannel1
        white_image = np.ones((image.height, image.width, 3), dtype=np.float32)

        # Create a PyOpenGL context
        if not glfw.init():
            raise RuntimeError("Failed to initialize GLFW")

        glfw.window_hint(glfw.VISIBLE, glfw.FALSE)
        window = glfw.create_window(image.width, image.height, "Hidden Window", None, None)
        if not window:
            glfw.terminate()
            raise RuntimeError("Failed to create GLFW window")

        glfw.make_context_current(window)

        # Compile the shader program
        vertex_shader = gl.glCreateShader(gl.GL_VERTEX_SHADER)
        gl.glShaderSource(vertex_shader, VERTEX_SHADER)
        gl.glCompileShader(vertex_shader)

        fragment_shader = gl.glCreateShader(gl.GL_FRAGMENT_SHADER)
        gl.glShaderSource(fragment_shader, FRAGMENT_SHADER)
        gl.glCompileShader(fragment_shader)

        shader_program = gl.glCreateProgram()
        gl.glAttachShader(shader_program, vertex_shader)
        gl.glAttachShader(shader_program, fragment_shader)
        gl.glLinkProgram(shader_program)

        gl.glUseProgram(shader_program)

        # Set up vertex buffer object (VBO) and vertex array object (VAO)
        vertices = np.array([
            -1.0, -1.0, 0.0, 0.0, 0.0,
            1.0, -1.0, 0.0, 1.0, 0.0,
            -1.0, 1.0, 0.0, 0.0, 1.0,
            1.0, 1.0, 0.0, 1.0, 1.0
        ], dtype=np.float32)

        vao = gl.glGenVertexArrays(1)
        gl.glBindVertexArray(vao)

        vbo = gl.glGenBuffers(1)
        gl.glBindBuffer(gl.GL_ARRAY_BUFFER, vbo)
        gl.glBufferData(gl.GL_ARRAY_BUFFER, vertices.nbytes, vertices, gl.GL_STATIC_DRAW)

        gl.glVertexAttribPointer(0, 3, gl.GL_FLOAT, gl.GL_FALSE, 5 * vertices.itemsize, None)
        gl.glEnableVertexAttribArray(0)

        gl.glVertexAttribPointer(1, 2, gl.GL_FLOAT, gl.GL_FALSE, 5 * vertices.itemsize, ctypes.c_void_p(3 * vertices.itemsize))
        gl.glEnableVertexAttribArray(1)

        # Set up textures
        texture0 = gl.glGenTextures(1)
        gl.glActiveTexture(gl.GL_TEXTURE0)
        gl.glBindTexture(gl.GL_TEXTURE_2D, texture0)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
        gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, image.width, image.height, 0, gl.GL_RGB, gl.GL_FLOAT, img_array)

        texture1 = gl.glGenTextures(1)
        gl.glActiveTexture(gl.GL_TEXTURE1)
        gl.glBindTexture(gl.GL_TEXTURE_2D, texture1)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_REPEAT)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
        gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, gl.GL_RGB, image.width, image.height, 0, gl.GL_RGB, gl.GL_FLOAT, white_image)

        # Set shader uniforms
        gl.glUniform1i(gl.glGetUniformLocation(shader_program, "iChannel0"), 0)
        gl.glUniform1i(gl.glGetUniformLocation(shader_program, "iChannel1"), 1)
        gl.glUniform3f(gl.glGetUniformLocation(shader_program, "iResolution"), image.width, image.height, 0.0)
        gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iTime"), time)
        gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iAngleNum"), angle_num)
        gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iSampNum"), samp_num)
        gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iLineWidth"), line_width)
        gl.glUniform1f(gl.glGetUniformLocation(shader_program, "iVignette"), vignette)

        # Render the shader
        gl.glClear(gl.GL_COLOR_BUFFER_BIT)
        gl.glDrawArrays(gl.GL_TRIANGLE_STRIP, 0, 4)

        # Read the rendered image from the framebuffer
        img_data = gl.glReadPixels(0, 0, image.width, image.height, gl.GL_RGB, gl.GL_FLOAT)
        img_array = np.frombuffer(img_data, dtype=np.float32).reshape((image.height, image.width, 3))

        # Clean up OpenGL resources
        gl.glDeleteTextures(2, [texture0, texture1])
        gl.glDeleteBuffers(1, [vbo])
        gl.glDeleteVertexArrays(1, [vao])
        gl.glDeleteProgram(shader_program)
        gl.glDeleteShader(vertex_shader)
        gl.glDeleteShader(fragment_shader)

        glfw.destroy_window(window)
        glfw.terminate()

        # Convert the processed image back to a PIL image
        processed_image = Image.fromarray((img_array * 255).astype(np.uint8))

        return processed_image

    def t2p(self, t):
        if t is not None:
            i = 255.0 * t.cpu().numpy().squeeze()
            p = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
        return p

    def p2t(self, p):
        if p is not None:
            i = np.array(p).astype(np.float32) / 255.0
            t = torch.from_numpy(i).unsqueeze(0)
        return t