// Copyright 2019 The Chromium OS 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 encoding

import (
	"context"
	"fmt"
	"io/ioutil"
	"math"
	"regexp"
	"strconv"
	"strings"

	"chromiumos/tast/common/perf"
	"chromiumos/tast/errors"
	"chromiumos/tast/testing"
)

// This file contains helper functions for parsing metric values from log files generated by video_encode_accelerator_unittest.

var regExpFPS = regexp.MustCompile(`(?m)^Measured encoder FPS: ([+\-]?[0-9.]+)$`)
var regExpEncodeLatency50 = regexp.MustCompile(`(?m)^Encode latency for the 50th percentile: (\d+) us$`)
var regExpEncodeLatency75 = regexp.MustCompile(`(?m)^Encode latency for the 75th percentile: (\d+) us$`)
var regExpEncodeLatency95 = regexp.MustCompile(`(?m)^Encode latency for the 95th percentile: (\d+) us$`)

// ReportFPS reports FPS info from log file and sets as the perf metric.
func ReportFPS(ctx context.Context, p *perf.Values, name, logPath string) error {
	b, err := ioutil.ReadFile(logPath)
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", logPath)
	}

	matches := regExpFPS.FindAllStringSubmatch(string(b), -1)
	if len(matches) != 1 {
		return errors.Errorf("found %d FPS matches in %q; want 1", len(matches), b)
	}

	fps, err := strconv.ParseFloat(matches[0][1], 64)
	if err != nil {
		return errors.Wrapf(err, "failed to parse FPS value %q", matches[0][1])
	}

	p.Set(perf.Metric{
		Name:      getMetricName(name, "fps"),
		Unit:      "fps",
		Direction: perf.BiggerIsBetter,
	}, fps)

	testing.ContextLogf(ctx, "> FPS: %.2f", fps)

	return nil
}

// ReportEncodeLatency reports encode latency from log file and sets as the perf metrics.
func ReportEncodeLatency(ctx context.Context, p *perf.Values, name, logPath string) error {
	b, err := ioutil.ReadFile(logPath)
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", logPath)
	}

	// Iterate over different latency measurements, extracting and reporting each.
	for _, st := range []struct {
		key string         // metric key
		re  *regexp.Regexp // matches latency stat
	}{
		{"encode_latency.50_percentile", regExpEncodeLatency50},
		{"encode_latency.75_percentile", regExpEncodeLatency75},
		{"encode_latency.95_percentile", regExpEncodeLatency95},
	} {
		match := st.re.FindStringSubmatch(string(b))
		if match == nil {
			return errors.Errorf("didn't find match for latency %q in %q", st.re, b)
		}
		val, err := strconv.Atoi(match[1])
		if err != nil {
			return errors.Wrapf(err, "failed converting %q latency %q", st.key, match[1])
		}
		p.Set(perf.Metric{
			Name:      getMetricName(name, st.key),
			Unit:      "us",
			Direction: perf.SmallerIsBetter,
		}, float64(val))

		testing.ContextLogf(ctx, "> "+getMetricName(name, st.key)+": %vus", val)
	}

	return nil
}

// ReportCPUUsage reports CPU usage from log file and sets as the perf metric.
func ReportCPUUsage(ctx context.Context, p *perf.Values, name, logPath string) error {
	b, err := ioutil.ReadFile(logPath)
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", logPath)
	}

	vstr := strings.TrimSuffix(string(b), "\n")
	v, err := strconv.ParseFloat(vstr, 64)
	if err != nil {
		return errors.Wrapf(err, "failed to parse %q to float", vstr)
	}

	p.Set(perf.Metric{
		Name:      getMetricName(name, "cpu_usage"),
		Unit:      "percent",
		Direction: perf.SmallerIsBetter,
	}, v)

	testing.ContextLogf(ctx, "> CPU usage: %.2f%%", v)

	return nil
}

// ReportPowerConsumption reports power consumption from log file and sets as the perf metric.
func ReportPowerConsumption(ctx context.Context, p *perf.Values, name, logPath string) error {
	b, err := ioutil.ReadFile(logPath)
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", logPath)
	}

	vstr := strings.TrimSuffix(string(b), "\n")
	v, err := strconv.ParseFloat(vstr, 64)
	if err != nil {
		return errors.Wrapf(err, "failed to parse %q to float", vstr)
	}

	p.Set(perf.Metric{
		Name:      getMetricName(name, "power_consumption"),
		Unit:      "watt",
		Direction: perf.SmallerIsBetter,
	}, v)

	testing.ContextLogf(ctx, "> Power consumption: %.2f", v)

	return nil
}

// channelStats records min, max, sum of statistics (e.g. PSNR) for a channel.
type channelStats struct {
	min, max, sum float64
	num           int
}

// report saves min, max, avg perf metrics to p.
func (cs *channelStats) report(p *perf.Values, prefix, unit string) {
	setPerf := func(stat string, val float64) {
		p.Set(perf.Metric{
			Name:      prefix + stat,
			Unit:      unit,
			Direction: perf.BiggerIsBetter,
		}, val)
	}
	setPerf("min", cs.min)
	setPerf("max", cs.max)
	if cs.num > 0 {
		setPerf("avg", cs.sum/float64(cs.num))
	}
}

// update updates statistics (e.g. PSNR) with the input value.
func (cs *channelStats) update(value float64) {
	cs.min = math.Min(value, cs.min)
	cs.max = math.Max(value, cs.max)
	cs.sum += value
	cs.num++
}

// newChannelStats returns a channelStats struct with default values.
func newChannelStats() channelStats { return channelStats{min: math.MaxFloat64} }

// frameStats records frame-wise statistics (e.g. PSNR) for YUV and combined channels.
type frameStats struct{ y, u, v, combined channelStats }

// report saves all stats to p.
// name is the initial arg passed to getMetricName.
// typ is the metric type, e.g. "ssim" or "psnr", and is as the unit of perf metric.
func (fs *frameStats) report(p *perf.Values, name, typ string) {
	// For "y", "u", "v" channels, the metric key is formatted as "quality.<typ>.<channel>.<stat>", such as "quality.ssim.y.max".
	// For "combined" channel, the metric key is formatted as "quality.<typ>.<stat>", such as "quality.psnr.avg".
	prefix := getMetricName(name, fmt.Sprintf("quality.%s.", typ))
	fs.y.report(p, prefix+"y.", typ)
	fs.u.report(p, prefix+"u.", typ)
	fs.v.report(p, prefix+"v.", typ)
	fs.combined.report(p, prefix, typ)
}

// newFrameStats returns a frameStats struct with default values.
func newFrameStats() frameStats {
	return frameStats{
		y:        newChannelStats(),
		u:        newChannelStats(),
		v:        newChannelStats(),
		combined: newChannelStats(),
	}
}

// ReportFrameStats reports quality from log file which assumes input is YUV420 (for MSE samples per channel), and sets as the perf metrics.
func ReportFrameStats(p *perf.Values, name, logPath string) error {
	b, err := ioutil.ReadFile(logPath)
	if err != nil {
		return errors.Wrapf(err, "failed to read file %s", logPath)
	}

	ssim := newFrameStats()
	psnr := newFrameStats()

	for i, line := range strings.Split(string(b), "\n") {
		if i == 0 || line == "" {
			// Skip the first CSV header line or blank line.
			continue
		}
		values := strings.Split(line, ",")
		if len(values) != 9 {
			return errors.Errorf("line %d does not contain 9 comma-separated values %q", i, line)
		}

		var index, width, height, ssimY, ssimU, ssimV, mseY, mseU, mseV float64
		for j, dst := range []*float64{&index, &width, &height, &ssimY, &ssimU, &ssimV, &mseY, &mseU, &mseV} {
			if *dst, err = strconv.ParseFloat(values[j], 64); err != nil {
				return errors.Wrapf(err, "failed to parse %q in field %d", values[j], j)
			}
		}

		ssim.y.update(ssimY)
		ssim.u.update(ssimU)
		ssim.v.update(ssimV)
		// Weighting of YUV channels for SSIM taken from libvpx.
		ssim.combined.update(0.8*ssimY + 0.1*(ssimU+ssimV))

		// Samples per MSE score assumes YUV420 subsampling.
		psnr.y.update(mseToPSNR(width*height*4/4, 255, mseY))
		psnr.u.update(mseToPSNR(width*height*4/4, 255, mseU))
		psnr.v.update(mseToPSNR(width*height*4/4, 255, mseV))
		psnr.combined.update(mseToPSNR(width*height*6/4, 255, mseY+mseU+mseV))
	}

	if ssim.y.num == 0 {
		return errors.New("frame statistics do not exist")
	}

	ssim.report(p, name, "ssim")
	psnr.report(p, name, "psnr")
	return nil
}

// mseToPSNR calculates PSNR from MSE for a frame.
func mseToPSNR(samples, peak, mse float64) float64 {
	const maxPSNR = 100.0
	// Prevent a divide-by-zero, MSE at 0 is perfect quality (no error).
	if mse == 0 {
		return maxPSNR
	}
	psnr := 10.0 * math.Log10(peak*peak*samples/mse)
	return math.Min(psnr, maxPSNR)
}

// getMetricName wraps the stream name and key into the metric name.
// For example, name should contain both stream and codec name such like "tulip2-1280x720_h264", key is the metric name such like "fps".
func getMetricName(name, key string) string {
	return fmt.Sprintf("%s.%s", name, key)
}