import TensorFlow

extension Array {  
/// A functional version of `Array.subscript.modify`.  
/// Differentiation does yet not support `Array.subscript.modify` because  
/// it is a coroutine.  
@differentiable(where Element: Differentiable)  
func updated(at index: Int, with newValue: Element) -> [Element] {  
var result = self  
result[index] = newValue  
return result  
}  
}

extension Array where Element: Differentiable {  
/// An inefficient derivative for a functional version of `Array.subscript.modify`.  
@derivative(of: updated)  
func vjpUpdated(at index: Int, with newValue: Element) -> (  
value: [Element], pullback: (TangentVector) -> (TangentVector, Element.TangentVector)  
) {  
let value = updated(at: index, with: newValue)  
return (value, { [count = count] v in  
// `dself`: `v` but with `v[index] = 0`.  
// `dnewValue`: `v[index]`.  
var dself = v  
dself.base[index] = .zero  
return (dself, v[index])  
})  
}  
}

struct PathCalc : Differentiable {  
@noDerivative let N: Int  
@noDerivative let Nmat: Int  
@noDerivative let Delta: Float  
@noDerivative let Z: Float

@differentiable(wrt: (Rates, Lambda))  
func applied(Rates:Array<Float>,Lambda:Array<Float> ) -> Array<Float> {

var L = Array<Float>(repeating: 0.0,count: N)  
let sq = Delta.squareRoot()

for n in 0...Nmat-1{  
let sqez = sq\*Z  
var v = Float(0)

for i in n+1...N-1{  
let lam = Lambda[i - n - 1 ]  
let con1 = Delta\*lam  
v = v + (con1\*Rates[i])/(1.0 + Delta\*L[i])  
let first = con1\*v  
let second = lam\*(sqez - 0.5\*con1)  
let vrat = exp(first + second)  
L[i] = Rates[i]\*vrat

}  
}  
return L  
}

}

struct Portfolio: Differentiable {

@noDerivative let N: Int  
@noDerivative let Nmat: Int  
@noDerivative let Nopt: Int  
@noDerivative let Delta: Float  
@noDerivative let Maturities: Array<Int>  
@noDerivative let SwapRates: Array<Float>  
@noDerivative let Z: Float

@differentiable(wrt: (self, Rates, Lambda))  
func calculate(Rates:Array<Float>,Lambda:Array<Float>) -> Float {

let pathCalcModel = PathCalc(N: N, Nmat: Nmat, Delta: Delta, Z:Z)  
let L = pathCalcModel.applied(Rates:Rates , Lambda:Lambda )

var S = Array<Float>(repeating: 0.0,count: N)  
var B = Array<Float>(repeating: 0.0,count: N)

var b = Float(1.0)  
var s = Float(0)

for n in Nmat...N-1{  
b = b / (1.0 + Delta\*L[n])  
s = s + Delta\*b  
B[n] = b  
S[n] = s  
}

var v = Float(0)  
var swapval = Float(0)

for i in 0...Nopt-1{  
let m = Int(Maturities[i] + Nmat - 1)  
swapval = B[m] + SwapRates[i]\*S[m] - 1.0  
if (swapval < 0.0){  
v = v + -100.0\*swapval  
}  
}  
// Apply discount  
for n in 0...Nmat-1{  
v = v / (1.0 + Delta\*L[n])

}

return v  
}

}

let Nmat = 40  
let N = Nmat + 40  
let Nopt = 15  
let maturities = Array<Int>([4,4,4,8,8,8,20,20,20,28,28,28,40,40,40])  
let swapRates = Array<Float>([0.045,0.05,0.055,0.045,0.05,0.055,0.045,0.05,  
0.055,0.045,0.05,0.055,0.045,0.05,0.055])

let rates = Array<Float>(repeating: 0.05,count: N)  
let lambdas = Array<Float>(repeating: 0.2,count: N)

let portfolioModel = Portfolio(N:N , Nmat:Nmat , Nopt: Nopt , Delta:0.25,  
Maturities: maturities , SwapRates:swapRates , Z:0.3 )  
let portfolioValue = portfolioModel.calculate(Rates:rates , Lambda:lambdas )  
print(portfolioValue)

let (g_model, g_rates,g_lambdas) = gradient(at: portfolioModel, rates, lambdas)  
{ portfolioModel, rates, lambdas in  
portfolioModel.calculate(Rates: rates, Lambda:lambdas)  
}  
print(g_model)  
print(g_rates)

    _modify {
      _makeMutableAndUnique() // makes the array native, too
      _checkSubscript_mutating(index)
      let address = _buffer.mutableFirstElementAddress + index
      defer { _endMutation() }
      yield &address.pointee
    }

    _modify {
      let (bucket, found) = _variant.mutatingFind(key)
      let native = _variant.asNative
      if !found {
        let value = defaultValue()
        native._insert(at: bucket, key: key, value: value)
      }
      let address = native._values + bucket.offset
      defer { _fixLifetime(self) }
      yield &address.pointee
    }

if (foo) {
  yield &something
} else
  yield &other
}

_modify {
  var y : Float;
  yield &y;
}