Machine Learning: SuperTrend Strategy TP/SL [YinYangAlgorithms]
// This source code is subject to the terms of the Mozilla Public License 2.0 at https://mozilla.org/MPL/2.0/
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@ @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@    ,@@@@@@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@      @@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@        @@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@         @@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@           @@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@        .@@@@@@@@@@@@@@@            @@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@          *@@@@@@@@@@@@@@             @@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@         @@@@@@@@@@@@@@@               @@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@     @@@@@@@@@@@@@@@@                 @@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                  @@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@.                    @@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                      @@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@.                         @
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                             @
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@,                                       @
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                                                @
// @@@@@@@@@@@@@@@@@@@@@@@@@@@                                                    @
// @@@@@@@@@@@@@@@@@@@@@@@@@                                                     @@
// @@@@@@@@@@@@@@@@@@@@@@@                                                       @@
// @@@@@@@@@@@@@@@@@@@@@@                                                       @@@
// @@@@@@@@@@@@@@@@@@@@@*                @@@@@                                 @@@@
// @@@@@@@@@@@@@@@@@@@@@               @@@@@@@@@                              @@@@@
// @@@@@@@@@@@@@@@@@@@@@              @@@@@@@@@@@                           @@@@@@@
// @@@@@@@@@@@@@@@@@@@@@               @@@@@@@@%                           @@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@                                                @@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@                                            @@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@                                        %@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@                                   @@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@                           @@@@@@@@@@@@@@@@@@@@@@@@
// @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@                @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
// © YinYangAlgorithms

//@version=5
strategy("Machine Learning: SuperTrend Strategy - Trailing Take Profit and Stop Loss [YinYangAlgorithms]", overlay=true, initial_capital=1000000, default_qty_type=strategy.percent_of_equity, default_qty_value=100)
// ~~~~~~~~~~~~ INPUTS ~~~~~~~~~~~~ //
//Supertrend
atrPeriod = input.int(4, "ATR Length", minval = 1, group="Supertrend", tooltip="ATR Length used to create the Original Supertrend.")
factor = input.float(2.94, "Factor", minval = 0.01, step = 0.01, group="Supertrend", tooltip="Multiplier used to create the Original Supertrend.")
stoplossMult = input.float(0.0000025, "Stop Loss Multiplier", group="Supertrend", tooltip="0 = Don't use Stop Loss. Stop loss can be useful for helping to prevent false signals but also may result in more loss when hit and less profit when switching trends.")
takeprofitMult = input.float(0.022, "Take Profit Multiplier", group="Supertrend", tooltip="Take Profits can be useful within the Supertrend Strategy to stop the price reverting all the way to the Stop Loss once it's been profitable.")
//Machine Learning
onlyUseSameTrendDirections = input.bool(true, "Only Factor Same Trend Direction", group="Machine Learning", tooltip="Very useful for ensuring that data used in KNN is not manipulated by different SuperTrend Directional data. Please note, it doesn't effect KNN Exponential.")
useRationalQudratics = input.string("All", "Rationalized Source Type", options=["All", "Close", "HL2", "ATR", "None"], group="Machine Learning", tooltip="Should we Rationalize only a specific source, All or None?")
useMachineLearning = input.string("KNN Average", "Machine Learning Type", options=["KNN Average", "KNN Exponential Average", "Simple Average", "None"], group="Machine Learning", tooltip="Are we using a Simple ML Average, KNN Mean Average, KNN Exponential Average or None?")
distanceType = input.string("Both", "KNN Distance Type", options=["Both", "Max", "Min"], group="Machine Learning", tooltip="We need to check if distance is within the KNN Min/Max distance, which distance checks are we using.")
mlType = input.string("SMA", "Machine Learning Smoothing Type", options=["SMA", "EMA", "VWMA"], group="Machine Learning" ,tooltip="How should we smooth our Fast and Slow ML Datas to be used in our KNN Distance calculation? SMA, EMA or VWMA?")
mlLength = input.int(100, "Machine Learning Length", group="Machine Learning", tooltip="How far back is our Machine Learning going to keep data for.")
knnLength = input.int(2, "k-Nearest Neighbour (KNN) Length", group="Machine Learning", tooltip="How many k-Nearest Neighbours will we account for?")
fastLength = input.int(3, "Fast ML Data Length", group="Machine Learning", tooltip="What is our Fast ML Length?? This is used with our Slow Length to create our KNN Distance.")
slowLength = input.int(50, "Slow ML Data Length", group="Machine Learning", tooltip="What is our Slow ML Length?? This is used with our Fast Length to create our KNN Distance.")

// ~~~~~~~~~~~~ VARIABLES ~~~~~~~~~~~~ //
rationalizeHL2 = useRationalQudratics == "HL2" or useRationalQudratics == "All"
rationalizeClose = useRationalQudratics == "Close" or useRationalQudratics == "All"
rationalizeATR = useRationalQudratics == "ATR" or useRationalQudratics == "All"
distanceTypeMin = distanceType == "Min" or distanceType == "Both"
distanceTypeMax = distanceType == "Max" or distanceType == "Both"

// ~~~~~~~~~~~~ FUNCTIONS ~~~~~~~~~~~~ //
//@jdehorty Kernel Function
//used to turn a source into a rational quadratic which performs better in ML calculations
rationalQuadratic(_src, _lookback, _relativeWeight, startAtBar) =>
	float _currentWeight = 0.
	float _cumulativeWeight = 0.
	_size = array.size(array.from(_src))
    for i = 0 to _size + startAtBar
        y = _src[i]
        w = math.pow(1 + (math.pow(i, 2) / ((math.pow(_lookback, 2) * 2 * _relativeWeight))), -_relativeWeight)
        _currentWeight += y*w
        _cumulativeWeight += w
    yhat = _currentWeight / _cumulativeWeight
    yhat

//Calculate a traditional Supertrend but with the ability to replace its sources with Rationalized versions as well as with a Stop loss
getSupertrend(factor, atrPeriod) =>
    src = rationalizeHL2 ? rationalQuadratic(hl2, 8, 8., 25) : hl2
    atr = rationalizeATR ? rationalQuadratic(ta.atr(atrPeriod), 8, 8., 25) : ta.atr(atrPeriod)
    upperBand = src + factor * atr
    lowerBand = src - factor * atr
    prevLowerBand = nz(lowerBand[1])
    prevUpperBand = nz(upperBand[1])

    closeSrc = rationalizeClose ? rationalQuadratic(close, 8, 8., 25) : close

    lowerBand := lowerBand > prevLowerBand or closeSrc[1] < prevLowerBand ? lowerBand : prevLowerBand
    upperBand := upperBand < prevUpperBand or closeSrc[1] > prevUpperBand ? upperBand : prevUpperBand
    int direction = na
    float superTrend = na
    float stoploss = direction[1] == -1 ? lowerBand - (lowerBand * stoplossMult) : upperBand + (upperBand * stoplossMult)
    prevSuperTrend = superTrend[1]
    if na(atr[1])
        direction := 1
    else if prevSuperTrend == prevUpperBand
        direction := closeSrc > stoploss ? -1 : 1
    else
        direction := closeSrc < stoploss ? 1 : -1
    superTrend := direction == -1 ? lowerBand : upperBand
    [superTrend, direction, stoploss]

//Get the exponential average of an array, where the exponential weight is focused on the first value of this array (current source)
getExponentialDataAverage(_data, _length) =>
    avg = 0.
    maxLen = math.min(_data.size(), _length)
    if maxLen > 0
        for i = 0 to maxLen - 1
            curData = _data.get(i)
            tempAvg = curData
            if i > 0
                for a = 0 to i
                    tempAvg += array.get(_data, a)
                tempAvg := math.avg(_data.get(0), tempAvg / i)
            avg += math.avg(_data.get(0), math.avg(curData, tempAvg))
        avg / _length
    else
        avg

//Uses KNN to sort distances with our ML fast and slow data
//This is a modified version that sorts distances but rather than saving and outputting the distance average it outputs the mean of the valid distance and averages the total
knnAverage_fromDistance(_currentDirection, _dataFast, _dataSlow, _directions, _minDist, _maxDist) =>
    //failsafe we need at least 1 distance
    maxDist = not _minDist ? true : _maxDist
    //Calculate the distance between slow and fast moving Datas
    distances = array.new_float(0)
    for i = 0 to _dataSlow.size() - 1
        if not onlyUseSameTrendDirections or _currentDirection == _directions.get(i)
            distance = _dataSlow.get(i) - _dataFast.get(i)
            distances.push(distance)
    //clone the array so it doesn't modify it
    clonedDistances = distances.copy()    
    //slice the length from the array and calculate the max value from this slice
    splicedDistances = clonedDistances.slice(0, math.min(knnLength, clonedDistances.size()))
    maxDistanceAllowed = splicedDistances.max()
    minDistanceAllowed = splicedDistances.min()
    //scan all distances and add any that are less than max distance
    validDistances = array.new_float(0)
    for i = 0 to distances.size() - 1
        if (not maxDist or distances.get(i) <= maxDistanceAllowed) and (not _minDist or distances.get(i) >= minDistanceAllowed)
         and (not onlyUseSameTrendDirections or _currentDirection == _directions.get(i))
            distAvg = (_dataSlow.get(i) + _dataFast.get(i)) / 2
            validDistances.push(distAvg)

    // Get exponential or regular average
    if useMachineLearning == "KNN Exponential Average"
        getExponentialDataAverage(validDistances, 1)
    else
        validDistances.avg()

// ~~~~~~~~~~~~ CALCULATIONS ~~~~~~~~~~~~ //
//Get our current supertrend
[supertrendOG, directionOG, stoplossOG] = getSupertrend(factor, atrPeriod)

//Create save variables based on our current Supertrend
float supertrend = supertrendOG
int direction = directionOG
float stoploss = stoplossOG

//Apply Machine Learning logic to our Supertrend if selected
if useMachineLearning != "None"
    //Save our directions so we can filter data based on directions (removes inconsistencies)
    directions = array.new_int(mlLength)
    //populate our directions
    for d = 0 to mlLength - 1
        directions.set(d, direction[d])

    if useMachineLearning == "Simple Average"
        // -- A simple machine Learning Average
            //essentially just a simple Average of supertrend data (potentially filtered to account for directions)
        supertrendData = array.new_float(0)
        for i = 0 to mlLength - 1
            if not onlyUseSameTrendDirections or direction == directions.get(i)
                supertrendData.push(supertrend[i])
        supertrend := supertrendData.avg()
    else
        // -- A still simple but more complex approach to Machine Learning using KNN to sort validDistances
        //Calculate our fast and slow MA's based on the current Supertrend
        float supertrendFast = mlType == "SMA" ? ta.sma(supertrend, fastLength) : mlType == "EMA" ? ta.ema(supertrend, fastLength) : ta.vwma(supertrend, fastLength)
        float supertrendSlow = mlType == "SMA" ? ta.sma(supertrend, slowLength) : mlType == "EMA" ? ta.ema(supertrend, slowLength) : ta.vwma(supertrend, slowLength)
        //create storage data for ML lookbacks
        supertrendFastData = array.new_float(mlLength)
        supertrendSlowData = array.new_float(mlLength)
        //populate our ML storage with lookbacks at our Fast and Slow Supertrend MAs
        for i = 0 to mlLength - 1
            supertrendFastData.set(i, supertrendFast[i])
            supertrendSlowData.set(i, supertrendSlow[i])
        //calculate our new Supertrend using KNN by using distances within KNN min/max and filtering by direction
        supertrend := knnAverage_fromDistance(direction, supertrendFastData, supertrendSlowData, directions, distanceTypeMin, distanceTypeMax)

    //Recalculate the Stop loss based on new supertrend (this is used in the direction recalculation below so needs to be done before)
    stoploss := direction == -1 ? supertrend - (supertrend * stoplossMult) : supertrend + (supertrend * stoplossMult)
    //Recalculate the direction based on the new supertrend
    direction := close > stoploss ? -1 : 1

//Calculate the take profit location
takeprofit = direction == -1 ? supertrend + (supertrend * takeprofitMult) : supertrend - (supertrend * takeprofitMult)

// ~~~~~~~~~~~~ PLOTS ~~~~~~~~~~~~ //
//Plot all of our supertrend data
supertrend := barstate.isfirst ? na : supertrend
bodyMiddle = plot(barstate.isfirst ? na : (open + close) / 2, display = display.none)
upTrend = plot(direction < 0 ? supertrend : na, "Up Trend", color = color.green, style = plot.style_linebr)
downTrend = plot(direction < 0 ? na : supertrend, "Down Trend", color = color.red, style = plot.style_linebr)
upTrendStoploss = plot(stoplossMult != 0 ? direction < 0 ? stoploss : na : na, "Stop Loss", color =color.green, style = plot.style_linebr)
downTrendStoploss = plot(stoplossMult != 0 ? direction < 0 ? na : stoploss : na, "Stop Loss", color = color.red, style = plot.style_linebr)
upTrendTakeprofit = plot(takeprofitMult != 0 ? direction < 0 ? takeprofit : na : na, "Stop Loss", color =color.rgb(16, 105, 179), style = plot.style_linebr)
downTrendTakeprofit = plot(takeprofitMult != 0 ? direction < 0 ? na : takeprofit : na, "Stop Loss", color = color.rgb(206, 113, 1), style = plot.style_linebr)
//Fill all of our supertrend data
fill(bodyMiddle, upTrend, color.new(color.green, 90))
fill(bodyMiddle, downTrend, color.new(color.red, 90))
fill(upTrendStoploss, upTrend, color.new(color.green, 95))
fill(downTrendStoploss, downTrend, color.new(color.red, 95))

// ~~~~~~~~~~~~ STRATEGY ~~~~~~~~~~~~ //
//Trend Change (includes Stop Loss)
if direction != direction[1]
    if direction == -1
        strategy.entry("Buy", strategy.long)
    else
        strategy.entry("Sell", strategy.short)

//Take Profit
if takeprofitMult != 0
    if direction == -1 and close > takeprofit
        strategy.close("Buy")
    else if direction == 1 and close < takeprofit
        strategy.close("Sell")

// ~~~~~~~~~~~~ ALERTS ~~~~~~~~~~~~ //
if (direction != direction[1]) or (takeprofitMult != 0 and ((direction == -1 and close > takeprofit) or (direction == 1 and close < takeprofit)))
    alert("{{strategy.order.action}} | {{ticker}} | {{close}}", alert.freq_once_per_bar)

// ~~~~~~~~~~~~ END ~~~~~~~~~~~~ //