open Microsoft.SolverFoundation.Services
open Microsoft.SolverFoundation.Common
open System.Collections.Generic
open Microsoft.SolverFoundation.SfsWrapper

// helper functions for grouping elements of the solution
let flatten (matrix:'a[,]) = matrix |> Seq.cast<'a> |> Seq.toArray

let getColumn c (matrix:_[,]) =
    flatten matrix.[*,c..c]

let getRow r (matrix:_[,]) =
    flatten matrix.[r..r,*]

let getSquare s (vars: 'a[,]) =
    let x = s % 3
    let y = s / 3
    let startx, endx = x * 3, (x + 1) * 3 - 1
    let starty, endy = y * 3, (y + 1) * 3 - 1
    flatten vars.[startx .. endx, starty .. endy]


// prints out a solution
let printSolution (sol: int[,]) =
    for y in 0 .. 8 do
        if y % 3 = 0 then
            for _ in 1 .. 35 do printf "-"
            printfn ""
        getRow y sol |>
        Seq.iteri(fun x value ->
                    if x % 3 = 0 then
                        printf " | "
                    printf " %i " value)
        printfn ""

let solveSudoku (input: seq<int>) =
    // sanity check on inputs
    let check setType i set =
        let set = Array.filter ((<>) 0) set
        let set' = Set.ofSeq set
        if set.Length <> set'.Count then
            failwith (sprintf "Error in %s[%i] contains non-unique number" setType i)

    let input' = Array.ofSeq input
    let squareInput = Array2D.init 9 9 (fun x y -> input'.[ x + (9 * y)])
    for i in 0 .. 8 do
        getRow i squareInput |> check "row" i
        getColumn i squareInput |> check "col" i
        getSquare i squareInput |> check "square" i


    // get context and create the model
    let context = SolverContext.GetContext()
    let model = new SfsModel(context)

    // create the variables that will represent each place in the puzzel
    let vars = Array2D.init 9 9 (fun _ _ -> model.CreateIntRangeVariable(1, 9))

    // helper to add contraint to model that terms are all different
    let addVarsAllDiff varArray =
        let termArray = varArray |> Seq.cast<SfsIntTerm<1>> |> Array.ofSeq
        let contraint = model.AllDifferent termArray
        model.AddConstraint contraint |> ignore

    // go though each column, rom, and square and add all diff contraint
    // i.e. these are the basic rules of the game
    for i in 0 .. 8 do
        addVarsAllDiff (getColumn i vars)
        addVarsAllDiff (getRow i vars)
        addVarsAllDiff (getSquare i vars)

    // add the contraints from the inputs to the model
    let varsFlat = flatten vars
    input |>
    Seq.iteri(fun i x ->
        if x <> 0 then
            model.AddConstraint (varsFlat.[i] ==== x) |> ignore)

    // tell the model to solve the problem, using a "Constraint Programming" solver
    let sol = model.Solve ConstraintProgramming

    // print the solution quality and if apporiate the solution itself
    printfn "Solution: %A" sol.Quality
    if sol.Quality = SolverQuality.Feasible then
        context.PropagateDecisions()
        let res = vars |> Array2D.map (fun x -> x.Value)
        printSolution res

// generate a blank input
let inputEmpty = [for _ in 1 .. 81 do yield 0]

// an explicit blank input, useful for creating new inputs from
let inputBlank =
    [ 0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;

      0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;

      0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0; ]

// an example puzzle
let input =
    [ 0; 9; 4;  2; 8; 0;  0; 0; 0;
      0; 3; 0;  1; 0; 6;  0; 0; 0;
      6; 0; 0;  0; 0; 0;  2; 0; 0;

      0; 8; 2;  4; 9; 0;  7; 0; 0;
      4; 0; 0;  0; 0; 7;  0; 9; 0;
      0; 0; 0;  0; 0; 0;  0; 0; 0;

      0; 6; 0;  0; 0; 0;  0; 5; 2;
      0; 0; 0;  0; 0; 0;  6; 4; 0;
      7; 0; 0;  9; 0; 0;  0; 0; 1; ]

// solve the example puzzel
solveSudoku input