using Combinatorics, PyPlot

SUITSIZE = 13
MAX_HAND_SIZE = 13

ALLOW_DUPLICATES_IF_UNAVOIDABLE = true

memo_good = Dict()
memo_goodable = Dict()
memo_hand = Dict()
deck = [ repeat([:c],SUITSIZE); repeat([:d],SUITSIZE); repeat([:h],SUITSIZE); repeat([:s],SUITSIZE) ]
;

function cuthand( hand, cut, target )
    # cut is a triplet (a,b,c) that says that (a,b) gets swapped with (b,c)
    target[:] = hand[:]
    target[ cut[1]:cut[3]-1 ] = [ hand[cut[2]:cut[3]-1] ; hand[cut[1]:cut[2]-1] ]
end

function shortenhand!( hand )
    # shorten a hand by removing adjacent duplicates
    markedind = []
    for i = 2:length(hand)
        if hand[i-1] == hand[i]
            push!(markedind,i)
        end
    end
    deleteat!(hand,markedind)
end

function isgood( hand )
    # is a hand in a legal sorted order?
    # assume hand has already been shortened as much as possible
    
    # memoisation: don't do more work than you have to!
    if get(memo_good,hand,"new_item") != "new_item"
        return memo_good[hand]
    else
        # suit duplicated (bad!)
        if count( hand .== :c ) > 1 || count( hand .== :d ) > 1 || count( hand .== :h ) > 1 || count( hand .== :s) > 1
            memo_good[hand] = false
            return false
        end
        
        if ALLOW_DUPLICATES_IF_UNAVOIDABLE
            # if you have a black card (spade or club) in hand, you can't have a heart next to a diamond.
            if :s in hand || :c in hand
                for i = 2:length(hand)
                    if hand[i-1:i] in ([:h,:d],[:d,:h])
                        memo_good[hand] = false
                        return false
                    end
                end
            end
            # if you have a red card (heart or diamond) in hand, you can't have a spade next to a club.
            if :h in hand || :d in hand
                for i = 2:length(hand)
                    if hand[i-1:i] in ([:s,:c],[:c,:s])
                        memo_good[hand] = false
                        return false
                    end
                end
            end
        else # heart next to diamond or spade next to club is ALWAYS bad in this case
            for i = 2:length(hand)
                if hand[i-1:i] in ([:h,:d],[:d,:h],[:s,:c],[:c,:s])
                    memo_good[hand] = false
                    return false
                end
            end
        end 
        memo_good[hand] = true
        return true
    end
end

function isgoodable( hand )
    if get(memo_goodable, hand, "new_item") != "new_item"
        return memo_goodable[hand]
    else    
        N = length(hand)
        # if the hand is already good, then it's clearly goodable
        if isgood(hand)
            memo_goodable[hand] = true
            return true
        end
        # no hope for a hand that has 8 or more cards
        if length(hand) >= 8
            memo_goodable[hand] = false
            return false
        end
        for cut in combinations(1:N+1,3)
            target = copy(hand)
            cuthand( hand, cut, target )
            shortenhand!(target)
            if isgood(target)
                memo_goodable[hand] = true
                return true
            end
        end
        memo_goodable[hand] = false
        return false
    end
end
            
permcount = [ length(permutations(1:SUITSIZE,ct)) for ct in 0:SUITSIZE ]

function handweight(hand)
    # compute the weight we should associate with a given hand
    s,h,c,d = count( hand .== :s ),count( hand .== :h ),count( hand .== :c ),count( hand .== :d )
    cts = [s,h,c,d]
    sort!(cts)
    if get(memo_hand, cts, "new_item") != "new_item"
        return memo_hand[cts]
    else
        memo_hand[cts] = permcount[s+1] * permcount[h+1] * permcount[c+1] * permcount[d+1]
        return memo_hand[cts]
    end
end      

function evaluategame(handsize)
    # assemble all possible hands and shorten them
    # for each hand, retrieve all possible moves. if any are good, mark hand as good.
    num,den = BigInt(0),BigInt(0)
    for hand in multiset_permutations(deck, handsize)
        wt = handweight(hand)
        shortenhand!(hand)
        den += wt
        if isgoodable(hand)
            num += wt
        end
    end
    num,den
end
;

nums = []
dens = []
props = []
handsizes = 1:MAX_HAND_SIZE
for handsize = handsizes
    println(handsize)
    num,den = @time evaluategame(handsize)
    push!(nums,num)
    push!(dens,den)
    push!(props,convert(Float64,num/den))
end

figure(figsize=(13,5))
bar(handsizes,props, tick_label=1:MAX_HAND_SIZE)
grid()
xlabel("hand size")
ylabel("probability")
if ALLOW_DUPLICATES_IF_UNAVOIDABLE
    title("Probability that a hand can be sorted in one move (adjacent same-color different-suit permitted only if unavoidable)")
else
    title("Probability that a hand can be sorted in one move (adjacent same-color different-suit is never permitted)")
end
for (i,p) in enumerate(props)
    annotate("$(round(p,digits=5))",(handsizes[i],p+0.01),ha="center")
end
xlim([0.4,13.6])
tight_layout()
savefig("hand_sort_updated.png")

# print out the probabilities and exact reduced fraction
[ handsizes props nums.//dens ]

# print out numerator and denominator
[ handsizes nums dens ]