{-# LANGUAGE ViewPatterns, BangPatterns, TupleSections #-}

import Data.List
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Char
import Data.Maybe
import Data.Ord
import Data.Tuple
import Data.Function

import Data.Monoid

import Debug.Trace

import Control.Monad

type Attack key = [Char] -> (key, String)
type Distribution = Char -> Rational

isLatin :: Char -> Bool
isLatin c = 'a' <= c && c <= 'z'

shift, unshift :: Integral n => Char -> n -> Char
shift c (fromIntegral -> n) = toEnum $ (fromEnum c - fromEnum 'a') + (n `mod` 26)
unshift c (fromIntegral -> n) = shift c (-n)

shift', unshift' :: Char -> Char -> Char
shift' c k = shift c $ fromEnum k - fromEnum 'a'
unshift' c k = shift c $ fromEnum 'a' - fromEnum k

-- shift :: Integral n => Char -> n -> Char
-- shift c@(fromIntegral . fromEnum -> c') (toInteger -> by)
--   = toEnum . fromIntegral $ a' + ((c' - a' + by) `mod` 26)
--   where
--     a = 'a'
--     a' = fromIntegral $ fromEnum a

main :: IO ()
main = do
  let
    attack :: Show key => [Char] -> ([Char] -> (key, String)) -> [Char] -> IO ()
    attack desc attack' cyphertext = do
      putStrLn $ desc ++ "..."
      let (k, plaintext) = attack' cyphertext
      putStrLn $ "Key: " ++ show k
      putStrLn plaintext

  attack "Substitution" substitution =<< (filter isLatin <$> readFile "substitution.txt")
  attack "Vigenere" vigenere =<< (filter isLatin <$> readFile "vigenere.txt")


vigenere :: Attack [Char]
vigenere cyphertext = (map key [0 .. keyLength - 1], zipWith unshift' [keyStream i | i <- [0..]] cyphertext)
  where
    statistic = multigraphs cyphertext
    multigraphs' = map (\(k, v) -> (k, diff v)) . take 3 . sortOn (Down . length . snd) . filter (\(k, _) -> length k >= 3) $ Map.toDescList statistic
    keyLengths = foldr count Map.empty . filter (not . null) . subsequences $ foldr (\(_, v) l -> v ++ l) [] multigraphs'
    count occs
      | gcd' == 1 = id
      | otherwise = Map.alter (Just . succ . fromMaybe 0) gcd'
      where
        gcd' = foldl gcd 0 occs
    keyLength = fst . head . sortOn (Down . snd) $ Map.assocs keyLengths
    keyStream = key . (`mod` keyLength)
    key n = fst . caesar $ do
      i <- [0..(length cyphertext `div` keyLength)]
      let i' = i * keyLength + n
      guard $ i' < length cyphertext
      return $ cyphertext !! i'

caesar :: Attack Char
caesar cyphertext = (closestKey, map (unshift' closestKey) cyphertext)
  where
    possibilities :: [(Char, Rational)]
    possibilities = [(k, char $ (-) <$> english <*> findDist (map (unshift' k) cyphertext)) | k <- ['a' .. 'z']]
    -- cDist = findDist cyphertext
    -- possibilities = [(k, sum $ map (\c -> square $ english c - cDist c) mostCommon) | k <- ['a' .. 'z']]
    -- mostCommon = take 3 . map fst $ sortOn (Down . snd) [(c, english c) | c <- ['a' .. 'z']]
    closestKey = undefined -- head . map fst $ sortOn snd possibilities

substitution :: Attack [(Char, Char)]
substitution cyphertext = (key, map (\c -> fromMaybe ' ' . lookup c $ map swap key) cyphertext)
  where
    cDist = invert $ findDist cyphertext
    english' = invert english
    keyStream :: Int -> Char
    keyStream i = (!! i) . map snd $ sortOn fst assoc
      where
        assoc = map (\((_, c1), (_, c2)) -> (c1, c2)) $ (zip `on` sortOn fst) english' cDist
    key = [(toEnum $ fromEnum 'a' + i, keyStream i) | i <- [0..25]]

multigraphs :: Ord a => [a] -> Map [a] [Int]
multigraphs = multigraphs' Map.empty 0
  where
    multigraphs' acc _ [] = acc
    multigraphs' acc i xs = let
      inits' = take 3 . filter (not . null) $ inits xs
      acc' = foldr (Map.alter (Just . (i :) . fromMaybe [])) acc inits'
      in multigraphs' acc' (succ i) $ tail xs

diff :: Num a => [a] -> [a]
diff [] = []
diff xs@(x:_) = tail $ diff' x xs
  where
    diff' _ [] = []
    diff' r (x:xs) = (r - x) : diff' x xs

english :: Distribution
english = fromMaybe 0 . flip lookup dist
  where
    dist = [('a', 8.2 / 100), ('b', 1.5 / 100),('c', 2.8 / 100), ('d', 4.2 / 100), ('e', 12.7 / 100), ('f', 2.2 / 100), ('g', 2 / 100), ('h', 6.1 / 100), ('i', 7 / 100), ('j', 0.1 / 100), ('k', 0.8 / 100), ('l', 4 / 100), ('m', 2.4 / 100), ('n', 6.7 / 100), ('o', 7.5 / 100), ('p', 1.9 / 100),('q', 0.1 / 100), ('r', 6 / 100), ('s', 6.3 / 100), ('t', 9 / 100), ('u', 2.8 / 100), ('v', 1 / 100), ('w', 2.4 / 100), ('x', 2 / 100), ('y', 0.1 / 100), ('z', 0.2 / 100)]

char :: Distribution -> Rational
char f = sum $ map (square . f) ['a'..'z']

square :: Num n => n -> n
square = (*) <$> id <*> id

findDist :: [Char] -> Distribution
findDist text = (/ fromIntegral (length text)) . fromMaybe 0 . flip Map.lookup counts
  where
    counts = foldr (Map.alter $ Just . succ . fromMaybe 0) Map.empty text

invert :: Distribution -> [(Rational, Char)]
invert f = [(f c, c) | c <- ['a' .. 'z']]