#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using DocOpt

const spot = 100.0
const strike = 110.0
const vol = 0.25
const maturity = 2.0

function init(numPaths::Int, numSteps::Int)
    # Pre-compute some values
    dt = maturity / numSteps
    volSqrtdt = vol * sqrt(dt)
    fwdFactor = vol^2 * dt * -0.5
    vsqrtdt_log2e = volSqrtdt * log2(exp(1.0))
    fwdFactor_log2e = fwdFactor * log2(exp(1.0))

    # Storage
    # Compute all per-path asset values for each time step
    # NOTE: hand-hoisting obscures the basic formula.
    # Can the compiler do this optimization?
    asset = Array(Array{Float64,1}, numSteps+1)
    asset[1] = Array(Float64, numPaths)
    fill!(asset[1], spot)
    for s=2:numSteps+1
        asset[s] = asset[s-1] .* exp2(fwdFactor_log2e .+ vsqrtdt_log2e .* randn(numPaths))
    end
    return asset
end

function model(strike, numPaths, numSteps, asset)
    avgPathFactor = 1.0/numPaths
    putpayoff = max(strike .- asset[numSteps+1], 0.0)
    cashFlowPut = putpayoff * avgPathFactor
    # Now go back in time using regression
    for s = numSteps:-1:2
        curasset = asset[s]
        vcond = curasset .<= strike
        valPut0 = curasset .* 0.0 # zeros(numPaths)
        valPut0[vcond] = 1.0
        valPut1 = valPut0 .* curasset # assetval
        valPut2 = valPut1 .* valPut1  # assetval ^ 2
        valPut3 = valPut2 .* valPut1  # assetval ^ 3
        valPut4 = valPut2 .* valPut2  # assetval ^ 4
        # compute the regression
        ridgecorrect = 0.01
        sum0 = sum(valPut0)
        sum1 = sum(valPut1)
        sum2 = sum(valPut2)
        sum3 = sum(valPut3)
        sum4 = sum(valPut4)
        sum0 += ridgecorrect
        sum1 += ridgecorrect
        sum2 += ridgecorrect
        sum3 += ridgecorrect
        sum4 += ridgecorrect
        sqrtest = Float64[sum0 sum1 sum2;
                          sum1 sum2 sum3;
                          sum2 sum3 sum4]
        invsqr = inv(sqrtest)
        vp0 = dot(cashFlowPut, valPut0)
        vp1 = dot(cashFlowPut, valPut1)
        vp2 = dot(cashFlowPut, valPut2)
        (betaPut0, betaPut1, betaPut2) = [vp0 vp1 vp2] * invsqr
        regImpliedCashFlow = valPut0 .* betaPut0 .+ valPut1 .* betaPut1 .+ valPut2 .* betaPut2
        payoff = valPut0 .* (strike .- curasset)
        pcond = payoff .> regImpliedCashFlow
        cashFlowPut[pcond] = payoff[pcond]
    end

    amerpayoffput = sum(cashFlowPut)
    finalputpayoff = sum(putpayoff)

    return amerpayoffput/numPaths, finalputpayoff/numPaths
end

function main()
    doc = """quant.jl

Quantitative option pricing model.

Usage:
  quant.jl -h | --help
  quant.jl [--assets=<assets>] [--iterations=<iterations>]

Options:
  -h --help                  Show this screen.
  --assets=<assets>          Specify the number of assets to simulate [default: 524288].
  --iterations=<iterations>  Specify a number of iterations [default: 256].
"""
    arguments = docopt(doc)

    paths = parse(Int, arguments["--assets"])
    steps = parse(Int, arguments["--iterations"])

    srand(0)

    println("assets = ", paths)
    println("iterations = ", steps)

    asset = init(paths, steps)
    tic()
    model(strike, 1, 1, asset)
    compiletime = toq()
    println("SELFPRIMED ", compiletime)

    tic()
    amerpayoffput, finalputpayoff = model(strike, paths, steps, asset)
    selftimed = toq()
    println("European Put Option Price: ", finalputpayoff)
    println("American Put Option Price: ", amerpayoffput)
    println("SELFTIMED ", selftimed)

end

main()


const J0_ROOTS(::Type{Float64}) = (
    (2.4048255576957730, -1.176691651530894e-16),
    (3.8317059702075125, -1.5269184090088067e-16),
    (5.5200781102863110, 8.088597146146722e-17),
    (7.0155866698156190, -9.414165653410389e-17),
    (8.6537279129110130, -2.92812607320779e-16),
    (10.173468135062722, 4.482162274768888e-16),
    (11.791534439014281, 2.812956912778735e-16),
    (13.323691936314223, 2.600408064718813e-16),
    (14.930917708487787, -7.070514505983074e-16),
    (16.470630050877634, -1.619019544798128e-15),
    (18.071063967910924, -9.658048089426209e-16),
    (19.615858510468243, -1.004445634526616e-15),
    (21.211636629879260, 4.947077428784068e-16),
    (22.760084380592772, -4.925749373614922e-16),
    (24.352471530749302, 9.169067133951066e-16),
    (25.903672087618382, 4.894530726419825e-16)
)
const J0_POLYS(::Type{Float64}) = (
    (0.0000000000000000000, -0.5191474972894669, 0.10793870175491979, 0.05660177443794914, -0.008657669593292222, -0.0021942003590739974, 0.0002643770365942964, 4.37291931443113e-5, -4.338825419759815e-6, -5.304927679598406e-7, 4.469819175606667e-8, 4.3284827345621115e-9, -3.135111000732148e-10, -2.628876489517534e-11),
    (-0.402759395702553, 2.476919088072758e-16, 0.20137969785127532, -0.017518715285670765, -0.013352611033152278, 0.0010359438492839443, 0.00037218755624680334, -2.4952042421113972e-5, -5.776086353844158e-6, 3.374317129436161e-7, 5.727482259215452e-8, -2.9561880489355444e-9, -3.905845672635605e-10, 1.971332566705736e-11),
    (0.000000000000000000, 0.34026480655836816, -0.030820651425593214, -0.05298855286760721, 0.004631042145890388, 0.002257440229081131, -0.00017518572879518415, -4.6521091062878115e-5, 3.199785909661533e-6, 5.716500268232186e-7, -3.5112898510841466e-8, -4.684643389757727e-9, 2.562685034682206e-10, 2.7958958795750104e-11),
    (0.30011575252613254, -1.6640272822046001e-16, -0.15005787626306408, 0.007129737603121546, 0.011742619737383848, -0.0006260583453094324, -0.00035093119008693753, 1.7929701912295164e-5, 5.6239324892485754e-6, -2.668437501970219e-7, -5.6648488273749086e-8, 2.48117399780498e-9, 3.8876537586241154e-10, -1.6657136713437192e-11),
    (0.000000000000000000, -0.27145229992838193, 0.015684124960953488, 0.044033774963413, -0.0025093022271948434, -0.0020603351551475315, 0.00011243486771159352, 4.482303558813413e-5, -2.288390108003442e-6, -5.679383588459768e-7, 2.693939234375692e-8, 4.737285529934781e-9, -2.0612709555352797e-10, -2.8163166483726606e-11),
    (-0.2497048770578432, 1.1807897766765572e-16, 0.12485243852891914, -0.0040907858517059345, -0.010102792347641438, 0.00038536375952213334, 0.00031859711440332953, -1.2373899600646271e-5, -5.3013932979548665e-6, 2.001098153528186e-7, 5.4711629662471434e-8, -1.9724572531751518e-9, -3.8121398193699247e-10, 1.3667679743782715e-11),
    (0.000000000000000000, 0.23245983136472478, -0.009857064513825458, -0.03818600911162367, 0.0016073972920762946, 0.0018420433388794816, -7.581358465623415e-5, -4.159284549011371e-5, 1.650645590334915e-6, 5.425453494592871e-7, -2.0556207467977526e-8, -4.620018928884712e-9, 1.642028058414746e-10, 2.7701605444102412e-11),
    (0.21835940724787295, -8.89726402965429e-17, -0.10917970362393398, 0.0027314677279632535, 0.008944552393700088, -0.00026391472261453583, -0.00028847875053074687, 8.858193371737123e-6, 4.9233776180403375e-6, -1.5077786827161215e-7, -5.190218733666561e-8, 1.5539413886301204e-9, 3.674809363354973e-10, -1.1113645791216594e-11),
    (0.00000000000000000, -0.20654643307799603, 0.006916736034268416, 0.034115572697347704, -0.001137276252948717, -0.0016680057255530109, 5.4841792064182565e-5, 3.837965853474541e-5, -1.2335804050962046e-6, -5.106259295634553e-7, 1.592333632709497e-8, 4.423517565793139e-9, -1.3138837384184105e-10, -2.6809397212536384e-11),
    (-0.19646537146865717, 6.979167865106427e-17, 0.09823268573432613, -0.001988037402152532, -0.008095530671166083, 0.00019440675128712672, 0.0002640383898036336, -6.666777683303928e-6, -4.5715696772304925e-6, 1.1666296153560847e-7, 4.8913639696764225e-8, -1.2379867207945651e-9, -3.508930968415813e-10, 9.07632091591013e-12),
    (0.00000000000000000, 0.18772880304043943, -0.005194182350684612, -0.031096513233785917, 0.0008577442641273341, 0.0015312251534677639, -4.184307585284775e-5, -3.5603170534217916e-5, 9.58002109234601e-7, 4.795250964600283e-7, -1.263434500625308e-8, -4.20550167685701e-9, 1.065791122326672e-10, 2.5727417675684577e-11),
    (0.18006337534431555, -5.638484737644332e-17, -0.09003168767215539, 0.0015299132863046024, 0.0074441453680234426, -0.00015060569680378768, -0.0002443398416353605, 5.227001519013193e-6, 4.267152607972633e-6, -9.295966007495808e-8, -4.610438011417262e-8, 1.0049092632275165e-9, 3.339442682325105e-10, -7.4991079301099e-12),
    (0.00000000000000000, -0.17326589422922986, 0.004084217951979124, 0.028749284970146657, -0.0006761643016121907, -0.0014215899173758441, 3.320978125391802e-5, 3.3264379323815026e-5, -7.684444100376941e-7, -4.515479818833484e-7, 1.0271599456634145e-8, 3.993903280527723e-9, -8.793975528824604e-11, -2.4610374225652004e-11),
    (-0.16718460047381803, 4.662597138876655e-17, 0.08359230023690671, -0.0012242529339116864, -0.006925682915280748, 0.00012100729852042988, 0.00022821854128498174, -4.230799709959437e-6, -4.007618360337058e-6, 7.601217108010105e-8, 4.358483157250522e-8, -8.317621452517008e-10, -3.178805429572483e-10, 6.284538399593043e-12),
    (0.000000000000000000, 0.16170155068925002, -0.0033200234006036146, -0.02685937038656159, 0.0005505380905909195, 0.0013316994659124243, -2.715683205388875e-5, -3.1289544794362e-5, 6.326900360777323e-7, 4.2697141781883964e-7, -8.532447325590315e-9, -3.799009445641295e-9, 7.379552811590934e-11, 2.353654960834717e-11),
    (0.15672498625285222, 1.1464880342445208e-19, -0.07836249312642612, 0.0010083833270351796, 0.006501011610557426, -9.993664895655577e-5, -0.00021478298462967253, 3.511086890959515e-6, 3.7857022791122945e-6, -6.350944888510364e-8, -4.135588012575766e-8, 7.135988717747828e-10, 3.2075281131621564e-10, 2.208506585533542e-12)
)
J0_POLY_PIO2(::Type{Float64}) = (
1.000000000000000, -0.25, 0.01562499999999994, -0.00043402777777725544, 6.781684026082576e-6,
-6.781683757550061e-8, 4.709479394601058e-10, -2.4016837144506874e-12, 9.104258208703104e-15
)

J1_ROOTS(::Type{Float64}) = (
        (1.8411837813406593, 4.7898393919093694e-18),
        (3.8317059702075125, -1.5269184090088067e-16),
        (5.3314427735250325, 1.5109105349471405e-16),
        (7.015586669815619, -9.414165653410389e-17),
        (8.536316366346286, -1.5433871213307537e-16),
        (10.173468135062722, 4.482162274768888e-16),
        (11.706004902592063, 7.121366942298246e-16),
        (13.323691936314223, 2.600408064718813e-16),
        (14.863588633909034, -6.265788988781879e-16),
        (16.470630050877634, -1.619019544798128e-15),
        (18.015527862681804, -1.1196999448424267e-16),
        (19.615858510468243, -1.004445634526616e-15),
        (21.16436985918879, 1.7024131380423588e-15),
        (22.760084380592772, -4.925749373614922e-16),
        (24.311326857210776, -2.614798558537172e-16),
        (25.903672087618382, 4.894530726419825e-16)
        )

@inline J1_POLYS(::Type{Float64}) = (
         (0.5818652242815964, 8.973834293036876e-17, -0.20511071214777496, 0.006058948324597746, 0.013801769808047723, -0.0003723170971916689, -0.0003949590750416506, 9.202953658715798e-6, 6.267292697927347e-6, -1.2681571675674473e-7, -6.319326736018798e-8, 1.057852721352679e-9, 4.839653511789343e-10, -2.1534683937074676e-11),
         (0.0, -0.402759395702553, 0.05255614585697725, 0.053410444132722765, -0.005179719245639115, -0.00223312533910116, 0.00017466429071755533, 4.6208700909637e-5, -3.0368632747366735e-6, -5.727790851165856e-7, 3.248220586398651e-8, 4.732693354840469e-9, -2.346792988667587e-10, -2.607612190048509e-11),
         (-0.34612620185379156, 2.6002760076991505e-17, 0.16697453550109428, -0.009678268542879972, -0.012099225779175138, 0.0006654009006841076, 0.0003541389011182841, -1.742720356958137e-5, -5.655294960447268e-6, 2.484316221112215e-7, 5.7101369594782676e-8, -2.2587104440771904e-9, -3.9854619701676605e-10, 1.854508641982197e-11),
         (0.0, 0.30011575252613254, -0.02138921280934146, -0.04697047894974007, 0.0031302917260424344, 0.0021055871432169443, -0.00012550790947431806, -4.499147505171648e-5, 2.401580334889493e-6, 5.665260872870135e-7, -2.7272058756924113e-8, -4.718921752439826e-9, 2.045523137027535e-10, 2.6923755198436e-11),
         (0.27329994163319987, -1.3511933704624036e-16, -0.13477468037992396, 0.005116340346495344, 0.010631861751990702, -0.000448743683849691, -0.0003268000185531749, 1.3382556752353113e-5, 5.363177082527157e-6, -2.0647461181851255e-7, -5.499892476485602e-8, 1.9780335410646812e-9, 3.8456897851545063e-10, -1.5543946040865833e-11),
         (0.0, -0.2497048770578432, 0.012272357555101388, 0.04041116939079181, -0.0019268187972545589, -0.0019115826893618788, 8.661729445925766e-5, 4.241116264477775e-5, -1.800978985430059e-6, -5.47159694982648e-7, 2.1682608288127194e-8, 4.629875997238198e-9, -1.6946465290961774e-10, -2.7001263147077087e-11),
         (-0.23330441717143408, 1.3347222236161232e-16, 0.11580092244607755, -0.0032489977328297747, -0.009372527206041913, 0.0003036138212749204, 0.00029804555521007126, -9.813819271763006e-6, -5.0242293237237345e-6, 1.6136482548346424e-7, 5.251755806094281e-8, -1.6213723779041528e-9, -3.7130918416720704e-10, 1.2840797195120732e-11),
         (0.0, 0.21835940724787295, -0.008194403183877394, -0.03577820957502993, 0.001319573612804675, 0.0017308725061602815, -6.200735153951766e-5, -3.938703730297954e-5, 1.3569938030086903e-6, 5.190652330962483e-7, -1.7082573736581586e-8, -4.464580500726025e-9, 1.3835552663370753e-10, 2.6508049073606443e-11),
         (0.20701265272531905, -1.1572365385200866e-16, -0.10303781563402706, 0.002289729546293356, 0.008432435058846974, -0.00021965941971569214, -0.000272721543802475, 7.36909520166486e-6, 4.682127415369e-6, -1.2611406291992532e-7, -4.9733737497256667e-8, 1.313333637862963e-9, 3.5584811981981077e-10, -1.0581432980865016e-11),
         (0.0, -0.19646537146865717, 0.005964112206447895, 0.032382122684890324, -0.0009720337562254532, -0.0015842303417456071, 4.6667442166602986e-5, 3.6572573450276185e-5, -1.0499611641790313e-6, -4.891787571238003e-7, 1.3609422867092438e-8, 4.2640592578446675e-9, -1.1321195999108497e-10, -2.5685037344174917e-11),
         (-0.18801748852581776, 9.806595982257605e-17, 0.09371909377243136, -0.0017233243363111925, -0.007714266760106843, 0.00016755862384546863, 0.000251827174474777, -5.7335177376628445e-6, -4.372425030172481e-6, 1.0045839852031556e-7, 4.698184959516747e-8, -1.0717704296017814e-9, -3.3952344362582674e-10, 8.773851724781324e-12),
         (0.0, 0.18006337534431555, -0.004589739858905968, -0.029776581472313136, 0.0007530284838491313, 0.0014660390526475475, -3.658900932751018e-5, -3.4137236390211275e-5, 8.366325091910564e-7, 4.6108475352995054e-7, -1.1047200388967438e-8, -4.058793425595929e-9, 9.363513647344476e-11, 2.4720051360744732e-11),
         (0.1734590492857464, -8.319859584484218e-17, -0.08653590193876039, 0.0013568189856187354, 0.007147216520714834, -0.00013295776598760658, -0.00023462959289705953, 4.6031437800947915e-6, 4.103196271673978e-6, -8.18373301294018e-8, -4.44403211204062e-8, 8.872617322304839e-10, 3.2359730989751983e-10, -7.354373969368619e-12),
         (0.0, -0.16718460047381806, 0.0036727588017285762, 0.02770273166133465, -0.0006050364924617241, -0.0013693112504440895, 2.9615596890234806e-5, 3.206096184277923e-5, -6.841058767433935e-7, -4.358877331737043e-7, 9.143763763164119e-9, 3.864037532317281e-9, -7.851667555762691e-11, -2.3728757802646805e-11),
         (-0.1618382095526585, 7.115722920799028e-17, 0.08078219522574913, -0.0011038527651152986, -0.006686444726790514, 0.00010870535083026479, 0.00022030197347228358, -3.791805814546456e-6, -3.871147931227952e-6, 6.80607764766372e-8, 4.2160129307780366e-8, -7.460427055677985e-10, -3.0872139357605097e-10, 6.24158293622982e-12),
         (0.0, 0.15672498625285222, -0.0030251499811055965, -0.026004046442225856, 0.0004996832448037306, 0.001288697907658241, -2.457760943647847e-5, -3.02856199504406e-5, 5.715983781675481e-7, 4.136275513873783e-7, -7.703846567878601e-9, -3.6853569596797437e-9, 6.676439988396374e-11, 2.277101500922319e-11))#for i in 2:2:18

J1_POLY_PIO2(::Type{Float64}) = (0.5, -0.0624999999999989, 0.002604166666657291, -5.42534721917933e-5, 6.781683542660179e-7, -5.651361336587487e-9, 3.36191211106159e-11, -1.4511302591871352e-13)

#for i in 1:2:18
#    display(((i+1)*pi/2-pts[i+1][1], (i+2)*pi/2-pts[i+1][1]))
#    poly = Float64.(Tuple(ratfn_minimax(x->r2(x,pts[i]), (i*pi/2-pts[i][1], (i+1)*pi/2-pts[i][1]), 13, 0)[1]))
#    push!(polys, poly)
#    poly = Float64.(Tuple(ratfn_minimax(x->r(x,pts[i+1]), ((i+1)*pi/2-pts[i+1][1], (i+2)*pi/2-pts[i+1][1]), 12, 0)[1]))
#    push!(polys, (0.0, poly...))
#end


@doc raw"""
    TabularMDP(T, R, discount)

Specify a discrete MDP in tabular form.

- The transition matrix is 3 dimensional of size ``|\mathcal{S}|\times |\mathcal{A}| \times |\mathcal{S}|``, then `T[sj, a, si]` corresponds to the probability of ending in `sj` while taking action `a` in `si`.
- The reward matrix is 2 dimensional of size ``|\mathcal{S}|\times |\mathcal{A}|``, where `R[s, a]` is the reward obtained when taking action `a` in state `s`.
"""
mutable struct TabularMDP <: MDP{Int64, Int64}
    T::Array{Float64, 3} # SPxAxS
    R::Matrix{Float64} # SxA
    discount::Float64
    InitialTable::Array{Float64}
end

@doc raw"""
    TabularPOMDP(T, R, O, discount)

Specify a discrete POMDP in tabular form.

- The transition matrix is 3 dimensional of size ``|\mathcal{S}|\times |\mathcal{A}| \times |\mathcal{S}|``, then `T[sj, a, si]` corresponds to the probability of ending in `sj` while taking action `a` in `si`.
- The observation matrix is also 3 dimensional of size ``|\mathcal{O}| \times |\mathcal{A}| \times |\mathcal{S}|``, `O[o, a, sp]` represents the probability of observing `o` in in state `sp` and action `a`.
- The reward matrix is 2 dimensional of size ``|\mathcal{S}|\times |\mathcal{A}|``, where `R[s, a]` is the reward obtained when taking action `a` in state `s`.
"""
mutable struct TabularPOMDP <: POMDP{Int64, Int64, Int64}
    T::Array{Float64, 3} # SPxAxS
    R::Matrix{Float64} # SxA
    O::Array{Float64, 3} # OxAxSP
    discount::Float64
    InitialTable::Array{Float64}
end

const TabularProblem = Union{TabularMDP, TabularPOMDP}

# Distribution Type and methods
# XXX: this should be replaced with Categorical when https://github.com/JuliaStats/Distributions.jl/issues/743 is fixed
struct DiscreteDistribution{P<:AbstractVector{Float64}}
    p::P
end

support(d::DiscreteDistribution) = 1:length(d.p)

pdf(d::DiscreteDistribution, sp::Int64) = d.p[sp] # T(s', a, s)

rand(rng::AbstractRNG, d::DiscreteDistribution) = sample(rng, Weights(d.p))

# MDP and POMDP common methods
states(p::TabularProblem) = 1:size(p.T, 1)
actions(p::TabularProblem) = 1:size(p.T, 2)

stateindex(::TabularProblem, s::Int64) = s
actionindex(::TabularProblem, a::Int64) = a

discount(p::TabularProblem) = p.discount

transition(p::TabularProblem, s::Int64, a::Int64) = DiscreteDistribution(view(p.T, :, a, s))

reward(prob::TabularProblem, s::Int64, a::Int64) = prob.R[s, a]

initialstate(p::TabularProblem) = DiscreteDistribution(p.InitialTable)

# POMDP only methods
n_observations(p::TabularProblem) = size(p.O, 1)

observations(p::TabularPOMDP) = 1:n_observations(p)

observation(p::TabularPOMDP, a::Int64, sp::Int64) = DiscreteDistribution(view(p.O, :, a, sp))

obsindex(p::TabularPOMDP, o::Int64) = o

# deprecated in POMDPs v0.9
POMDPs.initialstate_distribution(p::TabularProblem) = DiscreteDistribution(p.InitialTable)


#The key to the non-parametric method describedhere
#is that the sum of squared distances between pointsand
#their centroid is equal to (and can be calculateddirectly from)
#the sum of squared interpoint distancesdivided by the number of points
# https://onlinelibrary.wiley.com/doi/epdf/10.1111/j.1442-9993.2001.01070.pp.x

SS(D,N) = sum(D .^2)/N

function F(D,W,N,n,a)
    SST = SS(D,N)
    SSW = SS(W,n)
    SSA = SST - SSW
    return (SSA/(a-1))/(SSW/(N-a))
end

function F(SST ::Real,W,N,n,a)

    SSW = SS(W,n)
    SSA = SST - SSW
    return (SSA/(a-1))/(SSW/(N-a))
end

function getinds(group,N)
    W =Vector{CartesianIndex}(undef,0)
    notW = Vector{CartesianIndex}(undef,0)
    for i in 1:(N-1)
        for j in i+1:N
            if i !==1j
                if group[i] == group[j]
                    push!(W,(CartesianIndex(i,j)))
                else push!(notW,(CartesianIndex(i,j)))
                end
            end
        end
    end
    return W, notW
end

function permutest(D,group, n_perm = 1000)
    N =  length(group)
    a = length(unique(group))
    n = N/a
    Winds ,notWinds =getinds(group,N)

    @inbounds W = D[Winds]
    @inbounds notW= D[notWinds]

    Dvec = vcat(W,notW)

    SST = SS(Dvec,N)

    f = F(SST,W,N,n,a)
    Fs = Vector{Float64}(undef,n_perm)
    inds = 1:length(W)

    @inbounds for i in 1:n_perm
        shuffle!(Dvec)
        Fs[i] = F(SST,view(Dvec,inds),N,n,a)
    end
    P = sum(Fs .>= f)/n_perm

    g = unique(group)
    if a>2
        ppairs = NamedArray(zeros(a,a),( g,g ), ("group","group"))
        tpairs = NamedArray(zeros(a,a),( g,g ), ("group","group"))
        Threads.@threads for i in 1:a-1
            for j in i+1:a
                boolmask = (group .== g[i]) .| (group .== g[j])
                d = D[boolmask,boolmask]
                fstat,pstat =permutest(d,group[boolmask], n_perm)
                ppairs[j,i] = pstat
                tpairs[j,i] = sqrt(fstat)
            end
        end
    return (F = f,P = P, fpairs = tpairs,ppairs = ppairs)
    else
        return (F = f,P = P)
    end

end

X = rand(10,5)
D = pairwise(Euclidean(),X,X,dims = 1)
group = rand(1:2,10)
permutest(D,group)

X = vcat(rand(5,5),rand(5,5) .+3)
D = pairwise(Euclidean(),X,X,dims = 1)
group = [1,1,2,2,2,3,3,3,3,3]
permutest(D,group)

X = vcat(rand(800,500))
D = pairwise(Euclidean(),X,X,dims = 1)
group = vcat(fill(collect(1:10),80)...)
permutest(D,group)


using Jchemo

n = 6 ; p = 7 ; q = 2 ; m = 3
Xtrain = rand(n, p) ; Ytrain = rand(n, q)
Xtest = rand(m, p) ; Ytest = rand(m, q)

nlv = 5 ;
fm = plskern(Xtrain, Ytrain; nlv = nlv) ;
pnames(fm)
fm.T
fm.P

Jchemo.summary(fm, Xtrain).explvar

fm.T
Jchemo.transform(fm, Xtrain)
Jchemo.transform(fm, Xtrain; nlv = 1)
Jchemo.transform(fm, Xtest)

Jchemo.coef(fm).B
Jchemo.coef(fm).int
Jchemo.coef(fm; nlv = 2).B
Jchemo.coef(fm; nlv = 2).int
Jchemo.coef(fm; nlv = 0).B
Jchemo.coef(fm; nlv = 0).int

Jchemo.predict(fm, Xtest).pred
Jchemo.predict(fm, Xtest; nlv = nlv).pred

Jchemo.predict(fm, Xtest; nlv = 0:3).pred
Jchemo.predict(fm, Xtest; nlv = 0).pred

pred = Jchemo.predict(fm, Xtest).pred ;
msep(pred, Ytest)

### Weighted PLS

w = collect(1:n)
fm = plskern(Xtrain, Ytrain, w; nlv = nlv) ;


#Compute Threshold Functions

module ThresholdFunctions

using ..MathFuns

@doc raw"""
    Fb1(k,m,T)
Compute the one bosonic propagator threshold function without external momentum dependence.
```math
\tilde{\partial_k}\int \frac{d^3q}{2\pi^3}\sum_{q_0}G_b(q_0)
```
"""
function Fb1(k, m2, T)
    -1 / 24 * (
        k^4 * (
            2 * T * coth(Eb(k, m2) / (2 * T)) +
            Eb(k, m2) * csch(Eb(k, m2) / (2 * T))^2
        )
    ) / (Eb(k, m2)^3 * pi^2 * T)
end

@doc raw"""
    Fb2(k,m,T)
Compute the two bosonic propagator threshold function without external momentum dependence.
```math
\tilde{\partial_k}\int \frac{d^3q}{2\pi^3}\sum_{q_0}G_b(q_0)G_b(q_0)
```
If we want to start from the momentum dependent case and take the zero momentum limit, we need to
take the $p_0\to 0$ first then take the $p_s\to 0$, they are noncommutative, or we need to take the $p_0$
goes faster to $0$ than $p_s$
"""
function Fb2(k,m2,T)
    (
        k^4 * (
            -6 * T^2 * coth(Eb(k, m2) / (2 * T)) -
            csch(Eb(k, m2) / (2 * T))^2 *
            Eb(k, m2) *
            (3 * T + coth(Eb(k, m2) / (2 * T)) * Eb(k, m2))
        )
    ) / (48 * pi^2 * T^2 * Eb(k, m2)^5)
end

function ReFb2positive(p0, k, m2, T,ϵ=0.1)
    -1 / 24 * (
        k^4 *
        csch(Eb(k, m2) / (2 * T))^2 *
        (p0 - 2 * Eb(k, m2))^2 *
        sinh(Eb(k, m2) / T)
    ) / (pi^2 * (ϵ + (p0 - 2 * Eb(k, m2))^2)^2 * Eb(k, m2)^3) -
    (k^4 * csch(Eb(k, m2) / (2 * T))^2 * sinh(Eb(k, m2) / T)) /
    (24 * pi^2 * Eb(k, m2)^3 * (p0 + 2 * Eb(k, m2))^2) +
    (
        k^4 *
        csch(Eb(k, m2) / (2 * T))^2 *
        (p0 - 2 * Eb(k, m2)) *
        (Eb(k, m2) + 2 * T * sinh(Eb(k, m2) / T))
    ) / (48 * pi^2 * T * (ϵ + (p0 - 2 * Eb(k, m2))^2) * Eb(k, m2)^4) -
    (
        k^4 *
        csch(Eb(k, m2) / (2 * T))^2 *
        (Eb(k, m2) + 2 * T * sinh(Eb(k, m2) / T))
    ) / (48 * pi^2 * T * Eb(k, m2)^4 * (p0 + 2 * Eb(k, m2)))
end

#we add the ϵ by analitical continuation
function ReFb2positive_ac(ω, k, m2, T, ϵ = 0.1)
    -1 / 6 * (
        k^4 * (
            (ϵ^2 + (ω - 2 * Eb(k, m2))^2) *
            Eb(k, m2) *
            (ϵ^2 - ω^2 + 4 * Eb(k, m2)^2) *
            (ϵ^2 + (ω + 2 * Eb(k, m2))^2) +
            T *
            (
                (ϵ - ω) * (ϵ + ω) * (ϵ^2 + ω^2)^2 +
                4 * (5 * ϵ^4 - 14 * ϵ^2 * ω^2 + 5 * ω^4) * Eb(k, m2)^2 +
                112 * (ϵ - ω) * (ϵ + ω) * Eb(k, m2)^4 +
                192 * Eb(k, m2)^6
            ) *
            sinh(Eb(k, m2) / T)
        )
    ) / (
        pi^2 *
        T *
        (-1 + cosh(Eb(k, m2) / T)) *
        Eb(k, m2)^3 *
        (
            (ϵ^2 + ω^2)^2 +
            8 * (ϵ - ω) * (ϵ + ω) * Eb(k, m2)^2 +
            16 * Eb(k, m2)^4
        )^2
    )
end

function ReFb2positiveEu(p0, k, m2, T)
    -1 / 12 * (
        k^4 *
        csch(Eb(k, m2) / (2 * T))^2 *
        (
            Eb(k, m2) * (p0^2 + 4 * Eb(k, m2)^2) +
            T * (p0^2 + 12 * Eb(k, m2)^2) * sinh(Eb(k, m2) / T)
        )
    ) / (pi^2 * T * Eb(k, m2)^3 * (p0^2 + 4 * Eb(k, m2)^2)^2)
end

function ImFb2positive(p0, k, m2, T, ϵ = 0.1)
    -1 / 12 *
    (k^4 * coth(Eb(k, m2) / (2 * T)) * deltadxfun(p0 - 2 * Eb(k, m2), ϵ)) /
    (pi * Eb(k, m2)^3) -
    (
        k^4 *
        csch(Eb(k, m2) / (2 * T))^2 *
        deltafun(p0 - 2 * Eb(k, m2), ϵ) *
        (Eb(k, m2) + 2 * T * sinh(Eb(k, m2) / T))
    ) / (48 * pi * T * Eb(k, m2)^4)
end

function ImFb2(p0, k, m2, T, ϵ = 0.1)
    if p0>0.0
        ImFb2positive(p0, k, m2, T, ϵ)
    elseif p0<0.0
        -ImFb2positive(-p0, k, m2, T, ϵ)
    elseif p0==0.0
        0.0
    end
end

function ReFb2(p0, k, m2, T, ϵ = 0.1)
    if p0>=0.0
        ReFb2positive(p0, k, m2, T, ϵ)
    elseif p0<0.0
        ReFb2positive(-p0, k, m2, T, ϵ)
    end
end

function ReFb2_ac(p0, k, m2, T, ϵ = 0.1)
    if p0>=0.0
        ReFb2positive_ac(p0, k, m2, T, ϵ)
    elseif p0<0.0
        ReFb2positive_ac(-p0, k, m2, T, ϵ)
    end
end

function ReFb2Eu(p0, k, m2, T)
    if p0>=0.0
        ReFb2positiveEu(p0, k, m2, T)
    elseif p0<0.0
        ReFb2positiveEu(-p0, k, m2, T)
    end
end

end


using CrossEntropyMethod
using Distributions, Random
using Plots
# The goal of this example is to show how to estimate the optimal importance sampling distribution

# Actual distibution
base_dist = Categorical([0.74899, 0.2, 0.05, 0.001, 0.00001])

# starting sampling distribution
is_dist_0 = Dict{Symbol, Vector{Sampleable}}(:x => [Categorical(5)])

# Sample a random time series from the starting distribution
rand(is_dist_0)

# Define a loss function that identifies the values of 3,4,5 as those of intereste
function l(d, s)
    v = s[:x][1]
    -(v in [3,4,5])
end

# Define the likelihood ratio weighting function
function w(d, s)
    v = s[:x][1]
    exp(logpdf(base_dist, v) - logpdf(d, s))
end

# show what the weighting function does
s = Dict(:x => [5])
w(is_dist_0, s)

# Run the optimization
is_dist_opt = cross_entropy_method(l, is_dist_0; max_iter = 4, N=10000, weight_fn = w, verbose = true)

# See the final distribution.
is_dist_opt[:x][1].p


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

#=
Möller, T. & Trumbore, B., 1997.
(https://www.tandfonline.com/doi/abs/10.1080/10867651.1997.10487468)
=#
function intersection(f, r::Ray{3,T}, t::Triangle{3,T}) where {T}
  vs = vertices(t)
  o = origin(r)
  d = direction(r)

  e₁ = vs[3] - vs[1]
  e₂ = vs[2] - vs[1]
  p = d × e₂
  det = e₁ ⋅ p

  # keep det > 0, modify T accordingly
  if det > atol(T)
    τ = o - vs[1]
  else
    τ = vs[1] - o
    det = -det
  end

  if det < atol(T)
    # This ray is parallel to the plane of the triangle.
    return @IT NoIntersection nothing f
  end

  # calculate u parameter and test bounds
  u = τ ⋅ p
  if u < -atol(T) || u > det
    return @IT NoIntersection nothing f
  end

  q = τ × e₁

  # calculate v parameter and test bounds
  v = d ⋅ q
  if v < -atol(T) || u + v > det
    return @IT NoIntersection nothing f
  end

  λ = (e₂ ⋅ q) * (one(T) / det)

  if λ < -atol(T)
    return @IT NoIntersection nothing f
  end

  λ = clamp(λ, zero(T), typemax(T))

  return @IT IntersectingRayTriangle r(λ) f
end


function SimpRulT(A,B,Q,R,U,bet,n1,n2,F,SigmaXX,x10,cutoff=1.0)
#SimpRulT   Solving RE model for given simple rule u(t)  = -F * x(t),
#           where   x(t)  = [x1(t)' x2(t)']'.
#
#
#
#  Usage:    (M,C,J0) = SimpRulT(A,B,Q,R,U,bet,n1,n2,F,SigmaXX,x10,cutoff)
#
#  Input:     A,B,Q,R,U,bet,n1,n2: see ComItAlg
#             F          kxn matrix, decision rule
#             SigmaXX    n1xn1 covariance matrix of shocks
#             x10        n1 vector, initial values for predetermined
#                        variables. Not n1x1 array.
#             cutoff     scalar, max modulus of backward looking variables
#
#  Output:    M          n1xn1 matrix, x1(t+1) = M*x1(t) + e(t+1)
#             C          n2xn1 matrix, x2(t) = C*x1(t)
#             J0         scalar, value of loss function
#
#
#
#  Paul Soderlind, Paul.Soderlind@unisg.ch, Aug 2000, to Julia Jan 2016
#-----------------------------------------------------------------------

  Q = (Q + Q')/2                  #to make symmetric
  R = (R + R')/2

  A = A - B*F                     #subst for u(t) in evolution
  n = n1 + n2

  GS = schur(Matrix(1.0I,n,n),A)
  logconA = abs.(GS.beta) .<= (abs.(GS.alpha)*cutoff) #selecting stable eigenvalues
  ordschur!(GS,logconA)                                 #reordering, stable first

  if sum(logconA) < n1
    @warn("Too few stable roots: no stable solution")
    return NaN,NaN,NaN
  elseif sum(logconA) > n1
    @warn("Too many stable roots: inifite number of stable solutions")
    return NaN,NaN,NaN
  end

  Stt = GS.S[1:n1,1:n1]
  Zkt = GS.Z[1:n1,1:n1]
  Zlt = GS.Z[n1+1:n,1:n1]
  Ttt = GS.T[1:n1,1:n1]

  if cond(Zkt) > 1e+14
    @warn("Zkt is singular: rank condition for solution not satisfied")
    return NaN,NaN,NaN
  end

  Zkt_1 = inv(Zkt)         #inverting
  Stt_1 = inv(Stt)

  M = real(Zkt*Stt_1*Ttt*Zkt_1)       #x1(t+1) = M*x1(t) + e(t+1)
  C = real(Zlt*Zkt_1)                 #x2(t) = C*x1(t)

  P = [ I;
        C;
        (-F*[I;C]) ]
  QUUR = [ Q   U;
           U'  R ]
  PQUURP = P'QUUR*P

  vecV = (I - kron(M',bet*M'))\vec(PQUURP)
  V    = reshape(vecV,n1,n1)                 #solves V = W + bet*M'V*M

  J0 =  x10'V*x10 + (bet/(1-bet)) * tr(V*SigmaXX)

  return M,C,J0

end
#----------------------------------------------------------------------


using Manifolds
using Base.Test

M = Sphere()
x = randn(3)
g = randn(3)
@test norm(retract!(M, x)) ≈ 1
@test abs(dot(project_tangent(M,g,x),x)) < 1e-14


@testset "Cycle Basis" begin

    function evaluate(x,y)
        @test length(x) == length(y)
        for i=1:length(x)
            @test isempty(setdiff(x[i],y[i]))
            @test isempty(setdiff(x[i],y[i]))
        end
    end

    # No Edges
    ex = Graph(1)
    expected_cyclebasis = Array{Int64,1}[]
    for g in testgraphs(ex)
        ex_cyclebasis = @inferred cycle_basis(g)
        @test isempty(ex_cyclebasis)
    end

    # Only one self-edge
    elist = [(1,1)]
    ex = Graph(SimpleEdge.(elist))
    expected_cyclebasis = Array{Int64,1}[[1]]
    for g in testgraphs(ex)
        ex_cyclebasis = cycle_basis(g)
        evaluate(ex_cyclebasis,expected_cyclebasis)
    end

    # Graph with one cycle
    elist = [(1,2),(2,3),(3,4),(4,1),(1,5)]
    ex = Graph(SimpleEdge.(elist))
    expected_cyclebasis = Array{Int64,1}[
        [1,2,3,4] ]
    for g in testgraphs(ex)
        ex_cyclebasis = cycle_basis(g)
        evaluate(ex_cyclebasis,expected_cyclebasis)
    end

    # Graph with 2 of 3 cycles forming a basis
    elist = [(1,2),(1,3),(2,3),(2,4),(3,4)]
    ex = Graph(SimpleEdge.(elist))
    expected_cyclebasis = Array{Int64,1}[
        [2,3,4],
        [2,1,3] ]
    for g in testgraphs(ex)
        ex_cyclebasis = cycle_basis(g)
        evaluate(ex_cyclebasis,expected_cyclebasis)
    end

    # Testing root argument
    elist = [(1,2),(1,3),(2,3),(2,4),(3,4),(1,5),(5,6),(6,4)]
    ex = Graph(SimpleEdge.(elist))
    expected_cyclebasis = Array{Int64,1}[
        [2, 4, 3],
        [1, 5, 6, 4, 3],
        [1, 2, 3] ]
    for g in testgraphs(ex)
        ex_cyclebasis = @inferred cycle_basis(g,3)
        evaluate(ex_cyclebasis,expected_cyclebasis)
    end
end


# Tests related to validated integration

using TaylorModels
# using LinearAlgebra: norm
using Test

# const _num_tests = 1000
# const α_mid = TaylorModels.α_mid

setformat(:full)

# This is based on the examples from Buenger, Numer Algor (2018) 78:1001–1017,
# showing efectivity of shrink-wrapping
@testset "Testing `shrink_wrapping` 1" begin
    _order = 2
    set_variables("x", numvars=1, order=2*_order)
    x0 = IntervalBox(0..0, 1)
    dom = IntervalBox(-1..1, 1)
    for δ in 1/16:1/16:1
        rem = δ * Interval(-1, 1)
        p1 = TaylorModelN(TaylorN(1, order=_order), rem, x0, dom)
        p2 = deepcopy(p1)
        res = p1 * p2
        bound_prod = dom[1]^2 + δ*(2+δ)*dom[1]
        @test res(dom) ⊆ polynomial(res)(dom) + bound_prod
        @test remainder(res) ⊆ bound_prod
        q = [ res ]
        TaylorModels.shrink_wrapping!(q)
        @test q[1](dom) ⊆ res.pol(dom) + bound_prod
        @test remainder(q[1]) ⊆ bound_prod
        @test remainder(q[1]) == 0..0
    end
end

# This is based on the example of Sect 4 of M. Berz, K. Makino,
# in https://bt.pa.msu.edu/pub/papers/VIShrink06/VIShrink06.pdf
@testset "Testing `shrink_wrapping` 2" begin
    function two_state1!(v)
        x, y = v
        x_2 = x^2
        y_2 = y^2
        r_2 = x_2 + y_2
        s = sqrt(1 + r_2)
        v[1] = x * s
        v[2] = y * s
        return v
    end

    function two_state2!(v)
        x1, y1 = v
        x1_2 = x1^2
        y1_2 = y1^2
        r1_2 = x1_2 + y1_2
        r1 = sqrt( 1 + 4*r1_2 )
        s1 = 2 / (1+r1)
        ss = sqrt( s1 )
        v[1] = x1 * ss
        v[2] = y1 * ss
        return v
    end

    local δ = 0.05
    local ib = IntervalBox( -δ .. δ, 2)
    local mib = IntervalBox( 0..0, 2)

    # Diverges using naive Interval arithmetic methods
    ib0 = [ib...]
    for iter = 1:210 # iter is twice the number of iterates
        two_state1!(ib0)
        two_state2!(ib0)
    end
    @test minimum(diam.(ib0)) > 1e6

    # Taylor model controls grow and behaves as identity for after two iterates
    order = 10
    set_variables("x y", order=2*order)
    vm = [TaylorModelN(TaylorN(i, order=order), 0..0, mib, ib) for i = 1:2]
    vm0 = [deepcopy(vm)...]
    for iter = 1:1_000
        two_state1!(vm0)
        two_state2!(vm0)
    end
    @test maximum(remainder.(vm0)) < 1e-6 # 2.2e-7
    @test all(polynomial.(vm0) .== polynomial.(vm)) # only remainder grows!

    # Taylor model with shrink-wrapping after two iterates
    vm0 .= [deepcopy(vm)...]
    for iter = 1:1_000
        two_state1!(vm0)
        two_state2!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
    end
    @test maximum(remainder.(vm0)) < 1e-8 #2.2e-13

    # Taylor model with shrink-wrapping after each iterate
    vm0 .= [deepcopy(vm)...]
    for iter = 1:1_000
        two_state1!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
        two_state2!(vm0)
        TaylorModels.shrink_wrapping!(vm0)
    end
    @test maximum(remainder.(vm0)) < 2.2e-13
end


immutable Cauchy <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function Cauchy(l::Real, s::Real)
	    if s > 0.0
	    	new(float64(l), float64(s))
	    else
	    	error("scale must be positive")
	    end
	end
end

Cauchy(l::Real) = Cauchy(l, 1.0)
Cauchy() = Cauchy(0.0, 1.0)

@_jl_dist_2p Cauchy cauchy

entropy(d::Cauchy) = log(d.scale) + log(4.0 * pi)

insupport(d::Cauchy, x::Number) = isreal(x) && isfinite(x)

kurtosis(d::Cauchy) = NaN

mean(d::Cauchy) = NaN

median(d::Cauchy) = d.location

mgf(d::Cauchy, t::Real) = NaN

function cf(d::Cauchy, t::Real)
	m, theta = d.location, d.scale
	return exp(im * t * m - theta * abs(t))
end

modes(d::Cauchy) = [d.location]

skewness(d::Cauchy) = NaN

var(d::Cauchy) = NaN