# ------------------------------------------------------------------------------
# Channel usage examples
# coding: utf-8
# Copyright (c) 2017 The Foundry Visionmongers Ltd.  All Rights Reserved.
# ------------------------------------------------------------------------------


import mari

# ------------------------------------------------------------------------------
# Register Optimized "ACES" GLSL Shader Transforms:
# ------------------------------------------------------------------------------

def registerOptimizedAcesTransforms(aces_config_path, debug=False):
    """ Registers opimized GLSL shader transform code for the 'aces' OCIO config.
        This method can not only significantly improve the perfromance within Mari
        but can also improve the accuracy of the transform itself due to the ability
        to remove the need for a LUT.

        @param aces_config_path : (str)
            Path to the aces config file. This is the value that would be returned
            from mari.ColorspaceConfig.resolveFileName()
        @param debug=False : (bool)
            When True, will set the colorspace transform from "Utility - sRGB - Texture"
            to "ACES - ACSEcg" to output 100% green for debugging purposes.
    """

    mari.ocio.setShaderTransformCode(aces_config_path,
                                     'ACES - ACEScg',
                                     'ACES - ACES2065-1',
                                     '    mat3 ACES_AP1_TO_AP0 = mat3( 0.6954522414, 0.1406786965, 0.1638690622,\n'
                                     '                                 0.0447945634, 0.8596711185, 0.0955343182,\n'
                                     '                                -0.0055258826, 0.0040252103, 1.0015006723);\n'
                                     '    vec3 v1 = #Input.rgb * ACES_AP1_TO_AP0;\n'
                                     '    #Output = vec4(v1, #Input.a);\n')
    mari.ocio.setShaderTransformCode(aces_config_path,
                                     'ACES - ACES2065-1',
                                     'ACES - ACEScg',
                                     '    mat3 ACES_AP0_TO_AP1 = mat3( 1.4514393161, -0.2365107469, -0.2149285693,\n'
                                     '                                -0.0765537734,  1.1762296998, -0.0996759264,\n'
                                     '                                 0.0083161484, -0.0060324498,  0.9977163014);\n'
                                     '    vec3 v1 = #Input.rgb * ACES_AP0_TO_AP1;\n'
                                     '    #Output = vec4(v1, #Input.a);\n')

    mari.ocio.setShaderTransformCode(aces_config_path,
                                     'Utility - sRGB - Texture',
                                     'ACES - ACES2065-1',
                                     '    vec3 v1 = #Input.rgb;\n'
                                     '    bvec3 mask = lessThan(v1, vec3(0.04045));\n'
                                     '    vec3 v2 = v1 / 12.92;\n'
                                     '    v1 = pow((v1 + vec3(0.055)) / 1.055, vec3(2.4));\n'
                                     '    v1 = mix(v1, v2, mask);\n'
                                     '\n'
                                     '    mat3 ACES_Rec709_to_AP0 = mat3(0.4396467854, 0.3829813768, 0.1773715911,\n'
                                     '                                   0.0897802147, 0.8134397885, 0.0967796810,\n'
                                     '                                   0.0175446202, 0.1115570308, 0.8708982513);\n'
                                     '    v1 = v1 * ACES_Rec709_to_AP0;\n'
                                     '\n'
                                     '    #Output = vec4(v1, #Input.a);\n')
    mari.ocio.setShaderTransformCode(aces_config_path,
                                     'ACES - ACES2065-1',
                                     'Utility - sRGB - Texture',
                                     '    mat3 ACES_AP0_TO_Rec709 = mat3( 2.5216061325, -1.1340673962, -0.3875385104,\n'
                                     '                                   -0.2764821365,  1.3727177396, -0.0962352472,\n'
                                     '                                   -0.0153830779, -0.1529909300,  1.1683740700);\n'
                                     '    vec3 v1 = #Input.rgb * ACES_AP0_TO_Rec709;\n'
                                     '\n'
                                     '    bvec3 mask = lessThan(v1, vec3(0.04045 / 12.92));\n'
                                     '    vec3 v2 = 12.92 * v1;\n'
                                     '    v1 = 1.055 * pow(v1, vec3(1.0 / 2.4)) - vec3(0.055);\n'
                                     '\n'
                                     '    #Output = vec4(mix(v1, v2, mask), #Input.a);\n')

    # This is a useful method of debugging that the registration has succeeded. By passing in the 'debug' argument as
    # True the most common transform in the 'aces' config will result in blue, meaning things like your diffuse
    # channel's will appear blue.
    if debug:
        mari.ocio.setShaderTransformCode(aces_config_path,
                                         'Utility - sRGB - Texture',
                                         'ACES - ACEScg',
                                         '    vec3 v1 = #Input.rgb;\n'
                                         '    v1 = vec3(0.3, 0.3, 1.0) * v1;\n'
                                         '    #Output = vec4(v1, #Input.a);\n')
    else:
        mari.ocio.setShaderTransformCode(aces_config_path,
                                         'Utility - sRGB - Texture',
                                         'ACES - ACEScg',
                                         '    vec3 v1 = #Input.rgb;\n'
                                         '    bvec3 mask = lessThan(v1, vec3(0.04045));\n'
                                         '    vec3 v2 = v1 / 12.92;\n'
                                         '    v1 = pow((v1 + vec3(0.055)) / 1.055, vec3(2.4));\n'
                                         '    v1 = mix(v1, v2, mask);\n'
                                         '\n'
                                         '    mat3 ACES_Rec709_TO_AP1 = mat3(0.61311580, 0.33951018, 0.04737375,\n'
                                         '                                   0.07019675, 0.91635382, 0.01344908,\n'
                                         '                                   0.02061912, 0.10958016, 0.86980061);\n'
                                         '    v1 = v1 * ACES_Rec709_TO_AP1;\n'
                                         '\n'
                                         '    #Output = vec4(v1, #Input.a);\n')
    mari.ocio.setShaderTransformCode(aces_config_path,
                                     'ACES - ACEScg',
                                     'Utility - sRGB - Texture',
                                     '    mat3 ACES_AP1_TO_Rec709 = mat3( 1.70499807, -0.62174865, -0.08324921,\n'
                                     '                                   -0.13025803,  1.14080327, -0.01054488,\n'
                                     '                                   -0.02400765, -0.12898300,  1.15299072);\n'
                                     '    vec3 v1 = #Input.rgb * ACES_AP1_TO_Rec709;\n'
                                     '\n'
                                     '    bvec3 mask = lessThan(v1, vec3(0.04045 / 12.92));\n'
                                     '    vec3 v2 = 12.92 * v1;\n'
                                     '    v1 = 1.055 * pow(v1, vec3(1.0 / 2.4)) - vec3(0.055);\n'
                                     '\n'
                                     '    #Output = vec4(mix(v1, v2, mask), #Input.a);\n')

# ------------------------------------------------------------------------------

if mari.app.isRunning():
    # We want to use the optimized transforms all the time and not just when the user executes the example. The action
    # registered with the menu is really only there to make the user aware of the technique.
    registerOptimizedAcesTransforms('/Path/To/OpenColorIO-Configs/aces_1.0.3/config.ocio', debug=True)