-- | Calculates Thom polynomial of Sigma^{ij} with localization 
-- and the substituion trick
--
-- Some example usages:
--
-- > sigma-ij -h                           # help
-- > sigma-ij -i3 -j2 -n7                  # compute @Tp(Sigma^{3,2}(7))@ in the default ring
-- > sigma-ij -i3 -j2 -n7  -rZp            # compute @Tp(Sigma^{3,2}(7))@ in the hard-coded prime field
-- > sigma-ij -i3 -j2 -n10 -rZp -b3 -B10   # compute the 3rd part of @Tp(Sigma^{3,2}(10))@ divided into 10 pieces
--
-- The task can be parallezied using the @-B@ and @-b@ options
--

{-# LANGUAGE ScopedTypeVariables, TypeFamilies, BangPatterns, PackageImports, PatternGuards #-}
module Main where

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

import Data.Char
import Data.List
import Data.Ratio
import Data.Monoid

import Control.Monad
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.MVar

import System.Environment
import System.Mem
import System.IO
import System.Exit

import "time" Data.Time.Clock.POSIX

import Math.Combinat.Numbers.Primes

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.ThomPoly.Shared
import Math.ThomPoly.SigmaI   as SigmaI
import Math.ThomPoly.SigmaIJ  as SigmaIJ
import Math.ThomPoly.Formulae as Formulae

import Options.Applicative
    
--------------------------------------------------------------------------------

data Config = Config
  { _problem   :: !AnyProblem  
  , _tgtDim    :: !(Maybe Int)
  , _ring      :: !Ring 
  , _outFile   :: !(Maybe FilePath)
  , _batch     :: !(Maybe Batch)
  , _printStat :: !Bool
  , _dry       :: !Bool
  , _timeout   :: !(Maybe Int)
  }
  deriving Show

--------------------------------------------------------------------------------
  
run :: Config -> IO ()
run config = do
  -- print config
  void $ mbTimeout (_timeout config) $ do

    let problem = _problem config
        batch   = maybe defaultBatch id (_batch config)
        ring    = selectRing (_ring config)

    when (_printStat config) $ do
      print $ case problem of
        PI  si  -> calcStats si 
        PIJ sij -> calcStats sij

    let fname = case _outFile config of
          Just fname -> fname
          Nothing    -> case problem of
            PI  si  -> fullFName batch si 
            PIJ sij -> fullFName batch sij        

    let answer = case problem of
          PI  si  -> solveAny ring batch si 
          PIJ sij -> solveAny ring batch sij

    let text   = pretty answer

    answer `seq` do
      unless (_dry config) $ writeFile fname text

--------------------------------------------------------------------------------
-- * configuration

data AnyProblem 
  = PI  !SigmaI
  | PIJ !SigmaIJ
  | PI1 !SigmaI1            -- ^ we have an explicit formula for @Sigma^{i,1}@
  deriving Show

-- | Coefficient ring 
data Ring
  = Integers
  | Rationals
  | HardCodedZp             -- ^ temporary
  | PrimeField !Integer
  | SpecPrime  !Int         -- ^ special primes just below @2^k@ for @k=7,15,31,63@
  deriving Show

selectRing :: Ring -> AnyRing
selectRing r = case r of
  Integers    -> ringZZ
  Rationals   -> ringQQ
  HardCodedZp -> ringZp

-- | Primes close to the bounds of (signed) machine words.
specPrimes :: [(Int,Integer)]
specPrimes = 
  [ ( 7   , 2^7   -  1 )
  , ( 15  , 2^15  - 19 )
  , ( 31  , 2^31  -  1 )
  , ( 63  , 2^63  - 25 )
--  , ( 127 , 2^127 -  1 )
--  , ( 255 , 2^255 - 19 )
  ]

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

maybeRead :: Read a => String -> Maybe a
maybeRead s = case reads s of 
  [(x,"")] -> Just x
  _        -> Nothing

parseRing :: String -> Either String Ring
parseRing str0 
  | str `elem` ["zz","integer" ,"integers"  ]   =  Right Integers
  | str `elem` ["qq","rational","rationals" ]   =  Right Rationals
  | str `elem` ["zp","primefield" ]             =  Right HardCodedZp
--  | take 2
  where
    str = map toLower str0 

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

class Validate a where
  isValid :: a -> Maybe String
  
instance Validate Batch where
  isValid (Batch a b) 
    | b < 1           =  Just "the number of batches B should be at least 1"
    | a < 1 || a > b  =  Just "batch index b should be between 1 and B"
    | otherwise       =  Nothing

instance Validate Ring where
  isValid r = case r of
    PrimeField p -> if isProbablyPrime p 
      then Nothing 
      else Just "order of the finite field should be a prime"
    SpecPrime  q -> case lookup q specPrimes of 
      Nothing     -> Just "unimplemented special prime field (BITS should be one of 7, 15, 31 or 63)"
      _           -> Nothing
    _ -> Nothing       

instance Validate AnyProblem where

  isValid problem = case problem of
    PI (SigmaI i n)
      | i < 1     -> Just "the index I should be at least 1"
      | n < 1     -> Just "the source dimension N should be at least 1"
      | otherwise -> Nothing
    PIJ (SigmaIJ i j n)
      | i < 1          -> Just "the index I should be at least 1"
      | j < 1 || j > i -> Just "the index J should be between 1 and I"
      | n < 1          -> Just "the source dimension N should be at least 1"
      | otherwise      -> Nothing
    PI1 (SigmaI1 i n)
      | i < 1     -> Just "the index I should be at least 1"
      | n < 1     -> Just "the source dimension N should be at least 1"
      | otherwise -> Nothing

--------------------------------------------------------------------------------
-- * option parsing

configOpt :: Parser Config
configOpt = Config 
  <$> problemOpt 
  <*> mOpt 
  <*> ringNameOpt
  <*> outOpt 
  <*> batchOpt 
  <*> statFlag 
  <*> dryFlag 
  <*> timeoutOpt

problemOpt :: Parser AnyProblem
problemOpt = f <$> iOpt <*> jOpt <*> nOpt where
  f i mbj n = case mbj of
    Nothing -> PI $ SigmaI i n 
    Just j  -> case j of 
      0 -> PI  $ SigmaI  i   n
      _ -> PIJ $ SigmaIJ i j n

batchOpt :: Parser (Maybe Batch)
batchOpt = f <$> whichBatchOpt <*> nbatchOpt where
  f a b 
    | a >= 1 && a <= b  =  if b > 1 then Just (Batch a b) else Nothing
    | otherwise         =  error "the batch index should be between 1 and B"

ringNameOpt :: Parser Ring
ringNameOpt = option (eitherReader parseRing)
  (  long    "ring"
  <> short   'r'
  <> metavar "R"
  <> value   Rationals
  <> help    "The coefficient ring (or field) R we compute in, for example a prime field" 
  <> completeWith [ "Integers" , "Rationals" , "ZZ" , "QQ"
                  , "PrimeField" , "Zp"
--                  , "Zp7bit" , "Zp15bit" , "Zp31bit" , "Zp63bit"
                  ]
  <> showDefault
  )

timeoutOpt :: Parser (Maybe Int)
timeoutOpt = option (Just <$> auto)
  (  long    "timeout"
  <> short   't'
  <> metavar "TIMEOUT"
  <> value   Nothing
  <> help    "Timeout (specified in minutes)" 
  )

primeOpt :: Parser Integer
primeOpt = option auto
  (  long    "prime"
  <> short   'p'
  <> metavar "P"
  <> value   1000000007
  <> help    "The order of the prime field" 
  )

bitsOpt :: Parser Int
bitsOpt = option auto
  (  long    "bits"
  <> short   'q'
  <> metavar "BITS"
  <> value   63
  <> help    "Number of bits in the order of a special prime fields" 
  )

statFlag :: Parser Bool
statFlag = switch
  (  long  "stats"
  <> short 's'
  <> help  "print \"statistics\" (codimension, algebraic multiplicity)" 
  )

dryFlag :: Parser Bool
dryFlag = switch
  (  long  "dry"
  <> help  "do not write the result into a file" 
  )

outOpt :: Parser (Maybe FilePath)
outOpt = option (Just <$> str)
  (  long    "output"
  <> short   'o'
  <> metavar "FILE"
  <> value   Nothing
  <> help    "Write output to FILE (use --dry to skip)" 
  )

nbatchOpt :: Parser Int
nbatchOpt = option auto
  (  long    "nbatches"
  <> short   'B'
  <> metavar "B"
  <> value   1
  <> help    "number of batches" 
  )

whichBatchOpt :: Parser Int
whichBatchOpt = option auto
  (  long "batch"
  <> short   'b'
  <> metavar "b"
  <> value   1
  <> help    "which batch to run (from 1 to B)" 
  )

iOpt :: Parser Int
iOpt = option auto
  (  short   'i'
  <> metavar "I"
  <> help    "first Thom-Boardman index (I)"
  <> noArgError (ErrorMsg "specifying I is mandatory")
  )

jOpt :: Parser (Maybe Int)
jOpt = option (Just <$> auto)
  (  short   'j'
  <> metavar "J"
  <> value   Nothing
  <> help    "second Thom-Boardman index (J)"
  )

nOpt :: Parser Int
nOpt = option auto
  (  short   'n'
  <> metavar "N"
  <> help    "source dimension (N)"
  )

mOpt :: Parser (Maybe Int)
mOpt = option (Just <$> auto)
  (  short   'm'
  <> metavar "N"
  <> value   Nothing
  <> help    "target dimension (M, optional)"
  <> hidden
  )

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

main :: IO ()
main = execParser opts >>= run where
  opts = info (helper <*> configOpt)
    (  fullDesc
    <> progDesc shortDesc
    <> header   longDesc
    )
  shortDesc = "Thom polynomials of second order Thom-Boardman singularities"
  longDesc  = "A program computing Thom polynomials of second order Thom-Boardman singularities"

--------------------------------------------------------------------------------
-- * timeout

-- | argument: number of minutes
mbTimeout :: Maybe Int -> IO a -> IO (Maybe a)
mbTimeout mb action = 
  case mb of 
    Nothing      -> Just <$> action
    Just minutes -> do
      mv <- newEmptyMVar 
      t0 <- getPOSIXTime
      threadid <- forkIO $ do
        y <- action 
        putMVar mv $! y
      wait mv t0 minutes threadid

  where      
    wait mv t0 minutes threadid = do
      let seconds = minutes * 60
      let loop = do
            threadDelay 1000000     -- wait 1 sec
            mb <- tryTakeMVar mv
            case mb of
              Just y  -> return $ Just y
              Nothing -> do
                t <- getPOSIXTime
                if t - t0 < fromIntegral seconds
                  then loop
                  else do
                    putStrLn $ "timeout after " ++ show minutes ++ " minutes"
                    killThread threadid
                    return Nothing
      loop
      
--------------------------------------------------------------------------------