/*
    crt-pi - A Raspberry Pi friendly CRT shader.
    
    Metal Shading Language version converted from GLSL
    Copyright (C) 2015-2016 davej

    This program is free software; you can redistribute it and/or modify it
    under the terms of the GNU General Public License as published by the Free
    Software Foundation; either version 2 of the License, or (at your option)
    any later version.
*/

#include <metal_stdlib>
using namespace metal;

// Haven't put these as parameters as it would slow the code down.
#define SCANLINES
#define MULTISAMPLE  
#define GAMMA
//#define FAKE_GAMMA
//#define CURVATURE
//#define SHARPER
// MASK_TYPE: 0 = none, 1 = green/magenta, 2 = trinitron(ish)
#define MASK_TYPE 2

// Constants
#define CURVATURE_X 0.05
#define CURVATURE_Y 0.1
#define MASK_BRIGHTNESS 0.80
#define SCANLINE_WEIGHT 6.0
#define SCANLINE_GAP_BRIGHTNESS 0.12
#define BLOOM_FACTOR 3.5
#define INPUT_GAMMA 2.4
#define OUTPUT_GAMMA 2.2

struct VertexIn {
    float4 position [[attribute(0)]];
    float2 texcoord [[attribute(1)]];
};

struct VertexOut {
    float4 position [[position]];
    float2 TEX0;
#if defined(CURVATURE)
    float2 screenScale;
#endif
    float filterWidth;
};

vertex VertexOut vertex_main(VertexIn in [[stage_in]]) {
    VertexOut out;
    
#if defined(CURVATURE)
    out.screenScale = float2(1.0, 1.0);
#endif
    out.filterWidth = (768.0 / 256.0) / 3.0;
    out.TEX0 = float2(in.texcoord.x, 1.0 - in.texcoord.y) * 1.0001;
    out.position = in.position;
    
    return out;
}

#if defined(CURVATURE)
float2 Distort(float2 coord, float2 screenScale) {
    float2 CURVATURE_DISTORTION = float2(CURVATURE_X, CURVATURE_Y);
    // Barrel distortion shrinks the display area a bit, this will allow us to counteract that.
    float2 barrelScale = 1.0 - (0.23 * CURVATURE_DISTORTION);
    coord *= screenScale;
    coord -= float2(0.5);
    float rsq = coord.x * coord.x + coord.y * coord.y;
    coord += coord * (CURVATURE_DISTORTION * rsq);
    coord *= barrelScale;
    if (abs(coord.x) >= 0.5 || abs(coord.y) >= 0.5)
        coord = float2(-1.0);		// If out of bounds, return an invalid value.
    else {
        coord += float2(0.5);
        coord /= screenScale;
    }
    
    return coord;
}
#endif

float CalcScanLineWeight(float dist) {
    return max(1.0 - dist * dist * SCANLINE_WEIGHT, SCANLINE_GAP_BRIGHTNESS);
}

float CalcScanLine(float dy, float filterWidth) {
    float scanLineWeight = CalcScanLineWeight(dy);
#if defined(MULTISAMPLE)
    scanLineWeight += CalcScanLineWeight(dy - filterWidth);
    scanLineWeight += CalcScanLineWeight(dy + filterWidth);
    scanLineWeight *= 0.3333333;
#endif
    return scanLineWeight;
}

fragment float4 fragment_main(VertexOut in [[stage_in]], 
                              texture2d<float> Texture [[texture(0)]], 
                              sampler textureSampler [[sampler(0)]]) {
    float2 TextureSize = float2(320.0, 256.0);
#if defined(CURVATURE)
    float2 texcoord = Distort(in.TEX0, in.screenScale);
    if (texcoord.x < 0.0)
        return float4(0.0);
    else
#else
    float2 texcoord = in.TEX0;
#endif
    {
        float2 texcoordInPixels = texcoord * TextureSize;
#if defined(SHARPER)
        float2 tempCoord = floor(texcoordInPixels) + 0.5;
        float2 coord = tempCoord / TextureSize;
        float2 deltas = texcoordInPixels - tempCoord;
        float scanLineWeight = CalcScanLine(deltas.y, in.filterWidth);
        float2 signs = sign(deltas);
        deltas.x *= 2.0;
        deltas = deltas * deltas;
        deltas.y = deltas.y * deltas.y;
        deltas.x *= 0.5;
        deltas.y *= 8.0;
        deltas /= TextureSize;
        deltas *= signs;
        float2 tc = coord + deltas;
#else
        float tempY = floor(texcoordInPixels.y) + 0.5;
        float yCoord = tempY / TextureSize.y;
        float dy = texcoordInPixels.y - tempY;
        float scanLineWeight = CalcScanLine(dy, in.filterWidth);
        float signY = sign(dy);
        dy = dy * dy;
        dy = dy * dy;
        dy *= 8.0;
        dy /= TextureSize.y;
        dy *= signY;
        float2 tc = float2(texcoord.x, yCoord + dy);
#endif

        float3 colour = Texture.sample(textureSampler, tc).rgb;

#if defined(SCANLINES)
#if defined(GAMMA)
#if defined(FAKE_GAMMA)
        colour = colour * colour;
#else
        colour = pow(colour, float3(INPUT_GAMMA));
#endif
#endif
        scanLineWeight *= BLOOM_FACTOR;
        colour *= scanLineWeight;

#if defined(GAMMA)
#if defined(FAKE_GAMMA)
        colour = sqrt(colour);
#else
        colour = pow(colour, float3(1.0/OUTPUT_GAMMA));
#endif
#endif
#endif
#if MASK_TYPE == 0
        return float4(colour, 1.0);
#else
#if MASK_TYPE == 1
        float whichMask = fract((in.position.x * 1.0001) * 0.5);
        float3 mask;
        if (whichMask < 0.5)
            mask = float3(MASK_BRIGHTNESS, 1.0, MASK_BRIGHTNESS);
        else
            mask = float3(1.0, MASK_BRIGHTNESS, 1.0);
#elif MASK_TYPE == 2
        float whichMask = fract((in.position.x * 1.0001) * 0.3333333);
        float3 mask = float3(MASK_BRIGHTNESS, MASK_BRIGHTNESS, MASK_BRIGHTNESS);
        if (whichMask < 0.3333333)
            mask.x = 1.0;
        else if (whichMask < 0.6666666)
            mask.y = 1.0;
        else
            mask.z = 1.0;
#endif

        return float4(colour * mask, 1.0);
#endif
    }
}