import Foundation
import _Differentiation


//------------------------------------------------------------------------------------------
// Enable differentiable keyPath writes


//called 'writeSlowTo' because inout object wasn't working
@differentiable(where Object: Differentiable, Object == Object.TangentVector, Member: Differentiable, Member == Member.TangentVector)
public func writeSlowTo<Object, Member>(_ object: Object, at member: WritableKeyPath<Object, Member>, with value: Member) -> Object {
    var object = object
    object[keyPath: member] = value
    return object
}


@derivative(of: writeSlowTo)
public func vjpWriteSlowTo<Object, Member>(_ object: Object, at member: WritableKeyPath<Object, Member>, with value: Member) -> (value: Object, pullback: (Object.TangentVector) -> (Object.TangentVector, Member.TangentVector))
where Object: Differentiable, Object == Object.TangentVector, Member: Differentiable, Member == Member.TangentVector
{
    func pullback(_ dself: Object.TangentVector) -> (Object.TangentVector, Member.TangentVector) {
        var dself = dself
        let dWriteValue = dself[keyPath: member]
        dself[keyPath: member] = Member.zero
        return (dself, dWriteValue)
    }
    var object = object
    object = writeSlowTo(object, at: member, with: value)
    return (object, pullback)
}




//------------------------------------------------------------------------------------------
// Usage code
public struct Bar: Differentiable & AdditiveArithmetic {
    public var a: Double = 7
}


struct HoldsKeyPaths<Root, Value> {
    var keyPath: WritableKeyPath<Root, Value>


    @differentiable(where Root: Differentiable, Root == Root.TangentVector, Value: Differentiable, Value == Value.TangentVector)
    public func apply(to root: inout Root, using value: Value) {
        //---
        // force cast to print gradient
        var rootbar = root as! Bar
        rootbar = rootbar.withDerivative{ print("correct derivative", $0) }
        root = rootbar as! Root
        //---
        root = writeSlowTo(root, at: self.keyPath, with: value)
        for _ in 0..<1{} //useless loop, get rid of it and gradient is accurate
    }
}




let bar = Bar()
let keyPath = \Bar.a
let keyPathHolder = HoldsKeyPaths<Bar, Double>(keyPath: keyPath)


func writeToBar(bar: Bar, newValue: Double) -> Double {
    var bar = bar
    bar = bar.withDerivative { print("actual derivative", $0) }
    keyPathHolder.apply(to: &bar, using: newValue)
    return bar.a
}


let newValue: Double = 7
let valAndGrad = valueWithGradient(at: bar, newValue, in: writeToBar)


if valAndGrad.gradient.0.a != 0 {
    print("gradient is incorrect")
}
 
//prints:
//correct derivative Bar(a: 0.0)
//actual derivative Bar(a: 1.0)
//gradient is incorrect