import "lib/github.com/diku-dk/complex/complex"
import "lib/github.com/diku-dk/cpprandom/random"

module rand_engine = minstd_rand
module uniform_real = uniform_real_distribution f32 rand_engine

module cplx = mk_complex f32
type complex = cplx.complex

let N = 22i64
let PolyN = N + 1

type poly = [PolyN]f32

-- First element of pair steps fastest
let iota2 (n: i64) (m: i64): [](i64, i64) =
    flatten (map (\y -> map (\x -> (x, y)) (iota n)) (iota m))

let evaln_c (p: poly) (nterms: i64) (pt: complex): complex =
    foldr (\coef accum -> cplx.mk_re coef cplx.+ pt cplx.* accum)
          (cplx.mk_re p[nterms-1]) (take (nterms - 1) p)

let eval_c (p: poly) (pt: complex): complex = evaln_c p (length p) pt

let evaln_d (p: poly) (nterms: i64) (pt: f32): f32 =
    foldr (\coef accum -> coef + pt * accum)
          p[nterms-1] (take (nterms - 1) p)

let eval_d (p: poly) (pt: f32): f32 = evaln_d p (length p) pt

let derivative (p: poly): *poly = 
    map2 (\i v -> f32.i64 (i64.bool (i != PolyN - 1) * (i + 1)) * v)
         (0..<PolyN) (rotate 1 p)

-- Cauchy's bound: https://en.wikipedia.org/wiki/Geometrical_properties_of_polynomial_roots#Lagrange's_and_Cauchy's_bounds
let max_root_norm (p: poly): f32 =
    1 + f32.maximum (map (\coef -> f32.abs (coef / p[PolyN-1])) (init p))

module aberth = {
    type approx = [N]complex
    -- bound is 's' in the stop condition formulated at p.189-190 of
    -- https://link.springer.com/article/10.1007%2FBF02207694
    type context = {p: poly, deriv: poly, bound: poly, radius: f32}

    let gen_coord (r: f32) (rng: *rand_engine.rng): *(rand_engine.rng, f32) =
        uniform_real.rand (-r, r) rng

    let gen_coord_c (r: f32) (rng: *rand_engine.rng): (rand_engine.rng, complex) =
        let (rng, x) = gen_coord r rng
        let (rng, y) = gen_coord r rng
        in (rng, cplx.mk x y)

    let generate (ctx: context) (rng: *rand_engine.rng): *(rand_engine.rng, approx) =
        let rngs = rand_engine.split_rng N rng
        let (rngs, approx) = unzip (map (\rng -> gen_coord_c ctx.radius (copy rng)) rngs)
        let rng = rand_engine.join_rng rngs
        in (rng, approx)

    let compute_bound_poly (p: poly): *poly =
        map2 (\coef i -> f32.abs coef * f32.i64 (4 * i + 1)) p (0..<PolyN)

    let initialise (p: *poly): *context =
        let deriv = derivative p
        let bound = compute_bound_poly p
        let radius = max_root_norm p
        in {p, deriv, bound, radius}

    -- Jacobi-style step where the new elements are computed in parallel from
    -- the previous values
    let step (ctx: context) (approx: *approx): *(bool, approx) =
        let pvals = map (eval_c ctx.p) approx
        let derivvals = map (evaln_c ctx.deriv (PolyN - 1)) approx
        let quos = map2 (cplx./) pvals derivvals
        let sums = map (\i ->
                            reduce_comm (cplx.+) (cplx.mk_re 0.0)
                                (map (\j ->
                                        if i == j then cplx.mk_re 0.0
                                        else cplx.mk_re 1.0 cplx./
                                                (approx[i] cplx.- approx[j]))
                                     (0..<N)))
                       (0..<N)
        let offsets = map2 (\quo sum -> quo cplx./ (cplx.mk_re 1.0 cplx.- quo cplx.* sum))
                           quos sums
        let approx = map2 (cplx.-) approx offsets
        let svals = map (eval_d ctx.bound <-< cplx.mag) approx
        let conditions = map2 (\p s -> cplx.mag p <= 1e-5 * s) pvals svals
        let all_converged = all id conditions
        in (all_converged, approx)

    let iterate (ctx: context) (rng: *rand_engine.rng): (rand_engine.rng, *approx) =
        let (rng, approx) = generate ctx rng
        let (init_conv, approx) = step ctx approx
        let (rng, _, _, _, approx) =
            loop (rng, conv, tries, step_idx, approx) =
                    (rng, init_conv, 1, 1: i32, approx)
            while !conv
            do if step_idx + 1 > tries * 100
               then let (rng, approx) = generate ctx rng
                    let (conv, approx) = step ctx approx
                    in (rng, conv, tries + 1, 0, approx)
               else let (conv, approx) = step ctx approx
                    in (rng, conv, tries, step_idx + 1, approx)
        in (rng, approx)

    let aberth (p: *poly) (rng: *rand_engine.rng): *(rand_engine.rng, approx) =
        iterate (initialise p) rng
}

-- Set the constant coefficient to 1; nextDerbyshire will never change it
let init_derbyshire: poly =
    let res = [1] ++ replicate (PolyN - 1) (-1)
    in res :> poly

let next_derbyshire (p: *poly): *(bool, poly) =
    let (_, p, looped) =
        loop (i, p, cont) = (0, p, true)
        while cont && i < length p
        do if p[i] == -1
           then (i,     p with [i] =  1, false)
           else (i + 1, p with [i] = -1, true)
    in (looped, p)

let derbyshire_at_index (index: i32): *poly =
    let bitfield = (index << 1) + 1
    in tabulate PolyN (\i -> f32.i32 (i32.get_bit (i32.i64 i) bitfield * 2 - 1))


let calc_index (value: f32) (left: f32) (right: f32) (steps: i32): i32 =
    i32.f32 ((value - left) / (right - left) * (f32.i32 steps - 1) + 0.5)

let point_index
        (width: i32) (height: i32)
        (bottom_left: complex) (top_right: complex)
        (pt: complex)
        : i32 =
    -- Range for 'yi' is reversed because image coordinates go down in the y
    -- direction, while complex coordinates go up in the y direction
    let xi = calc_index (cplx.re pt) (cplx.re bottom_left) (cplx.re top_right) width
    let yi = calc_index (cplx.im pt) (cplx.im top_right) (cplx.im bottom_left) height
    in if 0 <= xi && xi < width && 0 <= yi && yi < height
        then width * yi + xi
        else -1

entry main_job
        (start_index: i32) (num_polys: i32)
        (width: i32) (height: i32)
        (left: f32) (top: f32) (right: f32) (bottom: f32)
        (seed: i32)
        : []i32 =
    -- Unnecessary to give each polynomial a different seed
    let rng = rand_engine.rng_from_seed [seed]
    let bottom_left = cplx.mk left bottom
    let top_right = cplx.mk right top
    let indices = flatten
            (map (\idx ->
                    let p = derbyshire_at_index idx
                    let (_, pts) = aberth.aberth p (copy rng)
                    in map (point_index width height bottom_left top_right) pts)
                 (start_index ..< start_index + num_polys))
    let filtered = filter (\i -> i != -1) indices
    in reduce_by_index (replicate (i64.i32 width * i64.i32 height) 0)
                       (+) 0
                       (map i64.i32 filtered)
                       (map (const 1) filtered)

entry main_all
        (width: i32) (height: i32)
        (left: f32) (top: f32) (right: f32) (bottom: f32)
        (seed: i32)
        : []i32 =
    main_job 0 (i32.i64 (1 << N)) width height left top right bottom seed

entry get_N: i32 = i32.i64 N