:- discontiguous(perform/4).
:- discontiguous(validate/4).
:- discontiguous(validate/3).
:- discontiguous(resolve_ref/3).
%%%%%%%%%%%%%
%%% Model %%%
%%%%%%%%%%%%%

%% Facts %%

block(b1).
block(b2).
block(b3).
block(b4).
block(b5).
block(b6).
block(b7).
block(b8).

size(b1, large).
size(b2, large).
size(b3, large).
size(b4, large).
size(b5, small).
size(b6, small).
size(b7, small).
size(b8, small).

color(b1, yellow).
color(b2, red).
color(b3, blue).
color(b4, green).
color(b5, yellow).
color(b6, red).
color(b7, blue).
color(b8, green).

initial_state([on(b1, table), on(b2, table), on(b3, table), on(b4, table), on(b5, table),
    on(b6, table), on(b7, table), on(b8, table)]).

% Insert final state here
final_state([on(b1, table), on(b2, b3), on(b3, b1), on(b4, b5), on(b5, b6),on(b6, table), 
            on(b7, b4), on(b8, b7)]).
    
%%%%%%%%%%%%%
%%% Rules %%%
%%%%%%%%%%%%%
    
clear(Block, State) :- \+ member(on(_, Block), State).
empty_hand(State) :- \+ member(hand(_), State).
holding(State, Block) :- member(hand(Block), State).
is_on(Block1, Block2, State) :- member(on(Block1, Block2), State).

validate(_, Block, State) :- 
    \+clear(Block, State), 
    write("I cannot do this because there is something on this block."), 
    nl, !, false.
validate(pickup, Block, State) :- 
    clear(Block, State), empty_hand(State).
validate(pickup, _, State) :-
    \+empty_hand(State),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
validate(pickup, table, _) :- 
    write("Oh no! Don't pick up the table!"), nl, !, false.
validate(pickup, Block1, Block2, State) :- clear(Block1, State), is_on(Block1, Block2, State),
    (empty_hand(State)).
validate(pickup, Block1, _, State) :- 
    \+clear(Block1, State), 
    write("I cannot do this because there is something on this block."), 
    nl, !, false.
validate(pickup, Block1, Block2, State) :-
    \+is_on(Block1, Block2, State),
    write("I cannot do this because this block is somewhere else."), 
    nl, !, false.
validate(pickup, _, _, State) :-
    \+empty_hand(State),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
validate(put, Block1, Block2, State) :- clear(Block1, State), clear(Block2, State),
    (empty_hand(State) ; holding(State, Block1)).
validate(put, Block1, Block2, State) :- 
    \+(clear(Block1, State); clear(Block2, State)),
    write("I cannot do this because one of the blocks has something on it."), 
    nl, !, false.
validate(put, Block, table, State) :- clear(Block, State), 
    (empty_hand(State) ; holding(State, Block)).
validate(put, Block, _, State) :- 
    \+(empty_hand(State) ; holding(State, Block)),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
validate(stack, Block1, Block2, State) :- clear(Block1, State), clear(Block2, State),
    (empty_hand(State) ; holding(State, Block1)).
validate(stack, Block1, Block2, State) :- 
    \+(clear(Block1, State); clear(Block2, State)),
    write("I cannot do this because one of the blocks has something on it."), 
    nl, !, false.
validate(stack, Block, _, State) :- 
    \+(empty_hand(State) ; holding(State, Block)),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
validate(unstack, Block, State) :- clear(Block, State), empty_hand(State).
validate(unstack, Block1, Block2, State) :- clear(Block1, State), is_on(Block1, Block2, State),
    (empty_hand(State)).
validate(unstack, _, State) :-
    \+empty_hand(State),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
validate(unstack, Block1, _, State) :- 
    \+clear(Block1, State), 
    write("I cannot do this because there is something on this block."),
    nl, !, false.
validate(unstack, Block1, Block2, State) :- 
    \+is_on(Block1, Block2, State), 
    write("I cannot do this because this block is somewhere else."), 
    nl, !, false.
validate(unstack, _, _, State) :- 
    \+empty_hand(State),
    write("I cannot do this because the hand is not empty."), 
    nl, !, false.
    
perform(pickup, Block, State, NewState) :- replace(on(Block,_), hand(Block), State, NewState).
perform(put, Block1, Target, State, NewState) :- perform(pickup, Block1, State, State1),
    replace(hand(Block1), on(Block1, Target), State1, NewState).
perform(stack, Block1, Target, State, NewState) :- perform(pickup, Block1, State, State1),
    replace(hand(Block1), on(Block1, Target), State1, NewState).
perform(unstack, Block, State, NewState) :- replace(on(Block,_), hand(Block), State, NewState).

%%%%%%%%%%%
%%% NLP %%%
%%%%%%%%%%%

%% Parsing %%

sentence(s(VP)) --> verb_phrase(VP), dot.
sentence(s(VP)) --> verb_phrase(VP).
noun_phrase(np(Det, Adjp, Noun)) --> det(Det), adj_phrase(Adjp), noun(Noun).
noun_phrase(np(Det, Noun)) --> det(Det), noun(Noun).
verb_phrase(vp(Verb, NP, PP)) --> verb(Verb), noun_phrase(NP), prep_phrase(PP).
verb_phrase(vp(Verb, NP)) --> verb(Verb), noun_phrase(NP).
prep_phrase(pp(Prep, NP)) --> prep(Prep), noun_phrase(NP).
adj_phrase(adjp(Adj)) --> adj(Adj).
adj_phrase(adjp(Adj, Adjp)) --> adj(Adj), adj_phrase(Adjp).
verb(v(put)) --> [put].
verb(v(pickup)) --> [pickup].
verb(v(stack)) --> [stack].
verb(v(unstack)) --> [unstack].
verb(v(V)) --> [V].
det(d(the)) --> [the].
adj(a(large)) --> [large].
adj(a(small)) --> [small].
adj(a(green)) --> [green].
adj(a(red)) --> [red].
adj(a(yellow)) --> [yellow].
adj(a(blue)) --> [blue].
adj(a(A)) --> [A].
noun(n(block)) --> [block].
noun(n(table)) --> [table].
noun(n(N)) --> [N].
prep(p(on)) --> [on].
prep(p(from)) --> [from].
dot --> ['.'].

%% Extracting Verbs and NPs %%

extract(s(vp(v(Verb), NounPhrase)), Verb, NounPhrase).
extract(s(vp(v(Verb), NounPhrase1, pp(_, NounPhrase2))), Verb, NounPhrase1, NounPhrase2).

%% Reference Resolution %%

%% 2-adjective
resolve_ref(NP, _, Block) :- NP = np(_, adjp(a(Size), adjp(a(Color))), n(block)), 
    size(Block, Size), color(Block, Color).
resolve_ref(NP, _, Block) :- NP = np(_, adjp(a(Color), adjp(a(Size))), n(block)), 
    size(Block, Size), color(Block, Color).
resolve_ref(NP, _, Block) :- NP = np(_, adjp(a(Adj1), adjp(a(_))), n(block)), 
    (   \+size(Block, Adj1), \+color(Block, Adj1)), 
    write("I cannot resolve "), 
    write(Adj1), write("."), nl, !, false.
resolve_ref(NP, _, Block) :- NP = np(_, adjp(a(_), adjp(a(Adj2))), n(block)), 
    (   \+size(Block, Adj2), \+color(Block, Adj2)), 
    write("I cannot resolve "), 
    write(Adj2), write("."), nl, !, false.
    
%% 1-adjective 
resolve_ref(NP, State, Block) :- NP = np(_, adjp(a(Adjective)), n(block)), 
    member(on(Block, _), State),
    (findall(Block, size(Block, Adjective), [Block]) ; 
     findall(Block, color(Block, Adjective), [Block])),
    clear(Block, State).
resolve_ref(NP, _, _) :- NP = np(_, adjp(a(Adjective)), n(block)), 
    write("I cannot resolve "), write(Adjective), 
    write("."), nl, !, false.

%% Special case for part 2
resolve_ref(NP1, NP2, State, Block1, Block2) :- 
    NP1 = np(_, adjp(a(Adjective1)), n(block)), 
    NP2 = np(_, adjp(a(Adjective2)), n(block)), 
    member(on(Block1, _), State),
    (findall(Block1, size(Block1, Adjective1), _) ; 
     findall(Block1, color(Block1, Adjective1), _)),
    clear(Block1, State),
    member(on(Block1, Block2), State),
    (findall(Block2, size(Block2, Adjective2), [Block2]) ; 
     findall(Block2, color(Block2, Adjective2), [Block2])).

%% 0-adjective  
resolve_ref(NP, State, Block) :- NP = np(d(the), n(block)), member(hand(Block), State).
resolve_ref(NP, _, table) :- NP = np(d(the), n(table)).

%% Natural Language Generation %%

nlg(Block) :- size(Block, S), color(Block, C), write("the "), write(S), write(" "), 
    write(C), write(" block"). 
nlg(table) :- write("the table"). 
nlg(on(Block, Target)) :- nlg(Block), write(" is on "), nlg(Target), nl.
nlg(hand(Block)) :- nlg(Block), write(" is in the hand "), nl.

nlg([StateItem]) :- nlg(StateItem).
nlg([StateItem|Rest]) :- nlg(StateItem), nlg(Rest).

%%%%%%%%%%%%%%%%%
%%% Main Loop %%%
%%%%%%%%%%%%%%%%%

perform_act(Parse, State, NewState) :- extract(Parse, Verb, NP), 
    resolve_ref(NP, State, Block), 
    validate(Verb, Block, State), 
    perform(Verb, Block, State, NewState).
perform_act(Parse, State, NewState) :- extract(Parse, Verb, NP1, NP2), 
    resolve_ref(NP1, NP2, State, Block1, Block2), 
    validate(Verb, Block1, Block2, State), 
    (perform(Verb, Block1, State, NewState) ; perform(Verb, Block1, Block2, State, NewState)).
perform_act(Parse, State, NewState) :- extract(Parse, Verb, NP1, NP2), 
    resolve_ref(NP1, State, Block1), 
    resolve_ref(NP2, State, Block2), 
    validate(Verb, Block1, Block2, State), 
    (perform(Verb, Block1, State, NewState) ; perform(Verb, Block1, Block2, State, NewState)).

process_input(State) :- read_word_list(Input), 
    (   Input == [exit] -> write("Bye!"), nl, ! ;
    process_input(State, Input)).
process_input(State) :- write("I do not understand this."), nl, talktome(State).
process_input(State, Input) :- 
    sentence(Parse, Input, []), 
    perform_act(Parse, State, NewState), 
    write("Currently..."), nl, 
    nlg(NewState), 
    talktome(NewState).
process_input(State, Input) :- 
    sentence(Parse, Input, []), 
    \+perform_act(Parse, State, _), 
    talktome(State).

talktome :- initial_state(State), write("Currently..."), nl, nlg(State), talktome(State).
talktome(State) :- write("What would you like to do next?"), nl, process_input(State).
        
%%%%%%%%%%%%%%%%
%%% Planning %%%
%%%%%%%%%%%%%%%%

%% My heurisric: form the stacks one by one
% start with blocks that are already at where they are supposed to be
% in this case probably the ones that are on the table to start with
% then for each of those blocks form its own stack
% keep looking for the block right on top until we reach the top-most block
% then we do the next stack
% I chose this heuristic since this can ensure that we do not go into
% a situation where we are trying to move a block
% that is under some other blocks
plan(CurrentState, CurrentState, [], []).
plan(CurrentState, FinalState, Steps, [BlockDone | Rest]) :- 
    form_one_stack(CurrentState, NextState, FinalState, Steps1, BlockDone),
    plan(NextState, FinalState, Steps2, Rest), 
    append(Steps2, Steps1, Steps).

% method for forming one stack
% keep looking for blocks on top until there isn't any
form_one_stack(CurrentState, CurrentState, FinalState, [], BlockDone) :- 
    \+find_on_top(FinalState, BlockDone, _).
form_one_stack(CurrentState, SemiFinalState, FinalState, Steps, BlockDone) :-
    find_on_top(FinalState, BlockDone, TopBlock), 
    validate(stack, TopBlock, BlockDone, CurrentState), 
    perform(stack, TopBlock, BlockDone, CurrentState, NextState),
    form_one_stack(NextState, NextNextState, FinalState, Steps1, TopBlock), 
    SemiFinalState = NextNextState,
    Steps = [putOn(TopBlock,BlockDone)|Steps1], !.

% solve method for convenience
solve(Steps) :- initial_state(InitialState), final_state(FinalState), 
    get_blocks_done(InitialState, FinalState, [], BlocksDone), 
    plan(InitialState, FinalState, Steps, BlocksDone).
    
% Get blocks that are already at their desired positions at the beginning
get_blocks_done([on(Block1, Block2)], [on(Block1, Block3)], BlocksDone, NewBlocksDone) :-
    (   Block2 == Block3 ->  NewBlocksDone = [Block1 | BlocksDone];
    (   NewBlocksDone = BlocksDone)).
get_blocks_done([on(Block1, Block2) | RestInitial], [on(Block1, Block3) | RestFinal], BlocksDone, NewBlocksDone) :-
    (   Block2 == Block3 ->  get_blocks_done(RestInitial, RestFinal, [Block1|BlocksDone], NewBlocksDone);
    (   get_blocks_done(RestInitial, RestFinal, BlocksDone, NewBlocksDone))).    

% find the block that is right on top
find_on_top(GoalState, BlockDone, TopBlock) :-
    member(on(TopBlock, BlockDone), GoalState).

% Utilities

replace(_, _, [], []).
replace(O, R, [O|T], [R|T2]) :- replace(O, R, T, T2).
replace(O, R, [H|T], [H|T2]) :- dif(H,O), replace(O, R, T, T2).

find_in(Pattern, [H|_], Result) :- Pattern = H, Result = H.
find_in(Pattern, [_|R], Result) :- find_in(Pattern, R, Result).
find_in(_, [], _) :- fail.

read_word_list(Ws) :-
    read_line_to_codes(user_input, Cs),
    atom_codes(A, Cs),
    tokenize_atom(A, Ws).