package quantiles

import (
	"math/rand"
	"time"
)

// The frugal estimator computes quantiles using very few resources at
// the cost of accuracy. It also cannot compute arbitrary quantiles;
// the target quantile must be chosen in advance. The parameter q sets
// that quantile; thereafter, every call to QuantileEstimator.Quantile
// will return the qth quantile, regardless of the value passed in.
//
// Algorithmic details come from 'Frugal Streaming for Estimating
// Quantiles', Qiang Ma, S. Muthukrishnan, Mark Sandler, 2014
// (http://arxiv.org/abs/1407.1121). This is the 2-unit frugal
// estimator, in the terms of that paper.
func NewFrugalQuantileEstimator(q float64) QuantileEstimator {
	if q == 0.5 {
		return &frugalMedian{
			estimate:     0,
			step:         0,
			lastWasAbove: false,
			initialized:  false,
		}
	} else {
		return &frugal{
			target:       q,
			estimate:     0,
			step:         0,
			lastWasAbove: false,
			initialized:  false,
			rng:          rand.New(rand.NewSource(time.Now().UnixNano())),
		}
	}
}

func InitializedFrugalQuantileEstimator(q float64, est, step float64) QuantileEstimator {
	if q == 0.5 {
		return &frugalMedian{
			estimate:     est,
			step:         step,
			lastWasAbove: false,
			initialized:  true,
		}
	} else {
		return &frugal{
			target:       q,
			estimate:     est,
			step:         step,
			lastWasAbove: false,
			initialized:  true,
			rng:          rand.New(rand.NewSource(time.Now().UnixNano())),
		}
	}
}

type frugal struct {
	// The quantile being estimated (eg '0.5' for median)
	target float64

	// The current estimated value of the target quantile
	estimate float64

	// How dramatically should we change our estimate when we receive new data?
	step float64

	// was the last added value above our previous estimate?
	lastWasAbove bool

	// Have we seen at least one value?
	initialized bool

	rng *rand.Rand
}

func (f *frugal) Add(val float64, weight int) {
	// Initialize with the first value we ever see.
	if !f.initialized {
		f.step = (val - f.estimate) / 2
		f.estimate = val
		f.initialized = true
		return
	}

	r := f.rng.Float64()

	if val > f.estimate && r > 1-f.target {
		if f.lastWasAbove {
			// We're continuing to see data above our old estimate. Increase
			// the step size.
			f.step += 1
		} else {
			// We've reversed direction! slow down our steps.
			f.step -= 1
		}

		f.lastWasAbove = true

		// Adjust by our step size, with a minimum change of 1.0.
		if f.step > 0 {
			f.estimate += f.step
		} else {
			f.estimate += 1
		}

		// Did we overshoot our mark?
		delta := f.estimate - val
		if delta > 0 {
			// Yep. Let's slow down a bit.
			f.step -= delta
			f.estimate = val
		}
		// And if the value is below our old estimate, follow all the same
		// instructions, but with the sign reversed.
	} else if val < f.estimate && r > f.target {
		if !f.lastWasAbove {
			f.step += 1
		} else {
			f.step -= 1
		}

		f.lastWasAbove = false

		if f.step > 0 {
			f.estimate -= f.step
		} else {
			f.estimate -= 1
		}

		delta := f.estimate - val
		if delta < 0 {
			f.step += delta
			f.estimate = val
		}
	}
}

func (f *frugal) Quantile(_ float64) float64 {
	return f.estimate
}

// frugalMedian is a special case of frugal when target is 0.5. This
// permits some optimization, since an RNG is no longer needed.
type frugalMedian struct {
	// The current estimated value of the target quantile
	estimate float64

	// How dramatically should we change our estimate when we receive new data?
	step float64

	// was the last added value above our previous estimate?
	lastWasAbove bool

	initialized bool
}

func (f *frugalMedian) Add(val float64, weight int) {
	// Initialize with the first value we ever see.
	if !f.initialized {
		f.estimate = val
		f.initialized = true
		return
	}

	if val > f.estimate {
		if f.lastWasAbove {
			// We're continuing to see data above our old estimate. Increase
			// the step size.
			f.step += 1
		} else {
			// We've reversed direction! slow down our steps.
			f.step -= 1
		}

		f.lastWasAbove = true

		// Adjust by our step size, with a minimum change of 1.0.
		if f.step > 0 {
			f.estimate += f.step
		} else {
			f.estimate += 1
		}

		// Did we overshoot our mark?
		delta := f.estimate - val
		if delta > 0 {
			// Yep. Let's slow down a bit.
			f.step -= delta
			f.estimate = val
		}
	} else if val < f.estimate {
		if !f.lastWasAbove {
			// We're continuing to see data below our old estimate. Increase
			// the step size.
			f.step += 1
		} else {
			// We've reversed direction! slow down our steps.
			f.step -= 1
		}

		f.lastWasAbove = false

		// Adjust by our step size, with a minimum change of 1.0.
		if f.step > 0 {
			f.estimate -= f.step
		} else {
			f.estimate -= 1
		}

		// Did we overshoot our mark?
		delta := f.estimate - val
		if delta < 0 {
			// Yep. Let's slow down a bit.
			f.step += delta
			f.estimate = val
		}
	}
}

func (f *frugalMedian) Quantile(_ float64) float64 {
	return f.estimate
}
