uniform sampler2D uRegionAtlas;
uniform vec2 uRegionAtlasSize,uRegionGridSize;
varying vec2 vUv;

float regionLuminance(vec3 c){return dot(c,vec3(.299,.587,.114));}
float distanceColor(vec3 a,vec3 b){vec3 d=a-b;float y=dot(d,vec3(.299,.587,.114));return dot(d,d)+.35*y*y;}
float sideColorCost(vec4 votes,vec3 target,vec3 c0,vec3 c1,vec3 c2,vec3 c3){return votes.x*distanceColor(c0,target)+votes.y*distanceColor(c1,target)+votes.z*distanceColor(c2,target)+votes.w*distanceColor(c3,target);}
vec4 atlasSampleData(vec2 cell,vec2 local){vec2 uv=(cell*8.+local+.5)/uRegionAtlasSize;return texture2D(uRegionAtlas,uv);}
vec3 atlasSample(vec2 cell,vec2 local){return atlasSampleData(cell,local).rgb;}
float nearestLabel(vec3 color,vec3 c0,vec3 c1,vec3 c2,vec3 c3){float d0=distanceColor(color,c0),d1=distanceColor(color,c1),d2=distanceColor(color,c2),d3=distanceColor(color,c3);if(d0<=d1&&d0<=d2&&d0<=d3)return 0.;if(d1<=d2&&d1<=d3)return 1.;if(d2<=d3)return 2.;return 3.;}
float majorityLabel(vec4 votes){float label=0.,count=votes.x;if(votes.y>count){label=1.;count=votes.y;}if(votes.z>count){label=2.;count=votes.z;}if(votes.w>count)label=3.;return label;}
void main(){
  vec2 pixel=floor(gl_FragCoord.xy);float paletteIndex=mod(pixel.x,5.);vec3 cell=vec3(floor(pixel.x/5.),pixel.y,0.);
  vec3 darkest=vec3(1.),lightest=vec3(0.);float minLum=2.,maxLum=-1.;
  for(int y=0;y<8;y++)for(int x=0;x<8;x++){
    vec3 c=atlasSample(cell.xy,vec2(float(x),float(y)));float l=regionLuminance(c);
    if(l<minLum){minLum=l;darkest=c;}if(l>maxLum){maxLum=l;lightest=c;}
  }
  vec3 third=darkest;float thirdDistance=-1.;
  for(int y=0;y<8;y++)for(int x=0;x<8;x++){
    vec3 c=atlasSample(cell.xy,vec2(float(x),float(y)));float d=min(distanceColor(c,darkest),distanceColor(c,lightest));
    if(d>thirdDistance){thirdDistance=d;third=c;}
  }
  vec3 fourth=darkest;float fourthDistance=-1.;
  for(int y=0;y<8;y++)for(int x=0;x<8;x++){
    vec3 c=atlasSample(cell.xy,vec2(float(x),float(y)));float d=min(distanceColor(c,darkest),min(distanceColor(c,lightest),distanceColor(c,third)));
    if(d>fourthDistance){fourthDistance=d;fourth=c;}
  }
  vec3 c0=darkest,c1=lightest,c2=third,c3=fourth;
  for(int iteration=0;iteration<2;iteration++){
    vec3 sum0=vec3(0.),sum1=vec3(0.),sum2=vec3(0.),sum3=vec3(0.);float n0=0.,n1=0.,n2=0.,n3=0.;
    for(int y=0;y<8;y++)for(int x=0;x<8;x++){
      vec3 c=atlasSample(cell.xy,vec2(float(x),float(y)));
      float d0=distanceColor(c,c0),d1=distanceColor(c,c1),d2=distanceColor(c,c2),d3=distanceColor(c,c3);
      if(d0<=d1&&d0<=d2&&d0<=d3){sum0+=c;n0+=1.;}
      else if(d1<=d2&&d1<=d3){sum1+=c;n1+=1.;}
      else if(d2<=d3){sum2+=c;n2+=1.;}
      else{sum3+=c;n3+=1.;}
    }
    if(n0>0.)c0=sum0/n0;if(n1>0.)c1=sum1/n1;if(n2>0.)c2=sum2/n2;if(n3>0.)c3=sum3/n3;
  }

  vec4 metadata=vec4(0.);
  if(paletteIndex>3.5){
    vec4 wholeVotes=vec4(0.);float cellEvidence=0.;
    for(int y=0;y<8;y++)for(int x=0;x<8;x++){
      vec4 sampleData=atlasSampleData(cell.xy,vec2(float(x),float(y)));vec3 c=sampleData.rgb;float label=nearestLabel(c,c0,c1,c2,c3);
      wholeVotes+=vec4(step(abs(label-vec4(0.,1.,2.,3.)),vec4(.1)));
      cellEvidence+=sampleData.a;
    }
    float singleMajority=max(max(wholeVotes.x,wholeVotes.y),max(wholeVotes.z,wholeVotes.w))/64.;
    float bestCost=1e20,bestConfidence=0.,bestAxis=0.,bestCut=4.,bestA=0.,bestB=0.;
    for(int axis=0;axis<2;axis++)for(int candidateIndex=0;candidateIndex<5;candidateIndex++){
      float cut=4.;
      if(candidateIndex==1)cut=3.;else if(candidateIndex==2)cut=5.;else if(candidateIndex==3)cut=2.;else if(candidateIndex==4)cut=6.;
      vec4 votesA=vec4(0.),votesB=vec4(0.),weightedVotesA=vec4(0.),weightedVotesB=vec4(0.);
      for(int y=0;y<8;y++)for(int x=0;x<8;x++){
        vec4 sampleData=atlasSampleData(cell.xy,vec2(float(x),float(y)));vec3 c=sampleData.rgb;float label=nearestLabel(c,c0,c1,c2,c3);
        vec4 oneHot=vec4(step(abs(label-vec4(0.,1.,2.,3.)),vec4(.1)));float sampleWeight=1.+3.*sampleData.a;
        float coordinate=axis==0?float(x)+.5:float(y)+.5;
        if(coordinate<cut){votesA+=oneHot;weightedVotesA+=oneHot*sampleWeight;}
        else{votesB+=oneHot;weightedVotesB+=oneHot*sampleWeight;}
      }
      float labelA=0.,labelB=0.,pairCost=1e20,pairVotes=0.;
      for(int a=0;a<4;a++)for(int b=0;b<4;b++){
        if(a!=b&&votesA[a]>0.&&votesB[b]>0.){
          vec3 colorA=a==0?c0:(a==1?c1:(a==2?c2:c3));
          vec3 colorB=b==0?c0:(b==1?c1:(b==2?c2:c3));
          float candidateCost=sideColorCost(weightedVotesA,colorA,c0,c1,c2,c3)+sideColorCost(weightedVotesB,colorB,c0,c1,c2,c3);
          if(candidateCost<pairCost){pairCost=candidateCost;labelA=float(a);labelB=float(b);pairVotes=votesA[a]+votesB[b];}
        }
      }
      if(pairCost<1e20){
        float accuracy=pairVotes/64.,gain=accuracy-singleMajority;
        float oldConfidence=smoothstep(.55,.75,accuracy)*smoothstep(.015625,.078125,gain);
        float confidence=.85+.15*oldConfidence;
        confidence=mix(confidence,1.,smoothstep(.02,.06,cellEvidence/64.));
        if(pairCost<bestCost){bestCost=pairCost;bestConfidence=confidence;bestAxis=float(axis);bestCut=cut;bestA=labelA;bestB=labelB;}
      }
    }
    metadata=vec4((bestAxis*8.+bestCut)/255.,bestA/255.,bestB/255.,bestConfidence);
  }

  vec3 result=c0;
  if(paletteIndex>.5&&paletteIndex<1.5)result=c1;
  else if(paletteIndex>1.5&&paletteIndex<2.5)result=c2;
  else if(paletteIndex>2.5&&paletteIndex<3.5)result=c3;
  gl_FragColor=paletteIndex>3.5?metadata:vec4(result,1.);
}