ExprJsonRoundtrip.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}

{- | Round-trip, pipeline, end-to-end, and wire-shape contract tests for the
expression/pipeline serializer.
-}
module IR.ExprJsonRoundtrip (tests) where

import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Char8 as BS8
import Data.Either (fromRight, isLeft)
import Test.HUnit

import qualified Data.Vector.Unboxed as VU
import qualified DataFrame as D
import DataFrame.Expression.Operators (ifThenElse, (.>.))
import qualified DataFrame.Functions as F
import DataFrame.Internal.Column (Columnable, TypedColumn (..), toVector)
import qualified DataFrame.Internal.Column as DI
import DataFrame.Internal.Expression (Expr, UExpr (..))
import DataFrame.Internal.Interpreter (interpret)

import DataFrame.LinearModel (defaultLinearConfig)
import DataFrame.Model (fit, predict)
import DataFrame.Transform (Transform (..), applyTransform)
import DataFrame.Transform.Serialize (
    loadTransformFromFile,
    saveTransformToFile,
 )

import DataFrame.Expr.Serialize (
    SomeExpr (..),
    decodeExprAny,
    decodeExprAt,
    decodeExprFromBytes,
    encodeExpr,
    loadExprAtFromFile,
    loadPipelineFromFile,
    saveExprToFile,
    savePipelineToFile,
 )
import System.IO.Temp (withSystemTempDirectory)

interpD :: D.DataFrame -> Expr Double -> [Double]
interpD df e = case interpret @Double df e of
    Right (TColumn c) -> either (const []) VU.toList (toVector @Double @VU.Vector c)
    Left err -> error (show err)

close :: [Double] -> [Double] -> Bool
close a b = length a == length b && and (zipWith (\x y -> abs (x - y) <= 1e-9) a b)

df3 :: D.DataFrame
df3 = D.fromNamedColumns [("x", DI.fromList [1, 2, 3 :: Double])]

regDF :: D.DataFrame
regDF =
    D.fromNamedColumns
        [ ("x1", DI.fromList xs1)
        , ("x2", DI.fromList xs2)
        , ("y", DI.fromList [2 * a - 3 * b + 1 | (a, b) <- zip xs1 xs2])
        ]
  where
    xs1 = [1, 2, 3, 4, 5, 6, 7, 8] :: [Double]
    xs2 = [2, 1, 4, 3, 6, 5, 8, 7] :: [Double]

{- | An expression round-trips when encode → decode → re-encode yields the same
JSON. Avoids needing an 'Eq' instance on the GADT.
-}
roundtrips :: (Columnable a) => Expr a -> Bool
roundtrips e = fromRight False $ do
    v1 <- encodeExpr e
    SomeExpr _ e' <- decodeExprAny v1
    v2 <- encodeExpr e'
    Right (v1 == v2)

testCase :: String -> Bool -> Test
testCase name ok = TestCase (assertBool name ok)

-- Row-wise expression node kinds.
exprRoundtripTests :: [Test]
exprRoundtripTests =
    [ testCase "col" (roundtrips (F.col @Double "x"))
    , testCase "lit" (roundtrips (F.lit @Int 42))
    , testCase "arithmetic" (roundtrips (F.col @Double "x" * F.lit 2 + F.lit 1))
    , testCase "division" (roundtrips (F.col @Double "x" / F.lit 2))
    , testCase "comparison" (roundtrips (F.col @Double "x" .>. F.lit 0))
    , testCase "unary sqrt" (roundtrips (sqrt (F.col @Double "x")))
    , testCase
        "if/then/else"
        ( roundtrips
            ( ifThenElse (F.col @Double "x" .>. F.lit 0) (F.lit 1.0) (F.lit (-1.0)) ::
                Expr Double
            )
        )
    ]

-- Aggregation and window node kinds (the decoder cases added for full parity).
aggRoundtripTests :: [Test]
aggRoundtripTests =
    [ testCase "sum" (roundtrips (F.sum (F.col @Double "x")))
    , testCase "count" (roundtrips (F.count (F.col @Double "x")))
    , testCase "minimum" (roundtrips (F.minimum (F.col @Double "x")))
    , testCase "maximum" (roundtrips (F.maximum (F.col @Int "x")))
    , testCase "mean" (roundtrips (F.mean (F.col @Double "x")))
    , testCase "variance" (roundtrips (F.variance (F.col @Double "x")))
    , testCase "median" (roundtrips (F.median (F.col @Double "x")))
    , testCase "over" (roundtrips (F.over ["g"] (F.sum (F.col @Double "x"))))
    ]

-- Documented limitation: collect's list-typed output has no type tag, so it
-- cannot be encoded.
limitationTests :: [Test]
limitationTests =
    [ testCase
        "collect cannot encode"
        (isLeft (encodeExpr (F.collect (F.col @Double "x"))))
    ]

-- Pin the wire shape Python decodes against: (x * 2) + 1 over doubles.
contractTests :: [Test]
contractTests =
    [ TestCase $ case Aeson.eitherDecodeStrict (BS8.pack contractJson) >>= decodeExprAt @Double of
        Left err -> assertFailure ("contract decode failed: " <> err)
        Right e ->
            assertBool "contract interprets to [3,5,7]" (close (interpD df3 e) [3, 5, 7])
    , testCase
        "collect agg node fails to decode"
        (isLeft (decodeExprFromBytes (BS8.pack collectJson)))
    ]

contractJson :: String
contractJson =
    "{\"node\":\"binary\",\"out_type\":\"double\",\"op\":\"add\",\"arg_type\":\"double\","
        <> "\"lhs\":{\"node\":\"binary\",\"out_type\":\"double\",\"op\":\"mult\",\"arg_type\":\"double\","
        <> "\"lhs\":{\"node\":\"col\",\"out_type\":\"double\",\"name\":\"x\"},"
        <> "\"rhs\":{\"node\":\"lit\",\"out_type\":\"double\",\"value\":2.0}},"
        <> "\"rhs\":{\"node\":\"lit\",\"out_type\":\"double\",\"value\":1.0}}"

collectJson :: String
collectJson =
    "{\"node\":\"agg\",\"out_type\":\"double\",\"agg\":\"collect\",\"arg_type\":\"double\","
        <> "\"arg\":{\"node\":\"col\",\"out_type\":\"double\",\"name\":\"x\"}}"

-- File round-trips for a single expression, a pipeline, and a fitted model.
ioTests :: [Test]
ioTests =
    [ TestCase $ withSystemTempDirectory "expr-ser" $ \dir -> do
        let fp = dir ++ "/expr.json"
            e = F.col @Double "x" * F.lit 2 + F.lit 1
        sr <- saveExprToFile fp e
        assertEqual "save expr ok" (Right ()) sr
        loaded <- loadExprAtFromFile @Double fp
        case loaded of
            Left err -> assertFailure ("load expr failed: " <> err)
            Right e' ->
                assertBool
                    "loaded expr interprets equal"
                    (close (interpD df3 e') (interpD df3 e))
    , TestCase $ withSystemTempDirectory "expr-ser" $ \dir -> do
        let fp = dir ++ "/pipeline.json"
            nes =
                [ ("z", UExpr (F.col @Double "x" * F.lit 2))
                , ("w", UExpr (F.col @Double "x" + F.lit 10))
                ]
        sr <- savePipelineToFile fp nes
        assertEqual "save pipeline ok" (Right ()) sr
        loaded <- loadPipelineFromFile fp
        case loaded of
            Left err -> assertFailure ("load pipeline failed: " <> err)
            Right nes' -> do
                let df' = applyTransform (Transform nes') df3
                assertBool "z = 2x" (close (interpD df' (F.col @Double "z")) [2, 4, 6])
                assertBool "w = x+10" (close (interpD df' (F.col @Double "w")) [11, 12, 13])
    , TestCase $ withSystemTempDirectory "expr-ser" $ \dir -> do
        let fp = dir ++ "/transform.json"
            t = Transform [("z", UExpr (F.col @Double "x" * F.lit 2))]
        sr <- saveTransformToFile fp t
        assertEqual "save transform ok" (Right ()) sr
        loaded <- loadTransformFromFile fp
        case loaded of
            Left err -> assertFailure ("load transform failed: " <> err)
            Right t' ->
                assertBool
                    "transform z = 2x"
                    (close (interpD (applyTransform t' df3) (F.col @Double "z")) [2, 4, 6])
    , TestCase $ withSystemTempDirectory "expr-ser" $ \dir -> do
        -- End-to-end: fit OLS, persist its prediction, reload, compare outputs.
        let fp = dir ++ "/model.json"
            m = fit defaultLinearConfig (F.col @Double "y") regDF
        sr <- saveExprToFile fp (predict m)
        assertEqual "save model ok" (Right ()) sr
        loaded <- loadExprAtFromFile @Double fp
        case loaded of
            Left err -> assertFailure ("load model failed: " <> err)
            Right e' ->
                assertBool
                    "reloaded prediction matches original"
                    (close (interpD regDF e') (interpD regDF (predict m)))
    ]

tests :: [Test]
tests =
    exprRoundtripTests
        ++ aggRoundtripTests
        ++ limitationTests
        ++ contractTests
        ++ ioTests