// Multistep Life patterns via Z3
// Inspired by https://www.youtube.com/watch?v=g4lhrVPDUG0

#r "bin/Microsoft.Z3.dll"

open System
open System.IO
open System.Runtime.InteropServices
open System.Collections.Generic
open Microsoft.Z3



let ctx = new Context([|("model", "true")|] |> dict |> Dictionary)

// helpers
let True : BoolExpr = ctx.MkTrue()
let False : BoolExpr = ctx.MkFalse()
let Int : int -> uint32 -> BitVecNum = fun v w -> ctx.MkBV(v, w)
let IntVar : string -> uint32 -> BitVecExpr = fun var w -> ctx.MkBVConst(var, w) 
let FreshVar : uint32 -> BitVecExpr = fun w -> ctx.MkBVConst(Guid.NewGuid().ToString(), w) 
let Eq : Expr -> Expr -> BoolExpr = fun l r -> ctx.MkEq(l, r) 
let Ite : BoolExpr -> Expr -> Expr -> Expr = fun p t e -> ctx.MkITE(p, t, e)
let And : BoolExpr[] -> BoolExpr = fun bools -> ctx.MkAnd(bools)
let Or : BoolExpr[] -> BoolExpr = fun bools -> ctx.MkOr(bools)
let Not : BoolExpr -> BoolExpr = fun bool -> ctx.MkNot(bool)
let Add : BitVecExpr -> BitVecExpr -> BitVecExpr = fun l r -> ctx.MkBVAdd(l, r)

let width = 15
let height = 7

let createBoard : int -> Expr[][] = fun index ->
    [| for i in {0..height - 1} ->
            [| for j in {0..width - 1} -> IntVar (sprintf "Board_%d_%d_%d" index i j) 1u :> _ |] |]

let pattern : int[][] = 
    [|[|0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0|];
      [|0; 1; 0; 0; 0; 1; 0; 1; 1; 1; 0; 1; 1; 1; 0|];
      [|0; 1; 0; 0; 0; 1; 0; 1; 0; 0; 0; 1; 0; 0; 0|]
      [|0; 1; 0; 0; 0; 1; 0; 1; 1; 0; 0; 1; 1; 0; 0|]
      [|0; 1; 0; 0; 0; 1; 0; 1; 0; 0; 0; 1; 0; 0; 0|]
      [|0; 1; 1; 1; 0; 1; 0; 1; 0; 0; 0; 1; 1; 1; 0|]
      [|0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0; 0|]|]

let validValues : Expr[][] -> BoolExpr = fun board -> board |> Array.collect id |> Array.map (fun x -> Or [|Eq x (Int 0 1u); Eq x (Int 1 1u)|] ) |> And
let validPattern : Expr[][] -> int[][] -> BoolExpr = fun board pattern ->  
    And [| for i in {0..height - 1} ->
            And [| for j in {0..width - 1} -> Eq board.[i].[j] (Int pattern.[i].[j] 1u) |] |]

let getNeighborhoods : Expr[][] -> int -> int -> Expr list = fun board i j ->
    match i, j with
    | 0, 0 -> [board.[0].[1]; board.[1].[0]; board.[1].[1] ]
    | i, j when i = height - 1 && j = width - 1 -> 
        [board.[i].[j - 1]; board.[i - 1].[j]; board.[i - 1].[j - 1] ] 
    | 0, j when j = width - 1 -> 
        [board.[0].[j - 1]; board.[1].[j - 1]; board.[1].[j] ]
    | i, 0 when i = height - 1 -> 
        [board.[i - 1].[0]; board.[i - 1].[1]; board.[i].[1] ]
    | 0, j -> 
        [board.[0].[j - 1]; board.[0].[j + 1]; board.[1].[j + 1]; board.[1].[j - 1]; board.[1].[j] ]
    | i, 0 -> 
        [board.[i - 1].[0]; board.[i + 1].[0]; board.[i + 1].[1]; board.[i - 1].[1]; board.[i].[1] ]
    | i, j when i = height - 1 -> 
        [board.[i].[j - 1]; board.[i].[j + 1]; board.[i - 1].[j + 1]; board.[i - 1].[j - 1]; board.[i - 1].[j] ] 
    | i, j when j = width - 1 -> 
        [board.[i - 1].[j]; board.[i + 1].[j]; board.[i + 1].[j - 1]; board.[i - 1].[j - 1]; board.[i].[j - 1] ] 
    | i, j -> 
        [board.[i - 1].[j - 1]; board.[i - 1].[j]; board.[i - 1].[j + 1]; 
         board.[i + 1].[j - 1]; board.[i + 1].[j]; board.[i + 1].[j + 1];
         board.[i].[j - 1]; board.[i].[j + 1] ]

let rec count : Expr list -> uint32 -> BitVecExpr -> BoolExpr = fun exprs w c ->
    match exprs with
    | [] -> Eq c (Int 0 w)
    | expr :: exprs ->
        let c' = FreshVar w
        let ite = 
            Ite (Eq expr (Int 1 1u)) 
                (Eq c (Add c' (Int 1 w)))
                (Eq c c') :?> _
        And [|ite; count exprs w c'|]

let countNeighborhoods : Expr[][] -> int -> int -> BitVecExpr -> BoolExpr = fun board i j c ->        
    count (getNeighborhoods board i j) 4u c
    
    
let rules : Expr[][] -> Expr[][] -> BoolExpr = fun fromBoard toBoard ->
    And [| for i in {0..height - 1} ->
            And [| for j in {0..width - 1} do
                        let c = FreshVar 4u 
                        let b = countNeighborhoods fromBoard i j c
                        let ite = 
                            Ite (Eq fromBoard.[i].[j] (Int 1 1u))
                                (Ite (Or [|(Eq c (Int 0 4u)); (Eq c (Int 1 4u)) |])
                                     (Eq toBoard.[i].[j] (Int 0 1u))
                                     (Ite (Or [|(Eq c (Int 4 4u)); (Eq c (Int 5 4u)); 
                                                (Eq c (Int 6 4u)); (Eq c (Int 7 4u)); (Eq c (Int 8 4u)) |])
                                        (Eq toBoard.[i].[j] (Int 0 1u))
                                        (Ite (Or [|(Eq c (Int 2 4u)); (Eq c (Int 3 4u))|])
                                             (Eq toBoard.[i].[j] (Int 1 1u))
                                            False))) 
                                (Ite (Eq c (Int 3 4u)) 
                                    (Eq toBoard.[i].[j] (Int 1 1u))
                                    (Eq toBoard.[i].[j] (Int 0 1u))) :?> _
                        yield And [|b; ite|] |] |]

let rec steps : Expr[][] list -> BoolExpr = 
    fun boards ->
        match boards with
        | [_] -> True
        | board :: board' :: boards ->
            And [|rules board board'; steps  (board' :: boards)|]
        | _ -> failwith "oups"

let evalPrintBoard : Expr[][] -> Model -> unit = fun board model ->
    for i in {0..height - 1} do
        for j in {0..width - 1} do
            let value = string <| model.Evaluate(board.[i].[j])
            let c = getNeighborhoods board i j 
                    |> List.filter (fun x -> string <| model.Evaluate(x) = "1") 
                    |> List.length
            match value with
            | "1" -> 
                match c with
                | 0 | 1 -> printf " "
                | 4 | 5 | 6 | 7 | 8 -> printf " "
                | 2 | 3 -> printf "*"
                | _ -> failwith "oups"
            | "0" -> 
                match c with
                | 3 -> printf "*"
                | _ -> printf " "
            | _ -> failwith "oups"

        printfn ""
    
let initBoard = createBoard 0
let middleBoard = createBoard 1
let finalBoard = createBoard 2
let c = IntVar "c" 8u
let formula = And [|validValues initBoard; validValues middleBoard; validPattern finalBoard pattern; 
                    steps [initBoard; middleBoard; finalBoard];
                    count (initBoard |> Array.collect id |> Array.toList) 8u c|]
                    

let solver = ctx.MkSolver()
solver.Assert(formula)
let flag = solver.Check() = Status.SATISFIABLE

let model = solver.Model

string <| model.Evaluate(c)

evalPrintBoard initBoard model
evalPrintBoard middleBoard model