using System;
using Excel = Microsoft.Office.Interop.Excel;

namespace Binomial
{
    class SOEURecursiveUtility
    {
        static void Main() {
            double alpha = 40; //CRUA, utility will be computed for a range of values
            double beta = 0.03; //rate of impatience
            double gamma = 40; //CRRA
            double delta = 1/1.5; //inverse EIS
            double c0 = 100; //time-zero consumption
            double muHigh = 0.05; //high drift scenario
            double muLow = 0.00; //low drift scenario
            double sigma = 0.03; //known consumption volatility
            int nPeriods = 120; //number of periods
            double h = 10 / (double)nPeriods; //period time-length = time horizon / nPeriods
            double qHigh = (1 + (muHigh / sigma) * Math.Sqrt(h)) / 2; //optimistic up prob
            double qLow = (1 + (muLow / sigma) * Math.Sqrt(h)) / 2; //pessimistic up prob
            double upcgrowth = Math.Exp(sigma * Math.Sqrt(h)); //up consumption growth
            double downcgrowth = 1 / upcgrowth; //down consumption growth
            double discount = Math.Exp(-beta*h);
           
            //build the consumption tree
            BinomialTree c = c0.ConstVolTree(nPeriods, upcgrowth, downcgrowth);
            Console.WriteLine("There are {0} periods and {1} nodes.", c.Periods, c.Nodes);
            double cmin = c0 * Math.Pow(downcgrowth, nPeriods);
            Console.WriteLine("Worst-case terminal consumption is {0}.", cmin);
           
            //construct the CE and time aggregator
            var hi = new SIEU(qHigh, gamma); //optimistic SI EU
            var lo = new SIEU(qLow, gamma); //pessimistic SI EU
            var so = new SIEU(0.5, alpha); //second-order SI EU
            Aggregator CE = (x, y) => so.CE(hi.CE(x, y), lo.CE(x, y)); //overall CE
            var tagg = new SIEU(discount, delta); //the time aggregator is tagg.CE
            BinomialTree U = c.BackwardRecursiveTree(tagg.CE, CE);
            Console.WriteLine("For CRUA = {0}, the time-zero utility is {1}.", so.CRRA, U[0, 0]);
            //uncomment following line to output consumption and utility trees to spreadsheet
            //TreeBuilder.ToSpreadsheet(c, U);
            Console.ReadKey();
        }
    }

    public class SIEU //Scale-Invariant Expected Utility
    {
        double p;

        public SIEU(double upProbability, double CRRA) {
            this.upProbability = upProbability;
            this.CRRA = CRRA;
        }
        public double upProbability {
            get { return p; }
            set { if ((value < 0) || (value > 1)) 
                    throw new ArgumentException("probability must be in (0,1). ", "upProbability");
                  p = value;
            }
        }
        public double CRRA {get; set;}
        
        public double CE(double x, double y) {
            if (CRRA == 1)
                return Math.Exp(p * Math.Log(x) + (1 - p) * Math.Log(y));
            else { 
                var power = 1 - CRRA;
                return Math.Pow(p * Math.Pow(x, power) + (1 - p) * Math.Pow(y, power), 1 / power);
            }
        }
    }

    public delegate double Aggregator(double up, double down);

    public class BinomialTree
    {
        readonly int T; // the number of periods, time runs through 0,1,...,T
        readonly int N; // the number of nodes
        double[] values; // lexicographically by (time, ups): (0,0), (1,0), (1,1), ...
        static int node(int t, int ups) { //map (t,ups) to node index of the values array
            return (t * (t + 1) / 2) + ups;
        }
        public BinomialTree(int nPeriods) {
            T = nPeriods;
            N = 1 + node(T, T);
            values = new double[N];
        }
        public int Periods {
            get { return T; }
        }
        public int Nodes {
            get { return N; }
        }
        public double this[int t, int u] {
            set {
                if (u > t) throw new ArgumentException
                    ("the number of time t nodes cannot exceed t", "ups");
                values[node(t, u)] = value;
            }
            get {
                if (u > t) throw new ArgumentException
                        ("the number of time t nodes cannot exceed t", "ups");
                return values[node(t, u)];
            }
        }
    }

    public static class TreeBuilder
    {
        public static BinomialTree ConstVolTree(this double initial, int nPeriods, double upReturn, double downReturn) {
            var C = new BinomialTree(nPeriods);
            C[0, 0] = initial;
            for (int t = 1; t <= nPeriods; t++) {
                C[t, 0] = C[t - 1, 0] * downReturn;
                for (int ups = 1; ups <= t; ups++)
                    C[t, ups] = C[t - 1, ups - 1] * upReturn;
            }
            return C;
        }

        public static BinomialTree BackwardRecursiveTree(this BinomialTree cons, Aggregator tAgg, Aggregator CE) {
            int N = cons.Periods;
            var U = new BinomialTree(N);
            for (int ups = 0; ups <= N; ups++)
                U[N, ups] = cons[N, ups];
            for (int t = N - 1; t >= 0; t--) {
                for (int ups = 0; ups <= t; ups++)
                    U[t, ups] = tAgg(CE(U[t + 1, ups + 1], U[t + 1, ups]), cons[t, ups]);
            }
            return U;
        }

        public static void ToSpreadsheet(BinomialTree c, BinomialTree U) {
            int N = U.Periods;
            var excel = new Excel.Application();
            excel.Visible = true;
            var workBook = excel.Workbooks.Add();
            for (int t = 0; t <= N; t++)
                for (int ups = 0; ups <= t; ups++) {
                    int downs = t - ups;
                    excel.Cells[1 + downs, 1 + t].Value = U[t, ups];
                    excel.Cells[N + 3 + downs, 1 + t].Value = c[t, ups];
                }
            // workBook.SaveAs(@"C:\Users\...\RecursiveUtility.xlsx");
        }
    }

}