// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package privacy

import (
	"config"
	"fmt"
	"math"
)

type ErrorCalculator struct {
	ParamsCalc PrivacyEncodingParamsCalculator
}

// Public factory method for creating an ErrorCalculator given a
// PrivacyEncodingParamsCalculator.
func NewErrorCalculator(paramsCalc PrivacyEncodingParamsCalculator) *ErrorCalculator {
	return &ErrorCalculator{paramsCalc}
}

// Public factory method for creating an ErrorCalculator given the file path
// of the PrivacyEncodingParams.
func NewErrorCalculatorFromPrivacyParams(privacyParamsPath string) (*ErrorCalculator, error) {
	paramsCalculator, err := NewPrivacyEncodingParamsCalculator(privacyParamsPath)
	if err != nil {
		return nil, err
	}
	errorCalculator := NewErrorCalculator(*paramsCalculator)
	if err != nil {
		return nil, err
	}
	return errorCalculator, nil
}

// Given a |metric|, |report|, and |params|, estimates the report row error.
//
// TODO(jaredweinstein): implement estimates for other report types.
func (e *ErrorCalculator) Estimate(metric *config.MetricDefinition, report *config.ReportDefinition, epsilon float64, population uint64) (estimate float64, err error) {
	reportType := report.GetReportType()

	sparsity, err := getSparsityForReport(metric, report)
	if err != nil {
		return -1, err
	}

	populationConstant := e.ParamsCalc.Constants.population
	privacyEncodingParams, err := e.ParamsCalc.GetPrivacyEncodingParams(epsilon, populationConstant, sparsity)
	if err != nil {
		return -1, err
	}

	var errorEstimate float64
	switch reportType {
	case config.ReportDefinition_UNIQUE_DEVICE_HISTOGRAMS:
		fallthrough // Calculate RMSE Error for each bucket.
	case config.ReportDefinition_UNIQUE_DEVICE_COUNTS:
		errorEstimate = SingleContributionRapporRMSE(population, privacyEncodingParams.ProbBitFlip)
	case config.ReportDefinition_HOURLY_VALUE_HISTOGRAMS:
		// (24 * population) pseudo-users to account for hourly values.
		errorEstimate = SingleContributionRapporRMSE(24*population, privacyEncodingParams.ProbBitFlip)
	case config.ReportDefinition_FLEETWIDE_OCCURRENCE_COUNTS:
		errorEstimate = MultipleContributionRapporRMSE(population, privacyEncodingParams.ProbBitFlip, report.MaxCount, report.NumIndexPoints)
	case config.ReportDefinition_FLEETWIDE_HISTOGRAMS:
		errorEstimate = MultipleContributionRapporRMSE(population, privacyEncodingParams.ProbBitFlip, 1, report.NumIndexPoints)
	default:
		reportType := config.ReportDefinition_ReportType_name[int32(reportType)]
		return -1, fmt.Errorf("Error estimation is not supported for reports of type %s", reportType)
	}
	return errorEstimate, nil
}

// Compute the 2D-RAPPOR RMSE for reports with single contributions.
//
// See Proposition 1 and 2 of go/histogram-aggregation-privacy-guarantee.
func SingleContributionRapporRMSE(population uint64, probBitFlip float64) float64 {
	return math.Sqrt(float64(population)*probBitFlip*(1-probBitFlip)) / (1 - 2*probBitFlip)
}

// Compute 2D-RAPPOR RMSE for reports with multiple or real-number contributions.
//
// See Proposition 1, 2, and 3 of go/histogram-aggregation-privacy-guarantee.
func MultipleContributionRapporRMSE(population uint64, probBitFlip float64, maxUserContribution uint64, discretization uint32) float64 {
	var n = float64(population)
	var p = probBitFlip
	var m = float64(maxUserContribution)
	var r = float64(discretization)

	var estimate = float64(n) * p * (1 - p) / math.Pow((1-2*p), 2)
	estimate = estimate * (2*math.Pow(r, 3) + 3*math.Pow(r, 2) + r) / 6
	estimate = estimate + (n / 4)
	estimate = m / r * math.Sqrt(estimate)
	return estimate
}