package tileworld.agent;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.TreeSet;
import javax.swing.text.html.HTMLDocument;
import sim.app.woims.Vector2D;
import sim.engine.Schedule;
import sim.field.grid.ObjectGrid2D;
import sim.util.Bag;
import sim.util.Int2D;
import sim.util.IntBag;
import tileworld.environment.NeighbourSpiral;
import tileworld.Parameters;
import tileworld.environment.*;

/**
 * TWAgentMemory
 * 
 * @author michaellees
 * 
 *         Created: Apr 15, 2010 Copyright michaellees 2010
 * 
 *         Description:
 * 
 *         This class represents the memory of the TileWorld agents. It stores
 *         all objects which is has observed for a given period of time. You may
 *         want to develop an entirely new memory system or extend/modify this one.
 * 
 *         The memory is supposed to have a probabilistic decay, whereby an element is
 *         removed from memory with a probability proportional to the length of
 *         time the element has been in memory. The maximum length of time which
 *         the agent can remember is specified as MAX_TIME. Any memories beyond
 *         this are automatically removed.
 */
public class TWAgentWorkingMemory {

    /**
     * Access to Schedule (TWEnvironment) so that we can retrieve the current timestep of the simulation.
     */
    private Schedule schedule;
    private SimpleTWAgent me;
    private final static int MAX_TIME = 120;
    private final static float MEM_DECAY = 0.5f;
    
    protected double lifetimeEstimate=200;
    
     private ObjectGrid2D memoryGrid;
    
    /*
     * This was originally a queue ordered by the time at which the fact was observed.
     * However, when updating the memory a queue is very slow.
     * Here we trade off memory (in that we maintian a complete image of the map)
     * for speed of update. Updating the memory is a lot more straightfoward.
     */
    private TWAgentPercept[][] objects;
    /**
     * Number of items recorded in memory, currently doesn't decrease as memory
     * is not degraded - nothing is every removed!
     */
    private int memorySize;

    /**
     * Stores (for each TWObject type) the closest object within sensor range,
     * null if no objects are in sensor range
     */
    private HashMap<Class<?>, TWEntity> closestInSensorRange;
    
    //three levels of spirals
    static private List<Int2D> spiral = new NeighbourSpiral(Parameters.defaultSensorRange * 10).spiral();
    static private List<Int2D> spiral2 = new NeighbourSpiral(Parameters.defaultSensorRange * 15).spiral();
    static private List<Int2D> spiral3 = new NeighbourSpiral(4).spiral();
    
    private TWSector[][] sectors;
    //TreeSet<TWSector> orderedSectors;
    
//    private List<TWAgent> neighbouringAgents = new ArrayList<TWAgent>();

    public TWAgentWorkingMemory(TWAgent moi, Schedule schedule, int x, int y) {

        closestInSensorRange = new HashMap<Class<?>, TWEntity>(4);
        this.me = (SimpleTWAgent)moi;

        this.objects = new TWAgentPercept[x][y];

        this.schedule = schedule;
        this.memoryGrid = new ObjectGrid2D(me.getEnvironment().getxDimension(), me.getEnvironment().getyDimension());
        
        TWSector a=null;
        int xdim=(int)Math.ceil(me.getEnvironment().getxDimension()/Parameters.sectorXSize);
        int ydim=(int)Math.ceil(me.getEnvironment().getyDimension()/Parameters.sectorYSize);
        sectors=new TWSector[ydim][xdim];
      
        for (int i = 0; i < xdim; i++) {
            for (int j = 0; j < ydim; j++) {
               a=new TWSector(i*Parameters.sectorXSize+Math.round(Parameters.sectorXSize/2), j*Parameters.sectorYSize+Math.round(Parameters.sectorYSize/2));
                
               sectors[j][i]=a;
            }    
        }
    }


    /**
     * Called at each time step, updates the memory map of the agent.
     * Note that some objects may disappear or be moved, in which case part of
     * sensed may contain null objects
     *
     * Also note that currently the agent has no sense of moving objects, so
     * an agent may remember the same object at two locations simultaneously.
     * 
     * Other agents in the grid are sensed and passed to this function. But it
     * is currently not used for anything. Do remember that an agent sense itself
     * too.
     *
     * @param sensedObjects bag containing the sensed objects
     * @param objectXCoords bag containing x coordinates of objects
     * @param objectYCoords bag containing y coordinates of object
     * @param sensedAgents bag containing the sensed agents
     * @param agentXCoords bag containing x coordinates of agents
     * @param agentYCoords bag containing y coordinates of agents
    */
    public void updateMemory(Bag sensedObjects, IntBag objectXCoords, IntBag objectYCoords, Bag sensedAgents, IntBag agentXCoords, IntBag agentYCoords) {
        //reset the closest objects for new iteration of the loop (this is short
        //term observation memory if you like) It only lasts one timestep
        closestInSensorRange = new HashMap<Class<?>, TWEntity>(4);
        
        //must all be same size.
        assert (sensedObjects.size() == objectXCoords.size() && sensedObjects.size() == objectYCoords.size());

//        me.getEnvironment().getMemoryGrid().clear();  // THis is equivalent to only having sensed area in memory
       this.decayMemory();       // You might want to think about when to call the decay function as well.
        int sectorX,sectorY;
        for (int i = 0; i < sensedObjects.size(); i++) {
            TWEntity o = (TWEntity) sensedObjects.get(i);
            if (!(o instanceof TWObject)) {
                //the object is not there anymore
                //here we should implement the decay!
                
                if(lifetimeEstimate==200 & (memoryGrid.get(objectXCoords.get(i), objectYCoords.get(i)) instanceof TWObstacle)) {
                    lifetimeEstimate=schedule.getTime();
                   /* System.out.println(memoryGrid.get(objectXCoords.get(i), objectYCoords.get(i)));
                    System.out.println("lifetimeEstimate is "+lifetimeEstimate);*/
                }
                if((memoryGrid.get(objectXCoords.get(i), objectYCoords.get(i)) instanceof TWTile) || (memoryGrid.get(objectXCoords.get(i), objectYCoords.get(i)) instanceof TWHole)){
                    this.me.sendFloor(objectXCoords.get(i), objectYCoords.get(i));
                }
                
                 
                memoryGrid.set(objectXCoords.get(i), objectYCoords.get(i), null);
                objects[objectXCoords.get(i)][objectYCoords.get(i)] =null;
                continue;
            }
            sectorX=(int)Math.floor(objectXCoords.get(i)/Parameters.sectorXSize);
            sectorY=(int)Math.floor(objectYCoords.get(i)/Parameters.sectorYSize);
            //if nothing in memory currently, then were increasing the number 
            //of items we have in memory by 1
            if(objects[objectXCoords.get(i)][objectYCoords.get(i)] == null) {
                memorySize++;
                //sectors[sectorX][sectorY].decreaseUnknown();
            }
           // else {sectors[sectorY][sectorX].decr(objects[objectXCoords.get(i)][objectYCoords.get(i)].getClass());}
            

            //Add the object to memory
            objects[objectXCoords.get(i)][objectYCoords.get(i)] = new TWAgentPercept(o, this.getSimulationTime());
            //sectors[sectorX][sectorY].incr(o.getClass());
            
            memoryGrid.set(objectXCoords.get(i), objectYCoords.get(i), o);

            updateClosest(o);
            
        }
        
        
        
//       Agents are currently not added to working memory. Depending on how 
//       communication is modelled you might want to do this.
//        neighbouringAgents.clear();
//        for (int i = 0; i < sensedAgents.size(); i++) {
//            
//            
//            if (!(sensedAgents.get(i) instanceof TWAgent)) {
//                assert false;
//            }
//            TWAgent a = (TWAgent) sensedAgents.get(i);
//            if(a.equals(me)){
//                continue;
//            }
//            neighbouringAgents.add(a);
//        }
    }

//    public TWAgent getNeighbour(){
//        if(neighbouringAgents.isEmpty()){
//            return null;
//        }else{
//            return neighbouringAgents.get(0);
//        }
//    }
    
    
    public void addPercept(TWAgentPercept p){
            int x =p.getO().getX();
            int y =p.getO().getY();
            
            if( (p.getO() instanceof TWFloor) ){
                objects[x][y]=null;
                memoryGrid.set(x, y,null);
                return;
                
            }
            
            if(!p.sameObject(objects[x][y])) {
                objects[x][y]=p;
                memoryGrid.set(x, y, p.getO());
            }
        }
    
    /**
     * updates memory using 2d array of sensor range - currently not used
     * @see TWAgentWorkingMemory#updateMemory(sim.util.Bag, sim.util.IntBag, sim.util.IntBag)
     */
    public void updateMemory(TWEntity[][] sensed, int xOffset, int yOffset) {
        for (int x = 0; x < sensed.length; x++) {
            for (int y = 0; y < sensed[x].length; y++) {
                objects[x + xOffset][y + yOffset] = new TWAgentPercept(sensed[x][y], this.getSimulationTime());
            }
        }
    }

    /**
     * removes all facts earlier than now - max memory time. 
     * TODO: Other facts are
     * remove probabilistically (exponential decay of memory)
     */
    public void decayMemory() {
        // put some decay on other memory pieces (this will require complete
        // iteration over memory though, so expensive.
        //This is a simple example of how to do this.
        int sectorX,sectorY;
           for (int x = 0; x < this.objects.length; x++) {
            for (int y = 0; y < this.objects[x].length; y++) {
                TWAgentPercept currentMemory =  objects[x][y];
               if(currentMemory!=null && currentMemory.getT() < schedule.getTime()-lifetimeEstimate){
                   
                   memoryGrid.set(x, y, null);
                   objects[x][y]=null;
                   
                   
                    /*sectorX=(int)Math.floor(x/Parameters.sectorXSize);
                    sectorY=(int)Math.floor(y/Parameters.sectorYSize);
                    sectors[sectorY][sectorX].increaseUnknown();*/
                }
            }
        }
    }


    public void removeAgentPercept(int x, int y){
        objects[x][y] = null;
    }


    public void removeObject(TWEntity o){
        removeAgentPercept(o.getX(), o.getY());
    }


    /**
     * @return
     */
    private double getSimulationTime() {
        return schedule.getTime();
    }

    /**
     * Finds a nearby tile we have seen less than threshold timesteps ago
     *
     * @see TWAgentWorkingMemory#getNearbyObject(int, int, double, java.lang.Class)
     */
    public TWTile getNearbyTile(int x, int y, double threshold) {
        return (TWTile) this.getNearbyObject(x, y, threshold, TWTile.class);
    }


    /**
     * Finds a nearby hole we have seen less than threshold timesteps ago
     *
     * @see TWAgentWorkingMemory#getNearbyObject(int, int, double, java.lang.Class)
     */
    public TWHole getNearbyHole(int x, int y, double threshold) {
        return (TWHole) this.getNearbyObject(x, y, threshold, TWHole.class);
    }
    
    /**
     * Finds a nearby N tiles 
     *
     * @see TWAgentWorkingMemory#getNearbyNObjects(int, int, java.lang.Class, int)
     */
    public TreeSet<TWObject> getNearbyNTiles(int x, int y) {
        
        return this.getNearbyNObjects(x, y, TWTile.class,Parameters.numberOfSearchedObjects);
    }


    /**
     * Finds a nearby N holes
     *
     * @see TWAgentWorkingMemory#getNearbyNObjects(int, int, java.lang.Class, int)
     */
    public TreeSet<TWObject> getNearbyNHoles(int x, int y, double threshold) {
        
        return this.getNearbyNObjects(x, y, TWHole.class,Parameters.numberOfSearchedObjects);
    }


    /**
     * Returns the number of items currently in memory
     */
    public int getMemorySize() {
        return memorySize;
    }


    
    /**
     * Returns the nearest object that has been remembered recently where recently
     * is defined by a number of timesteps (threshold)
     *
     * If no Object is in memory which has been observed in the last threshold
     * timesteps it returns the most recently observed object. If there are no objects in
     * memory the method returns null. Note that specifying a threshold of one
     * will always return the most recently observed object. Specifying a threshold
     * of MAX_VALUE will always return the nearest remembered tile.
     *
     * Also note that it is likely that nearby objects are also the most recently observed
     *
     *
     * @param x coordinate from which to check for tiles
     * @param y coordinate from which to check for tiles
     * @param threshold how recently we want to have see the object
     * @param type the class of object we're looking for (Must inherit from TWObject, specifically tile or hole)
     * @return
     */
    private TWObject getNearbyObject(int sx, int sy, double threshold, Class<?> type) {

        //If we cannot find an object which we have seen recently then we want
        //the one with maxTimestamp
        double maxTimestamp = 0;
        TWObject o = null;
        double time = 0;
        TWObject ret = null;
        int x, y;
        for (Int2D offset : spiral) {
            x = offset.x + sx;
            y = offset.y + sy;

            if (me.getEnvironment().isInBounds(x, y) && objects[x][y] != null) {
                o = (TWObject) objects[x][y].getO();//get mem object
                if (type.isInstance(o)) {//if it's not the type we're looking for do nothing

                    time = objects[x][y].getT();//get time of memory

                    if (this.getSimulationTime() - time <= threshold) {
                        //if we found one satisfying time, then return
                        return o;
                    } else if (time > maxTimestamp) {
                        //otherwise record the timestamp and the item in case
                        //it's the most recent one we see
                        ret = o;
                        maxTimestamp = time;
                    }
                }
            }
        }

        //this will either be null or the object of Class type which we have
        //seen most recently but longer ago than now-threshold.
        return ret;
    }
    
    
    /**
     * Returns nearest percept of class type
     * 
     * @param sx
     * @param sy
     * @param threshold
     * @param type
     * @param ax
     * @param ay
     * @param limit
     * @return
     */
    public TWAgentPercept getNearbyPercept(int sx, int sy, double threshold, Class<?> type,int ax, int ay , int limit) {

        //If we cannot find an object which we have seen recently then we want
        //the one with maxTimestamp
        double maxTimestamp = 0;
        TWObject o = null;
        TWAgentPercept p =null;
        double time = 0;
        TWAgentPercept ret = null;
        int x, y;
        for (Int2D offset : spiral2) {
            x = offset.x + sx;
            y = offset.y + sy;

            if (me.getEnvironment().isInBounds(x, y) && objects[x][y] != null) {
                 p=objects[x][y];//get mem object
                if (type.isInstance(p.getO())) {//if it's not the type we're looking for do nothing
                   
                    time = p.getT();//get time of memory
                    if (this.getSimulationTime() - time <= threshold && !p.isSent() && (Math.abs(ax-x)+Math.abs(ay-y))>limit) {
                        //if we found one satisfying time, then return
                        p.setSent(true);
                        objects[x][y].setSent(true);
                        return p;
                    } 
                }
            }
        }

        return null;
    }

    
    /**
     * Returns N nearest objects that has been remembered recently where recently
     * is defined by a number of timesteps (threshold)
     *
     * If no Object is in memory which has been observed in the last threshold
     * timesteps it returns the most recently observed object. If there are no objects in
     * memory the method returns null. Note that specifying a threshold of one
     * will always return the most recently observed object. Specifying a threshold
     * of MAX_VALUE will always return the nearest remembered tile.
     *
     * Also note that it is likely that nearby objects are also the most recently observed
     *
     *
     * @param x coordinate from which to check for tiles
     * @param y coordinate from which to check for tiles     
     * @param type the class of object we're looking for (Must inherit from TWObject, specifically tile or hole)
     * @param N number of objects to return
     * @return
     */
    public TreeSet<TWObject> getNearbyNObjects(int sx, int sy, Class<?> type, int N) {


        
        TWObject o = null;
        
        TreeSet<TWObject> ret = new TreeSet<TWObject>();
        int x, y;
        int i=0;
        for (Int2D offset : spiral) {
            x = offset.x + sx;
            y = offset.y + sy;

            if (me.getEnvironment().isInBounds(x, y) && objects[x][y] != null) {
                o = (TWObject) objects[x][y].getO();//get mem object
                if (type.isInstance(o)) {//if it's not the type we're looking for do nothing
                  o.setDistance((float)o.getDistanceTo(sx, sy));
                    ret.add(o);
                  i++;
                  
                  if(i==N) return ret;
                   
                }
            }
        }
        return ret;
    }
    
    
    
    
    
    /**
     * Returns array with quantity of [obstacles, tiles, holes]
     *    
     *
     *
     * @param x coordinate from which to check for tiles
     * @param y coordinate from which to check for tiles
     * @param radius number of objects
     * @return
     */
    private int[] getNearbyQuantity(int sx, int sy,int radius) {

        TWObject o = null;
        
        int[] ret = new int[3];
        ret[0]=0; //obstacles
        ret[1]=0; //tiles
        ret[2]=0; //holes
        
        int x, y;
        int i=0;
        for (Int2D offset : spiral) {
            x = offset.x + sx;
            y = offset.y + sy;
            i++;
            if (me.getEnvironment().isInBounds(x, y) && objects[x][y] != null) {
                o = (TWObject) objects[x][y].getO();//get mem object
                if (TWObstacle.class.isInstance(o)) ret[0]++;                
                else if (TWTile.class.isInstance(o)) ret[1]++;
                else if (TWHole.class.isInstance(o)) ret[2]++;
            }
            if(i>=radius){
                return ret;
            }
        }

        return ret;
    }
    
    /**
     * 
     * @param sx
     * @param sy
     * @return
     */
    public int[] getNearbyQuantity(int sx, int sy) {
        return getNearbyQuantity(sx, sy,48);
    }
    
    
    /**
     * Used for invalidating the plan, returns the object of a particular type
     * (Tile or Hole) which is closest to the agent and within it's sensor range
     *
     * @param type
     * @return
     */
    public TWEntity getClosestObjectInSensorRange(Class<?> type) {
        return closestInSensorRange.get(type);
    }

    private void updateClosest(TWEntity o) {
        assert (o != null);
        if (closestInSensorRange.get(o.getClass()) == null || me.closerTo(o, closestInSensorRange.get(o.getClass()))) {
            closestInSensorRange.put(o.getClass(), o);
        }
    }

    /**
     * Is the cell blocked according to our memory?
     * 
     * @param tx x position of cell
     * @param ty y position of cell
     * @return true if the cell is blocked in our memory
     */
    public boolean isCellBlocked(int tx, int ty) {
        if(tx<0 || ty<0 || tx>=me.getEnvironment().getxDimension() || ty>=me.getEnvironment().getyDimension()) return true;
        //no memory at all, so assume not blocked
        if (objects[tx][ty] == null) {
            return false;
        }

        TWEntity e = (TWEntity) objects[tx][ty].getO();
        
       
        return (e instanceof TWObstacle);
    }
    
    public boolean isCellTile(int tx, int ty) {

        //no memory at all, so assume not blocked
        if (objects[tx][ty] == null) {
            return false;
        }

        TWEntity e = (TWEntity) objects[tx][ty].getO();
        //is it an obstacle?
        return (e instanceof TWTile);
    }
    
    
    public boolean isCellHole(int tx, int ty) {

        //no memory at all, so assume not blocked
        if (objects[tx][ty] == null) {
            return false;
        }

        TWEntity e = (TWEntity) objects[tx][ty].getO();
        //is it an obstacle?
        return (e instanceof TWHole);
    }

    public ObjectGrid2D getMemoryGrid() {
        return this.memoryGrid;
    }
    
    
    
    /**
     * 
     * @param id
     * @return
     */
    public Vector2D bestSector(int id){        
        TWSector sec=sectors[0][0];
        Vector2D ret=new Vector2D(sec.getX(),sec.getY());
        
        double max=me.getEnvironment().getxDimension()/Parameters.sectorXSize;
        double maxX=max;
        double maxY=me.getEnvironment().getyDimension()/Parameters.sectorYSize;
        double minX=0;
        
        
        int i,j;
        for (Int2D offset : spiral3) {
            i = offset.x + (int)Math.floor(me.getX()/Parameters.sectorXSize);
            j = offset.y + (int)Math.floor(me.getY()/Parameters.sectorYSize);
            /*edit
             * 
             * if(id==0)maxX=max/maxY*j;
            else minX=max/maxY*j;*/
            
            if(i<minX||j<0||i>=maxX||j>=maxY) continue;
                
                if(sectors[j][i].getVisits()<sec.getVisits()) {
                    sec=sectors[j][i];
                    ret=new Vector2D(sec.getX(),sec.getY());
                    if(sec.getVisits()==0) return ret;
             }                
         }
        return ret;
        
    }
    
    /**
     * Increases where agent travels
     * @param sx
     * @param sy
     */
    public void visitSector(int sx, int sy){
        int sectorX=(int)Math.floor(sx/Parameters.sectorXSize);
        int sectorY=(int)Math.floor(sy/Parameters.sectorYSize);
        sectors[sectorY][sectorX].increaseVisits();
        for (int i = 0; i < me.getEnvironment().getxDimension()/Parameters.sectorXSize; i++) {
            for (int j = 0; j < me.getEnvironment().getyDimension()/Parameters.sectorYSize; j++) {
                if(sectors[j][i].getVisits()>0) sectors[j][i].decreaseVisits();
                
            }
            
        }
        
    }

}
