--- Imports ---


import Goal.Core
import Goal.Geometry


--- Globals ---


-- Functions --

f p = let (x,y) = toPair p in x^2 + 2*y^2 + (x-y)^2

df p =
    let (x,y) = toPair p
        x' = 2*x + 2*(x-y)
        y' = 4*y - 2*(x-y)
     in fromList (Tangent p) [x',y']

-- Plot --

res = 400
mn = -4
mx = 4
niso = 10
cntrf x y = f $ euclideanPoint [x,y]
rng = (mn,mx,res)
clrs = rgbaGradient (0.9,0,0,1) (0,0,0,1) niso
axprms = LinearAxisParams (show . round) 5 5

-- Gradient Descent --

p0 = euclideanPoint [-4,2]
eps = 0.01
nstps = 200
grds = take nstps $ gradientDescent eps df p0


--- Main ---


main = do

    -- Contour plots
    let rnbl = toRenderable . execEC $ do

            let cntrs = contours rng rng niso cntrf

            sequence_ $ do

                ((_,cntr),clr) <- zip cntrs clrs

                return . plot . liftEC $ do

                    plot_lines_style .= solidLine 3 clr
                    plot_lines_values .= cntr

            layout_x_axis . laxis_generate .= scaledAxis axprms (mn,mx)
            layout_x_axis . laxis_override .= axisGridHide
            layout_x_axis . laxis_title .= "x"
            layout_y_axis . laxis_generate .= scaledAxis axprms (mn,mx)
            layout_y_axis . laxis_override .= axisGridHide
            layout_y_axis . laxis_title .= "y"

            plot . liftEC $ do
                plot_points_style .= filledCircles 5 (opaque red)
                plot_points_values .= [(0,0)]

            plot . liftEC $ do
                plot_lines_style .= solidLine 3 (opaque black)
                plot_lines_values .= [toPair <$> grds]

    --void $ renderableToAspectWindow False 800 800 rnbl
    void $ renderableToFile (FileOptions (200,200) PDF) "gradient-descent.pdf" rnbl