-- | Shared code

{-# LANGUAGE ScopedTypeVariables, TypeFamilies, BangPatterns, PackageImports,
             TypeSynonymInstances, FlexibleInstances, FlexibleContexts,
             ExistentialQuantification 
  #-}
module Math.ThomPoly.Shared where

--------------------------------------------------------------------------------

import Data.List
import Data.Ratio
import Data.Proxy

import Math.Combinat.Classes
import Math.Combinat.Partitions.Integer
import Math.Combinat.Sets

import Math.FreeModule.Symbol
import Math.FreeModule.SortedList
import Math.FreeModule.PrettyPrint
import Math.FreeModule.PP
-- import FreeModule.Parser

import Math.Algebra.ModP
import Math.Algebra.Schur
import Math.Algebra.Determinant

--------------------------------------------------------------------------------
-- * Rings and fields

data AnyRing = forall r. CoeffRing r => AnyRing (Proxy r)

solveAny :: Problem problem => AnyRing -> Batch -> problem -> FreeMod Schur Integer
solveAny anyring batch prob = case anyring of
  AnyRing pxy -> solveAndProject pxy batch prob

ringZZ, ringQQ, ringZp :: AnyRing
ringZZ = AnyRing (Proxy :: Proxy Integer )
ringQQ = AnyRing (Proxy :: Proxy Rational)
ringZp = AnyRing (Proxy :: Proxy Zp      )

--------------------------------------------------------------------------------

class 
  ( Eq a , Num a , Show a , Determinant a
  , Eq (FieldOfFractions a) , Show (FieldOfFractions a) , Fractional (FieldOfFractions a)
  , Pretty (Term a)
  ) => CoeffRing a 
  where
    type FieldOfFractions a :: *
    embed    :: a -> FieldOfFractions a
    project  :: Proxy a -> FieldOfFractions a -> Maybe a
    toBigInt :: Proxy a -> FieldOfFractions a -> Maybe Integer

ratToInt :: Rational -> Maybe Integer
ratToInt x = case denominator x of { 1 -> Just (numerator x) ; _ -> Nothing }

--------------------------------------------------------------------------------

instance CoeffRing Integer where
  type FieldOfFractions Integer = Rational
  embed      = fromInteger
  project  _ = ratToInt
  toBigInt _ = ratToInt

instance CoeffRing Rational where
  type FieldOfFractions Rational = Rational
  embed      = id
  project  _ = Just
  toBigInt _ = ratToInt

instance CoeffRing Zp where
  type FieldOfFractions Zp = Zp
  embed      = id
  project  _ = Just
  toBigInt _ = Just . fromIntegral . fromZp

unsafeProject :: CoeffRing c => Proxy c -> FieldOfFractions c -> Integer
unsafeProject pxy x = case toBigInt pxy x of
  Just y  -> y
  Nothing -> error "cannot project back result"  

--------------------------------------------------------------------------------
-- * Thom polynomial problems

class Problem problem where
  baseFName :: problem -> String
  calcStats :: problem -> Stats
  solve     :: CoeffRing coeff => Proxy coeff -> Batch -> problem -> FreeMod Schur (FieldOfFractions coeff)

fullFName :: Problem problem => Batch -> problem -> FilePath
fullFName batch prob = baseFName prob ++ batchSuffix batch ++ ".txt"

solveAndProject 
  :: forall problem coeff. (Problem problem, CoeffRing coeff)
  => Proxy coeff -> Batch -> problem -> FreeMod Schur Integer
solveAndProject pxy batch prob = coeffMap (unsafeProject pxy) $ solve pxy batch prob where

--------------------------------------------------------------------------------

-- | \"Statistics\" of a problem
data Stats  = Stats 
  { _codim0   :: !Int        -- ^ codimension (for @m=n@)
  , _mu       :: !Int        -- ^ algebraic multiplicity (minus 1)
  , _maxPairs :: !Int        -- ^ maximum number of possible non-zero coefficients
  }
  deriving Show

-------------------------------------------------------------------------------
-- * Batches

data Batch = Batch 
  { _whichBatch :: !Int 
  , _nBatches   :: !Int 
  }
  deriving Show

defaultBatch :: Batch
defaultBatch = Batch 1 1

selectBatch :: Batch -> [a] -> [a]
selectBatch (Batch a b) xs  
  | a < 1     = error "selectBatch: a<1"
  | a > b     = error "selectBatch: a>b"
  | b == 1    = xs
  | otherwise = take bsize $ drop ((a-1)*bsize) $ xs
  where
    n     = length xs
    (q,r) = divMod n b
    bsize = case r of
      0 -> q
      _ -> q+1

batchSuffix :: Batch -> String
batchSuffix (Batch a b)
  | b == 1    = ""
  | otherwise = "_batch" ++ show a ++ "of" ++ show b

{-
-- sanity test
testBatch'' b n = concat [ selectBatch (Batch i b) [1..n] | i<-[1..b] ] == [1..n]
testBatch'  b   = and [ testBatch'' b n | n<-[0..1000] ]
testBatch       = and [ testBatch'  b   | b<-[1..100 ] ] 
-}
 
--------------------------------------------------------------------------------
-- * Misc

-- type CoeffRing = Zp  -- Rational -- Integer

type Term coeff = FreeMod Symbol coeff

alpha :: CoeffRing coeff => Int -> Term coeff
alpha i = fromBase $ Symbol "alpha" (Just i) 

newtype Schur = Schur Partition deriving (Eq,Ord,Show)

instance Pretty Schur where 
  pretty (Schur part) = 's' : show (fromPartition part)

--------------------------------------------------------------------------------
-- * Evaluate

evaluate :: (Num a, FreeModule x) => (Base x -> Coeff x -> a) -> x -> a
evaluate f = sum . map (uncurry f) . toList
         
--------------------------------------------------------------------------------
-- * Signed partitions

-- | Pairs of partition with weights of fix difference, given by 
-- the third parameter, @ofs=|pos|-|neg|@, and complementary length;
-- the first giving the positive deviation compared to the box of (m-n+i)*i,
-- and the second giving the negative one.
-- Picture:
--
-- >         m-n+i                    n-i
-- >    +------------------+----------------+---------+
-- >    |                  |          _____/          |
-- >  i |      lambda      |  pos____/                |
-- >    |                  |    /                     |
-- >    |                  |   /                      |
-- > mu +................._|__/                       | mu
-- >    |              __/ |       C lambda ~         |
-- >    |          ___/    |                          |
-- >    |         /  neg   |                          |
-- >    +--------+---------+--------------------------+
-- >                       m
--
-- The length ("width" in /combinat-speak/, unfortunately) of the partitions
-- are less than mu; the "height" (first element) of @pos@ is at most @(n-i)@,
-- the height of @neg@ is unlimited (well, it is limited by @(mu-1)*(n-i)@ of course).
--
-- Actually, in the case of sigmaij, length(pos)<=i and length(neg)<=mu-i !
--
partitionPairs :: Int -> Int -> Int -> Int -> [(Partition,Partition)]
partitionPairs mu n i ofs = 
  [ (pos,neg) 
  | d <- [0..i*(n-i)] 
  , pos <- partitions' (n-i,i) d
  , let l = width pos
  , neg <- partitions' (d,mu-i) (d-ofs)
  ]

-- | Given the parameters @(m-n+i,mu) (pos,neg)@, this computes @lambda@
-- in the picture above. 
posnegPairToPartition :: (Int,Int) -> (Partition,Partition) -> Partition
posnegPairToPartition (h,w) (pos,neg) = toPartitionUnsafe xs where
  xs = zipWith (+) ys (replicate w h)
  ys = pos' ++ replicate (w - width pos - width neg) 0 ++ map negate (reverse neg')
  pos' = fromPartition pos
  neg' = fromPartition neg

--------------------------------------------------------------------------------