using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace StateMonad
{
    public class Unit
    {
        public static readonly Unit Value = new Unit();

        private Unit () { }
    }

    /// <summary>
    /// State-content pair
    /// </summary>
    /// <typeparam name="S">Type of State</typeparam>
    /// <typeparam name="C">Type of Content</typeparam>
    public class Scp<S,C>
    {
        public Scp(S state, C content)
        {
            State = state;
            Content = content;
        }

        public S State { get; private set; }
        public C Content { get; private set; }

        public override string ToString()
        {
            return "{" + State + ", " + Content + "}";
        }
    }

    /// <summary>
    /// State monad
    /// </summary>
    /// <typeparam name="S">Type of State</typeparam>
    /// <typeparam name="C">Type of Content</typeparam>
    public class StateMonad<S, C>
    {
        public StateMonad(Func<S, Scp<S, C>> s2Scp)
        {
            S2Scp = s2Scp;
        }

        /// <summary>
        /// State to state-content pair
        /// </summary>
        public Func<S,Scp<S,C>> S2Scp { get; private set; }

    }

    /// <summary>
    /// State monad operations
    /// </summary>
    public static class StateMonad
    {
        /// <summary>
        /// "return" monadic operator
        /// </summary>
        /// <typeparam name="S">Type of State</typeparam>
        /// <typeparam name="C">Type of Content</typeparam>
        /// <param name="content">Content</param>
        /// <returns>State monad</returns>
        public static StateMonad<S, C> Return<S, C>(C content)
        {
            return new StateMonad<S, C>(s => new Scp<S, C>(s, content));
        }

        /// <summary>
        /// "return" monadic operator
        /// </summary>
        /// <typeparam name="S">Type of State</typeparam>
        /// <typeparam name="C">Type of Content</typeparam>
        /// <param name="dummy">Dummy argument for type inference</param>
        /// <param name="content">Content</param>
        /// <returns>State monad</returns>
        public static StateMonad<S, C> Return<S, C>(S dummy, C content)
        {
            return new StateMonad<S, C>(s => new Scp<S, C>(s, content));
        }

        /// <summary>
        /// "bind" monadic operator
        /// </summary>
        /// <param name="m">State monad</param>
        /// <param name="selector">Selector function</param>
        /// <returns>State monad</returns>
        public static StateMonad<S, C2> SelectMany<S, C1, C2>(this StateMonad<S, C1> m, Func<C1, StateMonad<S, C2>> selector)
        {
            return new StateMonad<S, C2>(s =>
                {
                    var scp = m.S2Scp(s);
                    return selector(scp.Content).S2Scp(scp.State);
                });
        }

        /// <summary>
        /// "bind" monadic operator
        /// </summary>
        /// <param name="m">State monad</param>
        /// <param name="selector">Selector function</param>
        /// <param name="result">Map result function</param>
        /// <returns>State monad</returns>
        public static StateMonad<S, C2> SelectMany<S, C1, C2, K>(this StateMonad<S, C1> m, Func<C1, StateMonad<S, K>> selector, Func<C1, K, C2> result)
        {
            return m.SelectMany(x => selector(x).SelectMany(y => Return<S,C2>(result(x, y))));
        }

        /// <summary>
        /// Select method for Query Pattern
        /// </summary>
        /// <param name="m">State monad</param>
        /// <param name="selector">Selector function</param>
        /// <returns>State monad</returns>
        public static StateMonad<S, C2> Select<S, C1, C2>(this StateMonad<S, C1> m, Func<C1, C2> selector)
        {
            return new StateMonad<S, C2>(s => 
            {
                var scp = m.S2Scp(s);
                return new Scp<S, C2>(scp.State, selector(scp.Content));
            });
        }

        /// <summary>
        /// State monad "get state" function
        /// </summary>
        /// <typeparam name="S">Type of state</typeparam>
        public static StateMonad<S, S> Get<S>()
        {
            return new StateMonad<S, S>(s => new Scp<S, S>(s, s));
        }

        /// <summary>
        /// State monad "put state" function
        /// </summary>
        /// <typeparam name="S">Type of state</typeparam>
        /// <param name="s">New state</param>
        public static StateMonad<S, Unit> Put<S>(S s)
        {
            return new StateMonad<S, Unit>(_ => new Scp<S, Unit>(s, Unit.Value));
        }
    }

    /// <summary>
    /// Tree constructor
    /// </summary>
    public static class Tr
    {
        /// <summary>
        /// Creates new Leaf
        /// </summary>
        /// <typeparam name="T">Type of content</typeparam>
        /// <param name="content">Content of Leaf</param>
        /// <returns>Binary tree</returns>
        public static Tr<T> Lf<T>(T content)
        {
            return new Lf<T>(content);
        }

        /// <summary>
        /// Creates new Branch
        /// </summary>
        /// <typeparam name="T">Type of content</typeparam>
        /// <param name="left">Left subtree</param>
        /// <param name="right">Right subtree</param>
        /// <returns>Binary tree</returns>
        public static Tr<T> Br<T>(Tr<T> left, Tr<T> right)
        {
            return new Br<T>(left, right);
        }
    }

    /// <summary>
    /// Base class for binary tree node
    /// </summary>
    /// <typeparam name="T">Type of content</typeparam>
    public abstract class Tr<T>
    {
        protected const int Ident = 2;

        /// <summary>
        /// Print node
        /// </summary>
        /// <param name="level"></param>
        public abstract void Show(int level);
    }

    /// <summary>
    /// Branch of binary tree
    /// </summary>
    /// <typeparam name="T">Type of content</typeparam>
    public class Br<T> : Tr<T>
    {
        public Br(Tr<T> left, Tr<T> right)
        {
            Left = left;
            Right = right;
        }

        /// <summary>
        /// Left subtree
        /// </summary>
        public Tr<T> Left { get; private set; }

        /// <summary>
        /// Right subtree
        /// </summary>
        public Tr<T> Right { get; private set; }

        public override void Show(int level)
        {
            Console.Write(new String(' ', level * Ident));
            Console.WriteLine("Branch:");
            Left.Show(level + 1);
            Right.Show(level + 1);
        }
    }

    /// <summary>
    /// Leaf of binary tree
    /// </summary>
    /// <typeparam name="T">Type of content</typeparam>
    public class Lf<T> : Tr<T>
    {
        public Lf(T content)
        {
            Content = content;
        }

        /// <summary>
        /// Content
        /// </summary>
        public T Content { get; private set; }

        public override void Show(int level)
        {
            Console.Write(new String(' ', level * Ident));
            Console.Write("Leaf: ");
            Console.Write(Content);
            Console.WriteLine();
        }
    }

    class Program
    {
        /// <summary>
        /// Monadic tree labeling
        /// </summary>
        /// <typeparam name="T">Type of content</typeparam>
        /// <param name="tree">Binary tree</param>
        /// <returns>State monad with binary tree with label-contents pair</returns>
        static StateMonad<int, Tr<Scp<int,T>>> GetLabelMonad<T>(Tr<T> tree)
        {
            if (tree is Lf<T>)
            {
                var leaf = tree as Lf<T>;
                var m = from s in StateMonad.Get<int>()
                        from _ in StateMonad.Put(s + 1)
                        from r in StateMonad.Return(0 /*dummy arg*/, Tr.Lf(new Scp<int, T>(s, leaf.Content)))
                        select r;
                return m;
            }
            else if(tree is Br<T>)
            {
                var br = tree as Br<T>;
                var m = from l in GetLabelMonad(br.Left)
                        from r in GetLabelMonad(br.Right)
                        from res in StateMonad.Return(0 /*dummy arg*/, Tr.Br(l, r))
                        select res;
                return m;
            }
            else
            {
                throw new Exception("Invalid object type");
            }
        }

        static void Main(string[] args)
        {
            var tree = 
                Tr.Br
                (
                    Tr.Br(
                        Tr.Lf(1),
                        Tr.Br(Tr.Lf(2), Tr.Lf(3))
                         ),
                    Tr.Lf(4)
                );
            tree.Show(0);
            Console.WriteLine();

            var m = GetLabelMonad(tree);

            var res = m.S2Scp(0);
            
            Console.WriteLine(res.State);
            res.Content.Show(0);

            Console.ReadLine();
        }
    }
}