// Color Cast Correction tool for ART
// Copyright (c) 2023 Ohnishi, Yasuo

// @ART-colorspace: "rec2020"
// @ART-label: "相対色領域補正"


// @ART-param: ["param_choice", "対象色領域", ["赤", "緑", "青", "シアン", "マゼンタ", "黄"]]
// @ART-param: ["R", "Red (シアン <--> 赤)", -1.0, 1.0, 0.0, 0.001, "対象色領域のRGB補正"]
// @ART-param: ["G", "Green (マゼンタ <--> 緑)", -1.0, 1.0, 0.0, 0.001, "対象色領域のRGB補正"]
// @ART-param: ["B", "Blue  (黄 <--> 青)", -1.0, 1.0, 0.0, 0.001, "対象色領域のRGB補正"]
// @ART-param: ["L", "明るさ  (暗 <--> 明)", -1.0, 1.0, 0.0, 0.001, "対象色領域のRGB補正"]
// @ART-param: ["Pivot", "狭める <--> 拡げる", -0.25, 0.25, 0.0, 0.001, "対象色の閾値調整"]
// @ART-param: ["Skew", "", -0.5, 0.5, 0.0, 0.001, "対象色の傾き"]
// @ART-param: ["Strength", "", 0.0, 3.0, 1.0, 0.001, "補正強度"]
// @ART-param: ["param_bool", "対象領域表示", false]


float[3] rgb_in_area(varying float p_rgb[3], float adj_rgb[4], float adj_L,
            float amount, float Strength, bool param_bool)
{
    float rgb[3];
    if (param_bool == false){
        for (int i = 0; i < 3; i = i+1) {
            rgb[i] = p_rgb[i] + adj_rgb[i] * amount * 4 * Strength;
            rgb[i] = rgb[i] * adj_L;
        }
    }else{
        rgb[0] = amount * 4 * Strength;
        rgb[1] = amount * 4 * Strength;
        rgb[2] = 0.0;
    }
    return rgb;
}


float[3] rgb_out_area(varying float p_rgb[3], bool param_bool)
{
    float rgb[3];
    if (param_bool == false){
        for (int i = 0; i < 3; i = i+1) {
            rgb[i] = p_rgb[i];
        }
    }else{
        for (int i = 0; i < 3; i = i+1) {
            rgb[i] = 0.0;
        }
    }
    return rgb;
}



float gamma22(varying float v)
{
    float v2;
    v2 = pow(v, 1 / 2.2);
    return v2;
}

float[3] colorCastCr(varying float r, varying float g,varying float b,
                    float adj_rgb[4], int param_choice, bool primary,
                    float Pivot, float Skew, float Strength,
                    bool param_bool)

{
    float p_rgb[3] = {r, g, b};
    float rgb[3];
    float Skew1;
    float Skew2;
    float adj_L;
    Skew1 = Skew + 0.5;
    Skew2 = 1.0 - Skew1;
    adj_L = adj_rgb[3] + 1.0;
    
    if (param_choice == 0) { // Red
        if (gamma22(r + Pivot) > gamma22(g * Skew1 + b * Skew2)) {
            float amount = gamma22(r + Pivot) - gamma22(g * Skew1 + b * Skew2);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        }else  {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }else if  (param_choice == 1){ // Green
        if (gamma22(g + Pivot) > gamma22(r * Skew1 + b * Skew2)) {
            float amount = gamma22(g + Pivot) - gamma22(r * Skew1 + b * Skew2);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        } else {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }else if  (param_choice == 2){ // Blue
        if (gamma22(b + Pivot) > gamma22(r * Skew1 + g * Skew2)) {
            float amount = gamma22(b + Pivot) - gamma22(r * Skew1 + g * Skew2);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        } else {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }else if  (param_choice == 3){ // Cyan
        if (gamma22(r - Pivot) < gamma22(g * Skew2 + b * Skew1)) {
            float amount = gamma22(g * Skew2 + b * Skew1) - gamma22(r - Pivot);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        } else {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }else if  (param_choice == 4){ // Magenta
        if (gamma22(g - Pivot) < gamma22(r * Skew2 + b * Skew1)) {
            float amount = gamma22(r * Skew2 + b * Skew1) - gamma22(g - Pivot);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        } else {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }else if  (param_choice == 5){ // Yellow
        if (gamma22(b - Pivot) < gamma22(g * Skew2 + r * Skew1)) {
            float amount = gamma22(g * Skew2 + r * Skew1) - gamma22(b - Pivot);
            rgb = rgb_in_area(p_rgb, adj_rgb, adj_L, 
                            amount, Strength, param_bool);
        } else {
            rgb = rgb_out_area(p_rgb, param_bool);
        }
    }

    return rgb;
}


void ART_main(varying float r, varying float g, varying float b,
              float R, float G, float B, float L, 
              float Pivot, float Skew, float Strength,
              output varying float rout,
              output varying float gout,
              output varying float bout,
              int param_choice, bool param_bool)
{
    float rgb[3];
    bool primary;
    float adj_rgb[4] = {R, G, B, L};



    rgb = colorCastCr(r, g, b, adj_rgb, param_choice, 
                    primary, Pivot, Skew, Strength, param_bool);

    rout = rgb[0];
    gout = rgb[1];
    bout = rgb[2];
}