{-
   Copyright 2017, Matthew Danish, Dominic Orchard, Andrew Rice, Mistral Contrastin

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-}

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PatternGuards #-}

module Camfort.Specification.Units.BackendTypes
  ( UnitSet, Dim, Sub, identDim, isIdentDim, dimFromUnitInfo, dimFromUnitInfos, dimToUnitInfo, dimToUnitInfos
  , subFromList, subToList, identSub, applySub, composeSubs, prop_composition, freeDimVars, dimSimplify
  , dimToConstraint, constraintToDim, dimMultiply, dimRaisePow, dimParamEq, dimParamEqCon, normaliseDim, dimFromList )
where

import qualified Data.Map.Strict as M
import qualified Data.Set as S
import Camfort.Specification.Units.Environment (UnitInfo(..), Constraint(..), flattenUnits, foldUnits, unitParamEq)
import Data.List (partition, foldl', foldl1', sortBy, maximumBy)
import Data.Ord (comparing)
import Data.Maybe (fromMaybe)

-- | Set of UnitInfos
type UnitSet = S.Set UnitInfo

-- | Represents a dimension: collection of units raised to a power,
-- multiplied together. Implemented as a map of unit -> power.
type Dim = M.Map UnitInfo Integer

-- | Represents a substitution: map of unit to be replaced ->
-- dimension to replace it with.
type Sub = M.Map UnitInfo Dim

-- Maintain invariant of Dim whereby all units raised to the zeroth
-- power are eliminated from the map.
removeZeroes :: Dim -> Dim
removeZeroes = M.filterWithKey f
  where
    f _ 0           = False
    f UnitlessVar _ = False
    f UnitlessLit _ = False
    f _ _           = True

-- Handy function to strip away the UnitPow & return the essence.
getUnitPow :: UnitInfo -> (UnitInfo, Integer)
getUnitPow (UnitPow u p) = (u', floor p * p')
  where (u', p') = getUnitPow u
getUnitPow u = (u, 1)

-- | The identity Dim: 1.
identDim :: Dim
identDim = M.empty

-- | Test for identity.
isIdentDim :: Dim -> Bool
isIdentDim = M.null . removeZeroes

-- | Convert from list of implicitly multipled units into a Dim.
dimFromUnitInfos :: [UnitInfo] -> Dim
dimFromUnitInfos = dimFromList . map getUnitPow

-- | Convert a UnitInfo to a Dim.
dimFromUnitInfo :: UnitInfo -> Dim
dimFromUnitInfo = dimFromUnitInfos . flattenUnits

-- | Convert a Dim into an implicitly multipled list of units.
dimToUnitInfos :: Dim -> [UnitInfo]
dimToUnitInfos = map (\ (u, p) -> UnitPow u (fromInteger p)) . M.toList

-- | Convert a Dim into a UnitInfo.
dimToUnitInfo :: Dim -> UnitInfo
dimToUnitInfo = foldUnits . dimToUnitInfos

-- | Convert a Constraint into a Dim where lhs/rhs is implicitly equal
-- to 1. Also normalise the powers by dividing by the gcd and making
-- the largest absolute value power be positive.
constraintToDim :: Constraint -> Dim
constraintToDim (ConEq lhs rhs) = normaliseDim (dimFromUnitInfo lhs `dimMultiply` dimRaisePow (-1) (dimFromUnitInfo rhs))
constraintToDim (ConConj cons)  = foldl' dimMultiply identDim $ map constraintToDim cons

-- | Divide the powers of a dimension by their collective GCD.
normaliseDim :: Dim -> Dim
normaliseDim dim
  | M.null dim = dim
  | otherwise  = M.map (`div` divisor) dim
  where
    divisor = (maxPow `div` abs(maxPow)) * gcds (M.elems dim)
    maxPow  = maximumBy (comparing abs) (M.elems dim)

    gcds []  = 1
    gcds [x] = x
    gcds xs  = foldl1' gcd xs

-- | Multiply two Dims
dimMultiply :: Dim -> Dim -> Dim
dimMultiply d1 d2 = removeZeroes $ M.unionWith (+) d1 d2

-- | Raise the dimension to the given power
dimRaisePow :: Integer -> Dim -> Dim
dimRaisePow 0 _ = identDim
dimRaisePow k d = M.map (* k) d

-- | Compare two Dims, not minding the difference between
-- UnitParam*Abs and UnitParam*Use versions of the polymorphic
-- constructors. Varies from a 'constraint parametric equality'
-- operator because it doesn't assume that dimRaisePow can be used
-- arbitrarily.
dimParamEq :: Dim -> Dim -> Bool
dimParamEq d1 d2 = dimParamEq' (M.toList d1) (M.toList d2)

dimParamEq' :: [(UnitInfo, Integer)] -> [(UnitInfo, Integer)] -> Bool
dimParamEq' [] []             = True
dimParamEq' [] _              = False
dimParamEq' ((u1, p1):d1') d2 = case partition (unitParamEq u1 . fst) d2 of
  ((u2, p2):d2', d2'') -> dimParamEq' (rem1 ++ d1') (rem2 ++ d2' ++ d2'')
    where
      (rem1, rem2) | p1 == p2  = ([], [])
                   | p1 < p2   = ([], [(u2, p2 - p1)])
                   | p1 > p2   = ([(u1, p1 - p2)], [])
                   | otherwise = error "dimParamEq'"
  _ -> False

-- | Similar to dimParamEq but assume that dimRaisePow can be
-- arbitrarily applied to each of the parameters, because they now
-- represent the unit equation 'd = 1'. In practice this means
-- computing the GCD of the powers and dividing.
dimParamEqCon :: Dim -> Dim -> Bool
dimParamEqCon d1 d2 = normaliseDim d1 `dimParamEq` normaliseDim d2

-- | Create a constraint that the given Dim is equal to the identity unit.
dimToConstraint :: Dim -> Constraint
dimToConstraint d = ConEq (dimToUnitInfo positives) (dimToUnitInfo (M.map (* (-1)) negatives))
  where
    (negatives, positives) = M.partition (< 0) d

-- | Convert a list of units paired with their corresponding power into a Dim.
dimFromList :: [(UnitInfo, Integer)] -> Dim
dimFromList = removeZeroes . M.fromListWith (+)

-- | Convert a list of units paired with their corresponding
-- substitution (as a Dim) into a Sub. Note that this is equivalent to
-- repeatedly composing substitutions, so earlier substitutions will
-- affect later ones.
subFromList :: [(UnitInfo, Dim)] -> Sub
subFromList = foldl' composeSubs identSub . map (uncurry M.singleton)

-- | Convert a Sub into an association-list format of unit mapped to unit.
subToList :: Sub -> [(UnitInfo, UnitInfo)]
subToList = map (fmap dimToUnitInfo) . M.toList

-- | Identity substitution (empty).
identSub :: Sub
identSub = M.empty

-- | Identity substitution filled out with some identity entries for a
-- certain set of units (useful for when two substitutions have
-- different sets of keys).
identSubWith :: [UnitInfo] -> Sub
identSubWith = M.fromList . map (\ u -> (u, dimFromList [(u, 1)]))

-- | Apply a substitution to a dimension.
applySub :: Sub -> Dim -> Dim
applySub sub dim =
  removeZeroes $ M.unionsWith (+) [ M.map (*p) (M.singleton ui 1 `fromMaybe` M.lookup ui sub) | (ui, p) <- M.toList dim ]

-- | Compose two substitutions.
composeSubs :: Sub -> Sub -> Sub
composeSubs sub1 sub2 = M.map (applySub sub1) (M.unionWith (curry snd) ident1 sub2)
  where
    ident1 = identSubWith (M.keys sub1)

-- | Test the composition property: f (g x) == (f . g) x
prop_composition :: Dim -> Sub -> Sub -> Bool
prop_composition d s1 s2 = applySub s1 (applySub s2 d) == applySub (composeSubs s1 s2) d

-- | Extract a list of 'free dimension variables' from a given Dim.
freeDimVars :: Dim -> [UnitInfo]
freeDimVars = filter isVar . M.keys
  where
    isVar (UnitParamPosAbs {}) = True
    isVar (UnitParamPosUse {}) = True
    isVar (UnitParamVarAbs {}) = True
    isVar (UnitParamVarUse {}) = True
    isVar (UnitParamLitAbs {}) = True
    isVar (UnitParamLitUse {}) = True
    isVar (UnitLiteral {})     = True
    isVar (UnitVar {})         = True
    isVar _                    = False

-- | The 'dimSimplify' algorithm as shown in Kennedy's technical report Fig 3.4.
dimSimplify :: UnitSet -> Dim -> Sub
dimSimplify excludes dim
  | null valids = identSub

  | (u, x):_ <- valids, x < 0
  , sub1     <- M.singleton u (M.singleton u (-1))
  , sub2     <- dimSimplify excludes (applySub sub1 dim) = composeSubs sub2 sub1

  | (u, x):[] <- valids = M.singleton u (dimFromList ((u, 1):[(v, -div y x) | (v, y) <- invals]))

  | (u, x):_ <- valids
  , sub1     <- M.singleton u (dimFromList ((u, 1):[(v, -div y x) | (v, y) <- M.toList dim, v /= u]))
  , sub2     <- dimSimplify excludes (applySub sub1 dim) = composeSubs sub2 sub1
  | otherwise = error "dimSimplify"
  where
    valids   = sortBy (comparing (abs . snd)) . filter ((`S.notMember` excludes) . fst) $ M.toList dim
    validSet = S.fromList (map fst valids)
    invals   = filter ((`S.notMember` validSet) . fst) $ M.toList dim

-- testVar x = UnitVar (x, x)

-- u0 = testVar "u0"
-- u1 = testVar "u1"
-- u2 = testVar "u2"
-- u3 = testVar "u3"
-- u4 = testVar "u4"

-- dim1 = dimFromList [(u1, 6), (u2, 15), (u3, -7), (u4, 12)]
-- dim2 = dimFromList [(u1, 2), (u2, 15), (u3, -9)]

-- test1 = applySub (dimSimplify (S.fromList [u3,u4]) dim1) dim1 == dimFromList [(u2, 3), (u3, 2)]

-- test2 =  (dimSimplify (S.fromList [u0]) (dimFromList [(u0, 1), (u1, -2)]))