module Bergamot.Syntax exposing
    ( Term(..), Metavariable(..), UnificationVar(..)
    , instantiate, instantiateList, InstantiationState, emptyInstantiationState, resetVars
    , unify, unifyList, UnificationState, emptyUnificationState
    , reify
    )

import Set exposing (Set)
import Dict exposing (Dict)
import Maybe exposing (Maybe)
import Tuple
import Debug

type alias Name = String
type Metavariable = MkMetavariable String
type UnificationVar = MkUnificationVar String

unMetavariable : Metavariable -> String
unMetavariable (MkMetavariable s) = s

unUnificationVar : UnificationVar -> String
unUnificationVar (MkUnificationVar s) = s

type Term a
    = IntLit Int
    | StringLit String
    | Call Name (List (Term a))
    | Var a

type alias InstantiationState =
    { counter : Int
    , vars : Dict String UnificationVar
    }

emptyInstantiationState = { counter = 0, vars = Dict.empty }

resetVars : InstantiationState -> InstantiationState
resetVars is = { is | vars = Dict.empty }

metavariable : Metavariable -> InstantiationState -> (UnificationVar, InstantiationState)
metavariable mv is =
    case Dict.get (unMetavariable mv) is.vars of
        Just v -> (v, is)
        Nothing ->
            let
                v = MkUnificationVar ("var" ++ (String.fromInt is.counter))
                isp =
                    { counter = is.counter + 1
                    , vars = Dict.insert (unMetavariable mv) v is.vars
                    }
            in (v, isp)

instantiateList : List (Term Metavariable) -> InstantiationState -> (List (Term UnificationVar), InstantiationState)
instantiateList ml is =
    case ml of
        mt :: mts ->
            instantiate mt is
            |> \(t, isp) -> Tuple.mapFirst (\ts -> t :: ts) (instantiateList mts isp)
        [] -> ([], is)

instantiate : Term Metavariable -> InstantiationState -> (Term UnificationVar, InstantiationState)
instantiate mt is =
    case mt of
        IntLit i -> (IntLit i, is)
        StringLit s -> (StringLit s, is)
        Call n mts -> Tuple.mapFirst (Call n) (instantiateList mts is)
        Var mv -> Tuple.mapFirst Var (metavariable mv is)

type alias UnificationInfo =
    { equivalence : Set String
    , term : Maybe (Term UnificationVar)
    }

type alias UnificationState = Dict String UnificationInfo

emptyUnificationState = Dict.empty

reconcile : Set String -> Maybe (Term UnificationVar) -> UnificationState -> UnificationState
reconcile eq mt us =
    let newValue = { equivalence = eq, term = mt }
    in Set.foldl (\v -> Dict.insert v newValue) us eq

merge : UnificationVar -> UnificationVar -> UnificationState -> Maybe UnificationState
merge v1 v2 us =
    let
        v1s = unUnificationVar v1
        v2s = unUnificationVar v2
    in
        case (Dict.get v1s us, Dict.get v2s us) of
            (Just ui1, Just ui2) ->
                let
                    newEq = Set.union ui1.equivalence ui2.equivalence
                in
                    case (ui1.term, ui2.term) of
                        (Just t1, Just t2) ->
                            unify t1 t2 us
                            |> Maybe.map (\(t, usp) -> reconcile newEq (Just t) usp)
                        (Just t1, Nothing) ->
                            Just (reconcile newEq (Just t1) us)
                        (Nothing, Just t2) ->
                            Just (reconcile newEq (Just t2) us)
                        (Nothing, Nothing) ->
                            Just (reconcile newEq Nothing us)
            (Just ui1, Nothing) -> Just (reconcile (Set.insert v2s ui1.equivalence) ui1.term us)
            (Nothing, Just ui2) -> Just (reconcile (Set.insert v1s ui2.equivalence) ui2.term us)
            (Nothing, Nothing) -> Just (reconcile (Set.fromList [v1s,v2s]) Nothing us)

set : UnificationVar -> Term UnificationVar -> UnificationState -> Maybe (Term UnificationVar, UnificationState)
set v t us =
    let
        vs = unUnificationVar v
    in
        case Dict.get vs us of
            Just ui ->
                case ui.term of
                    Just tp ->
                        unify t tp us
                        |> Maybe.map (\(tpp, usp) -> (tpp, reconcile ui.equivalence (Just tpp) usp))
                    Nothing ->
                        Just (t, reconcile ui.equivalence (Just t) us)
            Nothing -> Just (t, Dict.insert vs { equivalence = Set.singleton vs, term = Just t } us)

unifyList : List (Term UnificationVar) -> List (Term UnificationVar) -> UnificationState -> Maybe (List (Term UnificationVar), UnificationState)
unifyList l1 l2 us =
    case (l1, l2) of
        (t1 :: ts1, t2 :: ts2) ->
            unify t1 t2 us
            |> Maybe.andThen (\(t, usp) -> Maybe.map (Tuple.mapFirst (\ts -> t :: ts)) (unifyList ts1 ts2 usp))
        ([], []) -> Just ([], us)
        _ -> Nothing

unify : Term UnificationVar -> Term UnificationVar -> UnificationState -> Maybe (Term UnificationVar, UnificationState)
unify t1 t2 us =
    case (t1, t2) of
        (IntLit i1, IntLit i2) -> if i1 == i2 then Just (t1, us) else Nothing
        (StringLit s1, StringLit s2) -> if s1 == s2 then Just (t1, us) else Nothing
        (Call n1 ts1, Call n2 ts2) ->
            if n1 == n2
            then Maybe.map (Tuple.mapFirst (Call n1)) (unifyList ts1 ts2 us)
            else Nothing
        (Var v1, Var v2) ->
            merge v1 v2 us
            |> Maybe.map (\usp -> (Var v1, usp))
        (Var v1, _) -> set v1 t2 us
        (_, Var v2) -> set v2 t1 us
        _ -> Nothing

reify : Term UnificationVar -> UnificationState -> Term UnificationVar
reify t us =
    case t of
        IntLit i -> IntLit i
        StringLit s -> StringLit s
        Call n ts -> Call n (List.map (\tp -> reify tp us) ts)
        Var v ->
            case Dict.get (unUnificationVar v) us of
                Just ui ->
                    case ui.term of
                        Just tp -> reify tp us
                        _ -> Var v
                Nothing -> Var v