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

public class Locus : IEnumerable<double>
{
   private readonly HashSet<double> mutations;

   public Locus()
   {
      this.mutations = new HashSet<double>();
   }

   public Locus(IEnumerable<double> mutations)
   {
      this.mutations = new HashSet<double>(mutations);
   }

   public static Locus operator |(Locus c1, Locus c2)
   {
      if (c1 == null)
      {
         throw new ArgumentNullException("c1");
      }
      if (c2 == null)
      {
         throw new ArgumentNullException("c2");
      }
      return new Locus(c1.mutations.Union(c2.mutations));
   }

   public void AddMutation(double site)
   {
      this.mutations.Add(site);
   }

   public bool IsMutant(double site)
   {
      return this.mutations.Contains(site);
   }

   public IEnumerator<double> GetEnumerator()
   {
      return this.mutations.GetEnumerator();
   }
   System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
   {
      return GetEnumerator();
   }
}

public class CoalescentTreeNode
{
   public string Name { get; set; }

   private CoalescentTree tree;
   public CoalescentTree Tree
   {
      get { return this.tree; }
      internal set { this.tree = value; }
   }

   private CoalescentTreeNode left;
   public CoalescentTreeNode DescendantLeft
   {
      get { return this.left; }
      internal set { this.left = value; }
   }

   private CoalescentTreeNode right;
   public CoalescentTreeNode DescendantRight
   {
      get { return this.right; }
      internal set { this.right = value; }
   }

   private CoalescentTreeNode ancestor;
   public CoalescentTreeNode Ancestor
   {
      get { return this.ancestor; }
      internal set { this.ancestor = value; }
   }

   private double time;
   public double WaitingTime
   {
      get { return this.time; }
      internal set { this.time = value; }
   }

   private Locus locus;
   public Locus Locus
   {
      get
      {
         if (this.locus == null)
         {
            this.locus = new Locus();
         }
         return this.locus;
      }
      internal set { this.locus = value; }
   }

   public bool IsLeaf
   {
      get
      {
         return DescendantLeft == null && DescendantRight == null;
      }
   }

   internal bool HasDescendant
   {
      get
      {
         return DescendantLeft != null && DescendantRight != null;
      }
   }

   public CoalescentTreeNode()
   {
   }

   public void AddMutation(double site)
   {
      foreach (var node in EnumerateAll(this))
      {
         node.Locus.AddMutation(site);
      }
   }

   internal static IEnumerable<CoalescentTreeNode> EnumerateAll(CoalescentTreeNode baseNode)
   {
      if (baseNode == null)
      {
         throw new ArgumentNullException("baseNode");
      }

      yield return baseNode;
      if (baseNode.DescendantLeft != null)
      {
         foreach (var descendant in EnumerateAll(baseNode.DescendantLeft))
         {
            yield return descendant;
         }
      }
      if (baseNode.DescendantRight != null)
      {
         foreach (var descendant in EnumerateAll(baseNode.DescendantRight))
         {
            yield return descendant;
         }
      }
   }

   internal static IEnumerable<CoalescentTreeNode> EnumerateLeaves(CoalescentTreeNode baseNode)
   {
      if (baseNode == null)
      {
         throw new ArgumentNullException("baseNode");
      }

      CoalescentTreeNode node = baseNode;
      while (node.DescendantLeft != null)
      {
         node = node.DescendantLeft;
      }
      yield return node;
      while (node != baseNode)
      {
         node = node.Ancestor;
         if (node.DescendantRight != null)
         {
            foreach (var value in EnumerateLeaves(node.DescendantRight))
            {
               yield return value;
            }
         }
      }
   }

   internal double GetTotalTime()
   {
      double sum = 0.0;
      CoalescentTreeNode current = DescendantLeft;
      while (current != null)
      {
         sum += current.WaitingTime;
         current = current.DescendantLeft;
      }
      return sum;
   }
}

public class CoalescentTree : IEnumerable<CoalescentTreeNode>, IFormattable
{
   private readonly List<CoalescentTreeNode> nodes;

   public CoalescentTreeNode Root
   {
      get { return IsIncomplete ? null : this.nodes[0]; }
   }

   public bool IsIncomplete
   {
      get { return this.nodes.Count != 1; }
   }

   public CoalescentTree()
   {
      this.nodes = new List<CoalescentTreeNode>();
   }

   internal CoalescentTree(IEnumerable<CoalescentTreeNode> nodes)
   {
      this.nodes = new List<CoalescentTreeNode>(nodes);
   }

   public CoalescentTreeNode Coalesce(CoalescentTreeNode left, CoalescentTreeNode right, double time)
   {
      if (!Contains(left) || left.Ancestor != null)
      {
         throw new ArgumentException("left");
      }
      if (!Contains(right) || right.Ancestor != null)
      {
         throw new ArgumentException("right");
      }
      double leftTime = left.GetTotalTime();
      if (leftTime > time)
      {
         throw new ArgumentException("left");
      }
      double rightTime = right.GetTotalTime();
      if (rightTime > time)
      {
         throw new ArgumentException("right");
      }
      
      left.WaitingTime = time - leftTime;
      right.WaitingTime = time - rightTime;

      CoalescentTreeNode ancestor = new CoalescentTreeNode()
      {
         DescendantLeft = left,
         DescendantRight = right,
         Locus = left.Locus | right.Locus,
      };
      left.Ancestor = ancestor;
      right.Ancestor = ancestor;

      Remove(left);
      Remove(right);
      Add(ancestor);
      return ancestor;
   }

   public void Add(CoalescentTreeNode node)
   {
      if (node == null)
      {
         throw new ArgumentNullException("node");
      }
      if (node.Tree != null)
      {
         throw new InvalidOperationException();
      }
      this.nodes.Add(node);
      node.Tree = this;
   }

   private bool Remove(CoalescentTreeNode node)
   {
      return this.nodes.Remove(node);
   }

   private bool Contains(CoalescentTreeNode node)
   {
      return this.nodes.Contains(node);
   }

   internal IList<CoalescentTreeNode> GetNodes()
   {
      return new ReadOnlyCollection<CoalescentTreeNode>(nodes);
   }

   public override string ToString()
   {
      return ToString(null, null);
   }

   public string ToString(string format, IFormatProvider formatProvider)
   {
      if (format != null && Char.ToUpper(format[0]) == 'A')
      {
         return ToAlignment();
      }
      else
      {
         if (Root.HasDescendant)
         {
            return string.Format("({0},{1});", ToString(Root.DescendantLeft, format, formatProvider), ToString(Root.DescendantRight, format, formatProvider));
         }
         return string.Format("{0}:{1};", Root.Name, Root.WaitingTime);
      }
   }

   private string ToString(CoalescentTreeNode node, string format, IFormatProvider formatProvider)
   {
      if (node == null)
      {
         return string.Empty;
      }
      if (node.HasDescendant)
      {
         return string.Format("({0},{1}):{2}", ToString(node.DescendantLeft, format, formatProvider), ToString(node.DescendantRight, format, formatProvider), node.WaitingTime.ToString(format, formatProvider));
      }
      return string.Format("{0}:{1}", node.Name, node.WaitingTime.ToString(format, formatProvider));
   }

   private string ToAlignment()
   {
      StringBuilder alignment = new StringBuilder();
      HashSet<double> allMutations = new HashSet<double>(this.SelectMany(node => node.Locus));
      foreach (var node in this)
      {
         if (alignment.Length != 0)
         {
            alignment.AppendLine();
         }
         StringBuilder sequence = new StringBuilder();
         foreach (var site in allMutations)
         {
            sequence.Append(node.Locus.IsMutant(site) ? '1' : '0');
         }
         alignment.Append(node.Name).Append(": ").Append(sequence.ToString());
      }
      return alignment.ToString();
   }

   public IEnumerator<CoalescentTreeNode> GetEnumerator()
   {
      if (IsIncomplete)
      {
         throw new InvalidOperationException();
      }
      foreach (var node in CoalescentTreeNode.EnumerateLeaves(Root))
      {
         yield return node;
      }
   }
   System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
   {
      return GetEnumerator();
   }
}

public class CoalescentSimulator
{
   private readonly Random random;

   public CoalescentSimulator()
   {
      this.random = new Random();
   }

   public CoalescentTree PerformSimulation(IEnumerable<CoalescentTreeNode> samples, double theta = 0.0)
   {
      if (theta < 0)
      {
         throw new ArgumentOutOfRangeException("theta");
      }
      CoalescentTree tree = new CoalescentTree(samples);

      double time = 0.0;
      while (tree.IsIncomplete)
      {
         IList<CoalescentTreeNode> nodes = tree.GetNodes();
         int count = nodes.Count;
         time = GenerateNextCoalescentEventTime(time, count);
         int leftIndex;
         int rightIndex;
         GetRandomTwoIndexes(count, out leftIndex, out rightIndex);
         CoalescentTreeNode leftNode = nodes[leftIndex];
         CoalescentTreeNode rightNode = nodes[rightIndex];
         tree.Coalesce(leftNode, rightNode, time);
         AddMutation(leftNode, theta);
         AddMutation(rightNode, theta);
      }

      return tree;
   }

   private void GetRandomTwoIndexes(int count, out int index1, out int index2)
   {
      index1 = random.Next(count);
      index2 = random.Next(count - 1);
      if (index2 >= index1)
      {
         ++index2;
      }
   }

   private double GenerateNextCoalescentEventTime(double currentTime, int sampleCount)
   {
      double rate = (double)sampleCount * (sampleCount - 1);
      double waitingTime = random.NextExponential(rate);
      return currentTime + waitingTime;
   }

   private void AddMutation(CoalescentTreeNode node, double theta)
   {
      int count = random.NextPoisson(node.WaitingTime * theta);
      while (--count >= 0)
      {
         double site = random.NextDouble();
         node.AddMutation(site);
      }
   }
}

public static class RandomExtension
{
   public static double NextExponential(this Random random, double rate)
   {
      double r;
      do
      {
         r = random.NextDouble();
      } while (r == 0.0);
      return -Math.Log(r) / rate;
   }

   public static int NextPoisson(this Random random, double mean)
   {
      double lambda = Math.Exp(mean);
      int k = 0;
      if (!double.IsInfinity(lambda))
      {
         while ((lambda *= random.NextDouble()) > 1.0)
         {
            k++;
         }
      }
      else
      {
         lambda = mean;
         while ((lambda += Math.Log(random.NextDouble())) > 0.0)
         {
            k++;
         }
      }
      return k; 
   }
}

class Program
{
   static void Main()
   {
      CoalescentSimulator simulator = new CoalescentSimulator();
      var nodes = Enumerable.Range(0, 6).Select(i => new CoalescentTreeNode() { Name = (i + 1).ToString() }).ToArray();
      CoalescentTree tree = simulator.PerformSimulation(nodes, 5.0);
      Console.WriteLine(tree.ToString("0.000", null));
      Console.WriteLine(tree.ToString("A", null));
   }
}