{-
Experiments around solving the Pell equation using continued
fractions.

Copyright 2020 Ken Takusagawa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or (at
your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

-}
{-# LANGUAGE ScopedTypeVariables, LambdaCase, PackageImports #-}
module Main where {
import System.Environment(getArgs);
import Control.Exception(assert);
import Debug.Trace(trace);
import Data.Function((&));
import Control.Category((>>>));
import Prelude hiding((.),(>>));
--import qualified Prelude;
import System.IO(hSetBuffering,stdout,BufferMode(LineBuffering));
--import System.IO(stderr,hPutStrLn);
import qualified Data.List as List;
--import qualified Control.Monad as Monad;
--import Control.Monad(guard);
import qualified Data.Maybe as Maybe;
--import qualified Data.Map as Map; import Data.Map(Map);
--import qualified Data.Set as Set; import Data.Set(Set);
--import qualified Data.Bifunctor as Bifunctor;
--import qualified Data.Tuple as Tuple;
--import Control.Monad.GenericReplicate(genericReplicateM); -- igenericReplicateM
--import Data.Functor((<&>)); --foreach = flip map
import qualified Num.ContinuedFraction as Cf;
import Data.List.NonEmpty(NonEmpty);
import qualified Data.List.NonEmpty as Nonempty;
import qualified Data.Ratio as Ratio;
import qualified Numeric.QuadraticIrrational as Quadraticirrational;

{-
cf package:
Math.ContinuedFraction is doing something far fancier than simple continued fractions.  it has sqrt but cannot get simple terms.
Math.ContinuedFraction.Simple seems to be doing simple continued fractions, but it does not have sqrt.
-}

main :: IO();
main = do{
hSetBuffering stdout LineBuffering;
getArgs >>= \case{
  [] -> nonsquares & map pellrecord & mapM_ print;
  ["do",n] -> pellrecord (read n) & print;
  ["increasingdifficulty"] -> nonsquares & map pellrecord & map extractlength & scanl (\bestsofar item -> if snd item > snd bestsofar then item else bestsofar) (0,0) & List.group & map head & tail & mapM_ print;
  ["increasingdifficultyandsharpen"] -> nonsquares & map pellrecord & map extractlength & scanl (\bestsofar item -> if snd item > snd bestsofar then item else bestsofar) (0,0) & List.group & map head & tail & map fst & map pellrecord & filter (\(_,(_,(_,m))) -> Maybe.isJust m) & mapM_ print;
  ["twoways"] -> nonsquares & countinggoby 100 & mapM_ (\d -> if not(twoways d) then error ("twoways did not match: " ++ show d) else return ());
  ["truncatetest"] -> nonsquares & countinggoby 100 & mapM_ (\d -> if not(truncatetest d) then error ("truncatetest failed: " ++ show d) else return ());
_ -> undefined;
}};

type Ii = Integer;

-- useful for (id :: Endo Integer) to assert a type in a pipeline
type Endo a = a->a;

-- to avoid the redundancy warning
trace_placeholder :: ();
trace_placeholder = (trace,assert) & (id >>> error) "trace_placeholder";
-- ghc -O2 -fno-ignore-asserts

newton :: Rational -> Rational -> Rational;
newton s x = (x+s/x)/2;

tononempty :: [a] -> NonEmpty a;
tononempty l = (head l) Nonempty.:| (tail l);

cftofraction :: [Ii] -> Rational;
cftofraction = tononempty >>> Cf.collapseFraction;

approxsqrt :: Ii -> Rational;
approxsqrt = fromInteger >>> (id :: Endo Double) >>> sqrt >>> round >>> fromInteger;

precision :: Ii;
precision = 200;

-- no longer used, we use sqrtinqi
mysqrt :: Ii -> Rational;
mysqrt x = iterate (newton (fromInteger x)) (approxsqrt x) & dropWhile (\y -> Ratio.denominator y < 10^(precision::Ii)) & head;

onestep :: Rational -> (Ii,Maybe Rational) -> Maybe (Maybe Rational,(Ii,Maybe Rational));
onestep target (d,mbestsofar) = Just $ let {
  n :: Ii;
  n = target*fromInteger d & round; -- what to do about rounding 1/2 ?
  q :: Rational;
  q = (fromInteger n)/(fromInteger d);
  diff :: Rational;
  diff = q-target & abs;
} in if (case mbestsofar of {Nothing -> True; Just old -> diff < old})
then (Just q,(d+1,Just diff))
else (Nothing,(d+1,mbestsofar));

modlast :: [Ii] -> [[Ii]];
modlast l = do {
s <- enumFromTo 1 (last l);
return $ init l ++ [s];
};

autofrac :: [Ii] -> [(Rational,[Ii])];
autofrac x = do {
n <- enumFrom 1;
y <- take n x & modlast;
return $ (cftofraction y,y);
};

evalpell :: Ii -> Rational -> Ii;
evalpell d f = let {
x = Ratio.numerator f;
y = Ratio.denominator f;
} in x*x-d*y*y;

plusminussolve :: Ii -> ([(Rational,[Ii])],[Ii]);
plusminussolve d = let{
cf :: [Ii];
--cf = mysqrt d & Cf.continuedFraction;
cf = sqrtinqi d & qitocf;
} in (autofrac cf & filter (\l -> 1 == abs(evalpell d (fst l))),cf);

pellrecord :: Ii -> (Ii,((Rational,[Ii]),(Bool,Maybe Rational)));
pellrecord d = let {
  allsols :: [(Rational,[Ii])];
  cf :: [Ii];
  (allsols,cf) = plusminussolve d;
  f1 :: (Rational,[Ii]);
  f1 = head allsols;
  cfused :: [Ii];
  cfused = snd f1;
  solution :: Rational;
  solution = fst f1;
  onemore :: Maybe Rational;
  onemore = if evalpell d solution < 0
    then Just $ newton (fromInteger d) solution
    else Nothing;
} in (d,(f1,(take (length cfused) cf == cfused,onemore)));

isperfectsquare :: Ii -> Bool;
isperfectsquare x = (approxsqrt x & round)^(2::Ii) == x;

extractlength :: (Ii,((Rational,[Ii]),a)) -> (Ii,Ii);
extractlength (i,((_,l),_)) = (i, List.genericLength l);

sqrtinqi :: Ii -> Quadraticirrational.QI;
sqrtinqi x = Quadraticirrational.qi 0 1 x 1;

qitocf :: Quadraticirrational.QI -> [Ii];
qitocf = Quadraticirrational.qiToContinuedFraction >>> \case {
(integerpart, Quadraticirrational.CycList start repeatend) -> integerpart:(start++cycle repeatend);};

wherediverge :: (Ord a) => [a] -> [a] -> Maybe Ii;
wherediverge lx ly = zipWith3 (\i x y -> (i,x/=y)) [0..] lx ly & filter snd & \case { [] -> Nothing ; l3 -> l3 & head & fst & Just};

testqi :: Ii -> (Ii,Maybe Ii);
testqi i = (i,wherediverge (sqrtinqi i & qitocf) (mysqrt i & Cf.continuedFraction));

nonsquares :: [Ii];
nonsquares = filter (isperfectsquare >>> not) [2..];

countinggoby :: Ii -> [a] -> [a];
countinggoby every = zipWith (\n -> if mod n every == 0 then trace (show n) else id) (enumFrom 0);

-- verify that one addition newton iteration of a negative pell solution yields the first positive solution.
-- verified to 34000.
twoways :: Ii -> Bool;
twoways d = let {
  allsols :: [Rational];
  allsols = plusminussolve d & fst & map fst;
  isnegative :: Bool;
  isnegative = evalpell d firstsolution < 0;
  firstsolution :: Rational;
  firstsolution = head allsols;
  onemore :: Rational;
  onemore = if isnegative
    then newton (fromInteger d) firstsolution
    else firstsolution;
  positivesolution :: Rational;
  positivesolution = allsols & filter (\x -> evalpell d x > 0) & head;
} in positivesolution == onemore;

-- verify that the continued fraction segment that yields a solution is a simple truncation of the infinite continued fraction.
-- verified to 50000.
truncatetest :: Ii -> Bool;
truncatetest = pellrecord >>> snd >>> snd >>> fst;
} --end