-- This file is part of Quipper. Copyright (C) 2011-2016. Please see the
-- file COPYRIGHT for a list of authors, copyright holders, licensing,
-- and other details. All rights reserved.
--
-- ======================================================================

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}

{-# OPTIONS -fcontext-stack=50 #-}

-- | This module contains the Quipper implementation of the Quantum
-- Linear Systems Algorithm.
--
-- The algorithm estimates the radar cross section for a FEM
-- scattering problem by using amplitude estimation to calculate
-- probability amplitudes.
--
-- The notations are based on the paper
--
-- * B. D. Clader, B. C. Jacobs, C. R. Sprouse. Quantum algorithm to
-- calculate electromagnetic scattering cross
-- sections. <http://arxiv.org/abs/1301.2340>.
module Algorithms.QLS.QLS where

import Quipper

import QuipperLib.QFT
import QuipperLib.Arith
import QuipperLib.Decompose

import Data.Complex
import qualified Data.Map as Map

import qualified Algorithms.QLS.TemplateOracle as Template
import Algorithms.QLS.QDouble
import Algorithms.QLS.QSignedInt
import Algorithms.QLS.CircLiftingImport
import Algorithms.QLS.Utils

import Libraries.Auxiliary(boollist_of_int_bh)

-- | The type of 'oracle_A' input arguments during runtime.
type OracleARunTime = Double  -- ^Value resolution.
       -> Int                 -- ^Band.
       -> Bool                -- ^Argflag.
       -> ([Qubit],[Qubit],[Qubit])  -- ^(x=index,y+node,z+value).
       -> Circ ([Qubit],[Qubit],[Qubit])

-- | The type of 'oracle_b' and 'oracle_r' input arguments during runtime.
type OracleBRRunTime = Double  -- ^Magnitude resolution.
        -> Double              -- ^Phase resolution.
        -> ([Qubit],[Qubit],[Qubit]) -- ^(x=index,m+magnitude,p+phase).
        -> Circ ([Qubit],[Qubit],[Qubit])

-- | A type to encapsulate all three oracles.
data Oracle = Oracle {
  oracle_A :: RunTimeParam -> OracleARunTime,
  oracle_b :: RunTimeParam -> OracleBRRunTime,
  oracle_r :: RunTimeParam -> OracleBRRunTime
}

-- | A set of oracles using only blackboxes.
dummy_oracle :: Oracle
dummy_oracle = Oracle {
  oracle_A = \d r i b x -> named_gate "Oracle A" x,
  oracle_b = \d1 d2 r x -> named_gate "Oracle b" x,
  oracle_r = \d1 d2 r x -> named_gate "Oracle r" x
}

-- | A type to hold the runtime parameters.
data RunTimeParam = RT_param {
  k :: Double,       -- ^ Wave number.
  theta :: Double,   -- ^ Incident wave angle.
  phi :: Double,     -- ^ Direction of desired far field radiation pattern.
  e0 :: Double,      -- ^ Incident wave amplitude.
  lambda :: Double,  -- ^ Wavelength.
  xlength :: Double, -- ^ /x/-length of square scattering region.
  ylength :: Double, -- ^ /y/-length of square scattering region.

  scatteringnodes :: [(Int,Int)], -- ^ Metallic region.

  nx :: Int,
  ny :: Int,
  lx :: Double,
  ly :: Double,

  kappa :: Double,
  epsilon :: Double,
  t0 :: Double,
  r :: Double,
  b_max :: Double,
  r_max :: Double,
  d  :: Int,
  nb :: Int,
  p2 :: Double,
  n0 :: Int,
  n1 :: Int,
  n2 :: Int,
  n4 :: Int,

  -- argflags for oracle_A
  magnitudeArgflag :: Bool,
  phaseArgflag :: Bool
} deriving (Show)

-- | A convenient set of runtime parameters for testing.
dummy_RT_param :: RunTimeParam
dummy_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k dummy_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda dummy_RT_param),
  ylength = 2.0*(lambda dummy_RT_param),

  scatteringnodes =
    let rt = dummy_RT_param in
    let xul = round((fromIntegral $ nx rt)/2)
              - round((xlength rt)/(2.0*(lx rt))) in -- Upper left x index
    let yul = round((fromIntegral $ ny rt)/2)
              - round((ylength rt)/(2.0*(ly rt))) in -- Upper left y index
    let xlr = round((fromIntegral $ nx rt)/2)
              + round((xlength rt)/(2.0*(lx rt))) in  -- Lower right x index
    let ylr = round((fromIntegral $ ny rt)/2)
              + round((ylength rt)/(2.0*(ly rt))) in -- Lower right y in
      [(xul, yul), (xlr, ylr)],

  nx = 12885,
  ny = 12885,
  lx = 0.1,
  ly = 0.1,

  kappa = 1.0,
  epsilon = 1.0,
  t0 = 1.0,
  r = 1.0,
  b_max = 1.0,
  r_max = 1.0,
  d  = 3,
  nb = 2,
  p2 = 3,
  n0 = 3,
  n1 = 3,
  n2 = 3,
  n4 = 3,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | A set of larger values, for testing scalability.
large_RT_param :: RunTimeParam
large_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k large_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda large_RT_param),
  ylength = 2.0*(lambda large_RT_param),

  scatteringnodes =
    let rt = large_RT_param in
    let xul = round((fromIntegral $ nx rt)/2)
              - round((xlength rt)/(2.0*(lx rt))) in -- Upper left x index
    let yul = round((fromIntegral $ ny rt)/2)
              - round((ylength rt)/(2.0*(ly rt))) in -- Upper left y index
    let xlr = round((fromIntegral $ nx rt)/2)
              + round((xlength rt)/(2.0*(lx rt))) in  -- Lower right x index
    let ylr = round((fromIntegral $ ny rt)/2)
              + round((ylength rt)/(2.0*(ly rt))) in -- Lower right y in
      [(xul, yul), (xlr, ylr)],

  nx = 12885,
  ny = 12885,
  lx = 0.1,
  ly = 0.1,

  kappa = 1e4,
  epsilon = 0.01,
  t0 = 7.0e6,
  r = 2.5e12,
  b_max = 5.0,
  r_max = 1.01,
  d  = 7,
  nb = 9,
  p2 = (  1.0 / (4-(4**(1/3)))  ),
  n0 = 14,
  n1 = 24,
  n2 = 30,
  n4 = 65,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | A set of smaller values, for manageable yet meaningful output.
small_RT_param :: RunTimeParam
small_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k small_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda small_RT_param),
  ylength = 2.0*(lambda small_RT_param),

  scatteringnodes =
    let xul = 2 in
    let yul = 2 in
    let xlr = 3 in
    let ylr = 3 in
      [(xul, yul), (xlr, ylr)],

  nx = 4,
  ny = 4,
  lx = 0.1,
  ly = 0.1,

  kappa = 1e4,
  epsilon = 0.01,
  t0 = 7.0e6,
  r = 2.5e12,
  b_max = 5.0,
  r_max = 1.01,
  d  = 7,
  nb = 9,
  p2 = (  1.0 / (4-(4**(1/3)))  ),
  n0 = 14,
  n1 = 24,
  n2 = 6,
  n4 = 65,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | Apply an [exp −/iYt/] gate. The timestep /t/ is a parameter.
expYt :: Timestep -> Qubit -> Circ Qubit
expYt = named_rotation "exp(-i%Y)"

-- | Apply an [exp −/iYt/] gate. The timestep /t/ is a parameter.
expYt_at :: Timestep -> Qubit -> Circ ()
expYt_at = named_rotation_at "exp(-i%Y)"

-- | Read a list of bits and make it into a 'Double', by multiplying
-- its integer value by the provided factor.
dynamic_lift_double :: Double -> [Bit] -> Circ Double
dynamic_lift_double factor cl = do
      cdiscard cl

-- Implementation note: removed dynamic_lift as for now it breaks all
-- output formats except ASCII.

--      bl <- dynamic_lift cl
      let sign = 1 -- if (head $ reverse bl) then 1 else -1
      let unsigned_value = 1 -- integer_of_intm_unsigned $
                              -- intm_of_boollist_bh (tail $ reverse bl)
      return (sign * factor * (fromIntegral unsigned_value))

-- | A black box gate to stand in as a replacement for QFT.
qft_for_show :: [Qubit] -> Circ [Qubit]
qft_for_show qs = named_gate "QFT" qs

-- | Main function: for estimating the radar cross section for a FEM
-- scattering problem. The problem can be reduced to the calculation
-- of four angles: φ[sub /b/], φ[sub /bx/], φ[sub /r/1] and φ[sub /r/0].
qlsa_FEM_main :: RunTimeParam -> Oracle -> Circ Double
qlsa_FEM_main param oracle = do
     comment "FEM_main"
     phi_b  <- qlsa_AmpEst_phi_b param oracle
     phi_bx <- qlsa_AmpEst_phi_bx param oracle
     phi_r1 <- qlsa_AmpEst_phi_bxr param oracle True
     phi_r0 <- qlsa_AmpEst_phi_bxr param oracle False
     let sigma = ((((fromIntegral $ nb param) ^ 2) * ((b_max param) ^ 2) * ((r_max param) ^ 2) * ((sin phi_b) ^ 2)) / ( 4 * pi))
     comment "FEM_main"
     return sigma

-- * Amplitude Estimation Functions

-- | Estimates φ[sub /b/], related to the probability of success for the
-- preparation of the known state /b/, using amplitude amplification.
qlsa_AmpEst_phi_b :: RunTimeParam -> Oracle -> Circ Double
qlsa_AmpEst_phi_b param oracle = do
    g <- qinit $ replicate (n0 param) False
    with_ancilla_init (replicate (n2 param) False) $ \x -> do
       label (g,x) ("g","x")
       with_ancilla $ \a -> do
           label (a) ("anc. a")
           with_ancilla $ \b -> do
               label (b) ("anc. b")
               g <- map_hadamard g
               u_b (x,b)
               loop g u_g (x,b,a)
               return ()
           return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
        loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
        loop [] f x = return ()
        loop (h:t) f x = do
            f x `controlled` h
            loop t f' x;
            where
               f' x = do f x; f x

        u_b :: ([Qubit],Qubit) -> Circ ()
        u_b xb = qlsa_StatePrep param xb (oracle_b oracle param) (1.0/(b_max param))

        u_g :: ([Qubit],Qubit,Qubit) -> Circ ()
        u_g (x,b,a) = do
            comment "U_g starts"
            gate_Z_at b
            -- For unitrary linear transformation, adjoint == inverse, hence reverse_...
            (reverse_generic_imp u_b) (x,b)
            qnot_at a `controlled` x .==. (map (\x -> False) x)
            gate_X_at a
            gate_Z_at a
            gate_X_at a
            qnot_at a `controlled` x .==. (map (\x -> False) x)
            u_b (x,b)
            comment "U_g ends"
            return ()

-- | Testing function for 'qlsa_AmpEst_phi_b'.
test_qlsa_AmpEst_phi_b :: Bool -> IO ()
test_qlsa_AmpEst_phi_b dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    print_simple GateCount (qlsa_AmpEst_phi_b param dummy_oracle)
    print_simple Preview (qlsa_AmpEst_phi_b param dummy_oracle)

-- | Estimates φ[sub /bx/], related to the probability of success in
-- computing solution value /x/.
qlsa_AmpEst_phi_bx :: RunTimeParam -> Oracle -> Circ Double
qlsa_AmpEst_phi_bx param oracle  = do
    g <- qinit (take (n0 param) (repeat False))
    with_ancilla_init (take (n2 param) (repeat False)) $ \x -> do
      with_ancilla $ \a -> do
          with_ancilla $ \b -> do
              with_ancilla $ \s -> do
                  g <- map_hadamard g
                  u_bx (x,b,s)
                  loop g u_g (x,b,s,a)
                  return ()
              return ()
          return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
        loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
        loop [] f x = return ()
        loop (h:t) f x = do
            f x `controlled` h
            loop t f' x
            where
            f' x = do f x; f x

        u_bx :: ([Qubit],Qubit,Qubit) -> Circ ()
        u_bx (x,b,s) = do
            qlsa_StatePrep param (x,b) (oracle_b oracle param) (1.0/(b_max param))
            qlsa_Solve_x param (x,s) oracle
            return ()

        u_g :: ([Qubit],Qubit,Qubit,Qubit) -> Circ ()
        u_g (x,b,s,a) = do --named_gate_at "Ug_phi_bx" (x,b,s,a)
            qnot_at a `controlled` b .&&. s
            gate_Z_at a
            qnot_at a `controlled` b .&&. s
            (reverse_generic_imp u_bx) (x,b,s)
            qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. b .==. 0 .&&. s .==. 0
            gate_X_at a
            gate_Z_at a
            gate_X_at a
            qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. b .==. 0 .&&. s .==. 0
            u_bx (x,b,s)
            return ()

-- | Estimates φ[sub /r/0] and φ[sub /r/1] (depending on the boolean
-- parameter), related to the overlap of the solution with the
-- arbitrary state /r/.
qlsa_AmpEst_phi_bxr :: RunTimeParam -> Oracle -> Bool -> Circ Double
qlsa_AmpEst_phi_bxr param oracle target = do
    g <- qinit (take (n0 param) (repeat False))
    with_ancilla_init (take (n2 param) (repeat False)) $ \x -> do
      with_ancilla_init (take (n2 param) (repeat False)) $ \y -> do
        with_ancilla $ \a -> do
          with_ancilla $ \b -> do
             with_ancilla $ \s -> do
                with_ancilla $ \r -> do
                    with_ancilla $ \c -> do
                        g <- map_hadamard g
                        u_r (x,y,b,s,r,c)
                        loop g u_g (x,y,b,s,r,c,a)
                        return ()
                    return ()
                return ()
             return ()
          return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
    loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
    loop [] f x = return ()
    loop (h:t) f x = do
       f x `controlled` h
       loop t f' x
       where
         f' x = do f x; f x

    u_r :: ([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit) -> Circ ()
    u_r (x,y,b,s,r,c) = do
        qlsa_Solve_xr param (x,y,b,s,r,c) oracle
        return ()

    u_g :: ([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit,Qubit) -> Circ ()
    u_g (x,y,b,s,r,c,a) = do --named_gate_at "Ug_phi_bxr" (x,y,b,s,r,c,a)
         qnot_at a `controlled` (b .==. 1 .&&. s .==. 1 .&&. r .==. 1 .&&. c .==. target)
         gate_Z_at a
         qnot_at a `controlled` (b .==. 1 .&&. s .==. 1 .&&. r .==. 1 .&&. c .==. target)
         (reverse_generic_imp u_r) (x,y,b,s,r,c)
         qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. [ q .==. 0 | q <- y ] .&&. b .==. 0 .&&. s .==. 0 .&&. r .==. 0 .&&. c .==. 0
         gate_X_at a
         gate_Z_at a
         gate_X_at a
         qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. [ q .==. 0 | q <- y ] .&&. b .==. 0 .&&. s .==. 0 .&&. r .==. 0 .&&. c .==. 0
         u_r (x,y,b,s,r,c)
         return ()

-- * State Preparation.

-- | Prepares a quantum state /x/, as specified by an oracle function,
-- entangled with a single qubit flag /q/ marking the desired state.
qlsa_StatePrep ::
    RunTimeParam
    -> ([Qubit], Qubit) -- x & qare handles to wires to be changed by StatePrep
    -> OracleBRRunTime -- Common type of (oracle_b oracle) and (oracle_r oracle), if oracle is typed Oracle
    -> Double
    -> Circ ()
qlsa_StatePrep param (x, q) oracle phi0 = do
  _ <- (flip $ box ("qlsa_StatePrep_" ++ (show phi0))) (x,q) $ \(x,q) -> do
    -- comment "StatePrep starts"
    label (x,q) ("x", "q")
    with_ancilla_list (n4 param) $ \m -> do
        label (m) ("anc. m")
        with_ancilla_list (n4 param) $ \p -> do
            label (p) ("anc. p")
            x <- map_hadamard x
            (x, m, p) <- oracle phi0 (epsilon param) (x, m, p)
            qlsa_ControlledPhase p (epsilon param) False
            qlsa_ControlledRotation (m, q) phi0 False
            (x, m, p) <- oracle phi0 (epsilon param) (x, m, p)
            return (x,q)
    -- comment "StatePrep ends"
  return ()

-- | Testing function for 'qlsa_StatePrep'.
test_qlsa_StatePrep :: Bool -> IO ()
test_qlsa_StatePrep dummyRTParamFlag = do
          let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
          let oraclebRunTime = (oracle_b dummy_oracle param)
          let testCirc = do
              x <- qinit $ replicate (n2 param) False
              q <- qinit False
              qlsa_StatePrep param (x,q) oraclebRunTime 1.0
          print_simple GateCount testCirc
          print_simple Preview testCirc

-- * Linear System Solver Functions

-- | Implements the QLSA procedure to generate the solution state |/x/〉.
qlsa_Solve_x :: RunTimeParam ->([Qubit],Qubit) -> Oracle -> Circ ()
qlsa_Solve_x param (x,s) oracle = do
  _ <- (flip $ box "qlsa_Solve_x") (x,s) $ \(x,s) -> do
    with_ancilla_list (n1 param) $ \t -> do
        with_ancilla_list (length t) $ \f -> do
            let phi0 = 2 * pi  / (2 ** (fromIntegral $ n1 param) - (epsilon param))
            t <- map_hadamard t
            u_hs (t,x)
            t <- qft_big_endian t
            integer_inverse (t,f)
            qlsa_ControlledRotation (f,s) phi0 False
            integer_inverse (t,f)
            (reverse_generic_endo qft_big_endian) t
            (reverse_generic_imp u_hs) (t,x)
            t <- map_hadamard t
            return()
        return ()
    return (x,s)
  return ()
    where
        u_hs :: ([Qubit],[Qubit]) -> Circ ()
        u_hs (t,x) = do
            qlsa_HamiltonianSimulation param (t,x) (oracle_A oracle param)
            return ()

-- | Implementation of the integer division. The two registers are
-- supposed to be of the same size and represent little-headian
-- unsigned integers, i.e., the head of the list holds the least
-- significant bit.
integer_inverse :: ([Qubit],[Qubit]) -> Circ ()
integer_inverse (t,f) = do
  _ <- (flip $ box "integer_inverse") (t,f) $ \(t,f) -> do
    -- sanity check
    if (length t /= length f)
      then error "integer_inverse: registers of distinct sizes"
      else return ()
    -- initialize an unsigned integer to 2^(length t) - 1
    with_ancilla_init (map (\_ -> True) t) $ \num -> do
      -- perform the division (encapsulated in a subroutine)
      let d = classical_to_reversible $ \(t,num) -> do
                let x = ((qdint_of_qulist_lh num),(qdint_of_qulist_lh t))
                (_,_,f') <- uncurry q_div_unsigned x
                return $ qulist_of_qdint_lh f'
      d ((t,num),f)
      return (t,f)
  return ()

-- | Implements the complete QLSA procedure to find the
-- solution state |/x/〉 and then implements the swap protocol
-- required for estimation of 〈/x/|/r/〉.
qlsa_Solve_xr :: RunTimeParam ->([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit) -> Oracle -> Circ ()
-- qlsa_Solve_xr param (x,y,b,s,r,c) oracle  = named_gate_at "Solve_xr" (x,y,b,s,r,c)
qlsa_Solve_xr param (x,y,b,s,r,c) oracle = do
  _ <- (flip $ box "qlsa_Solve_xr") (x,y,b,s,r,c) $ \(x,y,b,s,r,c) -> do
    qlsa_StatePrep param (x,b) (oracle_b oracle param) (1.0/(b_max param))
    qlsa_Solve_x param (x,s) oracle
    qlsa_StatePrep param (y,r) (oracle_r oracle param) (1.0/(r_max param))
    hadamard_at c
    swap_at y x  `controlled` c
    hadamard_at c
    return (x,y,b,s,r,c)
  return ()

-- * Hamiltonian Simulation Functions.

-- | Uses a quantum register |/t/〉 to control the
-- implementation of the Suzuki method for simulating a Hamiltonian
-- specified by an oracle function.
qlsa_HamiltonianSimulation ::
    RunTimeParam
    -> ([Qubit], [Qubit])
    -> OracleARunTime
    -> Circ ()
qlsa_HamiltonianSimulation param (t, x) oracleA = do
  _ <- (flip $ box "qlsa_HamiltonianSimulation") (t,x) $ \(t,x) -> do
    -- Code the first line in a way that depends on the length of t rather
    -- explicitly on (n1 param) which is supposed to be the length of t
    -- Hence replaced the line below by the line after it.
    -- let denom = 2 * (r param) * (fromIntegral ((n1 param) - 1))
    let denom = 2 * (r param) * ( 2^((length t) - 1) )
    let t1 = (p2 param) * (t0 param) / denom
    let t2 = (1.0 - 4.0 * (p2 param)) * (t0 param) / denom
    (t,x) <- box_loopM  "TrotterLoop" (round $ r param) (t,x) (hs_loop t1 t2)
    return (t,x)
  return ()
    where
          hs_loop :: Double -> Double -> ([Qubit], [Qubit]) -> Circ ([Qubit], [Qubit])
          hs_loop t1 t2 (t,x) = do
            u_z_at (t,x) t1
            u_z_at (t,x) t1
            u_z_at (t,x) t2
            u_z_at (t,x) t1
            u_z_at (t,x) t1
            return (t,x)

          u_z_at :: ([Qubit], [Qubit]) -> Double -> Circ ()
          u_z_at (t, x) timeStep = do
              for (nb param) 1 (-1) $ \jj -> do
                  qlsa_HsimKernel param (t, x) jj timeStep oracleA
                  endfor
              for 1 (nb param) 1 $ \jj -> do
                  qlsa_HsimKernel param (t, x) jj timeStep oracleA
                  endfor
              return ()

-- | Testing function for 'qlsa_HamiltonianSimulation'.
test_qlsa_HamiltonianSimulation :: Bool -> IO ()
test_qlsa_HamiltonianSimulation dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let oracleARunTime = (oracle_A dummy_oracle param)
    let testCirc = do
              t <- qinit (take (n0 param) (repeat False))
              x <- qinit (take (n4 param) (repeat False))
              label(t,x) ("t", "x")
              qlsa_HamiltonianSimulation param (t,x) oracleARunTime
    print_simple GateCount testCirc
--    print_simple Preview testCirc

-- | Uses an oracle function and timestep control register (/t/) to
-- apply 1-sparse Hamiltonian to the input state |/t/, /x/〉.
qlsa_HsimKernel ::
    RunTimeParam
    -> ([Qubit], [Qubit])
    -> Int
    -> Double
    -> OracleARunTime
    -> Circ ()
-- qlsa_HsimKernel param tx band timeStep oracleA = named_gate_at "HsimKernel" tx

qlsa_HsimKernel param (t, x) band timeStep oracleA = do
  _ <- (flip $ box ("qlsa_HsimKernel" ++ (show band) ++ (show timeStep))) (t,x) $ \(t,x) -> do
    let phiP = (epsilon param)
    with_ancilla_list (n2 param) $ \y -> do
       with_ancilla_list (n4 param) $ \m -> do
           with_ancilla_list (n4 param) $ \p -> do
               label (y,m,p) ("y","m","p")
               -- phases
               oracleA phiP band (phaseArgflag param) (x,y,p)
               qlsa_ControlledPhase p phiP False
               oracleA phiP band (phaseArgflag param) (x,y,p)
               -- magnitudes
               let phiMag = 2 ** (negate $ fromIntegral $ after_radix_length)
               oracleA phiMag band (magnitudeArgflag param) (x,y,m)
               for 0 ((length t) - 1) 1 $ \ii -> do
                   let phi_mt = timeStep * phiMag * (2^ii)
                   qlsa_ApplyHmag param (x,y,m) phi_mt `controlled` (t !! ii)
                   endfor
               oracleA phiMag band (magnitudeArgflag param) (x,y,m)
               -- phases again
               oracleA phiP band (phaseArgflag param) (x,y,p)
               qlsa_ControlledPhase p phiP True
               oracleA phiP band (phaseArgflag param) (x,y,p)
               return ()
           return ()
       return ()
    return (t,x)
  return ()

-- | Testing function for 'qlsa_HsimKernel'.
test_qlsa_HsimKernel :: Bool -> IO ()
test_qlsa_HsimKernel dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let oracleARunTime = (oracle_A dummy_oracle param)
    let testCirc = do
              t <- qinit (take (n0 param) (repeat False))
              x <- qinit (take (n4 param) (repeat False))
              label(t,x) ("t", "x")
              qlsa_HsimKernel param (t,x) 1 0.1 oracleARunTime
    print_simple GateCount testCirc
--    print_simple Preview testCirc

-- | Applies the magnitude component of coupling elements in a
-- 1-sparse Hamiltonian.
qlsa_ApplyHmag :: RunTimeParam -> ([Qubit], [Qubit], [Qubit]) -> Double -> Circ ()
qlsa_ApplyHmag param (x,y,m) phi0 = do
  _ <- (flip $ box ("qlsa_ApplyHmag " ++ (show phi0))) (x,y,m) $ \(x,y,m) -> do
    let (onOne, onZero) = (True, False)
    with_ancilla $ \a -> do -- Assume initialized to False (0)
        label (a) ("anc. a")
        if (length x /= length y)
           then error "qlsa_ApplyHmag: Input registers x and y have different lengths."
           else return ()
        let length_xy = length x
        for 0 (length_xy - 1) 1 $ \ii -> do
            let (xi, yi) = (x !! ii, y !! ii)
            w (xi, yi)
            qnot_at a `controlled` (xi .==. onOne .&&. yi .==. onZero)
            endfor
        qlsa_ControlledPhase (m ++ [a]) phi0 False
        for (length_xy - 1) 0 (-1) $ \ii -> do
            let (xi, yi) = (x !! ii, y !! ii)
            qnot_at a `controlled` (xi .==. onOne .&&. yi .==. onZero)
            w (xi, yi)
            endfor
        return ()
    return (x,y,m)
  return ()

-- | Testing function for 'qlsa_ApplyHmag'.
test_qlsa_ApplyHmag :: Bool -> IO ()
test_qlsa_ApplyHmag dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let testCirc = do
              x <- qinit (take (n2 param) (repeat False))
              y <- qinit (take (n2 param) (repeat False))
              m <- qinit (take (n4 param) (repeat False))
              label(x,y,m) ("x", "y", "m")
              qlsa_ApplyHmag param (x,y,m) 0.1
    print_simple GateCount testCirc
    print_simple Preview testCirc

-- | Auxiliary function: the /W/-gate.
w :: (Qubit, Qubit) -> Circ ()
-- w xy = named_gate_at "W" xy
w (xi,yi) = do
  _ <- box "w" (\(xi, yi) -> do
    label (xi,yi) ("x[i]","y[i]")
    gate_X_at xi `controlled` yi
    gate_X_at yi `controlled` xi
    hadamard_at yi `controlled` xi
    gate_X_at yi `controlled` xi
    gate_X_at xi `controlled` yi
    return (xi,yi)) (xi,yi)
  return ()

-- | Testing function for 'w'.
test_w :: IO ()
test_w = do
    let testCirc = do
        xi <- qinit False
        yi <- qinit False
        w (xi, yi)
    print_simple GateCount testCirc
    print_simple Preview testCirc

-- * Controlled Logic Operations

-- | Applies a phase shift of φ\/2 to the signed input register |φ〉.

-- For c, bit locations are counted starting from 0 (least significant) to (length c - 2) (most significant)
-- last c (or c !! (n-1)) is the sign bit
qlsa_ControlledPhase :: [Qubit] -> Double -> Bool -> Circ ()
-- qlsa_ControlledPhase c phi0 f = named_gate_at "CPhase" c
qlsa_ControlledPhase c phi0 f = do
  _ <- (flip $ box ("qlsa_ControlledPhase " ++ (show phi0) ++ " " ++ (show f))) c $ \c -> do
    with_ancilla $ \a -> do -- ancilla is initialized to False
        if f then (qnot_at a) else return ()
        let signQubit = last c
        qnot_at a `controlled` signQubit
        for 0 (length c - 2) 1 $ \ii -> do
            let theta = ( (2.0^(ii)) * phi0 / 2.0) -- Note the divide by 2
            -- The following three lines are equivalent
            expZt_at theta a `controlled` (c !! ii)
            -- a <- expZt theta a `controlled` (c !! ii); return ()
            -- qlsa_ControlledU (c !! ii, a) (expZt theta) False
            endfor
        qnot_at a `controlled` signQubit
        return c
  return ()

-- | Applies a rotation of φ\/2 to the signed input register |φ〉.
qlsa_ControlledRotation :: ([Qubit], Qubit) -> Double -> Bool -> Circ ()
-- qlsa_ControlledRotation ct phi0 f = named_gate_at "CRotate" ct
qlsa_ControlledRotation (c, t) phi0 f = do
  _ <- (flip $ box ("qlsa_ControlledRotation " ++ (show phi0) ++ " " ++ (show f))) (c,t) $ \(c,t) -> do
    if f then (qnot_at t) else return ()
    let signQubit = last c
    qnot_at t `controlled` signQubit
    for 0 (length c - 2) 1 $ \ii -> do
        let theta = (2.0^(ii)) * phi0 / 2.0
        expYt_at theta t `controlled` (c !! ii)
        endfor
    qnot_at t `controlled` signQubit
    if f then (qnot_at t) else return ()
    return (c,t)
  return ()

----------------------------------------------------------------------
-- * Oracles

-- | Map a 'QDouble' into an integer, understood as being scaled by
-- the given factor. Take the factor and the size of the output
-- register as parameter.
make_factor_rep :: Double -> Int -> QDouble -> Circ [Qubit]
make_factor_rep factor size p = do
     -- number of high bits of p
     let p_int_size = (xdouble_length p) - (xdouble_exponent p)
     -- get the required number of bits for doing the encoding
     let auxsize = max (size - 1) p_int_size
     -- do the operation:
     --    (1) build a QDouble from the factor: need (size-1) bits of integer part
     qfactor <- qinit $ fdouble_of_double (xdouble_exponent p)
                                          (auxsize + xdouble_exponent p) factor
     --    (2) scale p
     p_large <- qdouble_pad_to_extent (auxsize,- (xdouble_exponent p)) p
     --    (3) divide p by factor
     qreal_multiples <- unpack template_symb_slash_ p_large qfactor
     --    (4) get the floor of the result
     q_floor <- template_floor
     qmultiples <- q_floor qreal_multiples
     --    (5) set them in the right format (it has size elements)
     let (SInt tp bp) = qmultiples
     let new_p = (take (size - 1) $ reverse tp) ++ [bp]
     return new_p

-- | Implements the oracle for the arbitrary state |/r/〉, using the
-- Template Haskell implementation of 'calcRweights'.
inline_oracle_r :: RunTimeParam -> Double -> Double -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_r rt factor_m factor_p =  box ("Or " ++ (show factor_m) ++ " " ++ (show factor_p)) $ decompose_generic Toffoli $ \(x',m',p') ->
    with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx ->
     with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
      with_ancilla_init (fdouble $ lx rt) $ \qlx ->
       with_ancilla_init (fdouble $ ly rt) $ \qly ->
        with_ancilla_init (fdouble $ k rt) $ \qk ->
         with_ancilla_init (fdouble $ theta rt) $ \qtheta ->
          with_ancilla_init (fdouble $ phi rt) $ \qphi ->
                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx
                    let f = classical_to_reversible $
                              \(x,nx,ny,lx,ly,k,theta,phi) -> do
                                 (m,p) <- unpack Template.template_calcRweights
                                            x nx ny lx ly k theta phi
                                 new_p <- make_factor_rep factor_p (n4 rt) p
                                 new_m <- make_factor_rep factor_m (n4 rt) m
                                 return (new_m, new_p)
                    f ((x,qnx,qny,qlx,qly,qk,qtheta,qphi),(m',p'))
                    return (x',m',p')

-- | Implements the oracle for the known state |/b/〉, using the
-- Template Haskell implementation of 'getKnownWeights'.
inline_oracle_b :: RunTimeParam -> Double -> Double -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_b rt factor_m factor_p = box ("Ob " ++ (show factor_m) ++ " " ++ (show factor_p)) $ decompose_generic Toffoli $ \(x',m',p') ->
    -- Make ancillas for constant values

    with_ancilla_init (listpair_fmap fromIntegral $ scatteringnodes rt :: [(FDouble,FDouble)]) $ \qscatteringnodes ->
     with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
      with_ancilla_init (fdouble $ lx rt) $ \qlx ->
       with_ancilla_init (fdouble $ ly rt) $ \qly ->
        with_ancilla_init (fdouble $ k rt) $ \qk ->
         with_ancilla_init (fdouble $ theta rt) $ \qtheta ->
          with_ancilla_init (fdouble $ e0 rt) $ \qe0 ->
           with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx -> do

                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx

                    let f = classical_to_reversible $
                              \(y,nx,ny,scatteringnodes,lx,ly,k,theta,e0) -> do
                                  -- get some QSInt and QDouble
                                  (m,p) <- unpack Template.template_getKnownWeights
                                              y nx ny scatteringnodes lx ly k theta e0 7
                                  new_p <- make_factor_rep factor_p (n4 rt) p
                                  new_m <- make_factor_rep factor_m (n4 rt) m

                                  return (new_m, new_p)

                    f ((x,qnx,qny,qscatteringnodes,qlx,qly,qk,qtheta,qe0),(m',p'))
                    return (x',m',p')

-- | Implementation of the oracle calculating the matrix /A/
-- corresponding to the discretization of the scattering problem,
-- using the Template Haskell implementation of
-- 'getNodeValuesMoreOutputs'.
inline_oracle_A ::  RunTimeParam -> Double -> Int -> Bool -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_A rt factor band argflag =  box ("OA " ++ (show band) ++ " " ++ (show argflag)) $ decompose_generic Toffoli $ \(x',y',p') -> do

    let argflag' = if argflag then PTrue else PFalse

    -- Make ancillas for constant values

    with_ancilla_init (listpair_fmap fromIntegral $ scatteringnodes rt :: [(FDouble, FDouble)]) $ \qscatteringnodes ->
      with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
        with_ancilla_init (fromRational $ toRational $ lx rt) $ \qlx ->
          with_ancilla_init (fromRational $ toRational $ ly rt) $ \qly ->
            with_ancilla_init (fromRational $ toRational $ k rt) $ \qk ->
             with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx -> do

                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx

                    let f = classical_to_reversible $
                              \(v,nx,ny,scatteringnodes,lx,ly,k) -> do
                               -- get some QSInt and QDouble
                               (y,p) <- unpack Template.template_getNodeValuesMoreOutputs
                                           v band nx ny scatteringnodes lx ly k argflag' 7

                               new_p <- make_factor_rep factor (n4 rt) p

                               -- set y in the correct format (it has n2 elements)
                               -- we can assume that the sign is positive.
                               -- we also assume that n2 < size of QSInt register
                               let (SInt ty _) = y
                               let new_y = (take (n2 rt) $ reverse ty)

                               -- return the values in the right format.
                               return (new_y,new_p)

                    f ((x,qnx,qny,qscatteringnodes,qlx,qly,qk),(y',p'))

                    return (x',y',p')

-- | Encapsulate the inline oracles in Template Haskell into an object
-- of type 'Oracle'.
inline_oracle :: Oracle
inline_oracle = Oracle {
  oracle_A = inline_oracle_A,
  oracle_b = inline_oracle_b,
  oracle_r = inline_oracle_r
}


{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE QuasiQuotes      #-}

module Layers
  ( ForwardLayer(..)
  , OutputLayer(..)
  , ForwardNN(..)
  , BackwardLayer(..)
  , BackputLayer(..)
  , BackwardNN(..)
  , TrainBatch(..)
  , NElement(..)
  , (<~)
  , (~>)
  , forward
  , output
  , backward
  , backput
  , learnForward
  , learnBackward
  , learn
  , learnAll
  , predict
  , evaluate
  ) where

import           Control.Arrow
import           Control.Lens            hiding ((<~))
import           Data.String.Interpolate as S (i)
import           Debug.Trace
import           Numeric
import           Numeric.LinearAlgebra
import           Synapse

type Weight a = Matrix a

type Bias a = Vector a

type SignalX a = Matrix a

type SignalY a = Matrix a

type Diff a = Matrix a

type TeacherBatch a = Matrix a

type InputBatch a = Matrix a

newtype TrainBatch a =
  TrainBatch (TeacherBatch a, InputBatch a)
  deriving (Show)

data ForwardLayer a
  = AffineForward (Weight a) (Bias a)
  | SigmoidForward
  | ReLUForward
  | JoinedForwardLayer (ForwardLayer a) (ForwardLayer a)
  deriving (Show)

instance (Numeric a, Eq a) => Eq (ForwardLayer a) where
  AffineForward w b == AffineForward w' b' = w == w' && b == b'
  SigmoidForward == SigmoidForward = True
  ReLUForward == ReLUForward = True
  JoinedForwardLayer a b == JoinedForwardLayer a' b' = a == a' && b == b'
  _ == _ = False

infixl 4 ~>

(~>) :: ForwardLayer a -> ForwardLayer a -> ForwardLayer a
a ~> (JoinedForwardLayer x y) = (a ~> x) ~> y
a ~> b = JoinedForwardLayer a b

data OutputLayer a =
  SoftmaxWithCrossForward
  deriving (Show)

data ForwardNN a =
  ForwardNN (ForwardLayer a) (OutputLayer a)
  deriving (Show)

data BackwardLayer a
  = AffineBackward (Weight a) (Bias a) (SignalX a)
  | SigmoidBackward (SignalY a)
  | ReLUBackward (SignalX a)
  | JoinedBackwardLayer (BackwardLayer a) (BackwardLayer a)
  deriving (Show)

instance (Numeric a, Eq a) => Eq (BackwardLayer a) where
  AffineBackward w b d == AffineBackward w' b' d' =
    w == w' && b == b' && d == d'
  SigmoidBackward y == SigmoidBackward y' = y == y'
  ReLUBackward x == ReLUBackward x' = x == x'
  JoinedBackwardLayer a b == JoinedBackwardLayer a' b' = a == a' && b == b'
  _ == _ = False

infixr 4 <~

(<~) :: BackwardLayer a -> BackwardLayer a -> BackwardLayer a
(JoinedBackwardLayer x y) <~ b = x <~ (y <~ b)
a <~ b = JoinedBackwardLayer a b

data BackputLayer a =
  SoftmaxWithCrossBackward (TeacherBatch a) (SignalY a)

data BackwardNN a =
  BackwardNN (BackwardLayer a) (BackputLayer a)

type NElement a = (Ord a, Floating a, Numeric a, Num (Vector a), Show a)

forward ::
     NElement a => ForwardLayer a -> SignalX a -> (BackwardLayer a, SignalY a)
forward (AffineForward w b) x = (AffineBackward w b x, affinem w b x)
forward SigmoidForward x = (SigmoidBackward y, y)
  where
    y = sigmoidm x
forward ReLUForward x = (ReLUBackward x, relum x)
forward (JoinedForwardLayer a b) x0 = (a' <~ b', x2)
  where
    (a', x1) = forward a x0
    (b', x2) = forward b x1

backward ::
     NElement a => a -> BackwardLayer a -> Diff a -> (ForwardLayer a, Diff a)
backward rate (AffineBackward w b x) d =
  (AffineForward (w - scale rate w') (b - scale rate b'), x')
  where
    (x', w', b') = affinemBackward w x d
backward _ (SigmoidBackward y) d = (SigmoidForward, sigmoidBackward y d)
backward _ (ReLUBackward x) d = (ReLUForward, relumBackward x d)
backward r (JoinedBackwardLayer a b) d0 = (a' ~> b', d2)
  where
    (b', d1) = backward r b d0
    (a', d2) = backward r a d1

output ::
     NElement a
  => OutputLayer a
  -> TeacherBatch a
  -> SignalY a
  -> (BackputLayer a, a)
output SoftmaxWithCrossForward t y = (SoftmaxWithCrossBackward t y', loss)
  where
    (y', loss) = softmaxWithCross t y

backput :: NElement a => BackputLayer a -> (OutputLayer a, Diff a)
backput (SoftmaxWithCrossBackward t y) =
  (SoftmaxWithCrossForward, softmaxWithCrossBackward t y)

learnForward :: NElement a => ForwardNN a -> TrainBatch a -> (BackwardNN a, a)
learnForward (ForwardNN layers loss) (TrainBatch (t, x)) =
  result `seq` (BackwardNN layers' loss', result)
  where
    (layers', y) = forward layers x
    (loss', result) = output loss t y

learnBackward :: NElement a => a -> BackwardNN a -> ForwardNN a
learnBackward rate (BackwardNN layers loss) = ForwardNN layers' loss'
  where
    (loss', d) = backput loss
    (layers', _) = backward rate layers d

learn :: NElement a => a -> ForwardNN a -> TrainBatch a -> (ForwardNN a, a)
learn rate a = first (learnBackward rate) . learnForward a

learnAll ::
     NElement a => a -> ForwardNN a -> [TrainBatch a] -> (ForwardNN a, [a])
learnAll rate origin batches = ls `seq` (nn, ls)
  where
    (nn, ls) = foldr f (origin, []) batches
    f batch (a, ls) = ls `seq` second (~: ls) $ learn rate a batch
    x ~: xs = trace [i|[#{length xs + 1}/#{n}] #{show x}|] (x : xs)
    n = length batches

predict :: NElement a => ForwardLayer a -> Vector a -> Int
predict layers = maxIndex . flatten . snd . forward layers . asRow

evaluate :: NElement a => ForwardLayer a -> [(Int, Vector a)] -> Double
evaluate layers samples = fromIntegral nOk / fromIntegral (length samples)
  where
    nOk = length $ filter (uncurry (==)) results
    results = map (second $ predict layers) samples


{-# LANGUAGE OverloadedStrings #-}

module Main where

import qualified    Data.Text.IO as T
import qualified    Data.Text as T
import              Data.Text (Text)
import              Data.Text.Encoding (encodeUtf8)

import qualified    Data.ByteString.Lazy as B
import              Data.ByteString.Lazy (ByteString)
import qualified    Data.Map as Map
import              Data.Map (Map)
import qualified    Data.Vector as V
import              Data.Vector (Vector)
import              Data.Csv
import              Data.Monoid

import              Parser
import              Type

import              Control.Monad
import              Control.Monad.Primitive
import              System.Random.MWC
import              Statistics.Distribution
import              Statistics.Distribution.Laplace

type County a = Map Text a
type District a = Map Text (County a)
type FlatDistrict a = Map Text a

iteration :: Int
iteration = 1000

main :: IO ()
main = do
    raw <- B.readFile "asset/sim_2.csv"
    case parse raw of
        Left err -> putStrLn err
        Right rows -> do
            let nation = buildUp rows
            calculateNation nation
            calculateCounty nation
            calculateDistrict nation

calculateNation :: District (Vector Row) -> IO ()
calculateNation nation = do
    putStrLn "全國平均就醫距離"
    let header = V.fromList [encodeUtf8 "全國"] :: Vector Name
    let result = [NationResult (nationAvg nation)]
    B.writeFile "result/original/nationwide.csv" (encodeByName header result)

    putStrLn "全國平均就醫距離 (added noise)"
    forM_ [-9 .. 5] $ \level -> do
        results <- replicateM iteration (nationAddNoise level nation)
        B.writeFile ("result/pertubated/nationwide/2^" <> show level <> ".csv") (encodeByName header results)
--
calculateCounty :: District (Vector Row) -> IO ()
calculateCounty nation = do
    putStrLn "各縣市平均就醫距離"
    let header = V.fromList (map encodeUtf8 (Map.keys nation)) :: Vector Name
    let avgs = fmap countyAvg nation
    let result = [CountyResult avgs]
    B.writeFile "result/original/countywide.csv" (encodeByName header result)

    withSystemRandom . asGenIO $ \gen -> do
        putStrLn "各縣市平均就醫距離 (added noise)"
        forM_ [-4 .. 10] $ \level -> do
            results <- replicateM iteration (countyAddNoise level nation gen)
            B.writeFile ("result/pertubated/countywide/2^" <> show level <> ".csv") (encodeByName header results)

calculateDistrict :: District (Vector Row) -> IO ()
calculateDistrict nation = do
    let flattened = flatten nation
    putStrLn "各鄉鎮市區的平均就醫距離"
    let header = V.fromList $ map encodeUtf8 $ Map.keys flattened
    let result = [DistrictResult (flatten $ fmap (fmap districtAvg) nation)]
    B.writeFile "result/original/districtwide.csv" (encodeByName header result)

    putStrLn "各鄉鎮市區的平均就醫距離 (added noise)"

    withSystemRandom . asGenIO $ \gen -> do
        forM_ [5 .. 20] $ \level -> do
            results <- replicateM iteration (districtAddNoise level flattened gen)
            B.writeFile ("result/pertubated/districtwide/2^" <> show level <> ".csv") (encodeByName header results)

nationAddNoise :: Int -> District (Vector Row) -> IO NationResult
nationAddNoise level nation = do
    let avg = nationAvg nation
    pertubated <- addNoise scale avg
    return (NationResult pertubated)
    where
        size = Map.foldl (\s elem -> s + (Map.foldl (\s vec -> s + V.length vec) 0 elem)) 0 nation
        eps = 2.0 ^^ level
        scale = 400000.0 / (fromIntegral size * eps)
            -- pertubated <- addNoise' scale (replicate iteration (nationwideAvg mapping))

countyAddNoise :: Int -> District (Vector Row) -> Gen (PrimState IO) -> IO CountyResult
countyAddNoise level nation gen = do
    pertubated <- mapM (\county -> do
        let avg = countyAvg county
        noise <- genContVar (laplace 0 (scale county)) gen
        return (avg + noise)) nation
    return (CountyResult pertubated)

    where
        countySize = Map.foldl (\s vec -> s + V.length vec) 0
        eps = 2.0 ^^ level
        scale county = 400000.0 / (fromIntegral (countySize county) * eps)

districtAddNoise :: Int -> FlatDistrict (Vector Row) -> Gen (PrimState IO) -> IO DistrictResult
districtAddNoise level flattened gen = do
    pertubated <- mapM (\district -> do
        let avg = districtAvg district
        noise <- genContVar (laplace 0 (scale district)) gen
        return (avg + noise)) flattened
    return (DistrictResult pertubated)

    where
        districtSize = V.length -- Map.foldl (\s vec -> s + V.length vec) 0
        eps = 2.0 ^^ level
        scale district = 400000.0 / (fromIntegral (districtSize district) * eps)

data NationResult = NationResult Double deriving (Show)
data CountyResult = CountyResult (County Double) deriving (Show)
data DistrictResult = DistrictResult (FlatDistrict Double) deriving (Show)

instance ToNamedRecord NationResult where
    toNamedRecord (NationResult value) = namedRecord
        [ (encodeUtf8 "全國") .= value ]

instance ToNamedRecord CountyResult where
    toNamedRecord (CountyResult value)
        = namedRecord (map (\(key, val) -> (encodeUtf8 key) .= val) (Map.assocs value))

flatten :: District a -> FlatDistrict a
flatten = Map.fromList . concat . map fuseName . Map.assocs
    where
        fuseName :: (Text, County a) -> [(Text, a)]
        fuseName (countyName, county) = map (\(districtName, distrct) -> (countyName <> districtName, distrct)) (Map.assocs county)

instance ToNamedRecord DistrictResult where
    toNamedRecord (DistrictResult nation)
        = namedRecord
            $ Map.elems
            $ Map.mapWithKey (\name value -> encodeUtf8 name .= value)
            $ nation

districtAvg :: Vector Row -> Double
districtAvg = average . V.toList . fmap distance

countyAvg :: County (Vector Row) -> Double
countyAvg = average . concat . Map.elems . fmap (V.toList . fmap distance)
    -- average . Map.elems . fmap districtAvg
    -- Map.foldl (\ acc distrct -> acc + districtAvg distrct) 0.0

nationAvg :: District (Vector Row) -> Double
nationAvg = average . concat . Map.elems . fmap (concat . Map.elems . fmap (V.toList . fmap distance))

buildUp :: Vector Row -> District (Vector Row)
buildUp = foldl insertCounty Map.empty
    where

        districtSingleton :: Row -> County (Vector Row)
        districtSingleton row = Map.singleton (snd (key row)) (V.singleton row)

        insertCounty :: District (Vector Row) -> Row -> District (Vector Row)
        insertCounty m row = Map.insertWith
            (insertDistrict row)
            (fst (key row))         -- key for TierCounty
            (districtSingleton row)    -- key for TierCounty (which is TierDistrict)
            m

        insertDistrict :: Row -> County (Vector Row) -> County (Vector Row) -> County (Vector Row)
        insertDistrict row _ old = Map.insertWith
            (V.++)
            (snd (key row))         -- key for TierDistrict
            (V.singleton row)       -- value for TierDistrict
            old

        key :: Row -> (Text, Text)
        key = T.splitAt 3 . townshipName . township . from

        townshipName :: Township -> Text
        townshipName (Township name _) = name

--------------------------------------------------------------------------------
-- Queries
--------------------------------------------------------------------------------

average :: [Double] -> Double
average vec = sum vec / fromIntegral (length vec)

distance :: Row -> Double
distance row = sqrt ((fromX - toX) * (fromX - toX) + (fromY - toY) * (fromY - toY))
    where
        (fromX, fromY) = coord (from row)
        (toX  , toY  ) = coord (to row)

--------------------------------------------------------------------------------
-- Noise
--------------------------------------------------------------------------------

sampleLaplace :: Double -> IO Double
sampleLaplace scale = withSystemRandom . asGenST $ genContVar (laplace 0 scale)

addNoise :: Double -> Double -> IO Double
addNoise scale x = do
    noise <- sampleLaplace scale
    return (x + noise)

-- -- 研究問題
--
-- -- 1. 各鄉鎮市區的平均就醫距離
-- -- 2. 各縣市的平均就醫距離
-- -- 3. 全國的平均就醫距離
--
-- -- 距離上限: 400 公里


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-----------------------------------------------------------------------------

module Numeric.AD.Mode
  (
  -- * AD modes
    Mode(..)
  , pattern KnownZero
  , pattern Auto
  ) where

import Numeric.Natural
import Data.Complex
import Data.Int
import Data.Ratio
import Data.Word

infixr 7 *^
infixl 7 ^*
infixr 7 ^/

class (Num t, Num (Scalar t)) => Mode t where
  type Scalar t
  type Scalar t = t

  -- | allowed to return False for items with a zero derivative, but we'll give more NaNs than strictly necessary
  isKnownConstant :: t -> Bool
  isKnownConstant _ = False

  asKnownConstant :: t -> Maybe (Scalar t)
  asKnownConstant _ = Nothing

  -- | allowed to return False for zero, but we give more NaN's than strictly necessary
  isKnownZero :: t -> Bool
  isKnownZero _ = False

  -- | Embed a constant
  auto  :: Scalar t -> t
  default auto :: (Scalar t ~ t) => Scalar t -> t
  auto = id

  -- | Scalar-vector multiplication
  (*^) :: Scalar t -> t -> t
  a *^ b = auto a * b

  -- | Vector-scalar multiplication
  (^*) :: t -> Scalar t -> t
  a ^* b = a * auto b

  -- | Scalar division
  (^/) :: Fractional (Scalar t) => t -> Scalar t -> t
  a ^/ b = a ^* recip b

  -- |
  -- @'zero' = 'lift' 0@
  zero :: t
  zero = auto 0

pattern KnownZero :: Mode s => s
pattern KnownZero <- (isKnownZero -> True) where
  KnownZero = zero

pattern Auto :: Mode s => Scalar s -> s
pattern Auto n <- (asKnownConstant -> Just n) where
  Auto n = auto n

instance Mode Double where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Float where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Int where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Integer where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Int8 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Int16 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Int32 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Int64 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Natural where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Word where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Word8 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Word16 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Word32 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Mode Word64 where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance RealFloat a => Mode (Complex a) where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)

instance Integral a => Mode (Ratio a) where
  isKnownConstant _ = True
  asKnownConstant = Just
  isKnownZero x = 0 == x
  (^/) = (/)


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-|
Module      : Grenade.Layers.Relu
Description : Rectifying linear unit layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Relu (
    Relu (..)
  ) where

import           Data.Serialize

import           GHC.TypeLits
import           Grenade.Core

import qualified Numeric.LinearAlgebra.Static as LAS

-- | A rectifying linear unit.
--   A layer which can act between any shape of the same dimension, acting as a
--   diode on every neuron individually.
data Relu = Relu
  deriving Show

instance UpdateLayer Relu where
  type Gradient Relu = ()
  runUpdate _ _ _ = Relu
  createRandom = return Relu

instance Serialize Relu where
  put _ = return ()
  get = return Relu

instance ( KnownNat i) => Layer Relu ('D1 i) ('D1 i) where
  type Tape Relu ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (relu y))
    where
      relu = LAS.dvmap (\a -> if a <= 0 then 0 else a)
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (relu' y * dEdy))
    where
      relu' = LAS.dvmap (\a -> if a <= 0 then 0 else 1)

instance (KnownNat i, KnownNat j) => Layer Relu ('D2 i j) ('D2 i j) where
  type Tape Relu ('D2 i j) ('D2 i j) = S ('D2 i j)

  runForwards _ (S2D y) = (S2D y, S2D (relu y))
    where
      relu = LAS.dmmap (\a -> if a <= 0 then 0 else a)
  runBackwards _ (S2D y) (S2D dEdy) = ((), S2D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a <= 0 then 0 else 1)

instance (KnownNat i, KnownNat j, KnownNat k) => Layer Relu ('D3 i j k) ('D3 i j k) where

  type Tape Relu ('D3 i j k) ('D3 i j k) = S ('D3 i j k)

  runForwards _ (S3D y) = (S3D y, S3D (relu y))
    where
      relu = LAS.dmmap (\a -> if a <= 0 then 0 else a)
  runBackwards _ (S3D y) (S3D dEdy) = ((), S3D (relu' y * dEdy))
    where
      relu' = LAS.dmmap (\a -> if a <= 0 then 0 else 1)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilyDependencies #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -Wwarn #-}
module Language.Test where

import Control.Arrow
import Data.Complex
import Data.Proxy
import Type.Reflection

import GHC.TypeLits

data INat = Z | S INat

data family Sing :: k -> *

data instance Sing (n :: INat) where
  SZ :: Sing 'Z
  SS :: Sing n -> Sing ('S n)

class IsSing a where sing :: Sing a
instance IsSing 'Z where sing = SZ
instance IsSing n => IsSing ('S n) where sing = SS sing
type SINat (n :: INat) = Sing n

newtype Unit a = MkUnit () deriving Show

unit :: Unit a
unit = MkUnit ()

type family Prod (n :: INat) (a :: *) = r | r -> n a where
  Prod 'Z a = Unit a
  Prod ('S n) a = (a, Prod n a)

fmapP' :: SINat n -> (a -> b) -> Prod n a -> Prod n b
fmapP' SZ     _ = const unit
fmapP' (SS m) f = f{- ALLOCATE HERE -} *** fmapP' m f

fmapPIx :: SINat n -> (Int -> a -> b) -> Int -> Prod n a -> Prod n b
fmapPIx SZ _ _ = const unit
fmapPIx (SS m) f k = f k *** fmapPIx m f (k+1)

fmapP :: IsSing n => (a -> b) -> Prod n a -> Prod n b
fmapP = fmapP' sing

type family Div2 (n :: INat) where
  Div2 'Z = 'Z
  Div2 ('S 'Z) = 'Z
  Div2 ('S ('S n)) = 'S (Div2 n)

div2 :: SINat n -> SINat (Div2 n)
div2 SZ = SZ
div2 (SS SZ) = SZ
div2 (SS (SS n)) = SS (div2 n)

assocL :: (a, (b, c)) -> ((a, b), c)
assocL (x, (y, z)) = ((x, y), z)

red :: SINat n -> ((a, a) -> a) -> Prod n a -> Prod (Div2 n) a
red SZ _ = id
red (SS SZ) _ = snd
red (SS (SS n)) f = assocL >>> f{- ALLOCATE HERE -} *** red n f

fold :: SINat n -> ((a, a) -> a) -> a -> Prod n a -> a
fold SZ _ z = const z
fold (SS SZ) _ _ = fst
fold n f z = red n f >>> fold (div2 n) f z

reduce :: (IsSing n, Monoid a) => Prod n a -> a
reduce = fold sing (uncurry mappend) mempty

data Fin (m :: INat) where
  FZ :: Fin ('S m)
  FS :: Fin m -> Fin ('S m)

data Le (n :: INat) (m :: INat) where
  LZ :: Le 'Z m
  LS :: Le n m -> Le ('S n) ('S m)

class (IsSing n, IsSing m) => IsLe (n :: INat) (m :: INat) where
  le :: Sing n -> Le n m
instance IsSing n => IsLe 'Z n where
  le SZ = LZ
instance IsLe n m => IsLe ('S n) ('S m) where
  le (SS n) = LS (le n)

get' :: Le n m -> Prod ('S m) a -> a
get' LZ     = fst
get' (LS n) = snd >>> get' n

get :: IsLe n m => SINat n -> Prod ('S m) a -> a
get n = get' (le n)

zw' :: forall a b c (n :: INat). SINat n -> ((a, b) -> c) -> (Prod n a, Prod n b) -> Prod n c
zw' SZ     _ = const unit
zw' (SS m) f = swap >>> f *** zw' m f

toNum :: forall (k :: INat) a. Num a => SINat k -> a
toNum SZ = 0
toNum (SS sm) = 1 + toNum sm

expn :: Float -> Float -> Complex Float -> Complex Float
expn n k c = cis (-2 * pi * k /  n) * c {- multiply exponential -}

add :: (Complex Float, Complex Float) -> Complex Float
add = uncurry (+)

sub :: (Complex Float, Complex Float) -> Complex Float
sub = uncurry (-)

type family Add (n :: INat) (m :: INat) where
  Add 'Z n = n
  Add ('S n) m = 'S (Add n m)

addn :: SINat n -> SINat m -> SINat (Add n m)
addn SZ n = n
addn (SS n) m = SS (addn n m)

type family Mul (n :: INat) (m :: INat) where
  Mul 'Z n = 'Z
  Mul ('S n) m = Add m (Mul n m)

mul :: SINat n -> SINat m -> SINat (Mul n m)
mul SZ _ = SZ
mul (SS n) m = addn m (mul n m)

type family Pow2 (n :: INat) where
  Pow2 'Z = 'S 'Z
  Pow2 ('S n) = Add (Pow2 n) (Pow2 n)

pow2 :: SINat n -> SINat (Pow2 n)
pow2 SZ = SS SZ
pow2 (SS n) = addn (pow2 n) (pow2 n)

-- Hack: proper way would be to traverse 'n' and 'm', but it is quite inefficient

szr :: SINat n -> Add n 'Z :~: n
szr SZ = Refl
szr (SS n) = case szr n of
               Refl -> Refl

ssr :: SINat n -> SINat m -> Add n ('S m) :~: 'S (Add n m)
ssr SZ _ = Refl
ssr (SS n) m = case ssr n m of
                 Refl -> Refl

addComm :: SINat n -> SINat m -> Add n m :~: Add m n
addComm SZ m = case szr m of Refl -> Refl -- unsafeCoerce Refl
addComm (SS n) m = case (addComm n m, ssr m n) of
                     (Refl, Refl) -> Refl

--deinterleave
split :: forall a (n :: INat). SINat n -> Prod (Add n n) a -> (Prod n a, Prod n a)
split SZ = id &&& id
split (SS n) =
  case addComm n (SS n) of
    Refl -> assocL >>> id *** split n >>> swap

assocR :: ((a, b), c) -> (a, (b, c))
assocR ((a, b), c) = (a, (b, c))

cat :: SINat n -> (Prod n a, Prod m a) -> Prod (Add n m) a
cat SZ = snd
cat (SS x) = assocR >>> id *** cat x

fromINat :: forall a n. Num a => SINat n -> a
fromINat SZ = 0
fromINat (SS x) = 1 + fromINat x

fromTo :: forall a n. Num a => SINat n -> Prod n a
fromTo = go 0

go :: forall a n. Num a => a -> SINat n -> Prod n a
go _ SZ = unit
go n (SS m) = (n, go (1 + n) m)

dft :: SINat n -> Prod (Pow2 n) (Complex Float) -> Prod (Pow2 n) (Complex Float)
dft SZ = id
dft (SS x)
   = split p2x >>>
     dft x {-EVENS-} *** dft x {-ODDS-} >>>
     id *** fmapPIx p2x (expn p2sx . fromIntegral) 0 {- Multiply by exponential -} >>>
     zw' p2x add {- Left side -} &&& zw' p2x sub {- right side -} >>>
     cat p2x
  where
    p2x = pow2 x
    p2sx :: Float
    p2sx = fromIntegral $ ((2 ^ (fromINat (SS x) :: Integer)) :: Integer)

--------------------------------------------------------------------------------
-- FFT BELOW -------------------------------------------------------------------

-- PRIM FUNCTIONS --------------------------------------------------------------

-- baseFFT is one of your PRIM funcs: you can copy from some online FFT implementation: e.g. Rosetta Code

baseFFT :: RealFloat a => [Complex a] -> [Complex a]
baseFFT [] = []
baseFFT [x] = [x]
baseFFT xs = zipWith (+) ys ts ++ zipWith (-) ys ts
    where n = length xs
          ys = baseFFT evens
          zs = baseFFT odds
          (evens, odds) = splitList xs
          ts = zipWith (\z k -> (exp' k n) * z) zs [0::Integer ..]

exp' :: (Floating a1, Integral a2, Integral a3) => a2 -> a3 -> Complex a1
exp' k n = cis $ -2 * pi * (fromIntegral k) / (fromIntegral n)

splitList :: [a] -> ([a], [a])
splitList [] = ([], [])
splitList [x] = ([x], [])
splitList (x:y:xs) = (x:xt, y:yt) where (xt, yt) = splitList xs

addPadding :: Num a => SINat n -> [a] -> [a]
addPadding sz l = l ++ replicate (padding $ length l) 0
  where
    padding k = 2 ^ (max ((fromINat sz) :: Integer) (ceiling (logBase 2 $ fromIntegral k :: Double))) - k

concatenate :: ([a], [a]) -> [a]
concatenate = uncurry (++)

{- Below should be your prim functions -}
mulExp :: (RealFloat a1, Integral a3) =>
                a3 -> Int -> [Complex a1] -> [Complex a1]
mulExp p2sx i l = zipWith (\k z -> exp' k (p2sx * fromIntegral len) * z) [i * len ..] l
  where
    len = length l
addc :: ([Complex Float], [Complex Float]) -> [Complex Float]
addc = uncurry $ zipWith (+)
subc :: ([Complex Float], [Complex Float]) -> [Complex Float]
subc = uncurry $ zipWith (-)

-- END PRIM FUNCTIONS ----------------------------------------------------------

-- ARROW DEFINITIONS -----------------------------------------------------------
---  Should not change, except types

type family Tree (n :: INat) (a :: *)  where
  Tree 'Z a = a
  Tree ('S n) a = (Tree n a, Tree n a)

splitL :: SINat m -> [Complex Float] -> Tree m [Complex Float]
splitL SZ = id
splitL (SS n) = splitList >>> splitL n *** splitL n

merge :: forall (n :: INat). SINat n -> Tree n [Complex Float] -> [Complex Float]
merge SZ  = id {- same: just concatenate all using a Prim -}
merge (SS n) = merge n *** merge n >>> concatenate

fmapTIx :: SINat n -> ((Int, [Complex Float]) -> [Complex Float]) -> Int -> Tree n [Complex Float] -> Tree n [Complex Float]
fmapTIx SZ f k = (\v -> f (k, v))
fmapTIx (SS x) f k = fmapTIx x f k *** fmapTIx x f (k + (2 ^ (fromINat x :: Integer)))

swap :: ((a,b), (c, d)) -> ((a, c), (b, d))
swap = ((fst >>> fst) &&& (snd >>> fst)) &&&
       ((fst >>> snd) &&& (snd >>> snd))

zwT :: SINat n -> (([Complex Float],[Complex Float]) -> [Complex Float])
    -> (Tree n [Complex Float], Tree n [Complex Float]) -> Tree n [Complex Float]
zwT SZ f = f
zwT (SS x) f = swap >>> zwT x f *** zwT x f

fft :: SINat n -> Tree n [Complex Float] -> Tree n [Complex Float]
fft SZ = baseFFT
fft (SS x)
  = fft x {-EVENS-} *** fft x {-ODDS-} >>>
    id *** fmapTIx x (uncurry $ mulExp p2sx) 0 {- Multiply by exponential -} >>>
    zwT x addc {- Left side -} &&& zwT x subc {- right side -}
  where
    p2sx :: Integer
    p2sx = 2 ^ (fromINat (SS x) :: Integer)

fastFourierR :: forall (n :: INat). SINat n -> [Complex Float] -> [Complex Float]
fastFourierR cores = addPadding cores >>> splitL cores >>> fft cores >>> merge cores

---------------------------------------------------------------------------------
-- UTILITY FUNCTIONS

fastFourier :: forall n. (KnownNat n, IsSing (FromNat n))
            => [Complex Float] -> [Complex Float]
fastFourier = fastFourierR (sing :: SINat (FromNat n))

fft8core :: [Complex Float] -> [Complex Float]
fft8core = fastFourier @3
fft16core :: [Complex Float] -> [Complex Float]
fft16core = fastFourier @4

type family ToNat (i :: INat) :: Nat where
  ToNat 'Z = 0
  ToNat ('S n) = 1 + ToNat n

type family FromNat (i :: Nat) :: INat where
  FromNat 0 = 'Z
  FromNat n = 'S (FromNat (n-1))

fmapT :: SINat n -> (a -> b) -> Tree n a -> Tree n b
fmapT SZ f = f
fmapT (SS x) f = fmapT x f *** fmapT x f

data SDict a where
  SDict :: Show a => SDict a

getDict :: Show a => SINat n -> Proxy a -> SDict (Tree n a)
getDict SZ _ = SDict
getDict (SS n) p = case getDict n p of
                        SDict -> SDict

test :: SINat n -> Tree n Int -> String
test n = case getDict n (Proxy :: Proxy Int) of
           SDict -> show

test2 :: Show (Tree n Int) => SINat n -> Tree n Int -> Bool
test2 _ _ =  False

test3 :: SINat n -> Tree n Int -> Bool
test3 n t = case getDict n (Proxy :: Proxy Int) of
              SDict -> test2 n t