module STExt where

import ST
import MyState
-- import Control.Lens

g2_1 :: Pos' ((x, y), z) ((p, y), z) x p
g2_1 = \ ((x, y), z) -> (x, \ x1 -> ((x1, y), z))

g2_2 :: Pos' ((x, y), z) ((x, p), z) y p
g2_2 = \ ((x, y), z) -> (y, \ y1 -> ((x, y1), z))


g3_1 :: Pos' ((x, y, z), w) ((p, y, z), w) x p
g3_1 = \ ((x, y, z), w) -> (x, \ x1 -> ((x1, y, z), w))

g3_2 :: Pos' ((x, y, z), w) ((x, p, z), w) y p
g3_2 = \ ((x, y, z), w) -> (y, \ y1 -> ((x, y1, z), w))

g3_3 :: Pos' ((x, y, z), w) ((x, y, p), w) z p
g3_3 = \ ((x, y, z), w) -> (z, \ z1 -> ((x, y, z1), w))

r2_1 :: Pos' (w, (x, y)) (w, (p, y)) x p
r2_1 = \ (w, (x, y)) -> (x, \ x1 -> (w, (x1, y)))

r2_2 :: Pos' (w, (x, y)) (w, (x, p)) y p
r2_2 = \ (w, (x, y)) -> (y, \ y1 -> (w, (x, y1)))


{-
int a = 0;

// 変な階乗
void foo(int n) {
    if (n == 0) {
        a += 1;
    } else {
        foo(n - 1);
        n *= 2;
        a *= n;
    }
}

int main(void) {
    foo(3);
    printf("%d ", a);
    foo(5);
    printf("%d ", a);

    return 0;
}
-}

callST :: y -> ST (x, y) a -> ST (x, z) a
callST y (ST st) = ST $ \ (x, z) ->
   let (a, (x1, _)) = st (x, y)
     in (a, (x1, z))

startST :: x -> ST (x, y) a -> ST (w, y) a
startST x (ST st) = ST $ \ (w, y) ->
   let (a, (_, y1)) = st (x, y)
     in (a, (w, y1))

foo :: Integer -> ST ((Integer, b), a) ()
foo n = callST (n, 0) foo'

foo' :: ST ((Integer, b), (Integer, y)) ()
foo' = do
    n <- get r2_1
    if n == 0 then do
        a <- get g2_1
        set g2_1 (a + 1)
    else do
        foo (n - 1)
        set r2_1 (n * 2)
        n <- get r2_1
        a <- get g2_1
        set g2_1 (a * n)

fooMain :: ST ((Integer, b), (Integer, y)) (Integer, Integer)
fooMain = do
  foo 3
  a <- get g2_1
  foo 5
  b <- get g2_1
  return (a, b)

v = evalST fooMain ((0,0), (0,0))

{-

int a = 0;

/*
int bar(int n) {
   if (n > 0) {
       bar(n / 2);
       a += n;
       bar(n - 2);
   }
}
*/

int bar(int n) {
   while (n > 0) {
       bar(n / 2);
       a += n;
       n -= 2;
   }
}

int main(void) {
    bar(10);
    printf("%d", a); // 62
    a = 0;
    bar(16)
    printf("%d", a); // 189

    return 0;
}
-}


bar :: Integer -> ST ((Integer, b), a) ()
bar n = callST (n, 0) bar'

bar' :: ST ((Integer, b), (Integer, y)) ()
bar' = do
    n <- get r2_1
    if n > 0 then do
        bar (n `div` 2)
        a <- get g2_1
        set g2_1 (a + n)
        set r2_1 (n - 2)
        bar'
    else
        return ()
  
barMain :: ST ((Integer, b), (Integer, y)) (Integer, Integer)
barMain = do
    set g2_1 0
    bar 10
    a <- get g2_1

    set g2_1 0
    bar 16
    b <- get g2_1
    return (a, b)

w = evalST barMain ((0,0), (0,0))