{-# OPTIONS_GHC -O2 -fglasgow-exts #-}

module Codes where

import Prelude hiding (span)
import System.Random
import Data.List (intersect,genericTake,genericLength)
import Control.Monad 
import Data.Bits
import qualified Data.Set as Set

-- A useful function for the interpreter
-- eg   minDist . ntuples 10 10 $$ newStdGen
infixr 8 $$
($$) :: (Monad m) => (a1 -> r) -> m a1 -> m r
($$) = liftM

-- GENERATOR
{--
The library random generator is very robust and has an large period. 
Unfortunately that actually makes it difficult to work out when we have exhausted all possible codewords. 
Instead we use a simple linear generator.
--}
	
-- Linear congruential generator of period exactly 2^n
newIntervalGen n = do
	let mult = 1 + (if (mod (2^n) 4 == 0) then 4 else 2)
	start <- randomRIO (0,2^n-1) 
	constant <- head $$ filter ((/=0) . flip mod 2) $$ randomRs (0,2^n) $$ newStdGen
	return (start,mult,constant)
	 
-- Returns every codeword of length n in some random order
ntuples n =
	let	m = 2^n
		randoms (start,mult,constant) =
			let 	next = mod (mult * start + constant) m
			in	next : randoms (next,mult,constant)
	in	genericTake (2^n) . randoms $$ newIntervalGen n

-- SIMPLE CODES

weight :: Integer -> Integer
weight 0 = 0
weight n = (mod n 2) + (weight $ div n 2)

dist a b = weight $ xor a b

minDist [a,b] = dist a b
minDist (a:as) = min (minDist as) (minimum $ map (dist a) as)

-- Finds a subset of the given list such that the minimum distance is d
-- Every time an ntuple is accepted one copy is sent to the output stream 
-- and the other is kept in the memory list to be compared with new ntuples.
distFilter d (a:as) = 
	let 	filter' memory [] = []
		filter' memory (a:as) =
			if all (\b -> dist a b >= d) memory
			then a : filter' (a : memory) as
			else filter' memory as
	in	a : filter' [a] as
 
-- LINEAR CODES

span :: [Integer] -> [Integer]
span gs =
	let 	span' [g] = Set.fromList [g,0]
		span' (g:gs) = 
			let rest = span' gs
			in	-- Check if g independent of gs
				if 	Set.member g rest
				then	rest
				else	Set.union rest (Set.map (xor g) rest)
	in	Set.toList $ span' gs

minWeight = minimum . filter (/=0) . map weight

-- Finds an independent subset of as satisying f
linearFilter :: (Integer -> Set.Set Integer -> Bool) -> [Integer] -> [Integer]
linearFilter f as = 
	let 	filter' spanSoFar [] = []
		filter' spanSoFar (a:as) =
			let newSpan = Set.map (xor a) spanSoFar
			in	if (Set.notMember a spanSoFar) && f a newSpan
				then a : filter' (Set.union spanSoFar newSpan) as
				else filter' spanSoFar as
	in	filter' (Set.singleton 0) as

-- Find an independent subset of list of codes
indFilter = linearFilter (\ a newSpan -> True)

-- Find an independent subset whose span has minWeight d
linearDistFilter d = linearFilter (\ a newSpan -> (all (\b -> weight b >= d) $ Set.toList newSpan)) 
	
-- SAMPLING
	
-- Applies f to k different values of some random generator gen.
sample k gen = sequence $ replicate k gen

assertSize r as = if genericLength as == r then as else error "Ran out of codewords"

bestDist r n = maximum $$ sample 10 (minDist $$ genericTake r $$ ntuples n)
bestSize d n = maximum $$ sample 10 (length $$ distFilter d $$ ntuples n)
bestWeight r n = maximum $$ sample 10 ((minWeight . span . assertSize r . genericTake r . indFilter) $$ ntuples n)
bestRank d n = maximum $$ sample 10 (length $$ linearDistFilter d $$ ntuples n)

-- For pasting examples to spreadsheet
csv as = putStr $ unlines $ map (\ ((a,b),c) -> show a ++ "," ++ show b ++ "," ++ show c) as
