{-# OPTIONS_GHC -Wall #-}

module Main where

-- Libs
import System.Console.ParseArgs hiding (args) -- cabal install parseargs
import System.FilePath
import System.IO
import Data.Maybe (isJust, fromJust)
import Data.List (unzip4)
import Control.Monad (when)

-- Music stuff
import MIR.Run

-- Parallelism
import Control.Parallel.Strategies

--------------------------------------------------------------------------------
-- Command-line arguments
--------------------------------------------------------------------------------

data MyArgs = SourceInputString | SourceInputFile | TargetInputFile 
            | InputDir | OpMode | Print | MaxErrorRate | WriteGT
  deriving (Eq, Ord, Show)

myArgs :: [Arg MyArgs]
myArgs = [
          Arg { argIndex = MaxErrorRate,
                argAbbr  = Just 'e',
                argName  = Just "max-error",
                argData  = argDataOptional "float" ArgtypeFloat,
                argDesc  = "Ignore pieces with higher error rate for diff"
              },
          Arg { argIndex = SourceInputString,
                argAbbr  = Just 'c',
                argName  = Just "chords",
                argData  = argDataOptional "string" ArgtypeString,
                argDesc  = "Input Chord Sequence to parse"
              },
          Arg { argIndex = SourceInputFile,
                argAbbr  = Just '1',
                argName  = Just "sfile",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Input file (source for diff)"
              },
          Arg { argIndex = TargetInputFile,
                argAbbr  = Just '2',
                argName  = Just "tfile",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Input file (target for diff)"
              },
          Arg { argIndex = InputDir,
                argAbbr  = Just 'd',
                argName  = Just "dir",
                argData  = argDataOptional "directory" ArgtypeString,
                argDesc  = "Input directory (to process all files within)"
              },
          Arg { argIndex = OpMode,
                argAbbr  = Just 'm',
                argName  = Just "mode",
                argData  = argDataRequired "parse|gdiff|stdiff"
                             ArgtypeString,
                argDesc  = "One of \"parse\", \"diff\", or \"stdiff\""
              },
          Arg { argIndex = Print,
                argAbbr  = Just 'p',
                argName  = Just "print",
                argData  = Nothing,
                argDesc  = "Set this flag to generate a .png for an input file"
              },
          Arg { argIndex = WriteGT,
                argAbbr  = Just 'g',
                argName  = Just "writegt",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Write ground truth to file"
              }
         ]

--------------------------------------------------------------------------------
-- Main
--------------------------------------------------------------------------------


data MatchMode = STDiff | GDiff
  deriving (Eq, Ord, Show)

err1, err2, err3 :: String
err1 = "Use a source file, or a directory."
err2 = "Use a source file and a target file, or a directory."
err3 = "Use a source file and optionally a target file."

main :: IO ()
main = do args <- parseArgsIO ArgsComplete myArgs
          let mode      = getRequiredArg args OpMode
              prnt      = gotArg args Print
          case mode of
            "parse"  -> mainParse prnt
            "gdiff"  -> mainDiff  False GDiff    -- diffs cannot be printed
            "stdiff" -> mainDiff  False STDiff
            s        -> usageError args ("Unknown mode: " ++ s)

mainParse :: Bool -> IO ()
mainParse p = do args <- parseArgsIO ArgsComplete myArgs
                 -- cdir <- getCurrentDirectory
                 let cStr = getArgString args SourceInputString
                     mf1  = getArgString args SourceInputFile
                     mf2  = getArgString args TargetInputFile
                     mdir = getArgString args InputDir
                 case (cStr, mf1,mf2,mdir, p) of
                   -- parse a string of chords
                   (Just c, Nothing, Nothing, Nothing , False) 
                     -> parseTree c
                   -- Parse one file, show full output
                   (Nothing, Just f1, Nothing, Nothing , False) 
                     -> readFile' f1 >>= parseTree
                   -- Parse all files in one dir, show condensed output
                   (Nothing, Nothing, Nothing, Just dir, False) 
                     -> testDir dir
                   _ -> usageError args err1   

                   
                   
mainDiff :: Bool -> MatchMode -> IO ()
mainDiff p m =
     do args <- parseArgsIO ArgsComplete myArgs
        let mf1   = getArg args SourceInputFile
            mf2   = getArg args TargetInputFile
            mdir  = getArg args InputDir
            me    = getArg args MaxErrorRate
            gt    = getArg args WriteGT
        case (mf1,mf2,mdir,p) of
          -- Parse source and target file, show full output
          (Just f1, Just f2, Nothing, False) -> diff m f1 f2
          -- Diff all files in one dir, show condensed output
          (Nothing, Nothing, Just dir, False) -> testDirMatch m me gt dir 
          _                                   -> usageError args err2

diff :: MatchMode -> String -> String -> IO ()
diff m f1 f2 = do (toks1, ts1, _, _) <- readFile' f1 >>= return . string2PieceC
                  (toks2, ts2, _, _) <- readFile' f2 >>= return . string2PieceC
                  case m of
                    STDiff -> putStrLn (diffChords toks1 toks2)
                    GDiff  -> print (diffPiece (head ts1) (head ts2))
                    
                       
filterNonQuery :: [FilePath] -> [FilePath]                       
filterNonQuery = concat . map snd . filter ((>1) . length . snd) . getClassSizes

testDirMatch :: MatchMode -> Maybe Float -> Maybe FilePath -> FilePath -> IO ()
testDirMatch m me od fp = 
  do fs <- readDataDir fp
     let process f   = do s <- readFile' (fp </> f)
                          let (toks, ps, _, errs) = string2PieceC s
                          return (f, toks, head ps, errorRatio errs toks)
     (fs', toks, ps, errs) <- fmap unzip4 (mapM process fs)
     when (isJust od) $ writeGroundTruth fs' (fromJust od)
     let fsq = filterNonQuery fs'
     putStr "true\ntrue\n"
     mapM_ (putStr . (++ "\t")) fsq
     putChar '\n'
     mapM_ (putStr . (++ "\t"). getId) fs'
     putStrLn ""
     let runDiff df l = [ [ df me ex ey x y | (y,ey) <- l ]
                        | ((x,ex),i) <- zip l (map getId fs'), i `elem` fsq ]
         printLine l = putStrLn (foldr (\a b -> showFloat a ++ "\t" ++ b) "" l)
                         >> hFlush stdout
         undf = error "stdiff should not inspect parsing results"
     let r = case m of
               STDiff -> runDiff diffChordsLen (zip toks (repeat undf))
               GDiff  -> runDiff diffPieceLen  (zip ps   errs)
              `using` parList rdeepseq
     sequence_ [ printLine x | x <- r]