
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveGeneric #-}

-- 1. Don't parse Strings, parse Text -- using takes takes 
--    only 40% of the time of Strings.
-- 2. 

import Control.Monad
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Read as T
import Data.Text (Text)
import Test.QuickCheck
import System.Random
import System.IO
import qualified Data.List as L
import qualified Data.Map.Strict as M
import Data.Map.Strict (Map, (!) )

import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Maybe as Maybe
import qualified Text.Printf as Printf
import qualified Control.Monad as Monad
import qualified Data.Array as Array

import Control.Arrow
import Control.DeepSeq
import GHC.Generics -- (Generic, Generic1)

import qualified Data.Vector.Unboxed as V
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed.Mutable as VMut
import Data.Vector.Unboxed.Mutable (IOVector)
import qualified Data.Vector.Algorithms.Intro as VV
import qualified Data.Vector.Algorithms.Search as VV

import Data.Either as E

import qualified Data.Vector.Unboxed as VU
import System.IO.Unsafe


bisectLeft :: Int -> Vector Int -> Int
bisectLeft n xs = bisectLeft' 0 (V.length xs) n xs

bisectLeft' :: Int -> Int -> Int -> Vector Int -> Int
bisectLeft' lo hi n xs = if lo >= hi then lo else
                    if n <= (xs V.! mid) then bisectLeft' lo mid n xs
                      else bisectLeft' (mid + 1) hi n xs
                    where mid = (lo + hi) `div` 2

bisectRight :: Int -> Vector Int -> Int
bisectRight n xs = bisectRight' 0 (V.length xs) n xs

bisectRight' :: Int -> Int -> Int -> Vector Int -> Int
bisectRight' lo hi n xs = if lo >= hi then hi else
                    if n < (xs V.! mid) then bisectRight' lo mid n xs
                      else bisectRight' (mid + 1) hi n xs
                    where mid = (lo + hi) `div` 2

bisectLeftV :: Int -> Vector Int -> Int
bisectLeftV n xs = unsafePerformIO $ do 
  xs' <- V.unsafeThaw xs
  res <- VV.binarySearchL xs' n
  return res

bisectRightV :: Int -> Vector Int -> Int
bisectRightV n xs = unsafePerformIO $ do 
  xs' <- V.unsafeThaw xs
  res <- VV.binarySearchR xs' n
  return res





sameyX (Positive n) =
  forAll (vectorOf n (arbitrary :: Gen Int)) $ \as ->
    forAll arbitrary $ \k ->
      let vec = V.fromList $ L.sort as
      in
      bisectLeft k vec ==
        bisectLeftV k vec
      
bisect n xs =
    (***) fold fold $ M.split n xs
  where
    fold :: Map Int Int -> Int
    fold = M.foldl' (+) 0  

bisectLeftX n xs =
  fst $ bisect n xs

bisectRightX n xs =
  (bisectLeftX n xs) + (M.findWithDefault 0 n xs)




fhiyoProb
  :: [String]
     -> (Int, Vector Int, [Int], Vector Int)
fhiyoProb [nStr, asStr, bsStr, csStr] = 
  let n = (read :: String -> Int) nStr
      as = V.fromList $ 
              {-# SCC "sortAs" #-} List.sort $ 
                ({-# SCC "parseAs" #-} parseLine $  asStr)
      bs = ({-# SCC "parseAs" #-} (parseLine  bsStr))
      cs = V.fromList $ 
              {-# SCC "sortCs" #-} List.sort $ 
                ({-# SCC "parseCs" #-} parseLine $ csStr)
  in (n, as, bs, cs)


fhiyoProbT
  :: [Text]
     -> (Int, Vector Int, [Int], Vector Int)
fhiyoProbT [nStr, asStr, bsStr, csStr] = 
  let [n] = parseLineT nStr
      as = V.fromList $ 
              {-# SCC "sortAs" #-} List.sort $ 
                ({-# SCC "parseAs" #-} parseLineT $  asStr)
      bs = ({-# SCC "parseAs" #-} (parseLineT  bsStr))
      cs = V.fromList $ 
              {-# SCC "sortCs" #-} List.sort $ 
                ({-# SCC "parseCs" #-} parseLineT $ csStr)
  in (n, as, bs, cs)

sortedVec :: Int -> Text -> Vector Int 
sortedVec n txt = V.create $ do
  vec <- VMut.unsafeNew n
  mapM_ (uncurry (VMut.unsafeWrite vec)) $ zip [0..] (parseLineT txt)
  VV.sort vec
  return vec

sortedList :: t -> Text -> Vector Int
sortedList n txt = 
  V.fromList $ L.sort $ parseLineT txt


fhiyoProbTSrt
  :: [Text]
     -> (Int, Vector Int, [Int], Vector Int)
fhiyoProbTSrt [nStr, asStr, bsStr, csStr] = 
  let [n] = parseLineT nStr
      as =  sortedList n asStr
      bs = ({-# SCC "parseAs" #-} (parseLineT  bsStr))
      cs = sortedList n csStr
  in (n, as, bs, cs)


fhiyoSolve (!n, !as, !bs, !cs) = 
  --putStrLn $ "as: " ++ show as
  --putStrLn $ "bisect Left b as: " ++ show
  --  (map (\b -> (b, bisectLeft b as)) bs)
  --putStrLn $ "\ncs: " ++ show cs
  --putStrLn $ "bisect Right b cs: " ++ show
  --  (map (\b -> (b,bisectRight b cs)) bs)
  foldl (+) 0 $ map (\ b -> (bisectLeftV b as) * (n - bisectRightV b cs)) bs



fhiyoProb2 [nStr, asStr, bsStr, csStr] =
  let n = (read :: String -> Int) nStr
      as = freqs . map (read :: String -> Int) . words $  asStr
      bs = map (read :: String -> Int) . words $  bsStr
      cs = freqs . map (read :: String -> Int) . words $ csStr
  in  (n, as, bs, cs)

fhiyoSolve2 (n, as, bs, cs) = do
  --putStrLn $ "as: " ++ show as
  --putStrLn $ "bisect Left b as: " ++ show
  --  (map (\b -> (b, bisectLeftX b as)) bs)
  --putStrLn $ "\ncs: " ++ show cs
  --putStrLn $ "bisect Right b cs: " ++ show
  --  (map (\b -> (b,bisectRightX b cs)) bs)
  foldl (+) 0 $ map (\ b -> (bisectLeftX b as) * (n - bisectRightX b cs)) bs



--fhiyo2 :: [String] -> Int
--fhiyo2 [nStr, asStr, bsStr, csStr] = 
--  let n = (read :: String -> Int) nStr
--      as = freqs . map (read :: String -> Int) . words $  asStr
--      bs = map (read :: String -> Int) . words $  bsStr
--      cs = freqs . map (read :: String -> Int) . words $ csStr
--  --in foldl (+) 0 $ map (\ b -> (bisectLeftX b as) * (n - bisectRight b cs)) bs
--  --in (n, as, bs, cs)
--  in sum $ map (\ b -> (bisectLeftX b as) * (n - bisectRightX b cs)) bs
--
testF :: IO ()
testF = do
  ls <- getFileInputT "sample.txt"
  let problem = {-# SCC "F_fhiyoProb" #-} (fhiyoProbTSrt ls)
  print $ {-# SCC "F_fhiyoSolveProb" #-} (fhiyoSolve problem)

stdinF = do
  ls <- getStdInputT 
  let problem = (fhiyoProbTSrt ls)
  print $ (fhiyoSolve problem)


testF2 = 
  (fhiyoSolve2 . fhiyoProb2 . map T.unpack <$> getFileInputT "sample.txt") 
  >>= print


data Problem = Problem {
  n :: Int,
  as :: [Int],
  bs :: [Int],
  cs :: [Int]
  }
  deriving (Show, Generic)

instance NFData Problem

data ProblemX = ProblemX {
  nx :: Int,
  asx :: Map Int Int,
  bsx :: Map Int Int,
  csx :: Map Int Int
  }
  deriving Show

freqs :: Ord k => [k] -> Map k Int
freqs xs = M.fromListWith (+) $ zip xs (repeat 1)

problemToX 
  (Problem n as bs cs) =
  ProblemX n (freqs as) (freqs bs) (freqs cs)

sample1 = ([
    "2"
  , "1 5"
  , "2 4"
  , "3 6"
  ], 3) 

samples = 
  [first parse sample1,
  first parse sample2,
  first parse sample3]


sample2 = ([
    "3"
  , "1 1 1"
  , "2 2 2"
  , "3 3 3"
  ], 27)

sample3 = ([
    "6"
  , "3 14 159 2 6 53"
  , "58 9 79 323 84 6"
  , "2643 383 2 79 50 288"
  ], 87)


getStdInputT :: IO [Text]  
getStdInputT = do
  res@[n, as, bs, cs] <- replicateM 4 T.getLine
  return res

getFileInputT :: FilePath -> IO [Text]
getFileInputT filename = 
  withFile filename ReadMode $ \h -> do
    res@[n, as, bs, cs] <- replicateM 4 (T.hGetLine h)
    return res


getFileInput :: FilePath -> IO [String]
getFileInput filename = 
  withFile filename ReadMode $ \h -> do
    res@[n, as, bs, cs] <- replicateM 4 (hGetLine h)
    return res

forceEither :: Either a b -> b
forceEither (Right r) = r 

parseLine :: String -> [Int]
parseLine str =
  map (read :: String -> Int) $ words $ str

parseLineT :: Text -> [Int]
parseLineT txt =
  map (fst . forceEither . T.decimal) $ T.splitOn " " txt



parse :: [Text] -> Problem 
parse [n, as, bs, cs] =
  Problem (read $ T.unpack  n) (parseList as) (parseList bs) (parseList cs)
  where
    parseList :: Text -> [Int]
    parseList txt =
      let
        bits =  T.splitOn (T.pack " ") txt
      in 
        L.sort $ map (read . T.unpack) bits

parseX :: [Text] -> ProblemX 
parseX [n, as, bs, cs] =
  ProblemX (read $ T.unpack  n) (toMap as) (toMap bs) (toMap cs)
  where
    toMap :: Text -> Map Int Int
    toMap txt = freqs $ parseList txt

    parseList :: Text -> [Int]
    parseList txt =
      let
        bits =  T.splitOn (T.pack " ") txt
      in 
        L.sort $ map (read . T.unpack) bits



bruteSolutions problem = do
  -- we want a top, middle, bottom
  -- bot > middle > top
  c <- cs problem
  b <- bs problem
  guard (c > b)
  a <- as problem
  guard (b > a)
  return (a,b,c)   

countSolutions :: Problem -> Int
countSolutions problem =
  length $ bruteSolutions problem

brute2 problem = sum $ do
  c <- M.keys (csx problem)
  let smallerThanC = fst $ M.split c (bsx problem)
  b <- M.keys smallerThanC
  let smallerThanB = fst $ M.split b (asx problem)
  a <- M.keys smallerThanB
  let res =  csx problem ! c * bsx problem ! b * asx problem ! a 
  return res  

perhapsBetter !problem = {-# SCC "PERHAPSBETTER_len" #-} length $ {-# SCC "PERHAPSBETTER" #-} do
  c <- {-# SCC "PBcs" #-} (cs problem)
  b <- {-# SCC "PBbs" #-} (filter (<c) $ bs problem)
  a <- {-# SCC "PBbas" #-} (filter (<b) $ as problem)
  return ()

--analyticSolution (Problem _ as bs cs) = 
--  (length cs) * (length bs) * (length cs)
--
samey as bs cs =
  let problem = Problem 0 as bs cs
  in  countSolutions problem == perhapsBetter problem

--analyticSolution (Problem _ as bs cs) = 
--  (length cs) * (length bs) * (length cs)
--
samey2 as bs cs =
  let problem = Problem 0 as bs cs
  in  countSolutions problem == (brute2 $ problemToX problem)

samey3 (Positive n) =
  forAll (vectorOf n (arbitrary :: Gen Int)) $ \as ->
    forAll (vectorOf n (arbitrary :: Gen Int)) $ \bs ->
      forAll (vectorOf n (arbitrary :: Gen Int)) $ \cs ->
        let problem = Problem n as bs cs
            fProb = (n, V.fromList $ L.sort as, bs, V.fromList $ L.sort cs)
        in  perhapsBetter problem == fhiyoSolve fProb
 
-- $ \v ->
--    length v == n
-- `suchThat` (\v -> True)

testSamey =
  forAll (listOf1 arbitrary) $ \as ->
    forAll (listOf1 arbitrary) $ \bs ->
      forAll (listOf1 arbitrary) $ \cs ->
        samey2 as bs cs



writeRandomDataFile :: Int -> Int -> FilePath -> IO ()
writeRandomDataFile n maxVal filename = 
  do
    withFile filename WriteMode $ \h -> do
      writeLine h [n]
      randomLine n >>= writeLine h
      randomLine n >>= writeLine h
      randomLine n >>= writeLine h
  where
  randomLine :: Int -> IO [Int]
  randomLine len = replicateM len (randomRIO (1, maxVal))

  writeItem :: Handle -> Int -> IO ()
  writeItem h n =
    hPutStr h (show n ++ " ")

  writeLine :: Handle -> [Int] -> IO ()
  writeLine h ns = do
    mapM_ (\n -> writeItem h n) (init ns)
    hPutStrLn h (show $ last ns)


test = do
  !problem <- parse <$> getFileInputT "sample.txt"
  -- print $ countSolutions problem
  print $ perhapsBetter problem
  --print $ brute2 $ problemToX problem

mainOLD = do
  !problem <- force . parse <$> getFileInputT "sample.txt"
  --problem <- parse <$> getStdInputT
  print $! perhapsBetter problem

main = -- testF
  stdinF

--main = print 'a'


