2 files changed, 56 insertions, 22 deletions

diff --git a/src/Sequence/Contact/Types.hs b/src/Sequence/Contact/Types.hs
index 80d4360..9bb9903 100644
--- a/src/Sequence/Contact/Types.hs
+++ b/src/Sequence/Contact/Types.hs
@@ -1,4 +1,4 @@
-{-# LANGUAGE TemplateHaskell, RecordWildCards, OverloadedStrings, OverloadedLists, ViewPatterns, RankNTypes, ImpredicativeTypes, FlexibleContexts #-}
+{-# LANGUAGE TemplateHaskell, RecordWildCards, OverloadedStrings, OverloadedLists, ViewPatterns, RankNTypes, ImpredicativeTypes, FlexibleContexts, TypeSynonymInstances, FlexibleInstances #-}
 
 module Sequence.Contact.Types
   ( module Sequence.Contact.Types
@@ -38,18 +38,10 @@ instance IsString Hitzone where
 
 makePrisms ''Hitzone
 
-makeLenses ''Armor
-
-instance Default Armor where
-  def = Armor
-        { _aBallistic = 0
-        , _aPiercing = 0
-        , _aBlunt = 0
-        , _aHeat = 0
-        , _aCold = 0
-        , _aToxic = 0
-        , _aExplosive = 0
-        }
+makePrisms ''DamageType
+
+instance {-# OVERLAPS #-} Default Armor where
+  def = const 0
 
 makeLenses ''TestResult
 
@@ -188,3 +180,41 @@ sUnconscious = folding $ do
   maxVitality <- preview sMaxVitality
   damage <- view sFatigue
   return $ liftM2 (>=) <$> Just (return damage) <*> maxVitality
+
+absorb :: Armor -> DamageType -> Int -> FormulaM Stats (Map DamageType Int)
+absorb armor = absorb' Map.empty
+  where
+    absorb' :: Map DamageType Int -> DamageType -> Int -> FormulaM Stats (Map DamageType Int)
+    absorb' old dType dmg = do
+      mass <- maybe (return (Nothing :: Maybe Int)) (fmap Just) =<< preview (ctx . sAMass)
+      armor' <- wArmor dType
+      let
+        current = transmit mass dType $ min dmg armor'
+      new <- Map.unionsWith (+) <$> (sequence . Map.elems $ Map.mapWithKey (absorb' current) current)
+      return . Map.unionWith (+) (leftover mass dType . clamp $ dmg - armor') $ if old /= current then new else current
+
+    wArmor :: DamageType -> Formula Stats
+    wArmor Passthrough = return $ 0
+    wArmor Electric = (2 *) . maximum <$> mapM wArmor ([minBound..maxBound] \\ [Electric] :: [DamageType])
+    wArmor x = armor x
+
+    clamp n
+      | n <= 0 = 0
+      | otherwise = n
+
+    transmit, leftover :: Maybe Int -> DamageType -> Int -> Map DamageType Int
+    transmit  _       Ballistic n = [ (Blunt, n) ]
+    transmit  _       Piercing  n = [ (Blunt, n) ]
+    transmit  _       _         _ = []
+
+
+    leftover (Just m) Blunt     n = [ (Fatigue, min n m)
+                                    , (Blunt, clamp $ n - m)
+                                    ]
+    leftover (Just m) Cold      n = [ (Fatigue, min n $ m * 2)
+                                    , (Cold, clamp $ n - m * 2)
+                                    ]
+    leftover _        Explosive n = [ (Fatigue, n)
+                                    , (Explosive, n)
+                                    ]
+    leftover _        dType     n = [ (dType, n) ]
diff --git a/src/Sequence/Contact/Types/Internal.hs b/src/Sequence/Contact/Types/Internal.hs
index 0fe6266..6fde09d 100644
--- a/src/Sequence/Contact/Types/Internal.hs
+++ b/src/Sequence/Contact/Types/Internal.hs
@@ -17,15 +17,19 @@ import Data.CaseInsensitive (CI)
 newtype Hitzone = Hitzone { _hitzone :: CI String }
   deriving (Eq, Ord)
 
-data Armor = Armor
-             { _aBallistic
-             , _aPiercing
-             , _aBlunt
-             , _aHeat
-             , _aCold
-             , _aToxic
-             , _aExplosive :: Int
-             }
+type Armor = DamageType -> Formula Stats
+
+data DamageType = Ballistic
+                | Piercing
+                | Blunt
+                | Heat
+                | Cold
+                | Toxic
+                | Explosive
+                | Fatigue
+                | Electric
+                | Passthrough
+  deriving (Eq, Ord, Enum, Bounded, Show)
 
 data TestResult = CritSuccess { _rWith, _rBy :: Int }
                 | Success { _rWith, _rBy :: Int }