classdef MCTSNode2 < handle
% A vertex of the MCTS implemenation. Holds the nodes state and 
% current time as well as the trajectory from the parent node to this node.
    properties
%         simulation_distance = 20
%         dwell_time_mult = .5
%         Ndwell_actions = 1;
%         action_cnt_max = inf; 
%         C = sqrt(2);  % MCTS param
%         beta_reward_multiplier = 10;
%         init_score_multiplier = 1; %1.0
%         gamma = .5;

        simulation_distance = 20
        dwell_time_mult = .5
        Ndwell_actions = 1;
        action_cnt_max = inf; 
        C = sqrt(2);  % MCTS param
        beta_reward_multiplier = 7;%7;
        init_score_multiplier = 1; %1.0
        gamma = .5;

        x0
        t0
        
        children
        parent
        num_actions
        num_parents
        
        % state trajectory and constant control to get to this node from
        % parent
        ts 
        xs 
        u
       
        % from BuildLattice
        lattice
        
        % index of the lattice vertex
        lattice_vertex
        
        % lattice vertices that this vertex can connect to.
        % corresponds to children---i.e. children(i) is reached through
        % action(i)
        actions
        
        % For MCTS:
        % Number of simulations for a given action (respective to actions)
        Nsa
        % Total number of simulations
        Ns
        % Total reward
        Wsa
        % Vertex UCB Score
        Zsa
        
        NLTL_automatas
        LTL_automatas
        LTL_states
        proposition_handlers
        internal_proposition_states
        constraints
        
        scene
        
        reward_stats
        reward_handlers
        driving_rewards 
        goal_rewards
    end
    
    methods
%         bl = BuildLattice([0,0,0,5], [.5, 1], [40, 10], 4, 10 , 3, [10,0; 40,0])
%         root = MCTSNode(MCTSNode.empty(0), [0, 0], [x0; x0], 0, bl.state_knots_index(1,:), bl)
        function obj = MCTSNode2(parent, ts, xs, u, lattice_vertex, lattice, scenario, LTL_automatas, proposition_handlers, reward_handlers, reward_stats, LTL_states, internal_proposition_states)
            obj.parent = parent;
            
            obj.ts = ts;
            obj.xs = xs;
            obj.x0 = xs(end,:);
            obj.t0 = ts(end);
            obj.u = u;
            obj.lattice_vertex = lattice_vertex;
            obj.lattice = lattice;
            
            obj.scene = scenario;
            
            obj.actions = obj.getActions(obj.lattice_vertex(1), obj.lattice_vertex(2));
            obj.num_actions = size(obj.actions,1);
            obj.children = cell(1,obj.num_actions);
%             assert(obj.num_actions ~= 0);
            
            cs.xBounds = [-inf*ones(4,1), inf*ones(4,1)];
            cs.uBounds = [-inf*ones(2,1), inf*ones(2,1)];
            
            if nargin >= 8
                % do LTL stuff.
                assert(numel(LTL_automatas)==numel(proposition_handlers));
                obj.NLTL_automatas = numel(LTL_automatas);
                obj.LTL_automatas = LTL_automatas;
                obj.proposition_handlers = proposition_handlers;
                obj.LTL_states = LTL_states;
                obj.reward_handlers = reward_handlers;
                obj.reward_stats = reward_stats;
                if nargin >= 13
                    obj.internal_proposition_states = internal_proposition_states;
                end

                obj.constraints = ConstraintFunctions(obj.scene, cs, obj.LTL_automatas, proposition_handlers);
            else
                obj.constraints = ConstraintFunctions(obj.scene, cs);
            end
            
            
            obj.driving_rewards = DrivingRewards(obj.scene, obj.reward_stats);
            obj.goal_rewards = GoalRewards(obj.scene, obj.reward_stats);
            
            obj.initializeScoring() 
            
        end

        
        function initializeScoring(obj)
            obj.Nsa = ones(1,size(obj.actions,1));
            obj.Ns = obj.num_actions;
            obj.Wsa = zeros(1,size(obj.actions,1));
            obj.Zsa = inf * ones(1,size(obj.actions,1));
            
            index_X1 = obj.lattice_vertex(1);
            index_v1 = obj.lattice_vertex(2);

            x1_X = obj.lattice.X_grid(index_X1);
            x1_v = obj.lattice.v_grid(index_v1);
            
            for i = 1:obj.Ns
                index_X2 = obj.actions(i,1);
                index_v2 = obj.actions(i,2);
            
                x2_X = obj.lattice.X_grid(index_X2);
                x2_v = obj.lattice.v_grid(index_v2);

                if obj.nearEq(x1_X, x2_X)
                    assert(obj.nearEq(x1_v, x2_v))
                    % the first Ndwell_actions are dwell actions each with
                    % linearly more time dwelling.
%                         th = obj.dwell_time_mult*rand_index;
                    a = 0;
                else
                    [~, a] = obj.lattice.dtAndAccForStateTransfer([x1_X, x1_v], [x2_X, x2_v]);
                end

                
                x1 = [x1_X, 0, 0, x1_v];
                x2 = [x2_X, 0, 0, x2_v];
                obj.Wsa(i) = obj.Wsa(i) + obj.computeReward(obj.ts(end), x1, x2, a, obj.internal_proposition_states) * obj.init_score_multiplier;
                
            end
            obj.updateZsa();
%             disp(obj.Zsa)
%             
%             disp(obj.NLTL_automatas)
%             
%             for i = 1:obj.NLTL_automatas
%                 LTL_states_out{i} = obj.LTL_automatas{i}.cur_state_index;
%                 
%             end
%             
%             for i = 1:size(obj.actions,1)
%                 if isempty(obj.internal_proposition_states)
%                     reward = obj.children{i}.simulate(LTL_states_out);
%                 else
%                     internal_proposition_states_out = obj.constraints.internal_proposition_state_out;
%                     disp('b')
%                     internal_proposition_states_out
%                     reward = obj.children{i}.simulate(LTL_states_out, internal_proposition_states_out);
% 
%                 end
%                 obj.Wsa(i) = obj.Wsa(i) + reward;
%                 disp(obj.Wsa(i))
%             end
%             obj.updateZsa();
            
        end
        
        % Actions are a list of acceptable vertices given by
        % lattice_connections. obj.actions has the form
        % obj.actions = [to_vert1_X, to_vert2_v;
        %                  ...     ,     ...   ;
        %                to_vertn_X, to_vertn_v]
        function actions_out = getActions(obj, index_X, index_v)
            actions_out = [];
            % Can dwell at current node only if velocity at lattice_vertex
            % is 0 -- i.e. the index of the velocity component is 1.
            if index_v == 1 
                actions_out = [ones(obj.Ndwell_actions,1)*index_X,...
                    ones(obj.Ndwell_actions,1)*index_v];
            end
            for k = 1:size(obj.lattice.is_connections,3)
                for l = 1:size(obj.lattice.is_connections,4)
                    if obj.lattice.is_connections(index_X, index_v, k, l)
                        actions_out = [actions_out; index_X + l, k];
                    end
                end
            end
            
        end

        
        function [reward, final_node] = selection(obj)
            final_node = MCTSNode2.empty(0);
            if obj.stopCondition(obj.x0(1), obj.x0(2))
                reward = 0;
                return
            end
            


            index = obj.getGreatestChildIndex();
            if isempty(obj.children{index})
                % first visit. Create child 
                [reward, final_node] = obj.expansion(index);
            else
                % already visited---select from child
                [reward, final_node] = obj.children{index}.selection();
            end
            obj.Ns = obj.Ns + 1;
            obj.Nsa(index) = obj.Nsa(index)+1;

            obj.Wsa(index) = obj.Wsa(index)+reward;
            
%             for i = 1:obj.num_actions
%                 if obj.Nsa(i) ~= 0
%                     Bsa = obj.C*sqrt(log(obj.Ns)/obj.Nsa(index));
%                     obj.Zsa(index) = obj.Wsa(index)/obj.Nsa(index) + Bsa;
%                 end
%             end
%             for i = 1:obj.num_actions
%                 if obj.Nsa(i) ~= 0
%                     Bsa = obj.C*sqrt(log(obj.Ns)/obj.Nsa(i));
%                     obj.Zsa(i) = obj.Wsa(i)/obj.Nsa(i) + Bsa;
%                 end
%             end
            obj.updateZsa();
        end
        
        function updateZsa(obj)
            for i = 1:obj.num_actions
                if obj.Nsa(i) ~= 0
                    Bsa = obj.C*sqrt(log(obj.Ns)/obj.Nsa(i));
                    obj.Zsa(i) = obj.Wsa(i)/obj.Nsa(i) + Bsa;
                end
            end
        end
            
        function [reward, final_node] = expansion(obj, index)
            final_node = MCTSNode2.empty(0);
            
            
            
            x2_X = obj.lattice.X_grid(obj.actions(index,1));
            x2_v = obj.lattice.v_grid(obj.actions(index,2));

%             if obj.stopCondition(x2_X, x2_v)
%                 % XXXXXXXXXX  SCREWED WITH REWARD HERE!!
%                 disp('here!!')
%                 reward = 0;
%                 return
%             end

            
            if obj.nearEq(obj.x0(1), x2_X)
                assert(obj.nearEq(x2_v, obj.x0(4)))
                % the first Ndwell_actions are dwell actions each with
                % linearly more time dwelling.
                th = obj.dwell_time_mult*index;
                a = 0;
            else
                [th, a] = obj.lattice.dtAndAccForStateTransfer(obj.x0([1,4]), [x2_X, x2_v]);
            end
            u_child = a;

            
%             disp([th, a])
            [ts_child, xs_child, us_child] = obj.buildTrajectorySegment(obj.x0, u_child, th);
    
            if isempty(obj.internal_proposition_states)
                constraints_satisfied = obj.constraints.compute(ts_child, xs_child, us_child, obj.LTL_states);
            else
                constraints_satisfied = obj.constraints.compute(ts_child, xs_child, us_child, ...
                    obj.LTL_states, obj.internal_proposition_states);
            end
            
            if ~constraints_satisfied
                % This is an unacceptable action, never try again!
                obj.Wsa(index) = -inf;
                reward = 0;
                disp('UNACCEPTABLE!! Removing')
                return
            end
            
            % This action is acceptable. Add it to the tree.
            

            
            for i = 1:obj.NLTL_automatas
                LTL_states_out{i} = obj.LTL_automatas{i}.cur_state_index;
            end
            
            if isempty(obj.internal_proposition_states)
                obj.children{index} = MCTSNode2(obj, ts_child, xs_child, u_child, ...
                    obj.actions(index,:), obj.lattice, obj.scene, obj.LTL_automatas, ...
                    obj.proposition_handlers, obj.reward_handlers, obj.reward_stats, LTL_states_out);
                reward = obj.children{index}.simulate(LTL_states_out);
            else
                internal_proposition_states_out = obj.constraints.internal_proposition_state_out;
                obj.children{index} = MCTSNode2(obj, ts_child, xs_child, u_child, ...
                    obj.actions(index,:), obj.lattice, obj.scene, obj.LTL_automatas, ...
                    obj.proposition_handlers, obj.reward_handlers, obj.reward_stats, LTL_states_out, ...
                    internal_proposition_states_out);
                reward = obj.children{index}.simulate(LTL_states_out, internal_proposition_states_out);
            end
            
            if obj.stopCondition(x2_X, x2_v)
                % XXXXXXXXXX  SCREWED WITH REWARD HERE!!
                reward = 0;
                return
            end
            
            
% 
%             if obj.stopCondition(x2_X, x2_v)
%                 % Expanded to the end! Full solution found!
%                 final_node = obj.children{index};
%             end

            
            
        end
        
        function do_stop = stopCondition(obj, X, v)
            do_stop = obj.scene.distToGoal([X, 0]) < 0 && v == 0;
            
        end
        
        function [ts_seg, xs_seg, us_seg] = buildTrajectorySegment(~, x0_seg, u_seg, th_seg)
            ts_seg = 0:.05:th_seg;
            
            if ts_seg(end) ~= th_seg
                ts_seg = [ts_seg,th_seg];
            end
            
            xs_seg = [x0_seg(1) + x0_seg(4)*ts_seg + u_seg/2*ts_seg.^2;
                  0*ts_seg;
                  0*ts_seg;
                  x0_seg(4) + ts_seg*u_seg]';
            us_seg = u_seg*ones(length(ts_seg),1);
        end
        
        function reward = simulate(obj, LTL_states, internal_proposition_states)
            index_X = obj.lattice_vertex(1);
            index_v = obj.lattice_vertex(2);
            
            x1_X = obj.lattice.X_grid(index_X);
            x1_v = obj.lattice.v_grid(index_v);
            
            x1_X_start = x1_X;
            
            total_th = 0;
            failed_cnt = 0;
            action_cnt = 0;
            reward = 0;
            
%             while index_X < obj.lattice.grid_size(2) ...
%                     && obj.lattice.X_grid(index_X) - x1_X_start < obj.simulation_distance
%             internal_proposition_states = obj.internal_proposition_states;
            while action_cnt < obj.action_cnt_max ...
                    && ~obj.stopCondition(x1_X, x1_v) ...
                    && obj.lattice.X_grid(index_X) - x1_X_start < obj.simulation_distance
                
                candidate_actions = obj.getActions(index_X, index_v);

                action_size = size(candidate_actions,1);
                if action_size == 0
                    reward = 0;
                    return;
                end
               
                % Randomly choose action:
%                 rand_index = randi(action_size);
                
                candidate_rewards = zeros(1,size(candidate_actions,1));
                for i = 1:size(candidate_actions,1)
                    
                    cand_index_X = candidate_actions(i,1);
                    cand_index_v = candidate_actions(i,2);

                
                    % Check constraints and set reward
                    x2_X = obj.lattice.X_grid(cand_index_X);
                    x2_v = obj.lattice.v_grid(cand_index_v);
                
                    if obj.nearEq(x1_X, x2_X)
                        assert(obj.nearEq(x1_v, x2_v))
                        % the first Ndwell_actions are dwell actions each with
                        % linearly more time dwelling.
%                         th = obj.dwell_time_mult*rand_index;
                        a = 0;
                    else
                        [~, a] = obj.lattice.dtAndAccForStateTransfer([x1_X, x1_v], [x2_X, x2_v]);
                    end
                    
                    x1 = [x1_X, 0, 0, x1_v];
                    x2 = [x2_X, 0, 0, x2_v];
                
                    candidate_rewards(i) = obj.computeReward(obj.ts(end), x1, x2, a, internal_proposition_states);
                end
                
                sq_candidate_rewards = candidate_rewards.^2;
                
                n_candidate_rewards = (sq_candidate_rewards)/sum(sq_candidate_rewards);
                
                n_s_candidate_rewards = ((1-obj.gamma)*n_candidate_rewards + obj.gamma/length(candidate_rewards));
%                 'asdf'
%                 disp(sum(n_s_candidate_rewards));
                assert(abs(sum(n_s_candidate_rewards)-1) < .0001)
                n_s_c_candidate_rewards = cumsum(n_s_candidate_rewards);
                rand_num = rand;
                temp = find(n_s_c_candidate_rewards > rand_num);
                rand_index = temp(1);
%                 disp('blah')
%                 disp([x1_X, x1_v])
%                 disp(n_s_c_candidate_rewards)
% %                 disp(rand_num)
%                 disp(rand_index)
                
%                 reward_params.v_to = x2_v;
%                 reward_params.d_goal = obj.scene.distToGoal([x1_X, 0]);
%                 if reward_params.d_goal < 1.0
%                     reward_params.in_goal = true;
%                     disp('XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXin goal is true!!!!');
%                 end
%                 reward = reward + obj.rewarder.compute(x1_v, a, reward_params);
                
                
                try_index_X = candidate_actions(rand_index,1);
                try_index_v = candidate_actions(rand_index,2);
                
                
                % Check constraints and set reward
                x2_X = obj.lattice.X_grid(try_index_X);
                x2_v = obj.lattice.v_grid(try_index_v);
                
                
                if obj.nearEq(x1_X, x2_X)
                    assert(obj.nearEq(x1_v, x2_v))
                    % the first Ndwell_actions are dwell actions each with
                    % linearly more time dwelling.
                    th = obj.dwell_time_mult*rand_index;
%                     disp('here')
%                     disp(th)
                    a = 0;
                else
                    [th, a] = obj.lattice.dtAndAccForStateTransfer([x1_X, x1_v], [x2_X, x2_v]);
                end
                
                %%%% XXXXXXXX
                if nargin >= 2
                    [ts_child, xs_child, us_child] = obj.buildTrajectorySegment([x1_X, 0, 0, x1_v], a, th);
                    if nargin >= 3
                        constraints_satisfied = obj.constraints.compute(ts_child, xs_child, us_child, ...
                            LTL_states, internal_proposition_states);
                    else
                        constraints_satisfied = obj.constraints.compute(ts_child, xs_child, us_child, ...
                            LTL_states);
                    end

%                     disp(internal_proposition_states{1})
%                     disp(LTL_states{1})
%                     disp(th)
%                     disp([x1_X, x1_v,x2_X, x2_v])

                    if ~constraints_satisfied
                        % This is an unacceptable action!
                        if obj.nearEq(x1_v,0.0) && failed_cnt < 10
                            % at stop-->try again may just need to stop
                            % longer
                            failed_cnt = failed_cnt+1;
                            continue
                        else
                            % not at stop or failed too much-->prior actions are important 
%                             % 0 reward!!
                            disp('Simulation Failed!!')
                            
                            disp([x1_X_start, x1_X, x1_v,x2_X, x2_v])

                            % loses reward proportional to the distance the
                            % rollout could have advanced.
                            
                            lost_reward_mult = abs(x1_X - x1_X_start)/obj.simulation_distance; 
                            disp(lost_reward_mult)
                            reward = reward * obj.beta_reward_multiplier * lost_reward_mult;
                            return;
                        end
                    end

                    % This action is acceptable. continue simulation
                    internal_proposition_states = obj.constraints.internal_proposition_state_out;
                    for i = 1:obj.NLTL_automatas
                        LTL_states{i} = obj.LTL_automatas{i}.cur_state_index;
                    end
                    
                    
                end
                %%%% XXXXXXXX                
                
                % collect reward
                % driving
%                 clear reward_params
%                 reward_params.v_to = x2_v;
%                 
%                 % goal 
%                 reward_params.d_goal = obj.scene.distToGoal([x1_X, 0]);
%                 
%                 if reward_params.d_goal < 0.0
%                     reward_params.in_goal = true;
%                 end
%                 
%                 % stop sign
%                 d_forward_stop_sign = obj.scene.distToStopSign([x1_X, 0]);
%                 if d_forward_stop_sign ~= inf
%                     reward_params.d_stop_sign = d_forward_stop_sign;
%                 end
%                 % TODO: should be implemented in using predicates
% %                 if reward_params.d_forward_stop_sign < 0.0
% %                     reward_params.in_stop_sign = true;
% %                 end
%                 % TODO: Predicate HasStopped
%                 
% %                 advancement_multiplier = 0.1 + .9*abs(x2_X - obj.lattice.X_grid(end))/obj.lattice.X_grid(end);
% %                 reward = reward + advancement_multiplier*obj.rewarder.compute(x1_v, a, reward_params);
%                 reward = reward + obj.rewarder.compute(x1_v, a, reward_params);
                reward = reward + candidate_rewards(rand_index);

                action_cnt = action_cnt + 1;
                
                
                
                index_X = try_index_X;
                index_v = try_index_v;
               
                x1_X = x2_X;
                x1_v = x2_v;
                
                total_th = total_th + th;
            end
            
            disp('action_cnt')
            disp(action_cnt)
            
            if action_cnt <= 0
                reward = 0;
            else
                advancement_multiplier = 2;
                a_gamma = 1;%.3;
                X_end = obj.lattice.X_grid(end);
                d_end = X_end-x1_X;
                disp([reward, advancement_multiplier*(1-d_end/X_end)]);
                reward = reward*(a_gamma + (1-a_gamma)*advancement_multiplier*(1-d_end/X_end));

                % normalize by number of actions.
                reward = obj.beta_reward_multiplier*reward/action_cnt; %+ 500*1/total_th^2;
            end

            disp([reward, total_th, obj.lattice.X_grid(obj.lattice_vertex(1)), obj.lattice.v_grid(obj.lattice_vertex(2))])
        end
        
        function reward = computeReward(obj, t, x1, x2, a, internal_proposition_states)
%             x1_v, x2_X, x2_v, a)
            

            reward = 0;
            
            % driving rewards
            % should have a "lane symbol" for this maybe...
            reward = reward + obj.driving_rewards.compute(t, x1, a, x2);
            
            % goal rewards
            reward = reward + obj.goal_rewards.compute(t, x1, a, x2);
            
% YYYYY
            % symbol rewards
            for i = 1:obj.NLTL_automatas
                if isempty(internal_proposition_states)
                    reward = reward + obj.reward_handlers{i}.compute(t, x1, a, x2);
                else
                    reward = reward + obj.reward_handlers{i}.compute(t, x1, a, x2, ...
                        internal_proposition_states{i});
                end
            end
            
        end
        
        function index = getGreatestChildIndex(obj)
%             [~,index] = max(obj.Zsa);
            assert(obj.num_actions ~= 0);
            val = max(obj.Zsa);
            candidate_indexes = find(val==obj.Zsa);
            index = candidate_indexes(randi(length(candidate_indexes)));
        end
        function dispVertex(obj)
            fprintf(1,'index: [%d, %d], x0: [%3f, %3f]\n', obj.lattice_vertex(1), obj.lattice_vertex(2), obj.x0(1), obj.x0(4));
        end
        
        function dispTree(obj)
            for i=1:obj.num_actions
                if ~isempty(obj.children{i})
                    obj.children{i}.dispTree();
                end
            end
            obj.dispVertex();
        end
        function plotTree(obj, ax, offset)
            if nargin < 3
                offset = 0;
            end
            for i=1:obj.num_actions
                if ~isempty(obj.children{i})
                    obj.children{i}.plotTree(ax, offset);
                end
            end
            obj.plotEdge(ax, 'g', offset)
        end
        function plotEdge(obj, ax, color, offset)
            if nargin < 4
                offset = 0;
            end
            axes(ax);
            plot(obj.xs(:,1), obj.xs(:,4) + offset, color)
        end
        function plotBranch(obj, ax, offset)
            if nargin < 3
                offset = 0;
            end
            if ~isempty(obj.parent)
                obj.parent.plotBranch(ax, offset)
            end
            obj.plotEdge(ax, 'k', offset)
        end
        function plotMostVisited(obj, ax, offset)
            if nargin < 3
                offset = 0;
            end
            if obj.num_actions>0
                val = max(obj.Nsa);
                most_visited_child = find(val==obj.Nsa);
                if ~isempty(obj.children{most_visited_child(1)})
                    obj.children{most_visited_child(1)}.plotMostVisited(ax, offset)
                end
            end
            obj.plotEdge(ax, 'r', offset)
        end
        function plotHighestReward(obj, ax, offset)
            if nargin < 3
                offset = 0;
            end
            if obj.num_actions>0
                val = max(obj.Wsa);
                higest_reward = find(val==obj.Wsa);
                if ~isempty(obj.children{higest_reward(1)})
                    obj.children{higest_reward(1)}.plotHighestReward(ax, offset)
                end
            end
            obj.plotEdge(ax, 'b', offset)
        end
        function plotHighestAverageReward(obj, ax, offset)
            if nargin < 3
                offset = 0;
            end
            if obj.num_actions>0
                vals = obj.Wsa./obj.Nsa;
                vals(obj.Nsa < obj.Ns/100) = 0;
                val = max(vals);
                higest_average_reward = find(val==vals);
                if ~isempty(obj.children{higest_average_reward(1)})
                    obj.children{higest_average_reward(1)}.plotHighestAverageReward(ax, offset)
                end
            end
            obj.plotEdge(ax, 'b', offset)
        end
        function [new_ts, new_xs] = save_branch(obj)
            if ~isempty(obj.parent)
                [traj_t, traj_x] = obj.parent.save_branch();
            else
                traj_x = [];
                traj_t = [];
            end
            new_xs = [traj_x; obj.xs];
            prev_t = 0.0;
            if ~isempty(traj_t)
                prev_t = traj_t(end);
            end
            new_ts = [traj_t, obj.ts + prev_t];
            
        end
        function is_near_eq = nearEq(~, a, b)
            is_near_eq = abs(a-b) < .000001;
        end
      
    end
end