#define QUANT 500

#define TEMPORAL true
#define K_CNT 10
#define PART_CNT 3200
#define DITH_AMT 4.0
#define ITERS 1
#define PALETTE 3
#define CRAWL 0.01
#define DIST_HEUR 0
#define SEED 102u

uint dith_sc = 4u;


struct Particle {
    vec3 pos;
    vec3 col;
};

struct K{
    vec3 pos;
    vec3 col;
    ivec4 accum;
    uvec4 min_col;
    uvec4 max_col;
    vec4 faccum;
    vec4 fmin_col;
    vec4 fmax_col;
    uint accum_cnt;
};

//struct ParticleData {
//};

layout(std430, binding = 6) coherent buffer aaaaaaaaaa{
    uint iter;
    K k[300];
    Particle particles[];
} buff;

uniform int input_tex;

vec3 mul3( in mat3 m, in vec3 v ){return vec3(dot(v,m[0]),dot(v,m[1]),dot(v,m[2]));}

vec3 mul3( in vec3 v, in mat3 m ){return mul3(m,v);}

vec3 srgb2oklab(vec3 c) {

    mat3 m1 = mat3(
    0.4122214708,0.5363325363,0.0514459929,
    0.2119034982,0.6806995451,0.1073969566,
    0.0883024619,0.2817188376,0.6299787005
    );

    vec3 lms = mul3(m1,c);

    lms = pow(lms,vec3(1./3.));

    mat3 m2 = mat3(
    +0.2104542553,+0.7936177850,-0.0040720468,
    +1.9779984951,-2.4285922050,+0.4505937099,
    +0.0259040371,+0.7827717662,-0.8086757660
    );

    return mul3(m2,lms);
}

vec3 oklab2srgb(vec3 c)
{
    mat3 m1 = mat3(
    1.0000000000,+0.3963377774,+0.2158037573,
    1.0000000000,-0.1055613458,-0.0638541728,
    1.0000000000,-0.0894841775,-1.2914855480
    );

    vec3 lms = mul3(m1,c);

    lms = lms * lms * lms;

    mat3 m2 = mat3(
    +4.0767416621,-3.3077115913,+0.2309699292,
    -1.2684380046,+2.6097574011,-0.3413193965,
    -0.0041960863,-0.7034186147,+1.7076147010
    );
    return mul3(m2,lms);
}

vec3 lab2lch( in vec3 c ){return vec3(c.x,sqrt((c.y*c.y) + (c.z * c.z)),atan(c.z,c.y));}

vec3 lch2lab( in vec3 c ){return vec3(c.x,c.y*cos(c.z),c.y*sin(c.z));}

vec3 srgb_to_oklch( in vec3 c ) { return lab2lch(srgb2oklab(c)); }
vec3 oklch_to_srgb( in vec3 c ) { return oklab2srgb(lch2lab(c)); }


vec4 sample_input_image(vec2 uv){
    vec3 col = texture(textures[input_tex],uv).rgb;
//    col = 1.-exp(-col*1.5);
    col = max(col,0.);
    col = pow(col,vec3(1./0.454545));
//    col = pow(col,vec3(1./0.454545));
    col = srgb_to_oklch(col);
//    col.z += floor(T)*.3;
    col.z = floor(T)*.3;
//    col.z += dot(col.x,col.x)*10;
//    col.y*=1000.;
//    col.y = 0.1 - col.x*0.1;
    col.y = max(col.y,0.);
//    col.y = (col.y - 0.5)*(1. + 0.2) + 0.5;
//    col.y *= 13.4;
    // col.x += sin(floor(time.elapsed))*.3;
    col.z = mod(col.z, tau);
    col = oklch_to_srgb(col);
    col = max(col,0.);
    return vec4(col,1.0);
}


mat3 m1 = mat3(
    0.4122214708, 0.5363325363, 0.0514459929,
    0.2119034982, 0.6806995451, 0.1073969566,
    0.0883024619, 0.2817188376, 0.6299787005
);

mat3 inverse_m1 = mat3(
    4.0767416621, -3.3077115913, 0.2309699292,
    -1.2684380046, 2.6097574011, -0.3413193965,
    -0.0041960863, -0.7034186147, 1.7076147010
);

mat3 mm2 = mat3(
    0.2104542553, 0.7936177850, -0.0040720468,
    1.9779984951, -2.4285922050, 0.4505937099,
    0.0259040371, 0.7827717662, -0.8086757660
);

mat3 inverse_m2 = mat3(
    1.0, 0.3963377774, 0.2158037573,
    1.0, -0.1055613458, -0.0638541728,
    1.0, -0.0894841775, -1.2914855480
);

float cbrt(float x) {
    float y = sign(x) * uintBitsToFloat((floatBitsToUint(abs(x)) / 3u) + 0x2a514067u);

    for (int i = 0; i < 1; i = i + 1) {
        y = (2.0 * y + x / (y * y)) * 0.333333333;
    }

    for (int i = 0; i < 1; i = i + 1) {
        float y3 = y * y * y;
        y = y * (y3 + 2.0 * x) / (2.0 * y3 + x);
    }

    return y;
}

vec3 cbrt_vec3(vec3 xyz) {
    return vec3(cbrt(xyz.x), cbrt(xyz.y), cbrt(xyz.z));
}

vec3 rgb2oklab(vec3 rgb) {
    return cbrt_vec3(rgb * m1) * mm2;
}

vec3 oklab2rgb(vec3 oklab) {
    // todo: max
    return pow((oklab * inverse_m2), vec3(3.0)) * inverse_m1;
}

vec3 to_space(vec3 c){
    if(DIST_HEUR < 1){
        
        return srgb2oklab(c);
    } else if(DIST_HEUR < 2){
        return srgb_to_oklch(c);
    } else {
        return c;
    }
}
vec3 from_space(vec3 c){
    if(DIST_HEUR < 1){
        // return oklab2srgb(c);
        return oklab2srgb(c);
    } else if(DIST_HEUR < 2){
        return oklch_to_srgb(c);
    } else {
        return c;
    }
}



float colour_dist(vec3 a, vec3 b) {
    if(DIST_HEUR < 1 ){
        vec3 aa = a;
        vec3 bb = b;
        vec3 v = abs(aa - bb);
        return length(v*vec3(0.8,0.5,1.0));
    } else if(DIST_HEUR < 2 ){
        // var aa = srgb_to_oklch(a);
        // var bb = srgb_to_oklch(b);
        // var aa = oklch_to_srgb(a);
        // var bb = oklch_to_srgb(b);
        vec3 aa = a;
        vec3 bb = b;
        vec3 v = abs(aa - bb);
        if(true){
            float h_a = bb.z;
            float h_b = bb.z + tau;
            // h_b = h_b % tau;
            float dist_a = abs(aa.z - h_a);
            float dist_b = abs(aa.z - h_b);
            // var hlen 
            if(dist_a < dist_b){
                // col.z = mix(col.z, h_a, alpha);
                v.z = abs(h_a - aa.z);
            } else {
                v.z = abs(h_b - aa.z);
                // col.z = mix(col.z, h_b, alpha);
            }
            // return abs(v.z);
            // return length(v*vec3f(0.5,0.02,1./tau));
            // return length(v*vec3f(1.0,4.0,1./tau));
            return length(v*vec3(0.4,1.0,0.01/tau));
        } else {
            return length(v*vec3(1.,1,1));
        }
    } else {
        return length(a-b);
    }
}


void init_k(uint id){
    seed = id + 15125u + F;
    // part.col = clamp(part.col,vec3f(0),vec3f(1);
    vec3 col;
    if(PALETTE < 1){
        col = pow(hash_v3(),vec3(3.));
    } else if(PALETTE < 2) {
        col = sample_input_image(hash_v2()).rgb;
        col = srgb_to_oklch( col );
        // col.z += 5.5;
        col.z += hash_f()*15.;
        // col.z -= 2.;
        // col.x *= 0.;
        col.z = mod(col.z,tau);
        col = oklch_to_srgb(col);
        col = pow(col,vec3(.5));

        col = max(col,vec3(0));
        col = min(col,vec3(1));
    } else if(PALETTE < 3) {
        vec3[16] palAppleII = vec3[](
        vec3(217, 60, 240)/255.,
        vec3(64, 53, 120)/255.,
        vec3(108, 41, 64)/255.,
        vec3(0, 0, 0)/255. + 0.001,

        vec3(236, 168, 191)/255.,
        vec3(128, 128, 128)/255.,
        vec3(217, 104, 15)/255.,
        vec3(64, 75, 7)/255.,

        vec3(191, 180, 248)/255.,
        vec3(38, 151, 240)/255.,
        vec3(128, 128, 128)/255.,
        vec3(19, 87, 64)/255.,

        vec3(255, 255, 255)/255.,
        vec3(147, 214, 191)/255.,
        vec3(191, 202, 135)/255.,
        vec3(38, 195, 15)/255.
        );

        col = palAppleII[id.x%16];
        // col = min(col,vec3f(1.0));
        col = pow(col,vec3(0.4));

    } else {
        // col = oklch_to_srgb( vec3f(0.3,0.1,hash_f()*tau + 0.01) );
        // col = oklch2rgb(vec3f(0.5,1.0,hash_f()*tau*0.01 + 0.3));

        col = sample_input_image(hash_v2()).rgb;
        col = srgb_to_oklch(col);
        col = oklch_to_srgb(vec3(0.5 +col.x*0.0 - hash_f()*0.4,hash_f()*0.02,hash_f()*tau*1.0));
        if(luma(col) > 0.02){
            // col = vec3f(0.5);
        }
        // col = oklch2rgb(vec3f(0.5,1.0,hash_f()*tau*0.01 + 0.3));
        // col = srgb_to_oklch( col );
        // col.x = 0.5;
        // col.y = 2.;
        // col = oklch_to_srgb( col );


        col = clamp(col, vec3(0), vec3(1));
        // oklch_to_srgb( vec3f);
    }

    col = to_space(col);
//    col *= 0.;

    buff.k[id].pos = col;
    buff.k[id].col = col;
}

#define IS_OKLCH (DIST_HEUR > 1 && DIST_HEUR < 2)

 
vec3 projj_p(vec3 in_pos) {
    vec3 pos = in_pos;
    pos -= 0.5;
    // pos *= 0.5;
    // pos.y += 0.5;
//    pos = rotY(-screendata.mouse_input.x*10) * pos;
    
    pos.xz *= rot(-mouse_ndc.x*10);

    pos.yz *= rot(mouse_ndc.y*4.5);
    

    // pos -= screendata.cam_pos;
    pos += vec3(0,0,1)*1.6;

    pos.x /= pos.z;
    pos.y /= pos.z;
    pos.x += 0.5;
    pos.y += 0.5;
    if(pos.z < 0.){
        pos = -vec3(1000);
    }
    return pos;
}



void draw_part(vec3 pos, vec3 col){
    for(int i = 0; i < 40; i++){
        vec3 k = projj_p(pos);

        if(k.z > 0.){
            vec2 uv = k.xy;
            vec2 X = hash_v2();
            vec2 disk_samp = vec2(sin(X.x), cos(X.x))*sqrt(X.y);
            uv.x += disk_samp.x*0.002;
            uv.y += disk_samp.y*0.002;
            dbg_plot_ss(uv);
        }
    }
}



void do_k_iter(uint id){
    if(id.x >= uint(K_CNT)){
        return;
    }
    uint k_id = id.x;
//    int k = buff.k[k_id];

    float acc_cnt = float((buff.k[k_id].accum_cnt));
//    float r = float((buff.k[k_id].accum[0]))/QUANT;
//    float g = float((buff.k[k_id].accum[1]))/QUANT;
//    float b = float((buff.k[k_id].accum[2]))/QUANT;
//    float w = float((buff.k[k_id].accum[3]))/QUANT;
    float r = float((buff.k[k_id].faccum[0]));
    float g = float((buff.k[k_id].faccum[1]));
    float b = float((buff.k[k_id].faccum[2]));
    float w = float((buff.k[k_id].faccum[3]));


    vec3 avg_col = vec3(r,g,b)/acc_cnt;

    if(IS_OKLCH){
        avg_col.z = atan(w,b);
        // avg_col.z = atan(b/w);
    }
    vec3 prev_pos = (buff.k[k_id]).pos;
    vec3 next_pos = mix(prev_pos,avg_col,CRAWL);


    if(acc_cnt > 0.01){
        // (*(k)).pos = avg_col;
        buff.k[k_id].pos = next_pos;
        // let pos = avg_col;
    
        float iters = 50;
        for(float i =0; i < iters; i++){
            draw_part(mix(prev_pos, next_pos, float(i)/float(iters)), vec3(0.6));
        }
        
//        draw_part(buff.k[k_id].pos, vec3(0.2,1.,0)*0.05);
    }



//    print(next_pos);
    if(TEMPORAL){
         seed += F;
//         if(acc_cnt < 4.0 && hash_f() < 0.01){
        if(
            (acc_cnt < 1.0 && hash_f() < 0.01) 
            ||  (colour_dist(avg_col, prev_pos) > 0.5 && hash_f() < 0.01)
            || hash_f() < 0.01
//            ||true
        ){

            print(prev_pos);
        // if(acc_cnt < 100.0 && hash_f() < colour_dist(avg_col, prev_pos)*1000){
            init_k(id.x);
        }
    }


   
    buff.k[k_id].accum_cnt=0;
    buff.k[k_id].accum = ivec4(0);
    buff.k[k_id].faccum = vec4(0);

//    vec4 faccum;
//    vec4 fmin_col;
//    vec4 fmax_col;
//    buff.k[k_id].accum[1]=0;
//    buff.k[k_id].accum[2]=0;
//    buff.k[k_id].accum[3]=0;

}

void do_particles_iter(int id){
    int part_id = id.x;
    if (part_id >= PART_CNT){
        return;
    }

    Particle part = buff.particles[part_id];
    float min_dist = 1000000.0;
    int k_idx = 0;
    for (int i = 0; i < K_CNT; i++){
        vec3 kpos = buff.k[i].pos;

        // let dist_to_k = length(part.col - k.pos);
        float dist_to_k = colour_dist(part.col, kpos);
        if (dist_to_k < min_dist){
            k_idx = i;
            min_dist = dist_to_k;
        }
    }
    //    let k = &buff.k[k_idx];
    #define IS_OKLCH (DIST_HEUR > 1 && DIST_HEUR < 2)


    draw_part(part.pos, vec3(0.2,1.,0)*0.05);

    //    if(custom.DIST_HEUR > 1 && custom.DIST_HEUR < 2){
//    }


//    buff.k[k_id].faccum = vec4(0);
    
    atomicAdd(buff.k[k_idx].accum[0], int(part.col[0]*QUANT));
    atomicAdd(buff.k[k_idx].accum[1], int(part.col[1]*QUANT));
    atomicAdd(buff.k[k_idx].faccum[0], part.col[0]);
    atomicAdd(buff.k[k_idx].faccum[1], part.col[1]);
    if(IS_OKLCH){
        atomicAdd(buff.k[k_idx].accum[2], int(cos(part.col[2])*QUANT));
        atomicAdd(buff.k[k_idx].accum[3], int(sin(part.col[2])*QUANT));
        
        atomicAdd(buff.k[k_idx].faccum[2], cos(part.col[2]));
        atomicAdd(buff.k[k_idx].faccum[3], sin(part.col[2]));
    } else {
        atomicAdd(buff.k[k_idx].accum[2], int(part.col[2]*QUANT));
        
        atomicAdd(buff.k[k_idx].faccum[2], part.col[2]);
    }
    atomicAdd(buff.k[k_idx].accum_cnt, uint(1));
}





















