{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE DefaultSignatures    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE OverloadedLists      #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Pretty printing for Generic data types

module Text.PrettyPrint.Generic (
  -- * Type classes
  Pretty(..),
  GPretty(..),

  -- * Utility functions
  pretty,  prettyShow,  prettyPrint,  hPrettyPrint,
  pretty', prettyShow', prettyPrint', hPrettyPrint',
  ) where

import           GHC.Exts                     (IsList (..))
import           GHC.Generics
import           System.IO
import           Text.PrettyPrint.ANSI.Leijen hiding (Pretty (..))

import           Control.Applicative          (Const, WrappedArrow,
                                               WrappedMonad, ZipList)
import qualified Data.ByteString.Char8        as S
import qualified Data.ByteString.Lazy.Char8   as L
import           Data.Complex
import           Data.Functor.Identity        (Identity)
import           Data.Int
import qualified Data.IntMap                  as IntMap
import qualified Data.IntSet                  as IntSet
import qualified Data.Map                     as Map
import           Data.Monoid                  (All, Alt, Any, First, Last,
                                               Product, Sum)
import           Data.Ratio
import qualified Data.Sequence                as Seq
import qualified Data.Set                     as Set
import qualified Data.Text                    as T
import qualified Data.Text.Encoding           as T
import qualified Data.Text.Lazy               as TL
import qualified Data.Vector                  as V
import qualified Data.Vector.Primitive        as VP
import qualified Data.Vector.Storable         as VS
import qualified Data.Vector.Unboxed          as VU
import           Data.Word

-- Color scheme

constructorS, numericS, stringS, operatorS, selectorS :: Doc -> Doc
constructorS = underline . bold . dullgreen
numericS     = magenta
stringS      = dullyellow
operatorS    = dullred
selectorS    = blue

-- Generic function

-- | Type class for generic representations
class GPretty f where
  -- | Pretty print a `Generic` value
  gprettyPrec :: Int -> f a -> [Doc]

instance GPretty V1 where
  gprettyPrec _ _ = error "this never happen"

instance GPretty U1 where
  gprettyPrec _ U1 = []

instance Pretty c => GPretty (Rec0 c) where
  gprettyPrec p (K1 c) = [prettyPrec p c]

instance GPretty f => GPretty (D1 d f) where
  gprettyPrec p (M1 a) = gprettyPrec p a

instance (GPretty f, Constructor c) => GPretty (C1 c f) where
  gprettyPrec p c@(M1 a)
    | conIsRecord c =
      [ con <+> encloseSep (lbrace <> space) (space <> rbrace) (comma <> space) es ]
    | null es      = [ con ]
    | p == 0       = [ con <+> align (sep es) ]
    | otherwise    = [ parens $ con <+> align (sep es) ]
    where
      con = constructorS $ text (conName c)
      es = gprettyPrec (p + 1) a

instance {-# OVERLAPPABLE #-} (GPretty f, Selector s) => GPretty (S1 s f) where
  gprettyPrec _ s@(M1 a) =
    [ selectorS (text (selName s)) <+> operatorS (text "=") <+> sep (gprettyPrec 0 a) ]
instance {-# OVERLAPPING #-} GPretty f => GPretty (S1 ('MetaSel 'Nothing su ss ds) f) where
  gprettyPrec p (M1 a) = gprettyPrec p a

instance (GPretty f, GPretty g) => GPretty (f :+: g) where
  gprettyPrec p (L1 a) = gprettyPrec p a
  gprettyPrec p (R1 a) = gprettyPrec p a

instance (GPretty f, GPretty g) => GPretty (f :*: g) where
  gprettyPrec p (a :*: b) = gprettyPrec p a ++ gprettyPrec p b

-- Wrapper function

-- | Type class for pretty printing
class Pretty a where
  -- | Pretty print a value to `Doc`
  prettyPrec :: Int -> a -> Doc
  default prettyPrec :: (Generic a, GPretty (Rep a)) => Int -> a -> Doc
  prettyPrec p = sep . gprettyPrec p . from

-- Utility functions

-- | Pretty print a value with decoration
pretty :: Pretty a => a -> Doc
pretty = prettyPrec 0

-- | Pretty print a value to `String`
prettyShow :: Pretty a => a -> String
prettyShow = show . pretty

-- | Pretty print a value to `stdout`
prettyPrint :: Pretty a => a -> IO ()
prettyPrint = hPrettyPrint stdout

-- | Pretty print a value
hPrettyPrint :: Pretty a => Handle -> a -> IO ()
hPrettyPrint h = hPutDoc h . (<> hardline) . pretty

-- | Pretty print a value without decoration
pretty' :: Pretty a => a -> Doc
pretty' = plain . pretty

-- | Plain version for `prettyShow`
prettyShow' :: Pretty a => a -> String
prettyShow' = show . pretty'

-- | Plain version for `prettyPrint`
prettyPrint' :: Pretty a => a -> IO ()
prettyPrint' = hPrettyPrint' stdout

-- | Plain version for `hPrettyPrint`
hPrettyPrint' :: Pretty a => Handle -> a -> IO ()
hPrettyPrint' h = hPutDoc h . (<> hardline) . pretty'

-- instances

instance Pretty ()      where prettyPrec _ = text . show
instance Pretty Char    where prettyPrec _ = stringS . text . show
instance Pretty Int     where prettyPrec _ = numericS . text . show
instance Pretty Integer where prettyPrec _ = numericS . text . show
instance Pretty Float   where prettyPrec _ = numericS . text . show
instance Pretty Double  where prettyPrec _ = numericS . text . show
instance Pretty Bool    where prettyPrec _ = numericS . text . show

instance Pretty Word    where prettyPrec _ = numericS . text . show
instance Pretty Word8   where prettyPrec _ = numericS . text . show
instance Pretty Word16  where prettyPrec _ = numericS . text . show
instance Pretty Word32  where prettyPrec _ = numericS . text . show
instance Pretty Word64  where prettyPrec _ = numericS . text . show

instance Pretty Int8    where prettyPrec _ = numericS . text . show
instance Pretty Int16   where prettyPrec _ = numericS . text . show
instance Pretty Int32   where prettyPrec _ = numericS . text . show
instance Pretty Int64   where prettyPrec _ = numericS . text . show

instance (Integral a, Pretty a) => Pretty (Ratio a) where
  prettyPrec _ r =
    pretty (numerator r) <+> operatorS (char '%') <+> pretty (denominator r)

instance (Pretty a) => Pretty (Complex a) where
  prettyPrec _ (r :+ i) =
    pretty r <+> operatorS (text ":+") <+> pretty i

instance {-# OVERLAPPABLE #-} Pretty a => Pretty [a] where
  prettyPrec _ = encloseSep (lbracket <> space) (space <> rbracket) (comma <> space) . map pretty

instance {-# OVERLAPPING #-} Pretty String where
  prettyPrec _ = stringS . dquotes . text . prettyString where
    prettyString cs = foldr ((.) . prettyChar) id cs ""
    prettyChar c
      | fromEnum c < 0x80 = showChar c
      | otherwise = (c:)

instance (Pretty a, Pretty b) => Pretty (a, b) where
  prettyPrec _ (a, b) = parens $ pretty a <> comma <+> pretty b

instance (Pretty a, Pretty b, Pretty c) => Pretty (a, b, c) where
  prettyPrec _ (a, b, c) = parens $ pretty a <> comma <+> pretty b <> comma <+> pretty c

instance (Pretty a, Pretty b, Pretty c, Pretty d) => Pretty (a, b, c, d) where
  prettyPrec _ (a, b, c, d) = parens $ pretty a <> comma <+> pretty b <> comma <+> pretty c <> comma <+> pretty d

instance (Pretty a, Pretty b, Pretty c, Pretty d, Pretty e) => Pretty (a, b, c, d, e) where
  prettyPrec _ (a, b, c, d, e) = parens $ pretty a <> comma <+> pretty b <> comma <+> pretty c <> comma <+> pretty d <> comma <+> pretty e

instance (Pretty a, Pretty b, Pretty c, Pretty d, Pretty e, Pretty f) => Pretty (a, b, c, d, e, f) where
  prettyPrec _ (a, b, c, d, e, f) = parens $ pretty a <> comma <+> pretty b <> comma <+> pretty c <> comma <+> pretty d <> comma <+> pretty e <> comma <+> pretty f

instance (Pretty a, Pretty b, Pretty c, Pretty d, Pretty e, Pretty f, Pretty g) => Pretty (a, b, c, d, e, f, g) where
  prettyPrec _ (a, b, c, d, e, f, g) = parens $ pretty a <> comma <+> pretty b <> comma <+> pretty c <> comma <+> pretty d <> comma <+> pretty e <> comma <+> pretty f <> comma <+> pretty g

instance Pretty a => Pretty (Maybe a)
instance (Pretty a, Pretty b) => Pretty (Either a b)
instance Pretty Ordering
instance Pretty Any
instance Pretty All
instance Pretty a => Pretty (First a)
instance Pretty a => Pretty (Last a)
instance Pretty a => Pretty (Sum a)
instance Pretty a => Pretty (Product a)
instance Pretty (f a) => Pretty (Alt f a)
instance Pretty a => Pretty (Identity a)
instance Pretty a => Pretty (Const a b)
instance Pretty a => Pretty (ZipList a)
instance Pretty (m a) => Pretty (WrappedMonad m a)
instance Pretty (a b c) => Pretty (WrappedArrow a b c)

-- bytestrings, texts

instance Pretty S.ByteString where
  prettyPrec _ bs = case T.decodeUtf8' bs of
    Left err -> pretty $ show err
    Right t -> pretty t

instance Pretty T.Text where
  prettyPrec _ = pretty . T.unpack

instance Pretty L.ByteString where
  prettyPrec _ = pretty . L.toStrict

instance Pretty TL.Text where
  prettyPrec _ = pretty . TL.toStrict

-- containers

instance (Pretty a, Ord a) => Pretty (Set.Set a) where
  prettyPrec _ = encloseSep (lbracket <> space) (space <> rbracket) (comma <> space) . map pretty . toList

instance Pretty IntSet.IntSet where
  prettyPrec _ = pretty . Set.fromList . toList

instance (Pretty a, Pretty b, Ord a) => Pretty (Map.Map a b) where
  prettyPrec _ = encloseSep (lbrace <> space) (space <> rbrace) (comma <> space) . map f . toList where
    f (key, val) = pretty key <> colon <+> pretty val

instance Pretty b => Pretty (IntMap.IntMap b) where
  prettyPrec _ = pretty . Map.fromList . toList

instance Pretty a => Pretty (Seq.Seq a) where
  prettyPrec _ = pretty . toList

-- vectors

instance Pretty a => Pretty (V.Vector a) where
  prettyPrec _ = pretty . toList

instance (Pretty a, VP.Prim a) => Pretty (VP.Vector a) where
  prettyPrec _ = pretty . toList

instance (Pretty a, VS.Storable a) => Pretty (VS.Vector a) where
  prettyPrec _ = pretty . toList

instance (Pretty a, VU.Unbox a) => Pretty (VU.Vector a) where
  prettyPrec _ = pretty . toList


module Syntax where

import qualified Data.Complex as ComplexData
import qualified Data.Ord as OrdData

type Identifier = String
type Package = Identifier
type SliceStart = Integer
type SliceEnd = Integer
type SliceStep = Integer
type Ellipsis = ()
type ArrayLength = Integer
type ReadChannel = ()
type WriteChannel = ()
type Receiver = [FunctionParameterDecl]
type ImportDecl = [ImportSpec]

data QualifiedIdentifier = Qualified Package Identifier
  deriving (Eq, Ord, Show)

data Declaration
  = ConstDecl [Identifier] (Maybe Type) [Expression]
  | TypeDecl Identifier Type
  | VarDecl [Identifier] (Maybe Type) [Expression]
  deriving (Eq, Ord, Show)

data TopLevelDecl
  = Declaration Declaration
  | FunctionDecl Identifier FunctionSignature (Maybe [Statement])
  | ExternFunctionDecl Identifier FunctionSignature    -- this is not part of go RFC
  | MethodDecl Receiver Identifier FunctionSignature (Maybe [Statement])
  deriving (Eq, Ord, Show)

data Expression
  = Unary UnaryExpr
  | Binary BinaryOp Expression Expression
  deriving (Eq, Ord, Show)

data UnaryExpr
  = PrimaryExpr PrimaryExpr
  -- FIXME: had to change this from 'UnaryExpr' to 'Expression'
  -- don't know if it will cause any problems
  | UnaryOpExpr UnaryOp Expression
  deriving (Eq, Ord, Show)

data PrimaryExpr
  = Operand Operand
  | Conversion Conversion
  | Selector PrimaryExpr Identifier
  | Index PrimaryExpr Expression
  | Slice PrimaryExpr (Maybe SliceStart) (Maybe SliceEnd) (Maybe SliceStep)
  | TypeAssertion PrimaryExpr Type
  | Arguments PrimaryExpr [Expression] (Maybe Type) (Maybe Ellipsis)
  deriving (Eq, Ord, Show)

data BinaryOp
  = Or        -- ||
  | And       -- &&
  | RelBinOp RelOp
  | AddBinOp AddOp
  | MulBinOp MulOp
  deriving (Eq, Ord, Show)

data RelOp
  = Equals    -- ==
  | NotEquals -- !=
  | Less      -- <
  | LessEq    -- <=
  | Greater   -- >
  | GreaterEq -- >=
  deriving (Eq, Ord, Show)

data AddOp
  = Plus      -- +
  | Minus     -- -
  | BitOr     -- |
  | BitXor    -- ^
  deriving (Eq, Ord, Show)

data MulOp
  = Mult      -- *
  | Div       -- /
  | Mod       -- %
  | LShift    -- <<
  | RShift    -- >>
  | BitAnd    -- &
  | AndXor    -- &^ - FIXME: no idea what is that
  deriving (Eq, Ord, Show)

data UnaryOp
  = UnPlus    -- +
  | UnMinus   -- -
  | Not       -- !
  | UnCaret   -- ^
  | Deref     -- *
  | Pointerof -- &
  | ChanRead  -- <-
  deriving (Eq, Ord, Show)

data TypeName
  = Identifier Identifier
  | QualifiedIdentifier QualifiedIdentifier
  deriving (Eq, Ord, Show)

data TypeWithEllipsis = TypeWithEllipsis Type (Maybe Ellipsis)

data Operand
  = Literal Literal
  | OperandName OperandName
  | MethodExpr Type Identifier
  | ExpressionInParens Expression
  deriving (Eq, Ord, Show)

data OperandName
  = OperandIdentifier Identifier
  | OperandQualified QualifiedIdentifier
  deriving (Eq, Ord, Show)

data Conversion = TypeConversion Type Expression
  deriving (Eq, Ord, Show)

data Type
  = TypeName TypeName
  | TypeLiteral TypeLiteral
  deriving (Eq, Ord, Show)

data TypeLiteral
  = ArrayType ArrayLength Type
  | StructType [FieldDecl]
  | PointerType Type
  | FunctionType FunctionSignature
  | InterfaceType [MethodSpec]
  | SliceType Type
  | MapType Type Type
  | ChannelType Type (Maybe ReadChannel) (Maybe WriteChannel)
  deriving (Eq, Ord, Show)

data FieldDecl
  = Field [Identifier] Type
  | AnonymousField TypeName
  deriving (Eq, Ord, Show)

data MethodSpec
  = Method Identifier FunctionSignature
  | Interface TypeName
  deriving (Eq, Ord, Show)

data FunctionSignature = Signature [FunctionParameterDecl] (Maybe FunctionResult)
  deriving (Eq, Ord, Show)

data FunctionParameterDecl = ParameterDecl [Identifier] Type (Maybe Ellipsis)
  deriving (Eq, Ord, Show)

data FunctionResult
  = FunctionTupleResult [FunctionParameterDecl]
  | FunctionSingleResult Type
  deriving (Eq, Ord, Show)

data Statement
  = DeclarationStmt Declaration
  | LabeledStmt Identifier Statement
  | SimpleStmt SimpleStmt
  | GoStmt Expression
  | ReturnStmt [Expression]
  | BreakStmt (Maybe Identifier)
  | ContinueStmt (Maybe Identifier)
  | GotoStmt Identifier
  | FallthroughStmt
  | Block [Statement]
  | IfStmt IfStmt
  | SwitchStmt SwitchStmt
  | SelectStmt [CommClauseStmt]
  | ForStmt ForStmt
  | DeferStmt Expression
  deriving (Eq, Ord, Show)

data SimpleStmt
  = Empty
  | ExpressionStmt Expression
  | SendStmt SendStmt
  | IncStmt Expression
  | DecStmt Expression
  | AssignmentStmt [Expression] AssignOp [Expression]
  | ShortVarDeclStmt [Identifier] [Expression]
  deriving (Eq, Ord, Show)

data IfStmt = If (Maybe SimpleStmt) Expression [Statement] (Maybe ElseStmt)
  deriving (Eq, Ord, Show)

data ElseStmt
  = ElseIf IfStmt
  | Else [Statement]
  deriving (Eq, Ord, Show)

data SwitchStmt
  = ExprSwitchStmt (Maybe SimpleStmt) (Maybe Expression) [ExprCaseClauseStmt]
  | TypeSwitchStmt (Maybe SimpleStmt) TypeSwitchGuardStmt [TypeCaseClauseStmt]
  deriving (Eq, Ord, Show)

data ExprCaseClauseStmt
  = ExprCaseClause [Expression] [Statement]
  | DefaultCaseClause [Statement]
  deriving (Eq, Ord, Show)

data TypeSwitchGuardStmt = TypeSwitchGuard (Maybe Identifier) PrimaryExpr
  deriving (Eq, Ord, Show)

data TypeCaseClauseStmt
  = TypeCaseClause [Type] [Statement]
  | DefaultTypeCaseClause [Statement]
  deriving (Eq, Ord, Show)

data CommClauseStmt = CommClause CommCaseStmt [Statement]
  deriving (Eq, Ord, Show)

data CommCaseStmt
  = SendCommClause SendStmt
  | RecvCommClause RecvStmt
  | DefaultCommClause
  deriving (Eq, Ord, Show)

data SendStmt = Send Expression Expression
  deriving (Eq, Ord, Show)

data RecvStmt = Recv RecvTarget Expression
  deriving (Eq, Ord, Show)

data RecvTarget
  = RecvAssignment [Expression]
  | RecvShortVarDecl [Identifier]
  deriving (Eq, Ord, Show)

data ForStmt = For ForDefinition [Statement]
  deriving (Eq, Ord, Show)

data ForDefinition
  = Condition Expression
  | ForClause (Maybe SimpleStmt) (Maybe Expression) (Maybe SimpleStmt)
  | RangeClause RangeTarget Expression
  deriving (Eq, Ord, Show)

data RangeTarget
  = RangeAssignment [Expression]
  | RangeShortVarDecl [Identifier]
  deriving (Eq, Ord, Show)

data AssignOp
  = AssignAddOp AddOp
  | AssignMulOp MulOp
  | Assign
  deriving (Eq, Ord, Show)

data Literal
  = IntLiteral Integer
  | FloatLiteral Double
  | ImaginaryLiteral ComplexValue
  | RuneLiteral Char
  | StringLiteral String
  | CompositeLiteral LiteralType LiteralValue
  | FunctionLiteral FunctionSignature [Statement]
  deriving (Eq, Ord, Show)

data LiteralType
  = StructTypeLiteral [FieldDecl]
  | ArrayTypeLiteral ArrayLength Type
  | AutoSizeArrayTypeLiteral Type
  | SliceTypeLiteral Type
  | MapTypeLiteral Type Type
  | TypeNameLiteral TypeName
  deriving (Eq, Ord, Show)

data LiteralValue = LiteralValue [LiteralElement]
  deriving (Eq, Ord, Show)

data LiteralElement = LiteralElement (Maybe LiteralKey) Value
  deriving (Eq, Ord, Show)

data LiteralKey
  = FieldName Identifier
  | KeyExpr Expression
  | KeyLiteralValue LiteralValue
  deriving (Eq, Ord, Show)

data Value
  = ValueExpr Expression
  | ValueLiteralValue LiteralValue
  deriving (Eq, Ord, Show)

data SourceFile = SourceFile Identifier [ImportDecl] [TopLevelDecl]
  deriving (Eq, Ord, Show)

data ImportSpec = ImportSpec Identifier String
  deriving (Eq, Ord, Show)

type ComplexValue = ComplexData.Complex Double

instance (Eq t, Ord t, RealFloat t) => Ord (ComplexData.Complex t) where
  compare a b =
    case byRealPart of
      OrdData.EQ -> byImagPart
      _ -> byRealPart
    where
      byRealPart = compare (ComplexData.realPart a) (ComplexData.realPart b)
      byImagPart = compare (ComplexData.imagPart a) (ComplexData.imagPart b)


module ParserSpec (spec) where

import Test.Hspec
import Data.Either (either)
import Data.Ratio ((%))
import Data.Complex (Complex((:+)))

import Language.Ruby.Parser
import Language.Ruby.Parser.Builder
import Language.Ruby.AST

run :: String -> Term
run = either error id . parseRuby

test title code term = it (title ++ " [ " ++ code ++ " ]") $ do
                          run code `shouldBe` term

xtest title code _ = it (title ++ " [ " ++ code ++ " ]") $ do
                          pending

spec :: Spec
spec = do
  describe "parseRuby" $ do
    test "empty stmt" "" Nil

    test "nil" "nil" Nil

    test "nil_expression" "()" (Begin [])
    test "nil_expression" "begin end" (KWBegin [])

    test "true" "true" RTrue

    test "false" "false" RFalse

    test "int" "42" (RInt 42)
    test "int" "+42" (RInt 42)
    test "int" "-42" (RInt $ -42)

    -- test "int___LINE__" "__LINE__" (RInt 1)
    test "int___LINE__" "__LINE__" Line

    test "float" "1.33" (RFloat 1.33)
    test "float" "-1.33" (RFloat $ -1.33)
    test "float" "-1.335" (RFloat $ -1.335)

    -- test "rational" "42r" (RRational $ 42 % 1)
    -- test "rational" "42.1r" (RRational $ 421 % 10)

    test "complex" "42i" (RComplex $ 0 :+ 42)
    -- test "complex" "42ri" (RComplex $ 0 :+ (42 % 1))
    -- test "complex" "42.1i" (RComplex $ 0 :+ 42.1)
    -- test "complex" "42.1ri" (RComplex $ 0 :+ (421 % 10))

    test "string_plain" "'foobar'" (Str "foobar")
    test "string_plain" "%q(foobar)" (Str "foobar")

    test "self" "self" Self

    test "lvar" "foo" (Lvar "foo")
    test "ivar" "@foo" (Ivar "@foo")
    test "cvar" "@@foo" (Cvar "@@foo")
    test "gvar" "$foo" (Gvar "$foo")

    test "back_ref" "$+" (BackRef "$+")
    test "nth_ref" "$10" (NthRef 10)
    test "const_toplevel" "::Foo" (Const Cbase "Foo")
    test "const_scoped" "Bar::Foo" (Const (Const Nil "Bar") "Foo")
    test "const_unscoped" "Foo" (Const Nil "Foo")
    test "const__ENCODING__" "__ENCODING__" Encoding

    test "defined?" "defined? foo" (Defined [Lvar "foo"])
    test "defined?" "defined?(foo)" (Defined [Lvar "foo"])
    test "defined?" "defined? @foo" (Defined [Ivar "@foo"])

    test "lvasgn" "var = 10; var" (Begin [lvasgn "var" (RInt 10), Lvar "var"])
    test "ivasgn" "@var = 10"  (ivasgn "@var" (RInt 10))
    test "cvasgn" "@@var = 10"  (cvasgn "@@var" (RInt 10))
    test "gvasgn" "$var = 10"  (gvasgn "$var" (RInt 10))

    test "asgn_cmd" "foo = m foo"  (lvasgn "foo" (Send Nil "m" [Lvar "foo"]))
    test "asgn_cmd" "foo = bar = m foo"  (lvasgn "foo" (lvasgn "bar" (Send Nil "m" [Lvar "foo"])))

    test "casgn_toplevel" "::Foo = 10" (casgn Cbase "Foo"  (RInt 10))
    test "casgn_scoped" "Bar::Foo = 10" (casgn (Const Nil "Bar") "Foo"  (RInt 10))
    test "casgn_unscoped" "Foo = 10" (casgn Nil "Foo" (RInt 10))

    test "masgn" "foo, bar = 1, 2"          (Masgn (Mlhs [Lvasgn "foo" Nothing, Lvasgn "bar" Nothing]) (RArray [RInt 1, RInt 2]))
    test "masgn" "(foo, bar) = 1, 2"        (Masgn (Mlhs [Lvasgn "foo" Nothing, Lvasgn "bar" Nothing]) (RArray [RInt 1, RInt 2]))
    test "masgn" "foo, bar, baz = 1, 2"     (Masgn (Mlhs [Lvasgn "foo" Nothing, Lvasgn "bar" Nothing, Lvasgn "baz" Nothing]) (RArray [RInt 1, RInt 2]))

    test "masgn_splat" "@foo, @@bar = *foo" (Masgn (Mlhs [Ivasgn "@foo" Nothing, Cvasgn "@@bar" Nothing]) (RArray [Splat (Just (Lvar "foo"))]))
    test "masgn_splat" "a, b = *foo, bar"   (Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (RArray [Splat (Just (Lvar "foo")), Lvar "bar"]))
    test "masgn_splat" "a, *b = bar"        (Masgn (Mlhs [Lvasgn "a" Nothing, Splat (Just (Lvasgn "b" Nothing))]) (Lvar "bar"))
    test "masgn_splat" "a, *b, c = bar"     (Masgn (Mlhs [Lvasgn "a" Nothing, Splat (Just (Lvasgn "b" Nothing)), Lvasgn "c" Nothing]) (Lvar "bar"))

    test "masgn_splat" "a, * = bar"         (Masgn (Mlhs [Lvasgn "a" Nothing, Splat Nothing]) (Lvar "bar"))
    test "masgn_splat" "a, *, c = bar"      (Masgn (Mlhs [Lvasgn "a" Nothing, Splat Nothing, Lvasgn "c" Nothing]) (Lvar "bar"))
    test "masgn_splat" "*b = bar"           (Masgn (Mlhs [Splat (Just (Lvasgn "b" Nothing))]) (Lvar "bar"))
    test "masgn_splat" "*b, c = bar"        (Masgn (Mlhs [Splat (Just (Lvasgn "b" Nothing)), Lvasgn "c" Nothing]) (Lvar "bar"))
    test "masgn_splat" "* = bar"            (Masgn (Mlhs [Splat Nothing]) (Lvar "bar"))
    test "masgn_splat" "*, c, d = bar"      (Masgn (Mlhs [Splat Nothing, Lvasgn "c" Nothing, Lvasgn "d" Nothing]) (Lvar "bar"))

  --  test "masgn_nested" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:lvasgn, :a),
  --         s(:mlhs,
  --           s(:lvasgn, :b),
  --           s(:lvasgn, :c))),
  --       s(:lvar, :foo)),
  --     %q{a, (b, c) = foo},

  -- test "masgn_nested" "((b, )) = foo" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:lvasgn, :b)),
  --       s(:lvar, :foo)),

  -- test "masgn_attr" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:send, s(:self), :a=),
  --         s(:indexasgn, s(:self), s(:int, 1), s(:int, 2))),
  --       s(:lvar, :foo)),
  --     %q{self.a, self[1, 2] = foo},
  --     %q{~~~~~~ expression (mlhs.send)
  --       |     ~ selector (mlhs.send)
  --       |            ^ begin (mlhs.indexasgn)
  --       |                 ^ end (mlhs.indexasgn)
  --       |        ~~~~~~~~~~ expression (mlhs.indexasgn)})

  -- test "masgn_attr" "self::a, foo = foo" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:send, s(:self), :a=),
  --         s(:lvasgn, :foo)),
  --       s(:lvar, :foo)),

  -- test "masgn_attr" "self.A, foo = foo" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:send, s(:self), :A=),
  --         s(:lvasgn, :foo)),
  --       s(:lvar, :foo)),

  -- test "masgn_const" "self::A, foo = foo" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:casgn, s(:self), :A),
  --         s(:lvasgn, :foo)),
  --       s(:lvar, :foo)),

  -- test "masgn_const" "::A, foo = foo" (
  --     s(:masgn,
  --       s(:mlhs,
  --         s(:casgn, s(:cbase), :A),
  --         s(:lvasgn, :foo)),
  --       s(:lvar, :foo)),

    test "masgn_cmd" "foo, bar = m foo" (Masgn (Mlhs [Lvasgn "foo" Nothing, Lvasgn "bar" Nothing]) (Send Nil "m" [Lvar "foo"]))

    test "asgn_mrhs" "foo = bar, 1"     (Lvasgn "foo" (Just (RArray [Lvar "bar", RInt 1])))
    test "asgn_mrhs" "foo = *bar"       (Lvasgn "foo" (Just (RArray [Splat (Just (Lvar "bar"))])))
    test "asgn_mrhs" "foo = baz, *bar"  (Lvasgn "foo" (Just (RArray [Lvar "baz", Splat (Just (Lvar "bar"))])))

    test "def" "def foo; end" (Def "foo" (Args []) Nil)
    test "def" "def String; end" (Def "String" (Args []) Nil)
    test "def" "def String=; end" (Def "String=" (Args []) Nil)
    test "def" "def until; end" (Def "until" (Args []) Nil)
    test "def" "def BEGIN; end" (Def "BEGIN" (Args []) Nil)
    test "def" "def END; end" (Def "END" (Args []) Nil)

    test "defs" "def self.foo; end" (Defs Self "foo" (Args []) Nil)
    test "defs" "def self::foo; end" (Defs Self "foo" (Args []) Nil)
    test "defs" "def (foo).foo; end" (Defs (Lvar "foo") "foo" (Args []) Nil)
    test "defs" "def String.foo; end" (Defs (Const Nil "String") "foo" (Args []) Nil)
    test "defs" "def String::foo; end" (Defs (Const Nil "String") "foo" (Args []) Nil)

    test "undef" "undef foo" (Undef [Sym "foo"])
    test "undef" "undef foo, :bar, :\"foo#{1}\"" (Undef [Sym "foo", Sym "bar", Dsym [Str "foo", Begin [RInt 1]]])

    test "alias" "alias :foo bar" (Alias (Sym "foo") (Sym "bar"))

    test "alias_gvar" "alias $a $b" (Alias (Gvar "$a") (Gvar "$b"))
    test "alias_gvar" "alias $a $+" (Alias (Gvar "$a") (BackRef "$+"))

    test "string_interp" "\"foo#{bar}baz\"" (Dstr [Str "foo", Begin [Lvar "bar"], Str "baz"])
    test "string_dvar" "\"#@a #@@a #$a\"" (Dstr [Ivar "@a", Str " ", Cvar "@@a", Str " ", Gvar "$a"])
    test "string_concat" "\"foo#@a\" \"bar\"" (Dstr [Dstr [Str "foo", Ivar "@a"], Str "bar"])
    test "string_concat" "'aaa' 'bbb'" (Dstr [Str "aaa", Str "bbb"])
    test "string_concat" "\"aaa\" \"bbb\"" (Dstr [Str "aaa", Str "bbb"])
  --  test "string___FILE__" "__FILE__" (Str "(assert_parses)")
    test "string___FILE__" "__FILE__" File
    test "character" "?a" (Str "a")

  --   test "heredoc" "<<HERE\nfoo\nbar\nHERE" (Dstr [Str "foo\n", Str "bar\n"])
  --   test "heredoc" "<<'HERE'\nfoo\nbar\nHERE" (Dstr [Str "foo\n", Str "bar\n"])
  --   test "heredoc" "<<`HERE`\nfoo\nbar\nHERE" (Dstr [Str "foo\n", Str "bar\n"])
  -- test "dedenting_heredoc" "p <<~E\nE" (Send Nil "p" (Dstr [])
  -- test "dedenting_heredoc" "p <<~E\n  E" (Send Nil "p" (Dstr [])
  -- test "dedenting_heredoc" "p <<~E\n  x\nE" (Send Nil "p" (Str "x\n")
  -- test "dedenting_heredoc" "p <<~E\n  x\n    y\nE" (Send Nil "p" (Dstr [Str "x\n", Str "  y\n"]))

  -- test "dedenting_heredoc" "p <<~E\n\tx\n    y\nE" (Send Nil "p" (Dstr [Str "x\n", Str "y\n"]))
  -- test "dedenting_heredoc" "p <<~E\n\tx\n        y\nE" (Send Nil "p" (Dstr [Str "x\n", Str "y\n"]))
  -- test "dedenting_heredoc" "p <<~E\n    \tx\n        y\nE" (Send Nil "p" (Dstr [Str "x\n", Str "y\n"]))
  -- test "dedenting_heredoc" "p <<~E\n        \tx\n\ty\nE" (Send Nil "p" (Dstr [Str "\tx\n", Str "y\n"]))
  -- test "dedenting_heredoc" "p <<~E\n  x\n\ny\nE" (Send Nil "p" (Dstr [Str "  x\n", Str "\n", Str "y\n"]))
  -- test "dedenting_heredoc" "p <<~E\n  x\n    \n  y\nE" (Send Nil "p" (Dstr [Str "x\n", Str "   \n", Str "y\n"]))

  --   assert_parses(
  --     (Send Nil, :p,
  --       (Dstr
  --         (Str "  x\n"),
  --         (Str "  y\n"))),
  --     %Q{p <<~E\n    x\n  \\  y\nE},
  --     %q{},

  --   assert_parses(
  --     (Send Nil, :p,
  --       (Dstr
  --         (Str "  x\n"),
  --         (Str "\ty\n"))),
  --     %Q{p <<~E\n    x\n  \\\ty\nE},
  --     %q{},

  --   assert_parses(
  --     (Send Nil, :p,
  --       (Dstr
  --         (Str "  x\n"),
  --         (Str ""),
  --         (Begin
  --           (Lvar :foo)),
  --         (Str "\n"))),
  --     %Q{p <<~"E"\n    x\n  \#{foo}\nE},
  --     %q{},

  --   assert_parses(
  --     (Send Nil, :p,
  --       s(:xstr,
  --         (Str "  x\n"),
  --         (Str ""),
  --         (Begin
  --           (Lvar :foo)),
  --         (Str "\n"))),
  --     %Q{p <<~`E`\n    x\n  \#{foo}\nE},
  --     %q{},

  --   assert_parses(
  --     (Send Nil, :p,
  --       (Dstr
  --         (Str "  x\n"),
  --         (Str ""),
  --         (Begin
  --           (Str "  y")),
  --         (Str "\n"))),
  --     %Q{p <<~"E"\n    x\n  \#{"  y"}\nE},
  --     %q{},
  -- end

  -- # Symbols

  -- test "symbol_plain" ":foo" (Sym :foo)
  -- test "symbol_plain" ":'foo'" (Sym :foo)

  -- test "symbol_interp"
  --   assert_parses(
  --     (Dsym
  --       (Str 'foo'),
  --       (Begin (Lvar :bar)),
  --       (Str 'baz')),
  --     %q{:"foo#{bar}baz"},

  -- end

  -- # Execute-strings

  -- test "xstring_plain"
  --   assert_parses(
  --     s(:xstr, (Str 'foobar')),
  --     %q{`foobar`},
  --     %q{^ begin
  --       |       ^ end
  --       |~~~~~~~~ expression})
  -- end

  -- test "xstring_interp"
  --   assert_parses(
  --     s(:xstr,
  --       (Str 'foo'),
  --       (Begin (Lvar :bar)),
  --       (Str 'baz')),
  --     %q{`foo#{bar}baz`},

  -- -- Regexp

  -- test "regex_plain"
  --   assert_parses(
  --     s(:regexp, (Str 'source'), s(:regopt, :i, :m)),
  --     %q{/source/im},

  -- test "regex_interp"
  --   assert_parses(
  --     s(:regexp,
  --       (Str 'foo'),
  --       (Begin (Lvar :bar)),
  --       (Str 'baz'),
  --       s(:regopt)),
  --     %q{/foo#{bar}baz/},

  -- end

  -- -- Arrays

  -- test "array_plain"
  --   assert_parses(
  --     (Array, (RInt 1), (RInt 2)),
  --     %q{[1, 2]},

  -- test "array_splat"
  --   assert_parses(
  --     (Array,
  --       (RInt 1),
  --       s(:splat, (Lvar :foo)),
  --       (RInt 2)),
  --     %q{[1, *foo, 2]},

  --   assert_parses(
  --     (Array,
  --       (RInt 1),
  --       s(:splat, (Lvar :foo))),
  --     %q{[1, *foo]},

  --   assert_parses(
  --     (Array,
  --       s(:splat, (Lvar :foo))),
  --     %q{[*foo]})
  -- end

  -- test "array_assocs"
  --   assert_parses(
  --     (Array,
  --       (Hash (Pair (RInt 1), (RInt 2)))),
  --     %q{[ 1 => 2 ]},

  --   assert_parses(
  --     (Array,
  --       (RInt 1),
  --       (Hash (Pair (RInt 2), (RInt 3)))),
  --     %q{[ 1, 2 => 3 ]},
  -- end

  -- test "array_words" "%w[foo bar]" (Array [Str 'foo', Str 'bar'])
  -- test "array_words_interp" "%W[foo #{bar}" (Array [Str 'foo', Dstr [Begin [Lvar "bar"]]])
  -- test "array_words_interp" "%W[foo #{bar}foo#@baz]" (Array [Str 'foo', Dstr [Begin [Lvar "bar"], Str "foo", Ivar "@baz"]])
  -- test "array_words_empty" "%w[]" (Array [])
  -- test "array_symbols" "%i[foo bar]" (Array [Sym "foo", Sym "bar"])
  -- test "array_symbols_interp" "%I[foo #{bar}]" (Array [Dsym [Str 'foo', Begin [Lvar "bar"]]])
  -- test "array_symbols_interp" "%I[foo#{bar}]" (Array [Dsym [Str 'foo', Begin [Lvar "bar"]]])

  -- test "array_symbols_empty" "%i[]" (Array [])
  -- test "array_symbols_empty" "%I()" (Array [])

  -- test "hash_label" "{ foo: 2 }" (Hash (Pair (Sym "foo"), (RInt 2)))

  -- test "hash_label_end" "{ 'foo': 2 }"
  --     (Hash (Pair (Sym "foo"), (RInt 2))),

  -- test "hash_label_end" "{ 'foo': 2, 'bar': {}}"
  --     (Hash
  --       (Pair (Sym :foo) (RInt 2)),
  --       (Pair (Sym :bar) (Hash))),

  -- test "hash_kwsplat" "{ foo: 2, **bar }" (Hash (Pair (Sym :foo) (RInt 2)) (Kwsplat (Lvar :bar)))

  -- -- Range

    test "range_inclusive" "1..2" (IRange (RInt 1) (RInt 2))
    test "range_inclusive" "1 .. 2" (IRange (RInt 1) (RInt 2))
    test "range_exclusive" "1...2"  (ERange (RInt 1) (RInt 2))
    test "range_exclusive" "1 ... 2"  (ERange (RInt 1) (RInt 2))

    test "range_endless" "1.." (IRange (RInt 1) Nil)
    test "range_endless" "1 .. " (IRange (RInt 1) Nil)
    test "range_endless" "1..." (ERange (RInt 1) Nil)
    test "range_endless" "1 ... " (ERange (RInt 1) Nil)

    test "if" "if foo then bar; end" (If (Lvar "foo") (Lvar "bar") Nil)
    test "if" "if foo; bar; end" (If (Lvar "foo") (Lvar "bar") Nil)

    test "if_nl_then" "if foo\nthen bar end" (If (Lvar "foo") (Lvar "bar") Nil)

    test "if_mod" "bar if foo" (If (Lvar "foo") (Lvar "bar") Nil)

    test "unless" "unless foo then bar; end" (If (Lvar "foo") Nil (Lvar "bar"))
    test "unless" "unless foo; bar; end" (If (Lvar "foo") Nil (Lvar "bar"))

    xtest "unless_mod" "bar unless foo" (If (Lvar "foo") Nil (Lvar "bar"))

    test "if_else" "if foo then bar; else baz; end" (If (Lvar "foo") (Lvar "bar") (Lvar "baz"))
    test "if_else" "if foo; bar; else baz; end" (If (Lvar "foo") (Lvar "bar") (Lvar "baz"))

    test "unless_else" "unless foo then bar; else baz; end" (If (Lvar "foo") (Lvar "baz") (Lvar "bar"))
    test "unless_else" "unless foo; bar; else baz; end" (If (Lvar "foo") (Lvar "baz") (Lvar "bar"))

    test "if_elsif" "if foo; bar; elsif baz; 1; else 2; end" (If (Lvar "foo") (Lvar "bar") (If (Lvar "baz") (RInt 1) (RInt 2)))

    xtest "ternary" "foo ? 1 : 2" (If (Lvar "foo") (RInt 1) (RInt 2))

    xtest "ternary_ambiguous_symbol" "t=1;(foo)?t:T" (Begin [Lvasgn "t" (Just $ RInt 1), If (Begin [Lvar "foo"]) (Lvar "t") (Const Nil "T")])

    test "if_masgn__24" "if (a, b = foo); end" (If (Begin [Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "foo")]) Nil Nil)

    test "not_masgn__24" "!(a, b = foo)" (Send (Begin [Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "foo")]) "!" [])

    test "cond_begin" "if (bar); foo; end" (If (Begin [Lvar "bar"]) (Lvar "foo") Nil)

    test "cond_begin_masgn" "if (bar; a, b = foo); end" (If (Begin [Lvar "bar", Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "foo")]) Nil Nil)

    --test "cond_match_current_line" "if /wat/; end" (If (MatchCurrentLine (Regexp (Str "wat") (Regopt))) Nil Nil)

    --test "cond_match_current_line" "!/wat/" (Send (MatchCurrentLine (Regexp (Str "wat") (Regopt))) "!")

    test "not" "not foo" (Send (Lvar "foo") "!" [])
    test "not" "not(foo)" (Send (Lvar "foo") "!" [])
    test "not" "not()" (Send (Begin []) "!" [])

    test "not_cmd" "not m foo" (Send (Send Nil "m" [Lvar "foo"]) "!" [])

    test "send_binary_op" "foo + 1"   (Send (Lvar "foo") "+"    [RInt 1])
    test "send_binary_op" "foo - 1"   (Send (Lvar "foo") "-"    [RInt 1])
    test "send_binary_op" "foo * 1"   (Send (Lvar "foo") "*"    [RInt 1])
    test "send_binary_op" "foo / 1"   (Send (Lvar "foo") "/"    [RInt 1])
    test "send_binary_op" "foo % 1"   (Send (Lvar "foo") "%"    [RInt 1])
    test "send_binary_op" "foo ** 1"  (Send (Lvar "foo") "**"   [RInt 1])
    test "send_binary_op" "foo | 1"   (Send (Lvar "foo") "|"    [RInt 1])
    test "send_binary_op" "foo ^ 1"   (Send (Lvar "foo") "^"    [RInt 1])
    test "send_binary_op" "foo & 1"   (Send (Lvar "foo") "&"    [RInt 1])
    test "send_binary_op" "foo <=> 1" (Send (Lvar "foo") "<=>"  [RInt 1])
    test "send_binary_op" "foo < 1"   (Send (Lvar "foo") "<"    [RInt 1])
    test "send_binary_op" "foo <= 1"  (Send (Lvar "foo") "<="   [RInt 1])
    test "send_binary_op" "foo > 1"   (Send (Lvar "foo") ">"    [RInt 1])
    test "send_binary_op" "foo >= 1"  (Send (Lvar "foo") ">="   [RInt 1])
    test "send_binary_op" "foo == 1"  (Send (Lvar "foo") "=="   [RInt 1])
    test "send_binary_op" "foo != 1"  (Send (Lvar "foo") "!="   [RInt 1])
    test "send_binary_op" "foo === 1" (Send (Lvar "foo") "==="  [RInt 1])
    test "send_binary_op" "foo =~ 1"  (Send (Lvar "foo") "=~"   [RInt 1])
    test "send_binary_op" "foo !~ 1"  (Send (Lvar "foo") "!~"   [RInt 1])
    test "send_binary_op" "foo << 1"  (Send (Lvar "foo") "<<"   [RInt 1])
    test "send_binary_op" "foo >> 1"  (Send (Lvar "foo") ">>"   [RInt 1])

    test "send_unary_op" "-foo" (Send (Lvar "foo") "-@" [])
    test "send_unary_op" "+foo" (Send (Lvar "foo") "+@" [])
    test "send_unary_op" "~foo" (Send (Lvar "foo") "~" [])

    test "and" "foo and bar" (And (Lvar "foo") (Lvar "bar"))
    test "and" "foo && bar" (And (Lvar "foo") (Lvar "bar"))

    test "or" "foo or bar" (Or (Lvar "foo") (Lvar "bar"))
    test "or" "foo || bar" (Or (Lvar "foo") (Lvar "bar"))

    test "and_or_masgn" "foo && (a, b = bar)" (And (Lvar "foo") (Begin [Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "bar")]))
    test "and_or_masgn" "foo || (a, b = bar)" (Or (Lvar "foo") (Begin [Masgn (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "bar")]))

    test "case_expr" "case foo; when 'bar'; bar; end" (Case [Lvar "foo", When [Str "bar", Lvar "bar"], Nil])

    test "case_expr_else" "case foo; when 'bar'; bar; else baz; end" (Case [Lvar "foo", When [Str "bar", Lvar "bar"], Lvar "baz"])

    test "case_cond" "case; when foo; 'foo'; end" (Case [Nil, When [Lvar "foo", Str "foo"], Nil])

    test "case_cond_else" "case; when foo; 'foo'; else 'bar'; end" (Case [Nil, When [Lvar "foo", Str "foo"], Str "bar"])

    test "case_cond_just_else" "case; else 'bar'; end" (Case [Nil, Str "bar"])

    test "when_then" "case foo; when 'bar' then bar; end" (Case [Lvar "foo", When [Str "bar", Lvar "bar"], Nil])

    test "when_multi" "case foo; when 'bar', 'baz'; bar; end" (Case [Lvar "foo", When [Str "bar", Str "baz", Lvar "bar"], Nil])

    test "when_splat" "case foo; when 1, *baz; bar; when *foo; end" (Case [Lvar "foo", When [RInt 1, Splat (Just (Lvar "baz")), Lvar "bar"], When [Splat (Just (Lvar "foo")), Nil], Nil])

    test "for" "for a in foo do p a; end" (For (Lvasgn "a" Nothing) (Lvar "foo") (Send Nil "p" [Lvar "a"]))

    test "for" "for a in foo; p a; end" (For (Lvasgn "a" Nothing) (Lvar "foo") (Send Nil "p" [Lvar "a"]))

    test "for_mlhs" "for a, b in foo; p a, b; end" (For (Mlhs [Lvasgn "a" Nothing, Lvasgn "b" Nothing]) (Lvar "foo") (Send Nil "p" [Lvar "a", Lvar "b"]))

    xtest "postexe" "END { 1 }" (Postexe (RInt 1))

    test "while" "while foo do meth end" (While (Lvar "foo") (Send Nil "meth" []))
    xtest "while" "while foo; meth end" (While (Lvar "foo") (Send Nil "meth" []))

    test "until" "until foo do meth end" (Until (Lvar "foo") (Send Nil "meth" []))
    test "until" "until foo; meth end" (Until (Lvar "foo") (Send Nil "meth" []))

    test "while_mod" "meth while foo" (While (Lvar "foo") (Send Nil "meth" []))

    test "until_mod" "meth until foo" (Until (Lvar "foo") (Send Nil "meth" []))

    test "while_post" "begin meth end while foo" (WhilePost (Lvar "foo") (KWBegin [Send Nil "meth" []]))

    test "until_post" "begin meth end until foo" (UntilPost (Lvar "foo") (KWBegin [Send Nil "meth" []]))

    test "rescue" "begin; meth; rescue; foo; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "foo"), Nil]])

    test "rescue_else" "begin; meth; rescue; foo; else; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "foo"), (Lvar "bar")]])

    test "ensure" "begin; meth; ensure; bar; end" (KWBegin [Ensure (Send Nil "meth" []) (Lvar "bar")])

    test "ensure_empty" "begin ensure end" (KWBegin [Ensure Nil Nil])

    test "rescue_ensure" "begin; meth; rescue; baz; ensure; bar; end" (KWBegin [Ensure (Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "baz"), Nil]) (Lvar "bar")])

    test "rescue_else_ensure" "begin; meth; rescue; baz; else foo; ensure; bar end" (KWBegin [Ensure (Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "baz"), Lvar "foo"]) (Lvar "bar")])

    test "rescue_mod" "meth rescue bar" (Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "bar"), Nil])

    xtest "rescue_mod_asgn" "foo = meth rescue bar" (Lvasgn "foo" (Just $ Rescue (Send Nil "meth" []) [Resbody Nil Nil (Lvar "bar"), Nil]))

    --test "rescue_mod_op_assign" "foo += meth rescue bar" (Op_asgn (Lvasgn "foo" Nothing), "+", (Rescue (Send Nil "meth" []), (Resbody Nil Nil (Lvar "bar")), Nil))

    --test "rescue_without_begin_end" "meth do; foo; rescue; bar; end" (Block, (Send Nil "meth" []), (Args), (Rescue (Lvar "foo"), (Resbody Nil Nil (Lvar "bar")), Nil))

    test "resbody_list" "begin; meth; rescue Exception; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody (RArray [Const Nil "Exception"]) Nil (Lvar "bar"), Nil]])

    test "resbody_list_mrhs" "begin; meth; rescue Exception, foo; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody (RArray [Const Nil "Exception", Lvar "foo"]) Nil (Lvar "bar"), Nil]])

    test "resbody_var" "begin; meth; rescue => ex; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody Nil (Lvasgn "ex" Nothing) (Lvar "bar"), Nil]])

    xtest "resbody_var" "begin; meth; rescue => @ex; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody Nil (Ivasgn "@ex" Nothing) (Lvar "bar"), Nil]])

    xtest "resbody_list_var" "begin; meth; rescue foo => ex; bar; end" (KWBegin [Rescue (Send Nil "meth" []) [Resbody (RArray [Lvar "foo"]) (Lvasgn "ex" Nothing) (Lvar "bar"), Nil]])

    test "retry" "retry" (Retry [])

    xtest "hash_empty" "{ }" (Hash [])

    xtest "hash_hashrocket" "{ 1 => 2 }" (Hash [Pair (RInt 1) (RInt 2)])

    xtest "hash_hashrocket" "{ 1 => 2, :foo => \"bar\" }]" (Hash [Pair (RInt 1) (RInt 2), Pair (Sym "foo") (Str "bar")])

    xtest "hash_label" "{ foo: 2 }" (Hash [Pair (Sym "foo") (RInt 2)])

    xtest "hash_label_end" "{ 'foo': 2 }" (Hash [Pair (Sym "foo") (RInt 2)])

    xtest "hash_label_end" "{ 'foo': 2, 'bar': {}}" (Hash [Pair (Sym "foo") (RInt 2), Pair (Sym "bar") (Hash [])])

    xtest "hash_label_end" "f(a ? \"a\":1)" (Send Nil "f" [If (Send Nil "a" []) (Str "a") (RInt 1)])

    xtest "hash_kwsplat" "{ foo: 2, **bar }" (Hash [Pair (Sym "foo") (RInt 2), KWSplat (Lvar "bar")])

    xtest "hash_no_hashrocket" "{ 1, 2 }" (Hash [Pair (RInt 1) (RInt 2)])

    test "args_cmd" "fun(f bar)" (Send Nil "fun" [Send Nil "f" [Lvar "bar"]])

    xtest "args_args_star" "fun(foo, *bar)" (Send Nil "fun" [Lvar "foo", Splat $ Just (Lvar "bar")])

    xtest "args_args_star" "fun(foo, *bar, &baz)" (Send Nil "fun" [Lvar "foo", Splat $ Just (Lvar "bar"), BlockPass (Lvar "baz")])

    xtest "args_star" "fun(*bar)" (Send Nil "fun" [Splat $ Just (Lvar "bar")])

    xtest "args_star" "fun(*bar, &baz)" (Send Nil "fun" [Splat $ Just (Lvar "bar"), BlockPass (Lvar "baz")])

    xtest "args_block_pass" "fun(&bar)" (Send Nil "fun" [BlockPass (Lvar "bar")])

    --test "args_args_comma" "foo[bar,]" (Index, (Lvar "foo"), (Lvar "bar"))

    xtest "args_assocs" "fun(:foo => 1)" (Send Nil "fun" [Hash [Pair (Sym "foo") (RInt 1)]])

    xtest "args_assocs" "fun(:foo => 1, &baz)" (Send Nil "fun" [Hash [Pair (Sym "foo") (RInt 1)], BlockPass (Lvar "baz")])

    --test "args_assocs_comma" "foo[:baz => 1,]" (Index, (Lvar "foo"), (Hash (Pair, (Sym "baz"), (RInt, 1))))

    xtest "args_args_assocs" "fun(foo, :foo => 1)" (Send Nil "fun" [Lvar "foo", Hash [Pair (Sym "foo") (RInt 1)]])

    xtest "args_args_assocs" "fun(foo, :foo => 1, &baz)" (Send Nil "fun" [Lvar "foo", Hash [Pair (Sym "foo") (RInt 1)], BlockPass (Lvar "baz")])

    --test "args_args_assocs_comma" "foo[bar, :baz => 1,]" (Index, (Lvar "foo"), (Lvar "bar"), (Hash (Pair, (Sym "baz"), (RInt, 1))))

    test "space_args_cmd" "fun (f bar)" (Send Nil "fun" [Begin [Send Nil "f" [Lvar "bar"]]])

    test "space_args_arg" "fun (1)" (Send Nil "fun" [Begin [RInt 1]])

    test "space_args_arg_newline" "fun (1\n)" (Send Nil "fun" [Begin [RInt 1]])

    --test "space_args_arg_block" "fun (1) {}" (Block (Send Nil "fun", (Begin (RInt, 1))), (Args), Nil)

    --test "space_args_arg_block" "foo.fun (1) {}" (Block (Send (Lvar "foo"), "fun", (Begin (RInt, 1))), (Args), Nil)

    xtest "space_args_arg_call" "fun (1).to_i" (Send Nil "fun" [Send (Begin [RInt 1]) "to_i" []])


{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# OPTIONS_GHC -Wall #-}
-- Show and Eq instances gives warnings.
{-# OPTIONS_GHC -Wno-orphans #-}

--
-- Copyright (c) 2009-2011, ERICSSON AB
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--
--     * Redistributions of source code must retain the above copyright notice,
--       this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of the ERICSSON AB nor the names of its contributors
--       may be used to endorse or promote products derived from this software
--       without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-- CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-- OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--

module Feldspar.Core.Types
       ( module Feldspar.Core.Types
       ) where

import Data.Array.IO (IOArray)
import Data.Bits
import Data.Complex
import Data.Hash (Hashable(..), combine, hashInt, Hash)
import Data.Int
import Data.IORef
import Data.List
import Data.Typeable (Typeable)
import Data.Word
import qualified Control.Monad.Par as MonadPar
import GHC.TypeLits

import Feldspar.Core.Tuple (Tuply(..))
import Feldspar.Lattice
import Feldspar.Range
import Feldspar.Core.NestedTuples

--------------------------------------------------------------------------------
-- * Heterogenous lists
--------------------------------------------------------------------------------

-- | Heterogeneous list
data a :> b = a :> b
  deriving (Eq, Ord, Show)

infixr 5 :>

instance (Lattice a, Lattice b) => Lattice (a :> b)
  where
    bot = bot :> bot
    top = top :> top
    (a1:>a2) \/ (b1:>b2) = (a1 \/ b1) :> (a2 \/ b2)
    (a1:>a2) /\ (b1:>b2) = (a1 /\ b1) :> (a2 /\ b2)

--------------------------------------------------------------------------------
-- * Integers
--------------------------------------------------------------------------------

-- | Convenience alias for bounded integers
type BoundedInt a = (Bounded a, Integral a, FiniteBits a,
                     Integral (UnsignedRep a), FiniteBits (UnsignedRep a))

-- | Witness for a type level literal. Native is encoded as 0.
data BitWidth (n :: Nat) where
    N8      :: BitWidth 8
    N16     :: BitWidth 16
    N32     :: BitWidth 32
    N64     :: BitWidth 64

instance KnownNat a => Hashable (BitWidth a) where
  hash = hashInt . fromIntegral . natVal

-- | Witness for 'U' or 'S'
data Signedness (s :: Symbol) where
    U :: Signedness "U"
    S :: Signedness "S"

instance KnownSymbol a => Hashable (Signedness a) where
  hash = hash . symbolVal

signedness :: Signedness s -> String
signedness U = "Word"
signedness S = "Int"

-- | A generalization of unsigned and signed integers. The first parameter
-- represents the signedness and the second parameter the number of bits.
type family GenericInt s n where
  GenericInt "U" 8      = Word8
  GenericInt "S" 8      = Int8
  GenericInt "U" 16     = Word16
  GenericInt "S" 16     = Int16
  GenericInt "U" 32     = Word32
  GenericInt "S" 32     = Int32
  GenericInt "U" 64     = Word64
  GenericInt "S" 64     = Int64

-- FIXME: IntN/WordN are hardcoded, not based on Options for architecture.

-- | Signed "natural" representation on an architecture
type IntN = Int32

-- | Unsigned "natural" representation on an architecture
type WordN = Word32

type Length = WordN
type Index  = WordN

--------------------------------------------------------------------------------
-- * Mutable data
--------------------------------------------------------------------------------

-- TODO Make newtypes?

-- | Monad for manipulation of mutable data
type Mut = IO

-- | Mutable references
instance Show (IORef a)
  where
    show _ = "IORef"

-- | Mutable arrays
type MArr a = IOArray Integer a

instance Show (MArr a)
  where
    show _ = "MArr"

--------------------------------------------------------------------------------
-- * Par Monad
--------------------------------------------------------------------------------

-- | Monad for parallel constructs
type Par = MonadPar.Par

-- | Immutable references
type IV = MonadPar.IVar

instance Show (IV a)
  where
    show _ = "IVar"

--------------------------------------------------------------------------------
-- * Elements Language
--------------------------------------------------------------------------------

newtype Elements a = Elements { unE :: [(Index, a)] }

instance Show a => Show (Elements a)
  where
    show (Elements a) = "Elements " ++ show a

instance Eq a => Eq (Elements a)
  where
    _ == _ = False

--------------------------------------------------------------------------------
-- * Future values
--------------------------------------------------------------------------------

newtype FVal a = FVal {unFVal :: a}

instance Show (FVal a)
  where
    show _ = "future"

instance Eq a => Eq (FVal a)
  where
    (FVal a) == (FVal b) = a == b

--------------------------------------------------------------------------------
-- * Type representation
--------------------------------------------------------------------------------

-- | Representation of supported types
data TypeRep a
  where
    UnitType      :: TypeRep ()
    BoolType      :: TypeRep Bool
    IntType       :: ( KnownNat n, KnownSymbol s, BoundedInt (GenericInt s n)
                     , Size (GenericInt s n) ~ Range (GenericInt s n)
                     ) =>
                       Signedness s -> BitWidth n -> TypeRep (GenericInt s n)
    FloatType     :: TypeRep Float
    DoubleType    :: TypeRep Double
    ComplexType   :: RealFloat a => TypeRep a -> TypeRep (Complex a)
    ArrayType     :: TypeRep a -> TypeRep [a]
    Tup2Type      :: TypeRep (Tuple '[a, b]) -> TypeRep (a,b)
    Tup3Type      :: TypeRep (Tuple '[a, b, c]) -> TypeRep (a,b,c)
    Tup4Type      :: TypeRep (Tuple '[a, b, c, d]) -> TypeRep (a, b, c, d)
    Tup5Type      :: TypeRep (Tuple '[a, b, c, d, e]) -> TypeRep (a, b, c, d, e)
    Tup6Type      :: TypeRep (Tuple '[a, b, c, d, e, f])
                  -> TypeRep (a, b, c, d, e, f)
    Tup7Type      :: TypeRep (Tuple '[a, b, c, d, e, f, g])
                  -> TypeRep (a, b, c, d, e, f, g)
    Tup8Type      :: TypeRep (Tuple '[a, b, c, d, e, f, g, h])
                  -> TypeRep (a, b, c, d, e, f, g, h)
    Tup9Type      :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i])
                  -> TypeRep (a, b, c, d, e, f, g, h, i)
    Tup10Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j)
    Tup11Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j, k])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j, k)
    Tup12Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j, k, l)
    Tup13Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j, k, l, m)
    Tup14Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m, n])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j, k, l, m, n)
    Tup15Type     :: TypeRep (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m, n, o])
                  -> TypeRep (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o)
    FunType       :: TypeRep a -> TypeRep b -> TypeRep (a -> b)
    MutType       :: TypeRep a -> TypeRep (Mut a)
    RefType       :: TypeRep a -> TypeRep (IORef a)
    MArrType      :: TypeRep a -> TypeRep (MArr a)
    ParType       :: TypeRep a -> TypeRep (Par a)
    ElementsType  :: TypeRep a -> TypeRep (Elements a)
    ConsType      :: TypeRep a -> TypeRep (Tuple b) -> TypeRep (Tuple (a ': b))
    NilType       :: TypeRep (Tuple '[])
    IVarType      :: TypeRep a -> TypeRep (IV a)
    FValType      :: TypeRep a -> TypeRep (FVal a)
      -- TODO `MArrType` Should have a target-specialized version. Or perhaps
      --      use a single type with a flag to distinguish between immutable and
      --      mutable arrays.

instance Show (TypeRep a)
  where
    show UnitType            = "()"
    show BoolType            = "Bool"
    show (IntType s n)       = signedness s ++ show (natVal n)
    show FloatType           = "Float"
    show DoubleType          = "Double"
    show (ComplexType t)     = "(Complex " ++ show t ++ ")"
    show (ArrayType t)       = "[" ++ show t ++ "]"
    show (Tup2Type t)        = show t
    show (Tup3Type t)        = show t
    show (Tup4Type t)        = show t
    show (Tup5Type t)        = show t
    show (Tup6Type t)        = show t
    show (Tup7Type t)        = show t
    show (Tup8Type t)        = show t
    show (Tup9Type t)        = show t
    show (Tup10Type t)       = show t
    show (Tup11Type t)       = show t
    show (Tup12Type t)       = show t
    show (Tup13Type t)       = show t
    show (Tup14Type t)       = show t
    show (Tup15Type t)       = show t
    show (FunType ta tb)                 = show ta ++ " -> " ++ show tb
    show (MutType ta)                    = unwords ["Mut", show ta]
    show (RefType ta)                    = unwords ["Ref", show ta]
    show (MArrType ta)                   = unwords ["MArr", show ta]
    show (ParType ta)                    = unwords ["Par", show ta]
    show (ElementsType ta)               = unwords ["Elements", show ta]
    show (ConsType ta tb)                = unwords [show ta, ":*", show tb]
    show NilType                         = "TNil"
    show (IVarType ta)                   = unwords ["IVar", show ta]
    show (FValType ta)                   = unwords ["FVal", show ta]

infixl 5 #
(#) :: Hashable a => Hash -> a -> Hash
h # x = h `combine` hash x

instance Hashable (TypeRep a) where
  hash UnitType              = hashInt 1
  hash BoolType              = hashInt 2
  hash (IntType sgn sz)      = hashInt 3 # sgn # sz
  hash FloatType             = hashInt 4
  hash DoubleType            = hashInt 5
  hash (ComplexType t)       = hashInt 6 # t
  hash (ArrayType t)         = hashInt 7 # t
  hash (Tup2Type t)          = hashInt 8 # t
  hash (Tup3Type t)          = hashInt 9 # t
  hash (Tup4Type t)          = hashInt 10 # t
  hash (Tup5Type t)          = hashInt 11 # t
  hash (Tup6Type t)          = hashInt 12 # t
  hash (Tup7Type t)          = hashInt 13 # t
  hash (Tup8Type t)          = hashInt 14 # t
  hash (Tup9Type t)          = hashInt 15 # t
  hash (Tup10Type t)         = hashInt 16 # t
  hash (Tup11Type t)         = hashInt 17 # t
  hash (Tup12Type t)         = hashInt 18 # t
  hash (Tup13Type t)         = hashInt 19 # t
  hash (Tup14Type t)         = hashInt 20 # t
  hash (Tup15Type t)         = hashInt 21 # t
  hash (FunType a b)         = hashInt 22 # a # b
  hash (MutType t)           = hashInt 23 # t
  hash (RefType t)           = hashInt 24 # t
  hash (MArrType t)          = hashInt 25 # t
  hash (ParType t)           = hashInt 26 # t
  hash (ElementsType t)      = hashInt 27 # t
  hash (ConsType a b)        = hashInt 28 # a # b
  hash NilType               = hashInt 29
  hash (IVarType t)          = hashInt 30 # t
  hash (FValType t)          = hashInt 31 # t

-- | The set of supported types
class (Eq a, Show a, Typeable a, Show (Size a), Lattice (Size a), TypeF a) => Type a
  where
    -- | Gives the type representation a value.
    typeRep  :: TypeRep a
    sizeOf   :: a -> Size a

instance Type ()      where typeRep = UnitType;          sizeOf = singletonRange
instance Type Bool    where typeRep = BoolType;          sizeOf = singletonRange
instance Type Word8   where typeRep = IntType U N8;      sizeOf = singletonRange
instance Type Int8    where typeRep = IntType S N8;      sizeOf = singletonRange
instance Type Word16  where typeRep = IntType U N16;     sizeOf = singletonRange
instance Type Int16   where typeRep = IntType S N16;     sizeOf = singletonRange
instance Type Word32  where typeRep = IntType U N32;     sizeOf = singletonRange
instance Type Int32   where typeRep = IntType S N32;     sizeOf = singletonRange
instance Type Word64  where typeRep = IntType U N64;     sizeOf = singletonRange
instance Type Int64   where typeRep = IntType S N64;     sizeOf = singletonRange
instance Type Float   where typeRep = FloatType;         sizeOf = singletonRange
instance Type Double  where typeRep = DoubleType;        sizeOf = singletonRange

instance Hashable a => Hashable (Complex a) where
  hash (re :+ im) = hash re `combine` hash im

instance (Type a, RealFloat a) => Type (Complex a)
  where
    typeRep             = ComplexType typeRep
    sizeOf _            = AnySize

instance Type a => Type [a]
  where
    typeRep       = ArrayType typeRep
    sizeOf as     = singletonRange (genericLength as) :> unions (map sizeOf as)

instance (Type a, Type b) => Type (a, b) where
  typeRep = Tup2Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c) => Type (a, b, c) where
  typeRep = Tup3Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d) => Type (a, b, c, d) where
  typeRep = Tup4Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e) => Type (a, b, c, d, e) where
  typeRep = Tup5Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f)
         => Type (a, b, c, d, e, f) where
  typeRep = Tup6Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g)
         => Type (a, b, c, d, e, f, g) where
  typeRep = Tup7Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h)
         => Type (a, b, c, d, e, f, g, h) where
  typeRep = Tup8Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i)
         => Type (a, b, c, d, e, f, g, h, i) where
  typeRep = Tup9Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j)
         => Type (a, b, c, d, e, f, g, h, i, j) where
  typeRep = Tup10Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j, Type k)
         => Type (a, b, c, d, e, f, g, h, i, j, k) where
  typeRep = Tup11Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j, Type k, Type l)
         => Type (a, b, c, d, e, f, g, h, i, j, k, l) where
  typeRep = Tup12Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j, Type k, Type l, Type m)
         => Type (a, b, c, d, e, f, g, h, i, j, k, l, m) where
  typeRep = Tup13Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j, Type k, Type l, Type m, Type n')
         => Type (a, b, c, d, e, f, g, h, i, j, k, l, m, n') where
  typeRep = Tup14Type typeRep
  sizeOf = sizeOf . unpack

instance (Type a, Type b, Type c, Type d, Type e, Type f, Type g, Type h,
          Type i, Type j, Type k, Type l, Type m, Type n', Type o)
         => Type (a, b, c, d, e, f, g, h, i, j, k, l, m, n', o) where
  typeRep = Tup15Type typeRep
  sizeOf = sizeOf . unpack

instance Type a => Type (IORef a)
  where
    typeRep = RefType typeRep

    sizeOf _ = universal

instance Type a => Type (MArr a)
  where
    typeRep = MArrType typeRep

    sizeOf _ = universal

instance Type a => Type (Elements a)
  where
    typeRep = ElementsType typeRep

    sizeOf _ = universal

instance (Type a, Typeable b, Type (Tuple b)) => Type (Tuple (a ': b)) where
    typeRep = ConsType typeRep typeRep
    sizeOf (x :* xs) = (sizeOf x, sizeOf xs)

instance Type (Tuple '[]) where
    typeRep = NilType
    sizeOf _ = universal

instance Type a => Type (IV a)
  where
    typeRep = IVarType typeRep

    sizeOf _ = universal

instance Type a => Type (FVal a)
  where
    typeRep = FValType typeRep

    sizeOf _ = universal

-- | Extend the class Type to higher order types
class (Typeable a, Show (Size a), Lattice (Size a)) => TypeF a where
  typeRepF :: TypeRep a
  sizeOfF :: a -> Size a

instance {-# OVERLAPPING #-} (TypeF a, TypeF b) => TypeF (a -> b) where
  typeRepF = FunType typeRepF typeRepF
  sizeOfF _ = universal

instance {-# OVERLAPPABLE #-} (Typeable a, Type a) => TypeF a where
  typeRepF = typeRep
  sizeOfF = sizeOf

--------------------------------------------------------------------------------
-- * Sized types
--------------------------------------------------------------------------------

data AnySize = AnySize
    deriving (Eq, Show, Ord)

anySizeFun :: AnySize -> AnySize
anySizeFun AnySize = AnySize

anySizeFun2 :: AnySize -> AnySize -> AnySize
anySizeFun2 AnySize AnySize = AnySize

instance Num AnySize
  where
    fromInteger _ = AnySize
    abs           = anySizeFun
    signum        = anySizeFun
    (+)           = anySizeFun2
    (-)           = anySizeFun2
    (*)           = anySizeFun2

instance Lattice AnySize
  where
    bot  = AnySize
    top  = AnySize
    (\/) = anySizeFun2
    (/\) = anySizeFun2

type family Size a where
  Size ()              = Range ()
  Size Bool            = Range Bool
  Size Word8           = Range Word8
  Size Int8            = Range Int8
  Size Word16          = Range Word16
  Size Int16           = Range Int16
  Size Word32          = Range Word32
  Size Int32           = Range Int32
  Size Word64          = Range Word64
  Size Int64           = Range Int64
  Size Float           = Range Float
  Size Double          = Range Double
  Size (Complex a)     = AnySize
  Size [a]             = Range Length :> Size a
  Size (a, b)          = Size (Tuple '[a, b])
  Size (a, b, c)       = Size (Tuple '[a, b, c])
  Size (a, b, c, d)    = Size (Tuple '[a, b, c, d])
  Size (a, b, c, d, e) = Size (Tuple '[a, b, c, d, e])
  Size (a, b, c, d, e, f) = Size (Tuple '[a, b, c, d, e, f])
  Size (a, b, c, d, e, f, g) = Size (Tuple '[a, b, c, d, e, f, g])
  Size (a, b, c, d, e, f, g, h) = Size (Tuple '[a, b, c, d, e, f, g, h])
  Size (a, b, c, d, e, f, g, h, i) = Size (Tuple '[a, b, c, d, e, f, g, h, i])
  Size (a, b, c, d, e, f, g, h, i, j)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j])
  Size (a, b, c, d, e, f, g, h, i, j, k)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j, k])
  Size (a, b, c, d, e, f, g, h, i, j, k, l)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l])
  Size (a, b, c, d, e, f, g, h, i, j, k, l, m)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m])
  Size (a, b, c, d, e, f, g, h, i, j, k, l, m, n)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m, n])
  Size (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o)
    = Size (Tuple '[a, b, c, d, e, f, g, h, i, j, k, l, m, n, o])
  Size (a -> b)        = (Size a, Size b)
  Size (Mut a)         = Size a
  Size (IORef a)       = Size a
  Size (MArr a)        = Range Length :> Size a
  Size (Par a)         = Size a
  Size (Elements a)    = Range Length :> Size a
  Size (Tuple (a ': b)) = (Size a, Size (Tuple b))
  Size (Tuple '[])     = Size ()
  Size (IV a)          = Size a
  Size (FVal a)        = Size a


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE CPP #-}

module Text.Namelist.Pretty
    ( DString, toString
    , Pretty(..)
    , PrettyConfig(..), Mode(..)
    ) where

import Data.Complex(Complex((:+)))
import Data.CaseInsensitive (CI, original)
import Data.List(intersperse)
import Data.Char(toUpper)

#if !MIN_VERSION_base(4,8,0)
import Data.Monoid(Monoid(..))
#endif

import Data.Monoid((<>))
import Data.Default.Class(Default(def))

import Text.Namelist.Types

newtype DString = DString (String -> String)

instance Monoid DString where
    mempty = DString id
    mappend (DString a) (DString b) = DString (a . b)
    {-# INLINABLE mempty #-}
    {-# INLINABLE mappend #-}

toString :: DString -> String
toString (DString d) = d []

fromString :: String -> DString
fromString s = DString (s ++)

fromShow :: Show a => a -> DString
fromShow = fromString . show

singleton :: Char -> DString
singleton c = DString (c:)

data Mode
    = Compact
    | Large { indent :: Int }

data PrettyConfig = PrettyConfig
    { prettyLogical :: String -> String
    , pairSeparator :: String
    , mode          :: Mode
    }

cfgCompact :: PrettyConfig -> Bool
cfgCompact (PrettyConfig _ _ Compact) = True
cfgCompact _ = False

instance Default PrettyConfig where
    def = PrettyConfig (map toUpper) " = " (Large 2)
    {-# INLINABLE def #-}

class Pretty a where
    ppr :: PrettyConfig -> a -> DString

instance Pretty Index where
    ppr _ (Index i)     = fromShow i
    ppr _ (Range l h s) = opt l <> colon <> opt h <> step
      where
        opt   = maybe mempty fromShow
        colon = singleton ':'
        step  = maybe mempty (\t -> colon <> fromShow t) s
    {-# INLINABLE ppr #-}

surround :: Char -> Char -> DString -> DString
surround a b c = singleton a <> c <> singleton b

instance Pretty [Index] where
    ppr _ = surround '(' ')' . mconcat . intersperse (singleton ',') . map (ppr def)
    {-# INLINABLE ppr #-}

ci :: CI String -> DString
ci = fromString . original

instance Pretty Key where
    ppr _ (Key k)       = ci k
    ppr _ (Indexed k i) = ci k <> ppr def i
    ppr _ (Sub k s)     = ci k <> singleton '%' <> ppr def s
    {-# INLINABLE ppr #-}

instance Pretty Value where
    ppr _   (Integer i) = fromShow i
    ppr _   (Real r)    = fromShow r

    ppr cfg (Complex (r :+ i)) = singleton '(' <> fromShow r <> singleton ',' <> sp <> fromShow i <> singleton ')'
      where
        sp = if cfgCompact cfg then mempty else singleton ' '

    ppr cfg (Logical True)  = fromString . prettyLogical cfg $ if cfgCompact cfg then "T" else ".True."
    ppr cfg (Logical False) = fromString . prettyLogical cfg $ if cfgCompact cfg then "F" else ".False."

    ppr _   (String s)
        | '\'' `notElem` s = singleton '\'' <> fromString s <> singleton '\''
        | '"'  `notElem` s = singleton '"'  <> fromString s <> singleton '"'
        | otherwise = singleton '\'' <> fromString (concatMap escape s) <> singleton '\''
      where
        escape '\'' = "''"
        escape a    = [a]

    ppr cfg (Array a) = mconcat . intersperse sep $ map (ppr cfg) a
      where
        sep = if cfgCompact cfg then singleton ',' else fromString ", "

    ppr cfg (r :* v)  = fromShow r <> singleton '*' <> ppr cfg v

    ppr _   Null = mempty
    {-# INLINABLE ppr #-}

instance Pretty Pair where
    ppr cfg (k := v) = ppr cfg k <> equal <> ppr cfg v
      where
        equal = if cfgCompact cfg then singleton '=' else fromString " = "
    {-# INLINABLE ppr #-}

instance Pretty [Pair] where
    ppr cfg@PrettyConfig{ mode = Compact } =
        mconcat . intersperse (fromString ", ") . map (ppr cfg)

    ppr cfg@PrettyConfig{ mode = Large i } =
        mconcat . intersperse (fromString ",\n") . map (\p -> ind <> ppr cfg p)
      where
        ind = fromString $ replicate i ' '
    {-# INLINABLE ppr #-}

instance Pretty Group where
    ppr cfg (Group g ps) = singleton '&' <> ci g <> gsp <> ppr cfg ps <> ged
      where
        gsp = if cfgCompact cfg then singleton ' ' else singleton '\n'
        ged = if cfgCompact cfg
              then fromString " /"
              else if null ps then singleton '/' else fromString "\n/"
    {-# INLINABLE ppr #-}

instance Pretty [Group] where
    ppr cfg@PrettyConfig{ mode = Compact } =
        mconcat . intersperse (singleton ' ') . map (ppr cfg)

    ppr cfg@PrettyConfig{ mode = Large _ } =
        mconcat . intersperse (singleton '\n') . map (ppr cfg)
    {-# INLINABLE ppr #-}


{-
#!/usr/bin/env stack
--stack --resolver lts-11.3 --install-ghc runghc --stack-yaml /home/aku/kanazawa/dev/cpvoh/stack.yaml
-}

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

--import CPVO.IO.Plot.Band

import CPVO.IO -- inshell2text
--import System.Environment (getArgs)
import Data.List.Split (splitOn)
import qualified Data.Text as T
import qualified Data.Text.Read as T
import Data.List -- intercalate
import Data.Either -- rights
import Numeric.LinearAlgebra.Data -- toLists, fromLists
import Data.Maybe -- fromJust
import Data.String  -- IsString
import System.Directory
import Options.Applicative as O
import Data.Semigroup -- <>

main :: IO ()
main = do
  (Opts fOut useOldBw judulUtama yr atomOs daftarLengkap) <- execParser withHelp
  plotBand fOut useOldBw judulUtama yr atomOs
    $ map (splitOn "#") daftarLengkap
  putStrLn "========DONE======="

debugLine :: Show a => Bool -> a -> IO ()
debugLine False _ = return ()
debugLine _ ss = putStrLn $ "===debug=== " ++ show ss

plotInit :: String
plotInit = unlines $ "set datafile missing '-'":
                     "plot 0 lt -1 lc rgb 'black' title '' , \\":[]

endMultiplot :: String
endMultiplot = unlines [ "unset multiplot" ]

plotBand :: String -> Bool -> String -> String -> String -> [[String]]
         -> IO ()
plotBand fOut useOldBw judulUtama yr atomOs daftarLengkap = do
  tempDir <- (T.unpack . head) <$> inshell2text "mktemp -d -p ./"

  putStrLn $ unlines $ concat daftarLengkap
  (xrangeatas,ticksbaru,arrow,plotplate) <- genSinglePic useOldBw atomOs daftarLengkap ("","",Nothing,[])

  putStrLn $ "============TEMPGLT"
  tempGLT <- head <$> inshell2text "mktemp -p ./"
  putStrLn $ show tempGLT

  writeFile (T.unpack tempGLT) $ genTEMPGLT tempDir judulUtama yr xrangeatas ticksbaru (fromMaybe "" arrow) plotplate

  system_ $ "gnuplot " ++ T.unpack tempGLT
  let target = map (\x -> unwords [ "mv "
                                  , "hasil" ++ x
                                  , "../" ++ fOut ++ x
                                  ]) [".eps",".png"]
  withCurrentDirectory tempDir $
    mapM_ system_ $
      "ps2eps --rotate=+ tmp.ps":
      "epstool --copy -b --quiet tmp.eps hasil.eps":
      "epstopdf hasil.eps":
      "pdftocairo -r 150 -singlefile -jpeg hasil.pdf tmp":
      "convert tmp.jpg -rotate 0 hasil.png":
      "rm -f tmp.jpg":
      target
  putStrLn "===end  : plotBand===="

genSinglePic :: Bool -> String -> [[String]]
              ->    (String, String, Maybe String, [String])
              -> IO (String, String, Maybe String, [String])
genSinglePic _        _      []            res              = return res
genSinglePic useOldBw atomOs (daftarLengkap:ds) (_,_,_,res) = do
  allPBAND@((xrangeatas,ticksbaru,arrow,_,_):_) <- plotSinglePic daftarLengkap useOldBw atomOs 1 []
  let generatedPFBAND = (map takePFBAND allPBAND)
  let generatedPBAND = T.intercalate "," $ filter (not . T.null) $ map fst generatedPFBAND
  let generatedFATBAND = T.intercalate "," $ filter (not . T.null) $ map snd generatedPFBAND
  genSinglePic useOldBw atomOs ds (xrangeatas, ticksbaru, arrow, (T.unpack $ T.intercalate "," $ filter (not . T.null) [generatedPBAND, generatedFATBAND]):res)
    where
      takePFBAND (_,_,_,p,f) = (p,f)

plotSinglePic :: [String]
               -> Bool
               -> String
               -> Int
               -> [(String, String, Maybe String, T.Text, T.Text)]
               -> IO [(String, String, Maybe String, T.Text, T.Text)]

plotSinglePic []                   _        _       _      res = return res
plotSinglePic (daftarLengkap:sisa) useOldBw atomOs colorId res = do
  let (foldernya:spinnya:legend:_) = splitOn ":" daftarLengkap
  bandFiles <- inshell2text $ unwords ["ls",concat[foldernya,"/bnd*spin",spinnya]]
  let daftarFolder = ":" ++ intercalate "@" [legend,spinnya,foldernya]

  let generatedPBAND = T.intercalate ", " $ map (genPlotPlate colorId) $ map (\x -> T.concat["'",x,"'"]) bandFiles
  valBand@(valBandX:valBandY:_) <- fmap (words . T.unpack . head) $ inshell2text $ unwords ["cat",concat [foldernya,"/bnd*spin*",spinnya],"|sed -e '/^#/d' |awk '{if ($3>0.1) print $2,$3}'|sort -k 2n|head -1"]
  (_:condBandY:_) <- fmap (words . T.unpack . head) $ inshell2text $ unwords ["cat", concat[foldernya,"/bnd*spin*",spinnya],"| sed -e '/^#/d' |awk '{if ($3<=0) print $2,$3}'|sort -k 2nr -u|sed -e '/^ *$/d'|head -1"]
  putStrLn $ "===valBand==" ++ show (valBand :: [String])
  gapCoordBandGap <- fmap runGAPband $ inshell2text $ "cat " ++ concat [foldernya,"/bnd*spin*",spinnya]
  let bandGap = last $ words gapCoordBandGap
  putStrLn $ "===allBand=== " ++ bandGap
  let arrow = genArrow bandGap valBandX valBandY condBandY
--  putStrLn $ show gapArrow
--  putStrLn $ "perintahDOS Removed == " ++ unwords [
--    "genPDOSvert.hs ", atomOs, fromJust gapnya, spinnya, daftarFolder]
  generatedFATBAND <- genPBAND useOldBw spinnya "0" atomOs [daftarFolder]
  putStrLn $ show generatedFATBAND
  (xrangeatas,ticksbaru) <- genBandTicks foldernya
  putStrLn $ "========" ++ show xrangeatas ++ "=====" ++ show ticksbaru
  putStrLn $ "=== FINISHED PROCESSING : " ++ daftarLengkap
  plotSinglePic sisa useOldBw atomOs (colorId+1) ((xrangeatas,ticksbaru,arrow,generatedPBAND,generatedFATBAND):res)

genArrow :: [Char] -> [Char] -> [Char] -> [Char] -> Maybe String
genArrow _ _ _ _ = Nothing
genArrow' :: [Char] -> [Char] -> [Char] -> [Char] -> Maybe String
genArrow' bandGap valBandX valBandY condBandY =
        if (bandGap == "0")
           then Nothing
           else --(,) (Just $ "0:" ++ bandGap ++ "@" ++ valBandY)
                    (Just $ unlines [ concat ["set label sprintf ('{/Symbol D}=%.2feV',"
                                      , bandGap
                                      , ") at (" ++ valBandX ++ "-0.55),"
                                      ,"(" ++ valBandY ++ "-0.4) font ',12'"
                                      ]
                             , concat [ "set arrow from "
                                      , valBandX ++ "," ++ condBandY
                                      , " to " ++ valBandX ++ "," ++ valBandY
                                      , " heads noborder lw 2 lc rgb 'blue'"
                                      ]
                             ]
                    )

genPBAND :: Bool -> String -> String -> String -> [String] -> IO T.Text
--genPBAND oldBw spin invStat atomNos foldernya = do
genPBAND   _      _   _       []      _         = return ""
genPBAND   useOldBw  _   invStat atomNos foldernya = do
    let daftaratomOs =  map (splitOn "@") $ splitOn "-" atomNos
        daftarJudulSpinFolders = filter (/=[""]) $ map (splitOn "@") $ splitOn ":" $ unwords foldernya
        daftarAOJSF = zip ([1..] :: [Integer]) $ concat $ map (\a -> zip daftaratomOs $ repeat a) daftarJudulSpinFolders

    let err'' = map ( \(_,([o,a],[_,s,folder'])) -> concat [ "cd ", folder', "; BandWeight.py PROCAR.",cekSpin s invStat "1" "UP" "DN" , " ",a, " ", o]) daftarAOJSF
    if (useOldBw == True) then mapM_ system_ err''
                      else return ()

{--
  (\$2>6.08991818?0.4+\$2:\$2):(\$2>6.08991818&&\$2<6.09991818?1/0:\$3)
--}
    return $ T.pack $ concat $ intercalate [", "]
      $ map ( \(i,([o,a],[j,s,folder])) -> ["'",folder,"/bw.",a,".",o,".PROCAR.",cekSpin s invStat "1" "UP" "DN" ,".dat' u ($1>6.09971818?0.4+$1:$1):2:($1>6.08991818&&$1<6.09991818?0:$3*3) ls ",show i," ps variable title '",concat $ intersperse "." [o,a,j],"'"] :: [String] )
      $ daftarAOJSF

genPlotPlate :: Int -> T.Text -> T.Text
genPlotPlate colorId bnd = T.concat [ bnd , T.pack $ unwords [" u ($2>6.08991818?0.4+$2:$2):($2>6.08991818&&$2<6.09991818?1/0:$3) with line lc", show colorId, "title ''"]]

runGAPband :: [T.Text] -> String
runGAPband bndSpin =
  let (valencePoint, conductionPoint) =
        (\(a,b) -> (last a,head b)) $
        break (\[_,a] -> a > 0) $
        Data.List.sortBy (\[_,a] [_,b] -> if a < b then LT else GT ) $
        toLists $
        (¿ [1,2]) . fromLists $
        filter (/=[]) $  -- kolom $2 x (k-point) -- $3 y Energy
        map ( (map fst) . rights . (map T.double) . T.words ) $
        filter (not . T.isPrefixOf "#") $
        bndSpin

  in if (last conductionPoint - last valencePoint < 0.01)  then "0 0 0"
                                                           else unwords $ map show  [head valencePoint, head conductionPoint, last conductionPoint - last valencePoint]

cekSpin :: (Eq a, IsString a) => a -> a -> a-> p -> p -> p
cekSpin s "0" ok sOk sNo = if s == ok then sOk else sNo
cekSpin s _ ok sOk sNo = if s == ok then sNo else sOk

genBandTicks :: String -> IO (String,String)
genBandTicks foldernya = do
    lTicks'' <- inshell2text $ unwords ["tail -n2  ", foldernya ++ "/bnd*spin1|sed -e '/^ $\\|==/d'|sort -u|awk '{print $2}' "]
    namaTicks'' <- inshell2text $ unwords ["cat ",foldernya ++ "/syml.* "]
    let lTicks = (:) "0.0" $ map T.unpack $ filter (not . T.null) lTicks''
    let namaTicks =  map (\a -> if a == "Gamma" then "{/Symbol G}" else a) $ (map head $ filter (/=[]) $ map (snd . (splitAt 7) . words . T.unpack)  namaTicks'')
    return $ (,) (last lTicks)
           $ unwords $ intersperse "," $ map (\(a,b) -> unwords ["'"++a++"'",b]) $ zip  namaTicks lTicks

genTEMPGLT :: String -> String -> String -> String -> String -> String
           -> [String]
           -> String
genTEMPGLT tempDir judulUtama yr xrangeatas ticksbaru arrow plotplate =
    let subTitles = splitOn "#" judulUtama
     in unlines [ "#!/home/aku/bin/gnuplot -persist"
                , "reset"
                , "set term post landscape enhanced color 'Times-Roman' 12"
                , "set output '" ++ tempDir ++ "/tmp.ps'"
                , "rydberg=13.605"
                , "if (!exists('MP_LEFT'))   MP_LEFT = .1"
                , "if (!exists('MP_RIGHT'))  MP_RIGHT = .95"
                , "if (!exists('MP_BOTTOM')) MP_BOTTOM = .1"
                , "if (!exists('MP_TOP'))    MP_TOP = .9"
                , "if (!exists('MP_GAP'))    MP_GAP = 0.05"
                , unwords [ "set multiplot layout 1," ++ (show $ length subTitles )
                          , "margins screen MP_LEFT, MP_RIGHT, MP_BOTTOM, MP_TOP"
                          , "spacing screen MP_GAP"
                          ]
                , "set xzeroaxis"
                , "set grid"
                , "set key right top Left"
                , "set key samplen 1 spacing 1"
                , "set size ratio 1.5"
                , "set mytics 10"
                , "set xlabel 'Wave Vector'"
                , "set ylabel 'Energy-E_F (eV)'"
--                , "set title '" ++ judulUtama ++ "'"
                , "set yrange ["++ yr ++ "]"
                , "set xrange [0.0:" ++ xrangeatas ++ "]"
                , "set grid noy"
                , "set grid xtics lt 0 lc rgb 'black'"
                , "set style line 1 lt 2 lw 1 lc rgb '#e41a1c'"
                , "set style line 2 lt 6 lw 1 lc rgb '#377eb8'"
                , "set style line 3 lt 6 lw 1 lc rgb '#4daf4a'"
                , "set style line 4 lt 6 lw 1 lc rgb '#984ea3'"
                , "set style line 5 lt 6 lw 1 lc rgb '#ff7f00'"
                , "set style line 8 lt 1 lw 2 lc rgb '#ffcc99'"
                , "set style line 9 lt 1 lw 2 lc rgb '#808080'"
                , "set style line 10 lt 1 lw 2 lc rgb '#94ffb5'"
                , "set style line 11 lt 1 lw 2 lc rgb '#8f7c0'"
                , "set style line 12 lt 1 lw 2 lc rgb '#9dcc0'"
                , "set style line 13 lt 1 lw 2 lc rgb '#c2088'"
                , "set style line 14 lt 1 lw 2 lc rgb '#03380'"
                , "set style line 15 lt 1 lw 2 lc rgb '#ffa45'"
                , "set style line 16 lt 1 lw 2 lc rgb '#ffa8bb'"
                , "set style line 17 lt 1 lw 2 lc rgb '#42660'"
                , "set style line 18 lt 1 lw 2 lc rgb '#ff010'"
                , "set style line 19 lt 1 lw 2 lc rgb '#5ef1f2'"
                , "set style line 20 lt 1 lw 2 lc rgb '#0998f'"
                , "set style line 21 lt 1 lw 2 lc rgb '#e0ff66'"
                , "set style line 22 lt 1 lw 2 lc rgb '#74aff'"
                , "set style line 23 lt 1 lw 2 lc rgb '#9900'"
                , "set style line 24 lt 1 lw 2 lc rgb '#ffff80'"
                , "set style line 25 lt 1 lw 2 lc rgb '#ffff0'"
                , "set style line 26 lt 1 lw 2 lc rgb '#ff505'"
                , "set style arrow 1 heads size screen 0.01,90 lw 2 lc rgb 'navy'"
                , "set key bottom left Left"
                , "set xtics (" ++ ticksbaru ++ ")"
                , (unlines
                     $ intersperse noMidLabel
                     $ zipWith (\s p -> unlines $ ("set title '" ++ s ++ "'"):
                                        unwords [arrow, plotInit ,p]:
                                        [])
                       subTitles plotplate
                  )
                , endMultiplot
                ]

noMidLabel :: String
noMidLabel = unlines $ "unset ylabel":
                       []

data Opts = Opts {
    _fOut       :: String,
    _useOldBw   :: Bool,
    _judulUtama :: String,
    _yr         :: String,
    _atomOs     :: String,
    _targetDir  :: [String]
                 } deriving Show

optsParser :: Parser Opts
optsParser = Opts
             <$> strOption (long "out" <> short 'o' <> metavar "OUTPUT" <>
                          help "target output file" <> value "test")
             <*> switch (long "oldbw" <> short 'b'
                            <> help "Using old BandWeight files")
             <*> strOption (long "titles" <> short 't' <>
                            metavar "SUBTITLES" <>
                            help "Titles of each Figures" <>
                            value "")
             <*> strOption  ( long "yrange" <> short 'y' <>
                            metavar "MIN:MAX" <>
                            help "Y-range of the figure, ex. \"-2.5:7.3\"")
             <*> strOption (long "orbitals" <> short 'r' <> metavar "ORBITALS"
                            <> help "List of atomic orbitals" <> value "")
             <*> (many $ argument str (metavar "TARGETDIRS"))

{-
         "~/kanazawa/dev/cpvoh/app/f2 -o outputfile -b -y \"-9.5:2\" ext.nico2o4.normal/nico2o4.6G10:1:\"11G20.Majority\"#ext.nico2o4.normal/nico2o4.6G10:2:\"11G20.Minority\" ext.nico2o4.normal/nico2o4.1G0:1:\"G0.Majority\"#ext.nico2o4.normal/nico2o4.1G0:2:\"G0.Minority\"":
-}
withHelp :: ParserInfo Opts
withHelp = info
  (helper <*> optsParser)
  (fullDesc <> (progDesc
    "Band Structure Generator for ecalj, please make sure that TARGETDIRS are the last arguments."
  ))


{-# LANGUAGE OverloadedStrings #-}
module ML.Gym.Data
  ( GymData(..)
  , getGymData
  , gymObservationToDoubleList
  , gymObservationToFloatList
  ) where

import           ML.Gym.DType
import           ML.Gym.Range
import           ML.Gym.Space
import           ML.Gym.Util
import           ML.Gym.Value

import           Control.Applicative
import qualified Control.Exception        as E
import           Control.Monad            (void, zipWithM, (>=>))
import qualified CPython                  as Py
import qualified CPython.Constants        as Py
import qualified CPython.Protocols.Object as Py
import qualified CPython.Types            as Py
import qualified CPython.Types.Exception  as Py
import qualified CPython.Types.Module     as Py
import qualified Data.ByteString          as BS
import           Data.Complex             (Complex)
import           Data.Maybe               (fromMaybe)
import qualified Data.Text                as T
import           System.Environment       (getArgs, getProgName)
import           System.Exit
import           System.IO
import           System.Random            (randomRIO)

type GymData = GymValue Double

gymObservationToDoubleList :: GymData -> [Double]
gymObservationToDoubleList = gymValueTo1D

gymObservationToFloatList :: GymData -> [Float]
gymObservationToFloatList = map realToFrac . gymObservationToDoubleList

getGymData :: GymSpace -> Py.SomeObject -> IO (Maybe GymData)
getGymData (GymSpaceDiscrete nr) obj  = fmap (GymScalar . fromIntegral) <$> pyToInteger obj
getGymData (GymSpaceBox range@(GymRange dt _ _) _) obj  = getGymValueWithRange range obj
getGymData (GymSpaceTuple spaces) obj = do
  list <- fromMaybe (error "Gym discrete space 'n' not recognized!") <$> python (Py.toUnicode "spaces" >>= Py.getAttribute obj >>= pyToList)
  fmap GymTuple . sequence <$> zipWithM getGymData spaces list

fromScalar :: GymDType -> Py.SomeObject -> IO (Maybe (GymValue Double))
fromScalar GymBool obj    = fmap (GymScalar . fromBool) <$> pyToBool obj
fromScalar GymInteger obj = fmap (GymScalar . fromIntegral) <$> pyToInteger obj
fromScalar GymFloat obj   = fmap GymScalar <$> pyToDouble obj
fromScalar dt _           = error $ "missing implemenations in fromScalar: " ++ show dt

from1D :: GymDType -> Py.SomeObject -> IO (Maybe (GymValue Double))
from1D GymBool obj    = fmap (Gym1D . map fromBool) <$> pyToListOf pyToBool obj
from1D GymInteger obj = fmap (Gym1D . map fromIntegral) <$> pyToListOf pyToInteger obj
from1D GymFloat obj   = fmap Gym1D <$> pyToListOf pyToDouble obj
from1D dt _           = error $ "missing implemenations in fromScalar: " ++ show dt

from2D :: GymDType -> Py.SomeObject -> IO (Maybe (GymValue Double))
from2D GymBool obj    = fmap (Gym2D . map (map fromBool)) <$> pyToListOfListOf pyToBool obj
from2D GymInteger obj = fmap (Gym2D . map (map fromIntegral)) <$> pyToListOfListOf  pyToInteger obj
from2D GymFloat obj   = fmap Gym2D <$> pyToListOfListOf pyToDouble obj
from2D dt _           = error $ "missing implemenations in fromScalar: " ++ show dt

from3D :: GymDType -> Py.SomeObject -> IO (Maybe (GymValue Double))
from3D GymBool obj    = fmap (Gym3D . map (map (map fromBool))) <$> pyToListOfListOfListOf pyToBool obj
from3D GymInteger obj = fmap (Gym3D . map (map (map fromIntegral))) <$> pyToListOfListOfListOf  pyToInteger obj
from3D GymFloat obj   = fmap Gym3D <$> pyToListOfListOfListOf pyToDouble obj
from3D dt _           = error ("missing implementations in fromScalar: " ++ show dt)

getGymValueWithRange :: GymRange -> Py.SomeObject -> IO (Maybe (GymValue Double))
getGymValueWithRange (GymRange dt (GymScalar _) _) obj = fromScalar dt obj
getGymValueWithRange (GymRange dt (Gym1D _) _) obj     = numPyArray obj >>= from1D dt
getGymValueWithRange (GymRange dt (Gym2D _) _) obj     = numPyArray obj >>= from2D dt
getGymValueWithRange (GymRange dt (Gym3D _) _) obj     = numPyArray obj >>= from3D dt
getGymValueWithRange (GymRange dt (GymTuple xs) _) obj = do
  list <- fromMaybe (error "Gym discrete space 'n' not recognized!") <$> python (Py.toUnicode "spaces" >>= Py.getAttribute obj >>= pyToList)
  fmap GymTuple . sequence <$> zipWithM (\v o -> getGymValueWithRange (GymRange dt v (error "unused")) o) xs list

numPyArray :: Py.SomeObject -> IO Py.SomeObject
numPyArray obj = python $ Py.callMethodArgs obj "tolist" []


{-# LANGUAGE CApiFFI #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module LatticeSymmetries.CSR where

-- import Data.Binary (Binary (..))

-- import Data.Vector.Binary

-- import Control.Exception.Safe (bracket, impureThrow, throwIO)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST
-- import qualified Control.Monad.ST.Unsafe (unsafeIOToST)
-- import Data.Aeson
-- import Data.Aeson.Types (typeMismatch)

-- import Data.Scientific (toRealFloat)

import Data.Bits (Bits, toIntegralSized)
import Data.Complex
import qualified Data.List
import Data.Type.Equality
-- import qualified Data.List.NonEmpty as NonEmpty

-- import qualified Data.Vector.Fusion.Stream.Monadic as Stream

-- import qualified Data.Vector.Storable as S
import qualified Data.Vector
import qualified Data.Vector as B
import qualified Data.Vector.Algorithms.Intro as Intro
import Data.Vector.Fusion.Bundle (Bundle)
import qualified Data.Vector.Fusion.Bundle as Bundle (inplace)
import qualified Data.Vector.Fusion.Bundle.Monadic as Bundle
import Data.Vector.Fusion.Bundle.Size (Size (..), toMax)
import Data.Vector.Fusion.Stream.Monadic (Step (..), Stream (..))
import qualified Data.Vector.Fusion.Util (unId)
import Data.Vector.Generic ((!))
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as GM
import qualified Data.Vector.Storable as S
import qualified Data.Vector.Unboxed
-- import qualified Data.Vector.Storable.Mutable as SM
-- import Data.Vector.Unboxed (Unbox)
-- import qualified Data.Vector.Unboxed as U
-- import Data.Yaml (decodeFileWithWarnings)
-- import Foreign.C.String (CString, peekCString)
import Foreign.C.Types (CInt (..), CUInt (..))
-- import Foreign.ForeignPtr
-- import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
-- import Foreign.Marshal.Alloc (alloca, free, malloc, mallocBytes)
-- import Foreign.Marshal.Array (newArray, withArrayLen)
import Foreign.Marshal.Utils (with)
import Foreign.Ptr (FunPtr, Ptr, castPtr)
-- import Foreign.StablePtr
import Foreign.Storable (Storable (..))
import qualified GHC.Exts as GHC (IsList (..))
import GHC.TypeLits
import qualified GHC.TypeLits as GHC
-- import qualified GHC.ForeignPtr as GHC (Finalizers (..), ForeignPtr (..), ForeignPtrContents (..))
-- import GHC.Generics
-- import qualified GHC.IORef as GHC (atomicSwapIORef)
-- import GHC.Prim
-- import qualified GHC.Ptr as GHC (Ptr (..))
-- import qualified Language.C.Inline as C
-- import qualified Language.C.Inline.Unsafe as CU
import LatticeSymmetries.Dense
-- import LatticeSymmetries.IO
-- import LatticeSymmetries.Types
import qualified System.IO.Unsafe
import qualified Unsafe.Coerce
-- import qualified System.Mem.Weak
import Prelude hiding (group, product, sort)

-- typedef struct ls_csr_matrix {
--   unsigned *offsets;
--   unsigned *columns;
--   _Complex double *off_diag_elements;
--   _Complex double *diag_elements;
--   unsigned dimension;
--   unsigned number_nonzero;
-- } ls_csr_matrix;
data Ccsr_matrix = Ccsr_matrix
  { c_csr_matrix_offsets :: !(Ptr CUInt),
    c_csr_matrix_columns :: !(Ptr CUInt),
    c_csr_matrix_off_diag_elements :: !(Ptr (Complex Double)),
    c_csr_matrix_diag_elements :: !(Ptr (Complex Double)),
    c_csr_matrix_dimension :: !CUInt,
    c_csr_matrix_number_nonzero :: !CUInt
  }

instance Storable Ccsr_matrix where
  alignment _ = 8
  sizeOf _ = 40
  peek p =
    Ccsr_matrix
      <$> peekByteOff p 0
      <*> peekByteOff p 8
      <*> peekByteOff p 16
      <*> peekByteOff p 24
      <*> peekByteOff p 32
      <*> peekByteOff p 36
  poke p x = do
    pokeByteOff p 0 (c_csr_matrix_offsets x)
    pokeByteOff p 8 (c_csr_matrix_columns x)
    pokeByteOff p 16 (c_csr_matrix_off_diag_elements x)
    pokeByteOff p 24 (c_csr_matrix_diag_elements x)
    pokeByteOff p 32 (c_csr_matrix_dimension x)
    pokeByteOff p 36 (c_csr_matrix_number_nonzero x)

data CsrMatrix
  = CsrMatrix
      !(S.Vector CUInt)
      !(S.Vector CUInt)
      !(S.Vector (Complex Double))
      !(S.Vector (Complex Double))
  deriving stock (Show, Eq)

withCsrMatrix :: CsrMatrix -> (Ptr Ccsr_matrix -> IO a) -> IO a
withCsrMatrix m@(CsrMatrix offsets columns offDiagElems diagElems) action =
  S.unsafeWith offsets $ \offsetsPtr ->
    S.unsafeWith columns $ \columnsPtr ->
      S.unsafeWith offDiagElems $ \offDiagElemsPtr ->
        S.unsafeWith diagElems $ \diagElemsPtr ->
          with (Ccsr_matrix offsetsPtr columnsPtr offDiagElemsPtr diagElemsPtr c_dimension c_number_non_zero) action
  where
    c_dimension = fromIntegral (csrDim m)
    c_number_non_zero = fromIntegral (csrNumberNonZero m)

unsafeNewCsrMatrix :: Int -> Int -> IO CsrMatrix
unsafeNewCsrMatrix dimension nnz =
  pure $
    CsrMatrix
      (G.replicate (dimension + 1) 0)
      (G.replicate nnz 0)
      (G.replicate nnz 0)
      (G.replicate dimension 0)

csrMatrixShrink :: HasCallStack => Int -> CsrMatrix -> CsrMatrix
csrMatrixShrink nnz m@(CsrMatrix offsets columns offDiagElems diagElems)
  | nnz == csrNumberNonZero m = m
  | nnz < csrNumberNonZero m = CsrMatrix offsets (G.take nnz columns) (G.take nnz offDiagElems) diagElems
  | otherwise = error "general resizing not implemented"

csrDim :: CsrMatrix -> Int
csrDim (CsrMatrix _ _ _ diagElems) = G.length diagElems

csrNumberNonZero :: CsrMatrix -> Int
csrNumberNonZero (CsrMatrix _ _ offDiagElems _) = G.length offDiagElems

csrMatrixFromDense ::
  (HasCallStack, G.Vector v (Complex Double), G.Vector v Int) =>
  DenseMatrix v (Complex Double) ->
  CsrMatrix
csrMatrixFromDense m@(DenseMatrix nRows nCols v)
  | nRows == nCols = runST $ do
    let nnz = countOffDiagNonZero m
        diagElems = G.convert $ extractDiagonal m
    offsetsBuffer <- GM.new (nRows + 1)
    columnsBuffer <- GM.new nnz
    offDiagElemsBuffer <- GM.new nnz
    let go2 !i !j !offset
          | i == j = go2 i (j + 1) offset
          | j < nCols =
            let x = indexDenseMatrix m (i, j)
             in if x /= 0
                  then do
                    GM.write columnsBuffer offset (fromIntegral j)
                    GM.write offDiagElemsBuffer offset x
                    go2 i (j + 1) (offset + 1)
                  else go2 i (j + 1) offset
          | otherwise = pure offset
        go1 !i !offset
          | i < nRows = do
            offset' <- go2 i 0 offset
            GM.write offsetsBuffer (i + 1) (fromIntegral offset')
            go1 (i + 1) offset'
          | otherwise = pure offset
    GM.write offsetsBuffer 0 0
    offset <- go1 0 0
    unless (nnz == offset) $ error "this is probably a bug"
    CsrMatrix
      <$> G.unsafeFreeze offsetsBuffer
      <*> G.unsafeFreeze columnsBuffer
      <*> G.unsafeFreeze offDiagElemsBuffer
      <*> pure diagElems
  | otherwise = error "expected a square matrix"

foreign import capi unsafe "csr.h ls_csr_plus"
  ls_csr_plus :: Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> IO ()

foreign import capi unsafe "csr.h ls_csr_minus"
  ls_csr_minus :: Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> IO ()

foreign import capi unsafe "csr.h ls_csr_times"
  ls_csr_times :: Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> IO ()

foreign import capi unsafe "csr.h ls_csr_kron"
  ls_csr_kron :: Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> IO ()

binaryOp :: (Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> Ptr Ccsr_matrix -> IO ()) -> CsrMatrix -> CsrMatrix -> CsrMatrix
binaryOp opKernel a b = System.IO.Unsafe.unsafePerformIO $ do
  let pessimisticNNZ = csrNumberNonZero a + csrNumberNonZero b
  c <- unsafeNewCsrMatrix (csrDim a) pessimisticNNZ
  nnz <-
    withCsrMatrix a $ \aPtr ->
      withCsrMatrix b $ \bPtr ->
        withCsrMatrix c $ \cPtr -> do
          opKernel aPtr bPtr cPtr
          fromIntegral . c_csr_matrix_number_nonzero <$> peek cPtr
  pure $ csrMatrixShrink nnz c

instance Num CsrMatrix where
  (+) = binaryOp ls_csr_plus
  (-) = binaryOp ls_csr_minus
  (*) = binaryOp ls_csr_times
  abs (CsrMatrix offsets columns offDiagElems diagElems) = CsrMatrix offsets columns (G.map abs offDiagElems) (G.map abs diagElems)
  signum (CsrMatrix offsets columns offDiagElems diagElems) = CsrMatrix offsets columns (G.map signum offDiagElems) (G.map signum diagElems)
  fromInteger _ = error "Num instance of CsrMatrix does not implement fromInteger"

csrScale :: Complex Double -> CsrMatrix -> CsrMatrix
csrScale c (CsrMatrix offsets columns offDiagElems diagElems) =
  CsrMatrix offsets columns (G.map (c *) offDiagElems) (G.map (c *) diagElems)

csrKron :: CsrMatrix -> CsrMatrix -> CsrMatrix
csrKron a b = System.IO.Unsafe.unsafePerformIO $ do
  let dimension = csrDim a * csrDim b
      nnz = csrNumberNonZero a * csrNumberNonZero b
  c <- unsafeNewCsrMatrix dimension nnz
  nnz' <-
    withCsrMatrix a $ \aPtr ->
      withCsrMatrix b $ \bPtr ->
        withCsrMatrix c $ \cPtr -> do
          ls_csr_kron aPtr bPtr cPtr
          fromIntegral . c_csr_matrix_number_nonzero <$> peek cPtr
  unless (nnz == nnz') $
    error $ "this is probably a bug: " <> show nnz <> " /= " <> show nnz'
  pure c

csrKronMany :: HasCallStack => [CsrMatrix] -> CsrMatrix
csrKronMany [] = error "expected a non-empty list of matrices"
csrKronMany xs = Data.List.foldl1' csrKron xs