module Euler066Test (euler66suite) where

import Prelude

import Data.Maybe (Maybe(..))
import Effect.Aff (Milliseconds(..), delay)
import Euler066 (euler66)
import Test.Unit (TestSuite, suite, test)
import Test.Unit.Assert as Assert

euler66suite :: TestSuite
euler66suite =
  suite "Euler 66" do
    test "Warmup" do
      delay (Milliseconds 0.0)
      Assert.equal (Just 5) (euler66 7)

    test "Real" do
      delay (Milliseconds 0.0)
      Assert.equal (Just 661) (euler66 1000)

module Euler066 where

import Prelude

import Data.BigInt (BigInt, fromInt)
import Data.Foldable (maximumBy)
import Data.Int (floor, toNumber)
import Data.Lazy (defer)
import Data.List.Lazy (List(..), Step(..), catMaybes, filter, head, range, step)
import Data.Maybe (Maybe)
import Data.Number (sqrt)

irrational :: Int -> Boolean
irrational n = 
  let si = floor $ sqrt $ toNumber n
  in si * si /= n

continued_fraction :: Int -> List Int
continued_fraction d =
  let p0 = 0
      q0 = 1
      a0 = floor $ sqrt $ toNumber d
      go p q a =
        let pn = a * q - p
            qn = (d - pn * pn) / q
            an = (a0 + pn) / qn
        in List $ defer \_ ->
             Cons a $ go pn qn an
  in go p0 q0 a0

convergents :: Int -> List {d :: Int, p :: BigInt, q :: BigInt}
convergents d =
  case step $ continued_fraction d of
    Nil -> List $ defer \_ -> Nil
    Cons a0 as' ->
      case step as' of
        Nil -> List $ defer \_ -> Nil
        Cons a1 as'' ->
          let a0b = fromInt a0
              a1b = fromInt a1
              b1  = fromInt 1
              e0 = {d, p: a0b, q: b1}
              e1 = {d, p: a0b * a1b + b1, q: a1b}
              go _ _ Nil = List $ defer \_ -> Nil
              go ep ec (Cons an as''') =
                let anb = fromInt an
                    en = {d, p: anb * ec.p + ep.p, q: anb * ec.q + ep.q}
                in List $ defer \_ ->
                    Cons en $ go ec en (step as''')
          in List $ defer \_ ->
               Cons e0 $
                 List $ defer \_ ->
                  Cons e1 $ go e0 e1 (step as'')

solve_pell :: Int -> Maybe {d :: Int, p :: BigInt, q :: BigInt}
solve_pell d =
  let cs = convergents d
      f {p, q} = p * p - fromInt d * q * q == fromInt 1
  in head $ filter f cs

euler66 :: Int -> Maybe Int
euler66 = map (_.d) <<< maximumBy (comparing (_.p)) <<< catMaybes <<< map solve_pell <<< filter irrational <<< range 1