namespace buildtimefun;

using System;
using System.Linq;
using MathNet.Numerics;
using MathNet.Numerics.Distributions;
using MathNet.Numerics.LinearAlgebra;
using MathNet.Numerics.Random;
using MathNet.Numerics.Statistics;

class Program
{
    public const int N = 50;

    public static void Main(string[] args)
    {
        Diagnose();
        double[] buildTimes = Build(N);
        PrintRaw(buildTimes);
        PrintDensity(buildTimes);

        // Run 500 iterations of expectation maximisation, before printing
        // results.
        var components = new Iteration(buildTimes, 2);
        for (int i = 0; i < 500; i++)
        {
            components = EM(components);
        }

        Console.WriteLine($"p = {components.Fraction[1]}");
        Console.WriteLine("Mu:\n" + components.Mean);
        Console.WriteLine("Var:\n" + components.Variance);
    }

    // The input and output to one iteration of expectation maximisation.
    // The Values vector contains the underlying data, and the other vectors
    // contain parameter estimations for each of the estimated components.
    private struct Iteration
    {
        public Vector<double> Values;
        public Vector<double> Fraction;
        public Vector<double> Mean;
        public Vector<double> Variance;

        // Create initial data for the first iteration, given some values and a
        // desire to estimate n Gaussian components.
        public Iteration(double[] values, int n)
        {
            this.Values = Vector.Dense(values);
            this.Fraction = Vector.Dense(n, 1.0 / n);

            double[] means = values.SelectCombination(n).ToArray();
            this.Mean = Vector.Dense(means);

            double variance = Statistics.Variance(values)/20.0;
            this.Variance = Vector.Dense(n, variance);
        }
    }

    // Get the likelihoods of observing the values given this iteration of
    // estimated parameters for the Gaussian components.
    private static Matrix<Double> Evaluate(Iteration components)
    {
        return Matrix.Dense(
            components.Values.Count,
            components.Fraction.Count,
            (i, d) => Normal.PDF(
                components.Mean[d],
                Math.Sqrt(components.Variance[d]),
                components.Values[i]
            )
        );
    }

    private static Matrix<double> Rows(Vector<double> row, int n)
    {
        return Matrix.DenseOfRowVectors(
            Enumerable.Repeat(row, n)
        );
    }

    private static Matrix<double> Columns(Vector<double> column, int n)
    {
        return Matrix.DenseOfColumnVectors(
            Enumerable.Repeat(column, n)
        );
    }

    // Perform one iteration of expectation maximisation.
    private static Iteration EM(Iteration components)
    {
        var n = components.Values.Count;
        var d = components.Fraction.Count;

        var probability = Evaluate(components);
        var gamma = probability.PointwiseMultiply(
            Rows(components.Fraction, n)
        ).NormalizeRows(1);

        var fraction = gamma.ColumnSums() / n;

        var mean = gamma.Transpose().Multiply(
            components.Values.ToColumnMatrix()
        ).PointwiseDivide(
            gamma.ColumnSums().ToColumnMatrix()
        );

        var means = Columns(mean.Column(0), n);
        var remainders = Rows(components.Values, d)
            .Subtract(means)
            .PointwisePower(2);
        var variance = gamma.Transpose().PointwiseMultiply(
            remainders
        ).RowSums().ToColumnMatrix().PointwiseDivide(
            gamma.ColumnSums().ToColumnMatrix()
        );

        return new Iteration
        {
            Values = components.Values,
            Fraction = fraction,
            Mean = mean.Column(0),
            Variance = variance.Column(0)
        };
    }

    private static void PrintRaw(double[] times)
    {
        Console.WriteLine("====== RAW ======");
        foreach (var time in times)
        {
            Console.WriteLine($"{time:0.0}");
        }
    }

    private static void PrintDensity(double[] times)
    {
        Console.WriteLine("====== DENSITY ======");
        Array.Sort(times);
        var buckets = times.GroupBy(obs => (int) obs);
        foreach (var bucket in buckets)
        {
            Console.Write($"{bucket.Key}|");
            foreach (var time in bucket)
            {
                var tenth = (time - bucket.Key)*10;
                Console.Write($"{tenth:0}");
            }
            Console.WriteLine("");
        }
    }

    private static double[] Build(int samples)
    {
        return BuildSuccess.Samples()
            .Take(samples)
            .Select(d => BuildTime[d].Sample())
            .ToArray();
    }

    private static void Diagnose()
    {
        Console.WriteLine("====== DIAGNOSTICS ======");
        foreach (var distribution in BuildTime)
        {
            Console.WriteLine(
                $"Distribution = {distribution.GetType()}"
            );
            PrintDensity(
                distribution.Samples().Take(500).ToArray()
            );
        }
    }

    private static readonly
    IContinuousDistribution[] BuildTime =
        new IContinuousDistribution[]
    {
        new BetaScaled(a: 3.0, b: 5.0, location: 0.0, scale: 43.0),
        new Normal(mean: 37.0, stddev: 1.5)
    };

    private static readonly Bernoulli BuildSuccess =
        new Bernoulli(p: 0.72);

    private static readonly MatrixBuilder<double> Matrix =
        Matrix<Double>.Build;

    private static readonly VectorBuilder<double> Vector =
        Vector<Double>.Build;

}