{-# LANGUAGE ParallelListComp #-}

module Main where

import Text.PrettyPrint

import System.Environment ( getArgs )

main = do
         [s] <- getArgs
         let n = read s
         mapM_ (putStrLn . render . generate) [2..n]

generate :: Int -> Doc
generate n =
  vcat [ text "#ifdef DEFINE_INSTANCES"
       , data_instance "MVector s" "MV"
       , data_instance "Vector" "V"
       , class_instance "Unbox"
       , class_instance "M.MVector MVector" <+> text "where"
       , nest 2 $ vcat $ map method methods_MVector
       , class_instance "G.Vector Vector" <+> text "where"
       , nest 2 $ vcat $ map method methods_Vector
       , text "#endif"
       , text "#ifdef DEFINE_MUTABLE"
       , define_zip "MVector s" "MV"
       , define_unzip "MVector s" "MV"
       , text "#endif"
       , text "#ifdef DEFINE_IMMUTABLE"
       , define_zip "Vector" "V"
       , define_zip_rule
       , define_unzip "Vector" "V"
       , text "#endif"
       ]

  where
    vars  = map (\c -> text ['_',c]) $ take n ['a'..]
    varss = map (<> char 's') vars
    tuple xs = parens $ hsep $ punctuate comma xs
    vtuple xs = parens $ sep $ punctuate comma xs
    con s = text s <> char '_' <> int n
    var c = text ('_' : c : "_")

    data_instance ty c
      = hang (hsep [text "data instance", text ty, tuple vars])
             4
             (hsep [char '=', con c, text "{-# UNPACK #-} !Int"
                   , vcat $ map (\v -> char '!' <> parens (text ty <+> v)) vars])

    class_instance cls
      = text "instance" <+> vtuple [text "Unbox" <+> v | v <- vars]
                        <+> text "=>" <+> text cls <+> tuple vars


    define_zip ty c
      = sep [text "-- | /O(1)/ Zip" <+> int n <+> text "vectors."
            ,name <+> text "::"
                  <+> vtuple [text "Unbox" <+> v | v <- vars]
                  <+> text "=>"
                  <+> sep (punctuate (text " ->") [text ty <+> v | v <- vars])
                  <+> text "->"
                  <+> text ty <+> tuple vars
             ,text "{-# INLINE_FUSED"  <+> name <+> text "#-}"
             ,name <+> sep varss
                   <+> text "="
                   <+> con c
                   <+> text "len"
                   <+> sep [parens $ text "unsafeSlice"
                                     <+> char '0'
                                     <+> text "len"
                                     <+> vs | vs <- varss]
             ,nest 2 $ hang (text "where")
                            2
                     $ text "len ="
                       <+> sep (punctuate (text " `delayed_min`")
                                          [text "length" <+> vs | vs <- varss])
             ]
      where
        name | n == 2    = text "zip"
             | otherwise = text "zip" <> int n

    define_zip_rule
      = hang (text "{-# RULES" <+> text "\"stream/" <> name "zip"
              <> text " [Vector.Unboxed]\" forall" <+> sep varss <+> char '.')
             2 $
             text "G.stream" <+> parens (name "zip" <+> sep varss)
             <+> char '='
             <+> text "Bundle." <> name "zipWith" <+> tuple (replicate n empty)
             <+> sep [parens $ text "G.stream" <+> vs | vs <- varss]
             $$ text "#-}"
     where
       name s | n == 2    = text s
              | otherwise = text s <> int n


    define_unzip ty c
      = sep [text "-- | /O(1)/ Unzip" <+> int n <+> text "vectors."
            ,name <+> text "::"
                  <+> vtuple [text "Unbox" <+> v | v <- vars]
                  <+> text "=>"
                  <+> text ty <+> tuple vars
                  <+> text "->" <+> vtuple [text ty <+> v | v <- vars]
            ,text "{-# INLINE" <+> name <+> text "#-}"
            ,name <+> pat c <+> text "="
                  <+> vtuple varss
            ]
      where
        name | n == 2    = text "unzip"
             | otherwise = text "unzip" <> int n

    pat c = parens $ con c <+> var 'n' <+> sep varss
    patn c n = parens $ con c <+> (var 'n' <> int n)
                              <+> sep [v <> int n | v <- varss]

    qM s = text "M." <> text s
    qG s = text "G." <> text s

    gen_length c _ = (pat c, var 'n')

    gen_unsafeSlice mod c rec
      = (var 'i' <+> var 'm' <+> pat c,
         con c <+> var 'm'
               <+> vcat [parens
                         $ text mod <> char '.' <> text rec
                                    <+> var 'i' <+> var 'm' <+> vs
                                        | vs <- varss])


    gen_overlaps rec = (patn "MV" 1 <+> patn "MV" 2,
                        vcat $ r : [text "||" <+> r | r <- rs])
      where
        r : rs = [qM rec <+> v <> char '1' <+> v <> char '2' | v <- varss]

    gen_unsafeNew rec
      = (var 'n',
         mk_do [v <+> text "<-" <+> qM rec <+> var 'n' | v <- varss]
               $ text "return $" <+> con "MV" <+> var 'n' <+> sep varss)

    gen_unsafeReplicate rec
      = (var 'n' <+> tuple vars,
         mk_do [vs <+> text "<-" <+> qM rec <+> var 'n' <+> v
                        | v  <- vars | vs <- varss]
               $ text "return $" <+> con "MV" <+> var 'n' <+> sep varss)

    gen_unsafeRead rec
      = (pat "MV" <+> var 'i',
         mk_do [v <+> text "<-" <+> qM rec <+> vs <+> var 'i' | v  <- vars
                                                              | vs <- varss]
               $ text "return" <+> tuple vars)

    gen_unsafeWrite rec
      = (pat "MV" <+> var 'i' <+> tuple vars,
         mk_do [qM rec <+> vs <+> var 'i' <+> v | v  <- vars | vs <- varss]
               empty)

    gen_clear rec
      = (pat "MV", mk_do [qM rec <+> vs | vs <- varss] empty)

    gen_set rec
      = (pat "MV" <+> tuple vars,
         mk_do [qM rec <+> vs <+> v | vs <- varss | v <- vars] empty)

    gen_unsafeCopy c q rec
      = (patn "MV" 1 <+> patn c 2,
         mk_do [q rec <+> vs <> char '1' <+> vs <> char '2' | vs <- varss]
               empty)

    gen_unsafeMove rec
      = (patn "MV" 1 <+> patn "MV" 2,
         mk_do [qM rec <+> vs <> char '1' <+> vs <> char '2' | vs <- varss]
               empty)

    gen_unsafeGrow rec
      = (pat "MV" <+> var 'm',
         mk_do [vs <> char '\'' <+> text "<-"
                                <+> qM rec <+> vs <+> var 'm' | vs <- varss]
               $ text "return $" <+> con "MV"
                                 <+> parens (var 'm' <> char '+' <> var 'n')
                                 <+> sep (map (<> char '\'') varss))

    gen_initialize rec
      = (pat "MV", mk_do [qM rec <+> vs | vs <- varss] empty)

    gen_unsafeFreeze rec
      = (pat "MV",
         mk_do [vs <> char '\'' <+> text "<-" <+> qG rec <+> vs | vs <- varss]
               $ text "return $" <+> con "V" <+> var 'n'
                                 <+> sep [vs <> char '\'' | vs <- varss])

    gen_unsafeThaw rec
      = (pat "V",
         mk_do [vs <> char '\'' <+> text "<-" <+> qG rec <+> vs | vs <- varss]
               $ text "return $" <+> con "MV" <+> var 'n'
                                 <+> sep [vs <> char '\'' | vs <- varss])

    gen_basicUnsafeIndexM rec
      = (pat "V" <+> var 'i',
         mk_do [v <+> text "<-" <+> qG rec <+> vs <+> var 'i'
                        | vs <- varss | v <- vars]
               $ text "return" <+> tuple vars)

    gen_elemseq rec
      = (char '_' <+> tuple vars,
         vcat $ r : [char '.' <+> r | r <- rs])
      where
        r : rs = [qG rec <+> parens (text "undefined :: Vector" <+> v)
                         <+> v | v <- vars]

    mk_do cmds ret = hang (text "do")
                          2
                          $ vcat $ cmds ++ [ret]

    method (s, f) = case f s of
                      (p,e) ->  text "{-# INLINE" <+> text s <+> text " #-}"
                                $$ hang (text s <+> p)
                                   4
                                   (char '=' <+> e)


    methods_MVector = [("basicLength",            gen_length "MV")
                      ,("basicUnsafeSlice",       gen_unsafeSlice "M" "MV")
                      ,("basicOverlaps",          gen_overlaps)
                      ,("basicUnsafeNew",         gen_unsafeNew)
                      ,("basicUnsafeReplicate",   gen_unsafeReplicate)
                      ,("basicUnsafeRead",        gen_unsafeRead)
                      ,("basicUnsafeWrite",       gen_unsafeWrite)
                      ,("basicClear",             gen_clear)
                      ,("basicSet",               gen_set)
                      ,("basicUnsafeCopy",        gen_unsafeCopy "MV" qM)
                      ,("basicUnsafeMove",        gen_unsafeMove)
                      ,("basicUnsafeGrow",        gen_unsafeGrow)
                      ,("basicInitialize",        gen_initialize)]

    methods_Vector  = [("basicUnsafeFreeze",      gen_unsafeFreeze)
                      ,("basicUnsafeThaw",        gen_unsafeThaw)
                      ,("basicLength",            gen_length "V")
                      ,("basicUnsafeSlice",       gen_unsafeSlice "G" "V")
                      ,("basicUnsafeIndexM",      gen_basicUnsafeIndexM)
                      ,("basicUnsafeCopy",        gen_unsafeCopy "V" qG)
                      ,("elemseq",                gen_elemseq)]
