using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.Boogie;
using System.Diagnostics;
using Graphing;

namespace Microsoft.Boogie
{

    public class UniformityAnalyser
    {
        private Program prog;

        private bool doAnalysis, unstructured;

        private ISet<Implementation> entryPoints;

        private IEnumerable<Variable> nonUniformVars;

        private bool ProcedureChanged;

        private Dictionary<string, KeyValuePair<bool, Dictionary<string, bool>>> uniformityInfo;

        private Dictionary<string, HashSet<int>> nonUniformLoops;

        private Dictionary<string, HashSet<Block>> nonUniformBlocks;

        private Dictionary<string, HashSet<int>> loopsWithNonuniformReturn;

        private Dictionary<string, List<string>> inParameters;

        private Dictionary<string, List<string>> outParameters;

        private List<WhileCmd> loopStack;

        private bool hitNonuniformReturn;

        public UniformityAnalyser(Program prog, bool doAnalysis, bool unstructured, ISet<Implementation> entryPoints, IEnumerable<Variable> nonUniformVars)
        {
            this.prog = prog;
            this.doAnalysis = doAnalysis;
            this.unstructured = unstructured;
            this.entryPoints = entryPoints;
            this.nonUniformVars = nonUniformVars;
            uniformityInfo = new Dictionary<string, KeyValuePair<bool, Dictionary<string, bool>>>();
            nonUniformLoops = new Dictionary<string, HashSet<int>>();
            nonUniformBlocks = new Dictionary<string, HashSet<Block>>();
            loopsWithNonuniformReturn = new Dictionary<string, HashSet<int>>();
            inParameters = new Dictionary<string, List<string>>();
            outParameters = new Dictionary<string, List<string>>();
            loopStack = new List<WhileCmd>();
        }

        public void Analyse()
        {
            var impls = prog.TopLevelDeclarations.OfType<Implementation>();

            foreach (var Impl in impls)
            {
                bool uniformProcedure = doAnalysis || entryPoints.Contains(Impl);

                uniformityInfo.Add(Impl.Name, new KeyValuePair<bool, Dictionary<string, bool>>
                    (uniformProcedure, new Dictionary<string, bool> ()));

                nonUniformLoops.Add(Impl.Name, new HashSet<int>());
                loopsWithNonuniformReturn.Add(Impl.Name, new HashSet<int>());

                foreach (var v in nonUniformVars)
                    SetNonUniform(Impl.Name, v.Name);

                foreach (Variable v in Impl.LocVars)
                {
                    if (doAnalysis)
                    {
                        SetUniform(Impl.Name, v.Name);
                    }
                    else
                    {
                        SetNonUniform(Impl.Name, v.Name);
                    }
                }

                inParameters[Impl.Name] = new List<string>();

                foreach (Variable v in Impl.InParams)
                {
                    inParameters[Impl.Name].Add(v.Name);
                    if (doAnalysis)
                    {
                        SetUniform(Impl.Name, v.Name);
                    }
                    else
                    {
                        SetNonUniform(Impl.Name, v.Name);
                    }
                }

                outParameters[Impl.Name] = new List<string>();
                foreach (Variable v in Impl.OutParams)
                {
                    outParameters[Impl.Name].Add(v.Name);
                    if (doAnalysis)
                    {
                        SetUniform(Impl.Name, v.Name);
                    }
                    else
                    {
                        SetNonUniform(Impl.Name, v.Name);
                    }
                }

                ProcedureChanged = true;
            }

            if (doAnalysis)
            {
                while (ProcedureChanged)
                {
                    ProcedureChanged = false;

                    foreach (var Impl in impls)
                    {
                        hitNonuniformReturn = false;
                        Analyse(Impl, uniformityInfo[Impl.Name].Key);
                    }
                }
            }

            foreach (var Impl in impls)
            {
                if (!IsUniform (Impl.Name))
                {
                    List<string> newIns = new List<String>();
                    newIns.Add("_P");
                    foreach (string s in inParameters[Impl.Name])
                    {
                        newIns.Add(s);
                    }
                    inParameters[Impl.Name] = newIns;
                }
            }
        }

        private void Analyse(Implementation Impl, bool ControlFlowIsUniform)
        {
            if (unstructured)
            {
                if (!ControlFlowIsUniform)
                {
                    nonUniformBlocks[Impl.Name] = new HashSet<Block>(Impl.Blocks);

                    foreach (Variable v in Impl.LocVars) {
                        if (IsUniform(Impl.Name, v.Name)) {
                            SetNonUniform(Impl.Name, v.Name);
                        }
                    }

                    foreach (Variable v in Impl.InParams) {
                        if (IsUniform(Impl.Name, v.Name)) {
                            SetNonUniform(Impl.Name, v.Name);
                        }
                    }

                    foreach (Variable v in Impl.OutParams) {
                      if (IsUniform(Impl.Name, v.Name)) {
                         SetNonUniform(Impl.Name, v.Name);
                      }
                    }
                        
                    return;
                }

                Graph<Block> blockGraph = prog.ProcessLoops(Impl);
                var ctrlDep = blockGraph.ControlDependence();

                // Compute transitive closure of control dependence info.
                bool changed;
                do
                {
                    changed = false;
                    foreach (var depEntry in ctrlDep)
                    {
                        var newDepSet = new HashSet<Block>(depEntry.Value);
                        foreach (var dep in depEntry.Value)
                        {
                            if (ctrlDep.ContainsKey(dep))
                                newDepSet.UnionWith(ctrlDep[dep]);
                        }
                        if (newDepSet.Count != depEntry.Value.Count)
                        {
                            depEntry.Value.UnionWith(newDepSet);
                            changed = true;
                        }
                    }
                } while (changed);

                var nonUniformBlockSet = new HashSet<Block>();
                nonUniformBlocks[Impl.Name] = nonUniformBlockSet;

                do {
                  changed = false;
                  foreach (var block in Impl.Blocks) {
                    bool uniform = !nonUniformBlockSet.Contains(block);
                    bool newUniform = Analyse(Impl, block.Cmds, uniform);
                    if (uniform && !newUniform) {
                      changed = true;
                      nonUniformBlockSet.Add(block);
                      Block pred = blockGraph.Predecessors(block).Single();
                      if (ctrlDep.ContainsKey(pred))
                        nonUniformBlockSet.UnionWith(ctrlDep[pred]);
                    }
                  }
                } while (changed);
            }
            else
            {
                Analyse(Impl, Impl.StructuredStmts, ControlFlowIsUniform);
            }
        }


        private void Analyse(Implementation impl, StmtList stmtList, bool ControlFlowIsUniform)
        {
            ControlFlowIsUniform &= !hitNonuniformReturn;
            foreach (BigBlock bb in stmtList.BigBlocks)
            {
                Analyse(impl, bb, ControlFlowIsUniform);
            }
        }

        private Implementation GetImplementation(string procedureName)
        {
            foreach (Declaration D in prog.TopLevelDeclarations)
            {
                if (D is Implementation && ((D as Implementation).Name == procedureName))
                {
                    return D as Implementation;
                }
            }
            return null;
        }

        private void Analyse(Implementation impl, BigBlock bb, bool ControlFlowIsUniform)
        {
            ControlFlowIsUniform &= !hitNonuniformReturn;
            Analyse(impl, bb.simpleCmds, ControlFlowIsUniform);

            if (bb.ec is WhileCmd)
            {
                WhileCmd wc = bb.ec as WhileCmd;
                loopStack.Add(wc);
                Analyse(impl, wc.Body, ControlFlowIsUniform && IsUniform(impl.Name, wc.Guard) && 
                    !nonUniformLoops[impl.Name].Contains(GetLoopId(wc)));
                loopStack.RemoveAt(loopStack.Count - 1);
            }
            else if (bb.ec is IfCmd)
            {
                IfCmd ifCmd = bb.ec as IfCmd;
                Analyse(impl, ifCmd.thn, ControlFlowIsUniform && IsUniform(impl.Name, ifCmd.Guard));
                if (ifCmd.elseBlock != null)
                {
                    Analyse(impl, ifCmd.elseBlock, ControlFlowIsUniform && IsUniform(impl.Name, ifCmd.Guard));
                }
                Debug.Assert(ifCmd.elseIf == null);
            }
            else if (bb.ec is BreakCmd)
            {
                if (!ControlFlowIsUniform && !nonUniformLoops[impl.Name].Contains(GetLoopId(loopStack[loopStack.Count - 1])))
                {
                    SetNonUniform(impl.Name, loopStack[loopStack.Count - 1]);
                }
            }

            if (bb.tc is ReturnCmd && !ControlFlowIsUniform)
            {
                hitNonuniformReturn = true;
                if (loopStack.Count > 0 && !nonUniformLoops[impl.Name].Contains(GetLoopId(loopStack[0])))
                {
                    SetNonUniform(impl.Name, loopStack[0]);
                    loopsWithNonuniformReturn[impl.Name].Add(GetLoopId(loopStack[0]));
                }
            }


        }

        private bool Analyse(Implementation impl, CmdSeq cmdSeq, bool ControlFlowIsUniform)
        {
            foreach (Cmd c in cmdSeq)
            {
                if (c is AssignCmd)
                {
                    AssignCmd assignCmd = c as AssignCmd;
                    foreach (var a in assignCmd.Lhss.Zip(assignCmd.Rhss))
                    {

                        if (a.Item1 is SimpleAssignLhs)
                        {
                            SimpleAssignLhs lhs = a.Item1 as SimpleAssignLhs;
                            Expr rhs = a.Item2;
                            if (IsUniform(impl.Name, lhs.AssignedVariable.Name) &&
                                (!ControlFlowIsUniform || !IsUniform(impl.Name, rhs)))
                            {
                                SetNonUniform(impl.Name, lhs.AssignedVariable.Name);
                            }

                        }
                    }
                }
                else if (c is CallCmd)
                {
                    CallCmd callCmd = c as CallCmd;
                    Implementation CalleeImplementation = GetImplementation(callCmd.callee);

                    if (CalleeImplementation != null)
                    {

                        if (!ControlFlowIsUniform)
                        {
                            if (IsUniform(callCmd.callee))
                            {
                                SetNonUniform(callCmd.callee);
                            }
                        }
                        for (int i = 0; i < CalleeImplementation.InParams.Length; i++)
                        {
                            if (IsUniform(callCmd.callee, CalleeImplementation.InParams[i].Name)
                                && !IsUniform(impl.Name, callCmd.Ins[i]))
                            {
                                SetNonUniform(callCmd.callee, CalleeImplementation.InParams[i].Name);
                            }
                        }

                        for (int i = 0; i < CalleeImplementation.OutParams.Length; i++)
                        {
                            if (IsUniform(impl.Name, callCmd.Outs[i].Name)
                                && !IsUniform(callCmd.callee, CalleeImplementation.OutParams[i].Name))
                            {
                                SetNonUniform(impl.Name, callCmd.Outs[i].Name);
                            }
                        }

                    }
                }
                else if (c is AssumeCmd)
                {
                    var ac = (AssumeCmd)c;
                    if (ControlFlowIsUniform && QKeyValue.FindBoolAttribute(ac.Attributes, "partition") &&
                        !IsUniform(impl.Name, ac.Expr))
                    {
                      ControlFlowIsUniform = false;
                    }
                }
            }

            return ControlFlowIsUniform;
        }

        private int GetLoopId(WhileCmd wc)
        {
            AssertCmd inv = wc.Invariants[0] as AssertCmd;
            Debug.Assert(inv.Attributes.Key.Contains("loophead_"));
            return Convert.ToInt32(inv.Attributes.Key.Substring("loophead_".Length));
        }

        private void SetNonUniform(string procedureName)
        {
            uniformityInfo[procedureName] = new KeyValuePair<bool,Dictionary<string,bool>>
                (false, uniformityInfo[procedureName].Value);
            RecordProcedureChanged();
        }

        private void SetNonUniform(string procedureName, WhileCmd wc)
        {
            nonUniformLoops[procedureName].Add(GetLoopId(wc));
            RecordProcedureChanged();
        }

        public bool IsUniform(string procedureName)
        {
            if (!uniformityInfo.ContainsKey(procedureName))
            {
                return false;
            }
            return uniformityInfo[procedureName].Key;
        }

        public bool IsUniform(string procedureName, Block b)
        {
            if (!nonUniformBlocks.ContainsKey(procedureName))
            {
                return false;
            }
            return !nonUniformBlocks[procedureName].Contains(b);
        }

        class UniformExpressionAnalysisVisitor : StandardVisitor {

          private bool isUniform = true;
          private Dictionary<string, bool> uniformityInfo;

          public UniformExpressionAnalysisVisitor(Dictionary<string, bool> uniformityInfo) {
            this.uniformityInfo = uniformityInfo;
          }

          public override Variable VisitVariable(Variable v) {
            if (!uniformityInfo.ContainsKey(v.Name)) {
              isUniform = isUniform && (v is Constant);
            } else if (!uniformityInfo[v.Name]) {
              isUniform = false;
            }

            return v;
          }

          internal bool IsUniform() {
            return isUniform;
          }
        }

        public bool IsUniform(string procedureName, Expr expr)
        {
            UniformExpressionAnalysisVisitor visitor = new UniformExpressionAnalysisVisitor(uniformityInfo[procedureName].Value);
            visitor.VisitExpr(expr);
            return visitor.IsUniform();
        }

        public bool IsUniform(string procedureName, string v)
        {
            if (!uniformityInfo.ContainsKey(procedureName))
            {
                return false;
            }

            if (!uniformityInfo[procedureName].Value.ContainsKey(v))
            {
                return false;
            }
            return uniformityInfo[procedureName].Value[v];
        }

        private void SetUniform(string procedureName, string v)
        {
            uniformityInfo[procedureName].Value[v] = true;
            RecordProcedureChanged();
        }

        private void RecordProcedureChanged()
        {
            ProcedureChanged = true;
        }

        private void SetNonUniform(string procedureName, string v)
        {
            uniformityInfo[procedureName].Value[v] = false;
            RecordProcedureChanged();
        }

        public void dump()
        {
            foreach (string p in uniformityInfo.Keys)
            {
                Console.WriteLine("Procedure " + p + ": "
                    + (uniformityInfo[p].Key ? "uniform" : "nonuniform"));
                foreach (string v in uniformityInfo[p].Value.Keys)
                {
                    Console.WriteLine("  " + v + ": " +
                        (uniformityInfo[p].Value[v] ? "uniform" : "nonuniform"));
                }
                Console.Write("Ins [");
                for (int i = 0; i < inParameters[p].Count; i++)
                {
                    Console.Write((i == 0 ? "" : ", ") + inParameters[p][i]);
                }
                Console.WriteLine("]");
                Console.Write("Outs [");
                for (int i = 0; i < outParameters[p].Count; i++)
                {
                    Console.Write((i == 0 ? "" : ", ") + outParameters[p][i]);
                }
                Console.WriteLine("]");
                Console.Write("Non-uniform loops:");
                foreach (int l in nonUniformLoops[p])
                {
                    Console.Write(" " + l);
                }
                Console.WriteLine();
                Console.Write("Non-uniform blocks:");
                foreach (Block b in nonUniformBlocks[p])
                {
                    Console.Write(" " + b.Label);
                }
                Console.WriteLine();
            }
        }


        public string GetInParameter(string procName, int i)
        {
            return inParameters[procName][i];
        }

        public string GetOutParameter(string procName, int i)
        {
            return outParameters[procName][i];
        }


        public bool knowsOf(string p)
        {
            return uniformityInfo.ContainsKey(p);
        }

        public void AddNonUniform(string proc, string v)
        {
            if (uniformityInfo.ContainsKey(proc))
            {
                Debug.Assert(!uniformityInfo[proc].Value.ContainsKey(v));
                uniformityInfo[proc].Value[v] = false;
            }
        }

        public bool IsUniform(string proc, WhileCmd wc)
        {
            if (unstructured)
                return false;

            return !nonUniformLoops[proc].Contains(GetLoopId(wc));
        }

        public bool HasNonuniformReturn(string proc, WhileCmd whileCmd)
        {
            return loopsWithNonuniformReturn[proc].Contains(GetLoopId(whileCmd));
        }
    }

}