content: module Symkell.SeriesSpec (spec) where
  
  import qualified Control.Exception as Exception
  import Data.Text (Text)
  import Test.Hspec
  import Test.QuickCheck
  
  import Symkell.Symbolic
  import Symkell.Series
  import Symkell.Symbolic.Simplify (simplify)
  
  -- Helper function to create a polynomial expression
  polynomial :: [(Integer, Integer)] -> Expression
  polynomial terms = foldr (:+:) (Number 0) $ map makeTerm terms
    where
      makeTerm (coef, power) 
        | power == 0 = Number coef
        | power == 1 = Number coef :*: Symbol "x"
        | otherwise = Number coef :*: (Symbol "x" :**: Number power)
  
  -- Parse a simplified expression to extract coefficients
  -- Only works for simple polynomials in standard form for testing purposes
  extractCoefficients :: Expression -> Int -> [Integer]
  extractCoefficients expr maxOrder = map (getCoef expr) [0..maxOrder]
    where
      getCoef (Number n) 0 = n
      getCoef (Number _ :*: (Symbol _ :**: Number p)) power 
        | p == fromIntegral power = 1
      getCoef (Number n :*: (Symbol _ :**: Number p)) power 
        | p == fromIntegral power = n
      getCoef (Number n :*: Symbol _) 1 = n
      getCoef (Symbol _) 1 = 1
      getCoef (x :+: y) power = getCoef x power + getCoef y power
      getCoef _ _ = 0
  
  spec :: Spec
  spec = do
    describe "Taylor Series Expansion" $ do
      it "expands a constant" $ do
        let expr = Number 5
            series = taylorSeries expr "x" 0 3
        -- A constant should just return the constant itself
        series `shouldBe` Number 5
      
      it "expands a linear function" $ do
        let expr = Number 2 :*: Symbol "x" :+: Number 3
            series = taylorSeries expr "x" 0 2
        -- f(x) = 2x + 3, expanded at x=0: 3 + 2x
        extractCoefficients series 2 `shouldBe` [3, 2, 0]
      
      it "expands a quadratic function around zero" $ do
        let expr = polynomial [(1, 2), (2, 1), (3, 0)]  -- x² + 2x + 3
            series = taylorSeries expr "x" 0 3
        -- f(x) = x² + 2x + 3, expanded at x=0: 3 + 2x + x²
        extractCoefficients series 3 `shouldBe` [3, 2, 1, 0]
      
      it "expands a quadratic function around non-zero point" $ do
        let expr = polynomial [(1, 2), (2, 1), (3, 0)]  -- x² + 2x + 3
            series = taylorSeries expr "x" 2 3
        -- f(x) = x² + 2x + 3, expanded at x=2: f(2) + f'(2)(x-2) + f''(2)(x-2)²/2
        -- f(2) = 11, f'(2) = 4, f''(2) = 2
        -- 11 + 4(x-2) + (x-2)²
        -- This should simplify to x² + 2x + 3 again (as it's a quadratic)
        let simplified = simplify series
        extractCoefficients simplified 2 `shouldBe` [3, 2, 1]
      
      it "expands e^x around zero" $ do
        let expr = UnaryApply Exp (Symbol "x")
            series = taylorSeries expr "x" 0 3
        -- e^x = 1 + x + x²/2 + x³/6 + ...
        let coeffs = extractCoefficients series 3
        coeffs !! 0 `shouldBe` 1
        coeffs !! 1 `shouldBe` 1
        coeffs !! 2 `shouldBe` 1  -- Note: division by 2 might not be represented directly as coefficient
        
      it "expands sin(x) around zero" $ do
        let expr = UnaryApply Sin (Symbol "x") 
            series = taylorSeries expr "x" 0 3
        -- sin(x) = x - x³/6 + ...
        let coeffs = extractCoefficients series 3
        coeffs !! 0 `shouldBe` 0
        coeffs !! 1 `shouldBe` 1
        coeffs !! 2 `shouldBe` 0
        
    describe "Laurent Series Expansion" $ do
      it "expands a function with no poles same as Taylor series" $ do
        let expr = polynomial [(1, 2), (2, 1), (3, 0)]  -- x² + 2x + 3
            laurent = laurentSeries expr "x" 0 3 0
            taylor = taylorSeries expr "x" 0 3
        -- No negative powers expected
        laurent `shouldBe` taylor
      
      it "expands 1/x around zero" $ do
        let expr = Number 1 :/: Symbol "x"
            series = laurentSeries expr "x" 0 0 1
        -- 1/x has a simple pole at x=0, with residue 1
        -- Laurent series is 1/x
        series `shouldBe` Number 1 :/: Symbol "x"
      
      it "expands 1/x² around zero" $ do
        let expr = Number 1 :/: (Symbol "x" :**: Number 2)
            series = laurentSeries expr "x" 0 0 2
        -- 1/x² has a double pole at x=0
        -- Laurent series is 1/x²
        series `shouldBe` Number 1 :/: (Symbol "x" :**: Number 2)
      
      it "expands a function with a removable singularity" $ do
        let expr = (Symbol "x" :**: Number 2) :/: Symbol "x"
            series = laurentSeries expr "x" 0 1 1
        -- x²/x = x, but treating it as having a potential pole
        -- Laurent series should simplify to x
        simplify series `shouldBe` Symbol "x"
      
      it "expands 1/(x-1) around x=0" $ do
        let expr = Number 1 :/: (Symbol "x" :-: Number 1)
            series = laurentSeries expr "x" 0 3 1
        -- 1/(x-1) = -1/1(1-x) = -1 - x - x² - x³ - ... for |x| < 1
        let simplified = simplify series
        let coeffs = extractCoefficients simplified 3
        coeffs !! 0 `shouldBe` -1
        coeffs !! 1 `shouldBe` -1
        coeffs !! 2 `shouldBe` -1
        coeffs !! 3 `shouldBe` -1 