package accrue.analysis.interprocanalysis;

import java.io.PrintStream;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

import polyglot.ast.CodeNode;
import polyglot.ast.ProcedureDecl;
import polyglot.frontend.ExtensionInfo;
import polyglot.frontend.Pass;
import polyglot.main.Report;
import polyglot.types.ProcedureInstance;
import polyglot.types.SemanticException;
import polyglot.util.InternalCompilerError;
import accrue.AccrueOptions;

/**
 * Manages the work queue of AnalysisUnits for interprocedural analyses. 
 * The work queue notices when analysis results change, and
 * adds dependent jobs back to the queue. Thus, the work queue finds fixed
 * points for the interprocedural analysis.
 * 
 * We could try to do some smart things, like figuring out the call graph,
 * and doing a topo sort on it, but we're being quick and dirty now.
 * 
 * @param <T> The type of the analysis results (not of the elements in the queue)
 */
public class WorkQueue<T extends Ordered<T>> {
    private static final int MAX_CURRENTLY_PROCESSING = Integer.MAX_VALUE;
    private final Queue<AnalysisUnit> queue;
    private final Set<AnalysisUnit> queueSet = new HashSet<AnalysisUnit>();
    private final Registrar registrar;
    private final ExtensionInfo extInfo;
    private final Pass pass;
    private final AnalysisFactory<T> factory;

    private Map<AnalysisUnit, AnalysisUnitRecord<T>> analysisUnitRecords =
            new LinkedHashMap<AnalysisUnit, AnalysisUnitRecord<T>>();

    /**
     * Current dependencies, i.e., when k is modified, need to recompute dependencies.get(k);
     */
    private Map<AnalysisUnit, Set<AnalysisUnit>> dependencies =
            new LinkedHashMap<AnalysisUnit, Set<AnalysisUnit>>();

    /**
     * Call graph between analysis units. May be different for different work queues due to
     * different AnalysisContexts
     */
    private CallGraph callGraph;

    private int processingCount = 0;

    public WorkQueue(Registrar registrar, ExtensionInfo extInfo,
            AnalysisFactory<T> factory, Pass pass) {
        this.registrar = registrar;
//        this.queue = new PriorityQueue<AnalysisUnit>(10, new AnalysisUnitComparator());
        this.queue = new LinkedList<AnalysisUnit>();
        this.extInfo = extInfo;
        this.factory = factory;
        this.pass = pass;
    }

    public AnalysisFactory<T> factory() {
        return this.factory;
    }

    public Registrar registrar() {
        return registrar;
    }

    public CallGraph callGraph() {
        return this.callGraph;
    }

    public Pass pass() {
        return this.pass;
    }

    public ExtensionInfo extensionInfo() {
        return this.extInfo;
    }

    public void processQueue() throws SemanticException {
        if (registrar.startProcs().isEmpty()) {
            if (((AccrueOptions) this.extInfo.getOptions()).helper().startProcClass == null) {
                throw new SemanticException("No start procedure found.");
            }
            else {
                throw new SemanticException("No start procedure found in class "
                        + ((AccrueOptions) this.extInfo.getOptions()).helper().startProcClass
                        + " (Is full class name correctly specified?)");
            }
        }

        Set<AnalysisUnit> startAnalysisUnits =
                new LinkedHashSet<AnalysisUnit>();
        AnalysisContext startContext = factory().startProcAnalysisContext();
        for (ProcedureInstance startPI : registrar.startProcs()) {
            AnalysisUnit sau =
                    new AnalysisUnit((ProcedureDecl) registrar.getCode(startPI),
                                     startContext);
            startAnalysisUnits.add(sau);
        }

        callGraph = new CallGraph(startAnalysisUnits);

        if (Report.should_report(this.factory.reportTopics(), 1))
            Report.report(1, "Starting to process " + factory.analysisName()
                    + " queue");
        long startTime = System.currentTimeMillis();

        for (AnalysisUnit pc : startAnalysisUnits) {
            AnalysisUtil<T> autil = factory().analysisUtil(this, pc);

            T initialInput = factory().startItem(autil, pc.procedureDecl());
            ExitMap<T> initialOutput =
                    initialAnalysisUnitOutput(pc, initialInput);
            analysisUnitRecords.put(pc,
                                    createAnalysisUnitRecord(initialInput,
                                                             initialOutput));

            addToQueue(pc);
        }

        while (!isEmpty()) {
            if (!currentlyProcessing.isEmpty()) {
                throw new InternalCompilerError("Should not be processing anything at this point, but it has "
                        + currentlyProcessing);
            }
            AnalysisUnit aunit = queue.remove();
            queueSet.remove(aunit);
            process(aunit);
        }

        long endTime = System.currentTimeMillis();
        if (Report.should_report(this.factory.reportTopics(), 1)) {
            Report.report(1, "Finished processing " + factory.analysisName()
                    + " queue: " + (endTime - startTime)
                    + " milliseconds; processed " + processingCount);
            final int HIST_SIZE = 100;
            int histogram[] = new int[HIST_SIZE];
            Set<Object> histogramSets[] = new Set[HIST_SIZE];
            int total = 0;
            Map<Integer, Integer> hist = new LinkedHashMap<Integer, Integer>();
            Report.report(1,
                          "  Histogram on number of times each analysis unit was run: number of records "
                                  + analysisUnitRecords.size());
            for (AnalysisUnit au : analysisUnitRecords.keySet()) {
                AnalysisUnitRecord<T> r = analysisUnitRecords.get(au);
                int count = r.runs;
                total += count;
                if (count < HIST_SIZE) {
                    histogram[count]++;
                    if (histogramSets[count] == null) {
                        histogramSets[count] = new LinkedHashSet<Object>();
                    }
                    histogramSets[count].add((au.cn instanceof ProcedureDecl) ? ((ProcedureDecl) au.cn).name()
                            : au.cn);
                }
                else {
                    Integer c = hist.get(count);
                    int d = c == null ? 0 : c;
                    hist.put(count, d + 1);
                }
            }
            for (int i = 0; i < HIST_SIZE; i++) {
                if (histogram[i] > 0) {
                    Report.report(1, "   " + i + "\t" + histogram[i]);
                    if (histogramSets[i].size() <= 10) {
                        Report.report(1, "           \t  " + histogramSets[i]);
                    }
                }
            }
            Report.report(1, hist.toString());
            Report.report(1, "total number of analysis units "
                    + analysisUnitRecords.keySet().size());
            Report.report(1, "total number of analyses " + total);

            int calleeCount = 0;
            int backEdges = 0;
            Set<Object>[] calledByHist = new Set[HIST_SIZE];
            int[] calledByHistA = new int[HIST_SIZE];
            for (AnalysisUnit au : analysisUnitRecords.keySet()) {
                Set<AnalysisUnit> callees = this.callGraph.callees(au);
                calleeCount += callees.size();
                for (AnalysisUnit callee : callees) {
                    boolean recursive =
                            this.callGraph.maybeRecursiveCall(au, callee);
                    if (recursive) {
                        backEdges++;
                        if (au.cn instanceof ProcedureDecl) {
                            System.out.println("Back edge "
                                    + au.procedureDecl().procedureInstance()
                                    + " -> " + callee);
                        }
                        else {
                            System.out.println("Back edge " + au.cn + " -> "
                                    + callee);
                        }
                    }
                }
                int calledByNum = this.callGraph.calledBy(au).size();
                if (calledByNum < HIST_SIZE) {
                    if (calledByHist[calledByNum] == null) {
                        calledByHist[calledByNum] = new LinkedHashSet<Object>();
                    }
                    calledByHist[calledByNum].add((au.cn instanceof ProcedureDecl) ? ((ProcedureDecl) au.cn).name()
                            : au.cn);
                    calledByHistA[calledByNum]++;
                }
                else {
                    System.out.println("   " + au + "  "
                            + this.callGraph.calledBy(au).size());
                }
            }
            Report.report(1, "Call graph: ");
            Report.report(1, "   total call edges " + calleeCount);
            Report.report(1, "   total back edges " + backEdges);
            Report.report(1, "   fan-in histogram: ");
            int faninsum = 0, faninwsum = 0;
            for (int i = 0; i < HIST_SIZE; i++) {
                if (calledByHist[i] != null && calledByHist[i].size() > 0) {
                    Report.report(1, "        " + i + "\t" + calledByHistA[i]);
                    if (calledByHist[i].size() <= 10) {
                        Report.report(1, "           \t  " + calledByHist[i]);
                    }
                    faninsum += calledByHistA[i];
                    faninwsum += i * calledByHistA[i];
                }
            }
            Report.report(1, "   fan in sum " + faninsum);
            Report.report(1, "   fan in weighted sum " + faninwsum);
        }
    }

    private Set<AnalysisUnit> currentlyProcessing =
            new LinkedHashSet<AnalysisUnit>();

    private void process(AnalysisUnit aUnit) throws SemanticException {
        this.processingCount++;

        if (currentlyProcessing.contains(aUnit)) {
            throw new InternalCompilerError("Should only be invoked if not currently processing");
        }

        if (queueSet.contains(aUnit)) {
            // we're processing it, so remove it from the queue.
            queue.remove(aUnit);
            queueSet.remove(aUnit);
        }

        AnalysisUnitRecord<T> rec = analysisUnitRecords.get(aUnit);
        AnalysisUtil<T> autil = factory().analysisUtil(this, aUnit);

        currentlyProcessing.add(aUnit);
        CodeNode cn = aUnit.codeNode();

        // analyze cd in the appropriate context.
        if (Report.should_report(factory.reportTopics(), 2))
            Report.report(1, "--- " + factory.analysisName() + " WorkQueue: "
                    + aUnit);

        if (rec == null || rec.input == null)
            throw new InternalCompilerError("No input!");
        T input = rec.input;
        ExitMap<T> output = autil.process(cn, input);
        if (Report.should_report(factory.reportTopics(), 4)) {
            Report.report(2, "WorkQueue (" + autil.analysisName()
                    + "): Just processed " + cn.codeInstance() + " and got "
                    + output);
        }

        recordResult(aUnit, autil, input, output);
        currentlyProcessing.remove(aUnit);

        if (Report.should_report(factory.reportTopics(), 2))
            Report.report(1, "--- " + factory.analysisName()
                    + " WorkQueue: done " + aUnit + "\truns: " + rec.runs);
    }

    /**
     * Request processing of an analysis unit with the given input. Returns true if
     * processing was actually performed
     */
    private boolean requestProcessing(AnalysisUnit requestor,
            AnalysisUnit requested, T input) throws SemanticException {
        // increase the input.
        AnalysisUnitRecord<T> rec = analysisUnitRecords.get(requested);
        // TODO BACKEDGE, use widen
        if (maybeRecursiveCall(requestor, requested)) {
            rec.input = input.widen(rec.input);
        }
        else {
            rec.input = input.upperBound(rec.input);
        }

        if (rec.output == null) {
            rec.output = initialAnalysisUnitOutput(requested, input);
        }

        // could be simple, and just call this.addToQueue(pc);
        // But we will be a little more complex to reduce the
        // the total amount of work.
        if (currentlyProcessing.size() >= MAX_CURRENTLY_PROCESSING
                || currentlyProcessing.contains(requested)) {
            // already processing this proc call, i.e., there is recursion in the calls...
            // We will finish processing this PC call, then do it again with the new input...

            if (Report.should_report(this.factory.reportTopics(), 4))
                Report.report(2, "Deferring processing of " + requested
                        + " as we are already processing it");

            this.addToQueue(requested);
            return false;
        }

        // at this point, could make a decision about whether to
        // add to the queue, or process now, say, based on the
        // size of currentlyProcessing. For now, we will
        // just process immediately.
        process(requested);
        return true;
    }

    public boolean isEmpty() {
        return queue.isEmpty();
    }

    public void recordResult(AnalysisUnit pc, AnalysisUtil<T> autil, T input,
            ExitMap<T> output) {
        AnalysisUnitRecord<T> rec = analysisUnitRecords.get(pc);
        if (isOutputChanged(output, rec, autil)) {
            // result changed.

            if (Report.should_report(factory.reportTopics(), 3)) {
                Report.report(2,
                              "WorkQueue.recordResult: recording changed result for "
                                      + pc);
            }
            // update result, and add anything that depended on it back to the queue
            updateOutput(rec, output, autil);

            this.addAllToQueue(getDependees(pc));

            // we have added to the queue any analysis units that may  
            // need to be re-run because the analysis results changed.
            // Clear the dependees: they will add themselves to the set
            // when they are re-analyzed, if needed. 
            clearDependees(pc);
        }
        else {
            // output didn't change, no need to recompute anything.            
            if (Report.should_report(factory.reportTopics(), 3)) {
                Report.report(2,
                              "WorkQueue.recordResult: no change in output of "
                                      + pc);
            }
        }
        rec.lastRunSound = autil.soundSoFarThisRun;
        rec.runs++;
    }

    protected void updateOutput(AnalysisUnitRecord<T> rec, ExitMap<T> output,
            AnalysisUtil<T> autil) {
        rec.output = output;
    }

    protected boolean isOutputChanged(ExitMap<T> output,
            AnalysisUnitRecord<T> rec, AnalysisUtil<T> autil) {
        return rec.output == null || !rec.output.equals(output);
    }

    public ExitMap<T> getAnalysisResult(AnalysisUnit callee, T input,
            AnalysisUtil<T> callerAUtil) throws SemanticException {
        AnalysisUnit caller = callerAUtil.currentAnalysisUnit();
        if (caller == null) {
            throw new IllegalArgumentException("Caller must be non-null");
        }
        if (input == null) {
            throw new IllegalArgumentException("Null input!");
        }

        callGraph.addEdge(caller, callee);

        // check if we have code for callee
        if (callee.cn instanceof ProcedureDecl
                && !registrar.hasCode(callee.procedureDecl()
                                            .procedureInstance())) {
            throw new SemanticException("This should only be called if there is code for "
                    + callee.procedureDecl());
        }

        AnalysisUnitRecord<T> rec;
        if (!analysisUnitRecords.containsKey(callee)) {
            rec = createAnalysisUnitRecord();
            analysisUnitRecords.put(callee, rec);
        }
        else {
            rec = analysisUnitRecords.get(callee);
        }

        // Is the existing analysis result suitable?
        ExitMap<T> output;
        boolean soundResultsUsed;
        if (existingResultSuitable(input, rec)) {
            // existing analysis result is OK.
            output = rec.output;
            soundResultsUsed = rec.lastRunSound;
        }
        else {

            // Request that callee be processed
            boolean processedNow = requestProcessing(caller, callee, input);

            // either the work item was processed immediately, in which case
            // the output is correct, or the work item will be processed
            // sometime, and (since there is a dependency from caller to callee)
            // caller will be reevaluated, but we are ok with returning the
            // most accurate result we currently have.
            AnalysisUnitRecord<T> calleeRec = analysisUnitRecords.get(callee);
            output = calleeRec.output;
            soundResultsUsed = processedNow && calleeRec.lastRunSound;
        }

        // update the read and write sets of the caller
        updateCallerRecordAfterCalleeProcessing(analysisUnitRecords.get(caller),
                                                rec);

        // modify the locations of the output to keep only the write set, and have the locations
        // of the input otherwise.               
        ExitMap<T> ret = restoreCallingFrame(output, input, rec);
        if (ret == null)
            throw new InternalCompilerError("null ExitMap returned");

        if (!soundResultsUsed) {
            // The results for the callee may not be sound (e.g., the callee
            // was a recursive call to something we are in the middle of
            // analyzing). Record a dependency from caller to callee
            // so that when the callee recomputes, the caller will get re-analyzed
            getDependees(callee).add(caller);
        }

        callerAUtil.soundSoFarThisRun =
                soundResultsUsed && callerAUtil.soundSoFarThisRun;

        return ret;
    }

    protected AnalysisUnitRecord<T> createAnalysisUnitRecord() {
        return new AnalysisUnitRecord<T>();
    }

    protected AnalysisUnitRecord<T> createAnalysisUnitRecord(T initialInput,
            ExitMap<T> initialOutput) {
        return new AnalysisUnitRecord<T>(initialInput, initialOutput);
    }

    protected void updateCallerRecordAfterCalleeProcessing(
            AnalysisUnitRecord<T> callerRec, AnalysisUnitRecord<T> rec) {
        // do nothing, this is a hook for subclasses
    }

    protected boolean existingResultSuitable(T input, AnalysisUnitRecord<T> rec) {
        return rec.input != null && input.leq(rec.input);
    }

    protected ExitMap<T> restoreCallingFrame(ExitMap<T> output, T input,
            AnalysisUnitRecord<T> rec) {
        return output;
    }

    public ExitMap<T> emptyExitMap() {
        return factory().emptyExitMap();
    }

    protected ExitMap<T> createExitMap(T input) {
        return emptyExitMap().normalTermination(input);
    }

    protected ExitMap<T> initialAnalysisUnitOutput(AnalysisUnit callee, T input) {
        ExitMap<T> f = createExitMap(input);
        return f;
    }

    public boolean maybeRecursiveCall(AnalysisUnit caller, AnalysisUnit callee) {
        return callGraph.maybeRecursiveCall(caller, callee);
    }

    protected Set<AnalysisUnit> getDependees(AnalysisUnit a) {
        Set<AnalysisUnit> dependees = dependencies.get(a);
        if (dependees == null) {
            dependees = new LinkedHashSet<AnalysisUnit>();
            dependencies.put(a, dependees);
        }
        return dependees;
    }

    /**
     * Clear the set of AnalysisUnits that depend on the
     * results of AnalysisUnit a.
     * @param a
     */
    protected void clearDependees(AnalysisUnit a) {
        Set<AnalysisUnit> dependees = dependencies.get(a);
        if (dependees != null) {
            dependees.clear();
        }
    }

    private void addToQueue(AnalysisUnit pc) {
        if (!queueSet.contains(pc)) {
            queue.add(pc);
            queueSet.add(pc);
        }
    }

    private void addAllToQueue(Collection<AnalysisUnit> pcs) {
        for (AnalysisUnit pc : pcs) {
            addToQueue(pc);
        }
    }

    public void dumpQueue(PrintStream out) {
        out.print("workqueue: ");
        out.println(queue);
        out.print("currently: ");
        out.println(currentlyProcessing);
    }

    public static class AnalysisUnitRecord<T extends Ordered<T>> {
        public AnalysisUnitRecord() {
            this(null, null);
        }

        public AnalysisUnitRecord(T input, ExitMap<T> output) {
            this.input = input;
            this.output = output;
        }

        private T input;
        private ExitMap<T> output;
        private int runs;

        /**
         * Did the last run have sound analysis results? This will be
         * false if the last run, say, had a recursive call that
         * needed to use unsound results.
         */
        private boolean lastRunSound;

        public ExitMap<T> output() {
            return this.output;
        }

        public T input() {
            return this.input;
        }
    }

    public AnalysisUnitRecord<T> getAnalysisUnitRecord(AnalysisUnit unit) {
        return this.analysisUnitRecords.get(unit);
    }
}
