module Diophantine where

import Data.Maybe

data PQa = PQa {getP :: Int, getQ :: Int, geta :: Int}

solvePQa :: Int -> Int -> Int -> ([PQa], [PQa])
-- assumption: q0 /= 0, D > 0, D /= k^2, p0^2 mod q0 = D mod q0
-- at j the 0-based position where second period starts so we need
-- elements as position 0,1..(j-1) i.e. j elements
solvePQa d p0 q0 = (uptoPeriodicPart, periodicPart)
  where
    periodicPart     = drop ir uptoPeriodicPart
    uptoPeriodicPart = take j series
    d'               = sqrt (fromIntegral d :: Double)
    getA p q = floor $ (fromIntegral p + d') / fromIntegral q
    -- x = (P + sqrt[D])/Q, x' = (P - sqrt[D])/Q is reduced if x > 1 and -1 < x' < 0
    isReduced p q = p' + d' > q' && -q' < p' - d' && p' - d' < 0
      where
        p' = fromIntegral p
        q' = fromIntegral q
    series = iterate
        (\(PQa p q a) ->
            let
                p' = a * q - p
                q' = (d - p' * p') `div` q
                a' = getA p' q'
            in PQa p' q' a'
        )
        (PQa p0 q0 (getA p0 q0))
    -- ir is the first 0-based position where (Pi + sqrt[D])/Qi is reduced
    (ir, PQa pir qir _) =
        head . filter ((\(PQa p q _) -> isReduced p q) . snd) . zip [0 ..] $ series
    -- j > ir is the 0-based positions where P(ir) = Pj and Q(ir) = Qj
    j =
        fst
            . head
            . filter ((\(PQa p q _) -> p == pir && q == qir) . snd)
            . drop (ir + 1)
            . zip [0 ..]
            $ series

solveLMM :: Int -> Int -> [(Int, Int)]
solveLMM d n = catMaybes $ concatMap solveF fs
  where
        -- solveLMM : fs =====================
    fs = go [1 .. floor (sqrt . fromIntegral $ n :: Double)]
      where
        go []       = []
        go (x : xs) = case n `quotRem` (x * x) of
            (m, 0) -> (x, m) : go xs
            _      -> go xs
    -- solveLMM : solveF ==================
    solveF (f, m) = map (getSolution f m . (\z -> solvePQa d z m'))
        -- find z so that -|m|/2 < z < |m|/2 and z^2 = D (mod |m|)
        $ filter (\z -> (z * z) `mod` m' == d `mod` m') [-m' `div` 2 .. m' `div` 2]
        where m' = abs m
    -- solveLMM : getSolution ==============
    getSolution f m (pqas, _)
        |
            -- if Qi = +-1 then we need G(i-1), B(i-1)
          v == m = Just (f * g, f * b)
        | v == -m = case getTU of
            (Just (t, u)) -> Just (f * (g * t + b * u * d), f * (g * u + b * t))
            _             -> Nothing
        | otherwise = Nothing
      where
            -- i is the 0-based first position where Qi = +-1
        i =
            fst
                . head
                . filter ((\q -> q == 1 || q == -1) . getQ . snd)
                . tail
                . zip [0 ..]
                $ pqas
        (g, b) = getGB (i - 1) pqas
        v      = g * g - d * b * b
    -- solveLMM : getGB ====================
    getGB k pqas = go (-p0, 1) (q0, 0) as !! k
      where
        p0 = getP . head $ pqas
        q0 = getQ . head $ pqas
        as = map geta pqas
        go _ _ [] = []
        go (g, b) (g', b') (a : as') =
            let
                g'' = a * g' + g
                b'' = a * b' + b
            in (g'', b'') : go (g', b') (g'', b'') as'
    -- solveLMM : getTU ===================
    getTU = -- solving x^2-dy^2=-1, p0=0, q0=1, l is odd then G(l-1), B(l-1)
        let
            (pqas', pqasPeriodic') = solvePQa d 0 1
            l                      = length pqasPeriodic'
        in if odd l then Just $ getGB (l - 1) (pqas' ++ pqasPeriodic') else Nothing
