meanReversion    = 0.1
volatilityTimes  = np.array([ 2.0,   5.0,   10.0,   20.0   ])
volatilityValues = np.array([ 0.005, 0.005,  0.005,  0.005 ])
hwModel          = HullWhiteModel(discCurve,meanReversion,volatilityTimes,volatilityValues)

# calibration targets
expiries  = np.array([ 2, 5, 10, 20 ])
swapterms = np.array([ 2, 5, 10, 20 ])
marketVols = np.zeros([len(expiries),len(swapterms)])
swaptions = []
for i in range(expiries.shape[0]):
    swaptionLine = []
    for j in range(swapterms.shape[0]):
        marketVols[i][j] = atmVols.loc[str(expiries[i])+'y',str(swapterms[j])+'y']
        swaptionLine.append(createSwaption(str(expiries[i])+'y',str(swapterms[j])+'y',discCurve,projCurve))
    swaptions.append(swaptionLine)

def objective(sigma):
    hwModel = HullWhiteModel(discCurve,meanReversion,volatilityTimes,sigma)
    modelVols = np.zeros(expiries.shape[0]*swapterms.shape[0])
    for i in range(expiries.shape[0]):
        for j in range(swapterms.shape[0]):
            modelVols[i*swapterms.shape[0]+j] = swaptions[i][j].npvHullWhite(hwModel,'v')
            modelVols[i*swapterms.shape[0]+j] -= marketVols[i][j]
    return modelVols        

#opt = least_squares(objective,volatilityValues,bounds=(0.1*volatilityValues,10*volatilityValues))
opt = least_squares(objective,volatilityValues,method='lm')
print(opt.x)
# swaption(s)

maturity = 20  # in years
swaptions = []
terms = []
inputVols = []
for k in range(1, maturity):
    expTerm = str(k) + "y"
    swpTerm = str(maturity - k) + "y"
    sigma = 0.01  # atmVols.loc[expTerm,swpTerm]
    terms.append(expTerm + "-" + swpTerm)
    inputVols.append(sigma)
    swaptions.append(
        createSwaption(
            expTerm, swpTerm, discCurve, projCurve, 0.03, ql.VanillaSwap.Receiver, sigma
        )
    )

bermudanSwaption = BermudanSwaption(swaptions, meanReversion)

table = pandas.DataFrame(
    [
        np.array(terms),
        np.array(inputVols),
        bermudanSwaption.model.volatilityTimes,
        bermudanSwaption.model.volatilityValues,
        bermudanSwaption.swaptionsNPV(),
        bermudanSwaption.bondOptionsNPV(),
    ]
).T
def objective(sigma):
    tmpModel = HullWhiteModel(discCurve, meanReversion, np.array([30.0]),
                              np.array([sigma]))
    return (createSwaption("10y", "10y", discCurve, projCurve).npvHullWhite(
        tmpModel, "v") - 0.01)  # we calibrate to 100bp 10y-10y vols
    4.30e-2,
    4.30e-2,
    4.30e-2,
]
rates2 = [r + 0.005 for r in rates]

discCurve = YieldCurve(terms, rates)
projCurve = YieldCurve(terms, rates2)

# Hull-White model mean reversion
meanReversion = 0.05

# first we have a look at the model-implied volatility smile

fig = plt.figure(figsize=(4, 6))
atmRate = createSwaption("10y", "10y", discCurve, projCurve).fairRate()
relStrikes = [-0.03 + 1e-3 * k for k in range(61)]
hwVols = [0.0050, 0.0075, 0.0100, 0.0125]
for hwVol in hwVols:
    hwModel = HullWhiteModel(discCurve, meanReversion, np.array([30.0]),
                             np.array([hwVol]))
    normalVols = [
        createSwaption("10y", "10y", discCurve, projCurve,
                       atmRate + strike).npvHullWhite(hwModel, "v") * 1e4
        for strike in relStrikes
    ]
    plt.plot(
        relStrikes,
        normalVols,
        label="a=" + str(meanReversion) + ", sigma_r=" + str(hwVol),
    )
flatCurve = YieldCurve(['30y'],[0.03])

terms = [    '1y',    '2y',    '3y',    '4y',    '5y',    '6y',    '7y',    '8y',    '9y',   '10y',   '12y',   '15y',   '20y',   '25y',   '30y', '50y'   ] 
rates = [ 2.70e-2, 2.75e-2, 2.80e-2, 3.00e-2, 3.36e-2, 3.68e-2, 3.97e-2, 4.24e-2, 4.50e-2, 4.75e-2, 4.75e-2, 4.70e-2, 4.50e-2, 4.30e-2, 4.30e-2, 4.30e-2 ] 
rates2 = [ r+0.005 for r in rates ]

discCurve = YieldCurve(terms,rates)
projCurve = YieldCurve(terms,rates2)

# Hull-White model mean reversion
meanReversion    = 0.05

# first we have a look at the model-implied volatility smile

fig = plt.figure(figsize=(4, 6))
atmRate = createSwaption('10y','10y',discCurve,projCurve).fairRate()
relStrikes = [ -0.03+1e-3*k for k in range(61)]
hwVols = [ 0.0050, 0.0075, 0.0100, 0.0125 ]
for hwVol in hwVols:
    hwModel = HullWhiteModel(discCurve,meanReversion,np.array([30.0]),np.array([hwVol]))
    normalVols = [ createSwaption('10y','10y',discCurve,projCurve,atmRate+strike).npvHullWhite(hwModel,'v')*1e+4
                   for strike in relStrikes ]
    plt.plot(relStrikes, normalVols, label='a='+str(meanReversion)+', sigma_r='+str(hwVol))
plt.ylim(0,250)
plt.legend()
plt.xlabel('Strike (relative to ATM)')
plt.ylabel('Model-implied normal volatility (bp)')
plt.show()

# since Hull White smile is essentially flat we now consider ATM swaptions
# we set up ATM swaptions on a grid of expiry and swap terms
    4.50e-2, 4.75e-2, 4.75e-2, 4.70e-2, 4.50e-2, 4.30e-2, 4.30e-2
]

terms = ['30y']  # flat curve
rates = [5.00e-2]

rates2 = [r + 0.005 for r in rates]

discCurve = YieldCurve(terms, rates)
projCurve = YieldCurve(terms, rates)

a = 0.03  # mean reversion
vol = 0.01
h = 1.0e-12

swaptionStrike = createSwaption('10y','10y',discCurve,projCurve).fairRate() \
                 + vol * np.sqrt(10.0) * np.sqrt(3.0)
swaptionStrike = 0.1048
print('swaptionStrike = ' + str(swaptionStrike))

res = []
for h in [10.0**(-k) for k in range(3, 18)]:
    print('h = ' + str(h))
    # swaptions
    s0 = createSwaption('10y',
                        '10y',
                        discCurve,
                        projCurve,
                        strike=swaptionStrike,
                        normalVolatility=vol)
    sp = createSwaption('10y',