/***************************************************************
*                                                      
*   Copyright (C) University of Southampton. All rights reserved
*                                                       
****************************************************************/

// Author: Matteo Venanzi

// Source code of the models for the Community-based Bayesian aggregation of Crowdsourced Estimates presented in the paper:
// 
// "Bayesian Modelling of Community-Based Multidimensional Trust in Participatory Sensing under Data Sparsity"
// by Venanzi, Matteo, Teacy, W. T. L., Rogers, Alex and Jennings, Nicholas R. (2015)
// In, 24th International Joint Conference on Artificial Intelligenge (IJCAI)

using MicrosoftResearch.Infer;
using MicrosoftResearch.Infer.Distributions;
using MicrosoftResearch.Infer.Models;
using MicrosoftResearch.Infer.Utils;
using System;
using System.Linq;

namespace IJCAI15
{
    /// <summary>
    /// The BACE model.
    /// 
    /// This model includes the core functionalities of the models CBACE and CBACEBias presented in Venanzi et al. (IJCAI 2015)
    /// 
    /// This model is the the MaxTrust model used as benchmark in Venanzi et al. (IJCAI 2015)
    /// </summary>
    public class BACE
    {

        public int ItemCount
        {
            get
            {
                return N == null ? 0 : N.SizeAsInt;
            }
        }
        public int UserCount
        {
            get
            {
                return K == null ? 0 : K.SizeAsInt;
            }
        }

        // Ranges
        protected Range N;
        protected Range K;
        protected Range KN;

        // Variables in the model
        protected VariableArray<double> ItemTrueValue;
        protected VariableArray<int> UserItemCount;
        protected VariableArray<double> UserTruePrecision;
        protected VariableArray<VariableArray<int>, int[][]> UserItemIndex;
        protected VariableArray<VariableArray<double>, double[][]> UserObservedValue;
        protected VariableArray<VariableArray<double>, double[][]> UserObservedPrecision;

        // Prior distributions
        protected VariableArray<Gamma> UserTrustPrior;
        protected VariableArray<Gaussian> ItemTrueValuePrior;

        // Model evidence
        protected Variable<bool> Evidence;

        // Inference engine
        protected InferenceEngine Engine;

        public int NumberOfIterations
        {
            get;
            set;
        }

        public BACE()
        {
            NumberOfIterations = 35;
        }

        public virtual void CreateModel(int itemCount, int userCount)
        {
            Evidence = Variable.Bernoulli(0.5).Named("Evidence");
            var evidenceBlock = Variable.If(Evidence);
            DefineVariablesAndRanges(itemCount, userCount);
            DefineGenerativeProcess();
            DefineInferenceEngine();
            evidenceBlock.CloseBlock();
        }

        protected virtual void DefineVariablesAndRanges(int taskCount, int userCount)
        {
            N = new Range(taskCount).Named("n");
            K = new Range(userCount).Named("k");

            UserItemCount = Variable.Array<int>(K).Named("UserItemCount");
            KN = new Range(UserItemCount[K]).Named("kn");
            UserItemIndex = Variable.Array(Variable.Array<int>(KN), K).Named("UserItemIndex");
            UserItemIndex.SetValueRange(N);
            UserObservedValue = Variable.Array(Variable.Array<double>(KN), K).Named("UserObservedValue");
            UserObservedPrecision = Variable.Array(Variable.Array<double>(KN), K).Named("UserObservedPrecision");
            ItemTrueValuePrior = Variable<Gaussian>.Array(N).Named("TrueItemValuePrior");
            ItemTrueValue = Variable<double>.Array(N).Named("TrueItemValue");
            ItemTrueValue[N] = Variable<double>.Random(ItemTrueValuePrior[N]);
            UserTrustPrior = Variable<Gamma>.Array(K).Named("UserTrustPrior");
            UserTruePrecision = Variable<double>.Array(K).Named("UserTruePrecision");
        }

        protected virtual void DefineGenerativeProcess()
        {
            // The process that generates the user's reports
            UserTruePrecision[K] = Variable<double>.Random(UserTrustPrior[K]);
            using (Variable.ForEach(K))
            {
                var trueValue = Variable.Subarray(ItemTrueValue, UserItemIndex[K]).Named("TrueItemValueSub");
                using (Variable.ForEach(KN))
                {
                    UserObservedValue[K][KN] = Variable.GaussianFromMeanAndPrecision(trueValue[KN], (UserObservedPrecision[K][KN] * UserTruePrecision[K]).Named("ScaledPrecision"));
                }
            }
        }

        protected virtual void DefineInferenceEngine()
        {
            Engine = new InferenceEngine(new VariationalMessagePassing());
            Engine.Compiler.UseParallelForLoops = true;
            Engine.ShowProgress = false;
            Engine.ShowFactorGraph = false;
            Engine.Compiler.WriteSourceFiles = true;
            Engine.ShowSchedule = false;
        }

        protected virtual void SetPriors()
        {
            ItemTrueValuePrior.ObservedValue = Util.ArrayInit(ItemCount, t => Gaussian.FromMeanAndPrecision(0, 1));
            UserTrustPrior.ObservedValue = Util.ArrayInit(UserCount, k => Gamma.FromMeanAndVariance(1, 1));
        }

        protected virtual void AttachData(int[][] itemIndices, double[][] userValues, double[][] userPrecisions)
        {
            UserItemCount.ObservedValue = itemIndices.Select(items => items.Length).ToArray();
            UserItemIndex.ObservedValue = itemIndices;
            UserObservedValue.ClearObservedValue();
        }

        public virtual BACEPosteriors Infer(int[][] itemIndices, double[][] userValues, double[][] userPrecisions)
        {
            SetPriors();
            AttachData(itemIndices, userValues, userPrecisions);
            var result = new BACEPosteriors();
            Engine.NumberOfIterations = NumberOfIterations;
            result.TrueItemValuePosterior = Engine.Infer<Gaussian[]>(ItemTrueValue);
            result.UserTrustPosterior = Engine.Infer<Gamma[]>(UserTruePrecision);
            result.LogEvidence = Engine.Infer<Bernoulli>(Evidence).LogOdds;
            return result;
        }
    }

    /// <summary>
    /// The BACE posteriors with accuracy metrics (Root mean square error and energy score).
    /// </summary>
    [Serializable]
    public class BACEPosteriors
    {
        public Gaussian[] TrueItemValuePosterior;
        public Gamma[] UserTrustPosterior;
        public Gaussian[][] UserPrediction;
        public double LogEvidence;

		
        public double ComputeRMSE(double[] TrueValue)
        {
            double[] Accuracy = new double[TrueValue.Length];
            for (int i = 0; i < TrueValue.Length; i++)
            {

                Accuracy[i] = Math.Abs(TrueItemValuePosterior[i].GetMean() - TrueValue[i]);
            }
            return Accuracy.Average();

        }

        public double ComputeEnergyScore(double[] TrueValue)
        {
            // Energy score
            double[] EnergyScore = new double[TrueValue.Length];
            int num_samples = 10000;
            for (int i = 0; i < TrueValue.Length; i++)
            {
                double[] samples = Util.ArrayInit(num_samples, t => TrueItemValuePosterior[i].Sample());
                double[] m_term_arr = new double[num_samples];
                for (int j = 0; j < num_samples; j++)
                {
                    m_term_arr[j] = Math.Sqrt(Math.Pow(samples[j] - TrueValue[i], 2)); // Norm
                }
                double m_term = m_term_arr.Average();

                double[] v_term_arr = new double[num_samples-1];
                for (int j=0; j < num_samples-1; j++)
                    v_term_arr[j] = Math.Abs(samples[j] - samples[j+1]);

                double v_term = v_term_arr.Sum() / (2 * (num_samples-1));
                EnergyScore[i] = m_term - v_term;
                    
            }
            return EnergyScore.Average();
        }
    }
	
	/// <summary>
    /// The CBACEBias model implemented as a sub-class of BACE
    /// </summary>
    class CBACEBias : BACE
    {
        // Additional ranges
        protected Range M;

        // Additional variables
        protected Variable<int> CommunityCount;
        protected VariableArray<double> UserBias;
        protected VariableArray<double> UserPrecisionLog;
        protected VariableArray<int> Community;
        protected VariableArray<Discrete> CommunityInitializer;
        protected Variable<Vector> CommunityProbs;
        protected VariableArray<double> CommunityBias;
        protected VariableArray<double> CommunityTrust;

        // Additional priors
        protected Variable<Dirichlet> CommunityProbsPrior;
        protected VariableArray<Gaussian> CommunityBiasPrior;
        protected VariableArray<Gaussian> CommunityTrustPrior;

        // Additional parameters
        protected int BiasPrecision;
        protected int TrustPrecision;
        protected int CommunityPseudoCount;

        public CBACEBias()
            : base()
        {
            BiasPrecision = 1;
            TrustPrecision = 1;
            CommunityPseudoCount = 10;
        }

        protected override void DefineVariablesAndRanges(int taskCount, int userCount)
        {
            base.DefineVariablesAndRanges(taskCount, userCount);

            CommunityCount = Variable.New<int>().Named("CommunityCount");
            M = new Range(CommunityCount).Named("m");

            // Community memberships
            CommunityProbsPrior = Variable.New<Dirichlet>().Named("CommunityProbPrior");
            CommunityProbs = Variable<Vector>.Random(CommunityProbsPrior).Named("CommunityProb");
            CommunityProbs.SetValueRange(M);
            Community = Variable.Array<int>(K).Named("Community");
            Community[K] = Variable.Discrete(CommunityProbs).ForEach(K);
            // Initialiser to break symmetry for community membership
            CommunityInitializer = Variable.Array<Discrete>(K).Named("CommunityInitializer");
            Community[K].InitialiseTo(CommunityInitializer[K]);

            // Community bias
            CommunityBiasPrior = Variable<Gaussian>.Array(M).Named("CommunityBiasPrior");
            CommunityBias = Variable<double>.Array(M).Named("CommunityBias");
            CommunityBias[M] = Variable<double>.Random(CommunityBiasPrior[M]);
            UserBias = Variable<double>.Array(K).Named("UserBias");

            // Community trust
            CommunityTrustPrior = Variable<Gaussian>.Array(M).Named("CommunityTrustPrior");
            CommunityTrust = Variable<double>.Array(M).Named("CommunityTrust");
            CommunityTrust[M] = Variable<double>.Random(CommunityTrustPrior[M]);
            
            // User log trust
            UserTrustLog = Variable<double>.Array(K).Named("UserTrustLog");
        }

        protected override void DefineGenerativeProcess()
        {
            // The process that generates the user's reports
            using (Variable.ForEach(K))
            {
                using (Variable.Switch(Community[K]))
                {
                    UserBias[K] = Variable.GaussianFromMeanAndPrecision(CommunityBias[Community[K]], BiasPrecision);
                    UserTrustLog[K] = Variable.GaussianFromMeanAndPrecision(CommunityTrust[Community[K]], TrustPrecision);
                }
                UserTruePrecision[K] = Variable.Exp(UserTrustLog[K]);

                var trueValue = Variable.Subarray(ItemTrueValue, UserItemIndex[K]).Named("ItemTrueValueSub");
                using (Variable.ForEach(KN))
                {
                    UserObservedValue[K][KN] = Variable.GaussianFromMeanAndPrecision((trueValue[KN] + UserBias[K]).Named("BiasedMean"), (UserObservedPrecision[K][KN] * UserTruePrecision[K]).Named("ScaledPrecision"));
                }
            }
        }

        protected void SetPriors(int userCount, int communityCount)
        {
            base.SetPriors();
            UserTrustPrior.ClearObservedValue();
            CommunityCount.ObservedValue = communityCount;
            CommunityBiasPrior.ObservedValue = Util.ArrayInit(communityCount, t => Gaussian.FromMeanAndPrecision(0, 1));
            CommunityTrustPrior.ObservedValue = Util.ArrayInit(communityCount, k => Gaussian.FromMeanAndPrecision(1, 1));
            CommunityProbsPrior.ObservedValue = Dirichlet.Symmetric(communityCount, CommunityPseudoCount);
            CommunityInitializer.ObservedValue = Util.ArrayInit(userCount, user => Discrete.PointMass(Rand.Int(communityCount), communityCount));
        }

        public BACEPosteriors Infer(int[][] itemIndices, double[][] userValues, double[][] userPrecisions, int communityCount)
        {
            int workerCount = userValues.Length;
            SetPriors(workerCount, communityCount);
            AttachData(itemIndices, userValues, userPrecisions);
            var result = new CBACEPosteriors();
            Engine.NumberOfIterations = NumberOfIterations;
            result.TrueItemValuePosterior = Engine.Infer<Gaussian[]>(ItemTrueValue);
            result.UserTrustLogPosterior = Engine.Infer<Gaussian[]>(UserTrustLog);
            result.UserBiasPosterior = Engine.Infer<Gaussian[]>(UserBias);
            result.CommunityBiasPosterior = Engine.Infer<Gaussian[]>(CommunityBias);
            result.CommunityTrustPosterior = Engine.Infer<Gaussian[]>(CommunityTrust);
            result.CommunityPosterior = Engine.Infer<Discrete[]>(Community);
            result.CommunityProbsPosterior = Engine.Infer<Dirichlet>(CommunityProbs);
            result.LogEvidence = Engine.Infer<Bernoulli>(Evidence).LogOdds;

            return result;
        }
    }

    /// <summary>
    /// The CBACE posteriors
    /// </summary>
    [Serializable]
    public class CBACEPosteriors : BACEPosteriors
    {
        public Gaussian[] UserBiasPosterior;
        public Gaussian[] UserTrustLogPosterior;
        public Gaussian[] CommunityBiasPosterior;
        public Gaussian[] CommunityTrustPosterior;
        public Discrete[] CommunityPosterior;
        public Dirichlet CommunityProbsPosterior;

    }
	
	/// <summary>
    /// The CBACE model implemented as a sub-class of CBACE bias
    /// </summary>
    class CBACE : CBACEBias
    {
        /// <summary>
        /// The generative process of CBACE
        /// </summary>
        protected override void DefineGenerativeProcess()
        {
            // The process that generates the user's reports
            using (Variable.ForEach(K))
            {
                using (Variable.Switch(Community[K]))
                {
                    UserBias[K] = Variable.GaussianFromMeanAndPrecision(CommunityBias[Community[K]], BiasPrecision);
                    UserTrustLog[K] = Variable.GaussianFromMeanAndPrecision(CommunityTrust[Community[K]], TrustPrecision);
                }
                UserTruePrecision[K] = Variable.Exp(UserTrustLog[K]);

                var trueValue = Variable.Subarray(ItemTrueValue, UserItemIndex[K]).Named("ItemTrueValueSub");
                using (Variable.ForEach(KN))
                {
                    UserObservedValue[K][KN] = Variable.GaussianFromMeanAndPrecision((trueValue[KN]), (UserObservedPrecision[K][KN] * UserTruePrecision[K]).Named("ScaledPrecision"));
                }
            }
        }  
    }
}