#!/usr/bin/python

import decimal
import os
import re
import string
import sys
import time

from optparse import OptionParser
from subprocess import check_call, CalledProcessError

class CPUScalingTest(object):

    def __init__(self):
        self.speedUpTolerance = 10.0 # percent
        self.retryLimit = 5
        self.retryTolerance = 5.0 # percent
        self.sysCPUDirectory = "/sys/devices/system/cpu"
        self.cpufreqDirectory = os.path.join(self.sysCPUDirectory, "cpu0", "cpufreq")
        self.idaFlag = "ida"
        self.idaSpeedupFactor = 8.0 # percent
        self.selectorExe = "cpufreq-selector"
        self.ifSelectorExe = None

    def getCPUFreqDirectories(self):
        if not os.path.exists(self.sysCPUDirectory):
            print "Error: no file %s" % self.sysCPUDirectory
            return None
        # look for cpu subdirectories
        pattern = re.compile("cpu(?P<cpuNumber>[0-9]+)")
        self.cpufreqDirectories = list()
        for subdirectory in os.listdir(self.sysCPUDirectory):
            match = pattern.search(subdirectory)
            if match and match.group("cpuNumber"):
                cpufreqDirectory = os.path.join(self.sysCPUDirectory,
                    subdirectory, "cpufreq")
                if not os.path.exists(cpufreqDirectory):
                    print "Error: cpu %s has no cpufreq directory %s" \
                        % (match.group("cpuNumber"), cpufreqDirectory)
                    return None
                # otherwise
                self.cpufreqDirectories.append(cpufreqDirectory)
        if len(self.cpufreqDirectories) is 0:
            return None
        # otherwise
        return self.cpufreqDirectories

    def checkParameters(self, file):
        current = None
        for cpufreqDirectory in self.cpufreqDirectories:
            parameters = self.getParameters(cpufreqDirectory, file)
            if not parameters:
                print "Error: could not determine cpu parameters from %s" \
                    % os.path.join(cpufreqDirectory, file)
                return None
            if not current:
                current = parameters
            elif not current == parameters:
                return None
        return current



    def getParameters(self, cpufreqDirectory, file):
        path = os.path.join(cpufreqDirectory, file)
        file = open(path)
        while 1:
            line = file.readline()
            if not line:
                break
            if len(line.strip()) > 0:
                return line.strip().split()
        return None

    def setParameter(self, setFile, readFile, value, skip=False, automatch=False):
        def findParameter(targetFile):
            for root, _, files in os.walk(self.sysCPUDirectory):
                for f in files:
                    rf = os.path.join(root,f)
                    if targetFile in rf:
                        return rf
            return None

        path = None
        if not skip:
            if automatch:
                path = findParameter(setFile)
            else:
                path = os.path.join(self.cpufreqDirectory,  setFile)

            try:
                check_call("echo \"%s\" > %s" % (value, path), shell=True)
            except CalledProcessError, exception:
                print "Error: command failed:"
                print exception
                return False

        # verify it has changed
        if automatch:
            path = findParameter(readFile)
        else:
            path = os.path.join(self.cpufreqDirectory, readFile)

        parameterFile = open(path)
        line = parameterFile.readline()
        if not line or line.strip() != str(value):
            print "Error: could not verify that %s was set to %s" % (path, value)
            if line:
                print "Actual Value: %s" % line
            else:
                print "parameter file was empty"
            return False

        return True

    def checkSelectorExecutable(self):
        def is_exe(fpath):
            return os.path.exists(fpath) and os.access(fpath, os.X_OK)

        if self.ifSelectorExe is None:
            # cpufreq-selector default path
            exe = os.path.join("/usr/bin/",self.selectorExe)
            if is_exe(exe):
                self.ifSelectorExe = True
                return True
            for path in os.environ["PATH"].split(os.pathsep):
                exe = os.path.join(path, self.selectorExe)
                if is_exe(exe):
                    self.ifSelectorExe = True
                    return True

            self.ifSelectorExe = False
            return False

    def setParameterWithSelector(self, switch, setFile, readFile, value):
        # Try the command for all CPUs
        skip = True
        if self.checkSelectorExecutable(): 
            try:
                check_call("cpufreq-selector -%s %s" % (switch, value), shell=True)
            except CalledProcessError, exception:
                print "Note: command failed: %s" % exception.cmd
                skip = False
        else:
            skip = False

        return self.setParameter(setFile, readFile, value, skip)

    def setFrequency(self, frequency):
        return self.setParameterWithSelector("f", "scaling_setspeed", "scaling_cur_freq", frequency)

    def setGovernor(self, governor):
        return self.setParameterWithSelector("g", "scaling_governor", "scaling_governor", governor)


    def getParameter(self, parameter):
        value = None
        parameterFilePath = os.path.join(self.cpufreqDirectory, parameter)
        try:
            parameterFile = open(parameterFilePath)
            line = parameterFile.readline()
            if not line:
                print "Error: failed to get %s for %s" % (parameter, self.cpufreqDirectory)
                return None
            value = line.strip()
            return value
        except IOError, exception:
            print "Error: could not open %s" % parameterFilePath
            print exception

        return None

    def getParameterList(self, parameter):
        values = list()
        for cpufreqDirectory in self.cpufreqDirectories:
            path = os.path.join(cpufreqDirectory, parameter)
            parameterFile = open(path)
            line = parameterFile.readline()
            if not line:
                print "Error: failed to get %s for %s" % (parameter, cpufreqDirectory)
                return None
            values.append(line.strip())
        return values

    def runLoadTest(self):
        print "Running CPU load test..."
        try:
            check_call("taskset -pc 0 %s" % os.getpid(), shell=True)
        except CalledProcessError, exception:
            print "Error: could not set task affinity"
            print exception
            return None

        runTime = None
        tries = 0
        while tries < self.retryLimit:
            sys.stdout.flush()
            (start_utime, start_stime, start_cutime, start_cstime, start_elapsed_time) = os.times()
            self.pi()
            (stop_utime, stop_stime, stop_cutime, stop_cstime, stop_elapsed_time) = os.times()
            if not runTime:
                runTime = stop_elapsed_time - start_elapsed_time
            else:
                thisTime = stop_elapsed_time - start_elapsed_time
                if (abs(thisTime-runTime)/runTime)*100 < self.retryTolerance:
                    return runTime
                else:
                    runTime = thisTime
            tries += 1

        print "Error: could not repeat load test times within %.1f%%" % self.retryTolerance
        return None

    def pi(self):
        decimal.getcontext().prec = 500
        s = decimal.Decimal(1)
        h = decimal.Decimal(3).sqrt()/2
        n = 6
        for i in range(170):
            s2 = ((1-h)**2+s**2/4)
            s = s2.sqrt()
            h = (1-s2/4).sqrt()
            n = 2*n

        return True

    def verifyMinimumFrequency(self, waitTime=5):
        sys.stdout.write("Waiting %d seconds..." % waitTime)
        sys.stdout.flush()
        time.sleep(waitTime)
        sys.stdout.write(" done.\n")
        minimumFrequency = self.getParameter("scaling_min_freq")
        currentFrequency = self.getParameter("scaling_cur_freq")
        if not minimumFrequency or not currentFrequency or (minimumFrequency != currentFrequency):
            return False

        # otherwise
        return True

    def getSystemCapabilities(self):
        print ""
        print "System Capabilites:"
        print "-------------------------------------------------"

        # Do the CPUs support scaling?
        if not self.getCPUFreqDirectories():
            return False
        if len (self.cpufreqDirectories) > 1:
            print "System has %u cpus"% len(self.cpufreqDirectories)

        # Ensure all CPUs support the same frequencies
        freqFileName = "scaling_available_frequencies"
        self.frequencies = self.checkParameters(freqFileName)
        if not self.frequencies:
            return False

        print ""
        print "Supported CPU Frequencies: "
        for freq in self.frequencies:
            f = string.atoi(freq)/1000
            print "    %u MHz" % f

        # Check governors to verify all CPUs support the same control methods
        governorFileName = "scaling_available_governors"
        self.governors = self.checkParameters(governorFileName)
        if not self.governors:
            return False

        print ""
        print "Supported Governors: "
        for governor in self.governors:
            print "    %s" % governor

        self.originalGovernors = self.getParameterList("scaling_governor")
        if self.originalGovernors:
            print ""
            print "Current governors:"
            i = 0
            for g in self.originalGovernors:
                print "    cpu%u: %s" % (i, g)
                i += 1
        else:
            print "Error: could not determine current governor settings"
            return False

        self.getCPUFlags()

        return True

    def getCPUFlags(self):
        self.cpuFlags = None
        try:
            cpuinfo_file = open('/proc/cpuinfo', 'r')
            cpuinfo = cpuinfo_file.read().split("\n")
            cpuinfo_file.close()

            for line in cpuinfo:
                if line.startswith('flags'):
                    pre, post = line.split(':')
                    self.cpuFlags = post.strip().split()
                    break
        except:
            print "Warning: could not read CPU flags"

    def runUserSpaceTests(self):
        print ""
        print "Userspace Governor Test:"
        print "-------------------------------------------------"
        self.minimumFrequencyTestTime = None
        self.maximumFrequencyTestTime = None

        success = True
        differenceSpeedUp = None
        governor = "userspace"
        if governor not in self.governors:
            print "Note: %s governor not supported" % governor
        else:

            # Set the governor to "userspace" and verify
            print "Setting governor to %s" % governor
            if not self.setGovernor(governor):
                success = False

            # Set the the CPU speed to it's lowest value
            frequency = self.frequencies[-1]
            print "Setting CPU frequency to %u MHz" % (string.atoi(frequency)/1000)
            if not self.setFrequency(frequency):
                success = False

            # Verify the speed is set to the lowest value
            minimumFrequency = self.getParameter("scaling_min_freq")
            currentFrequency = self.getParameter("scaling_cur_freq")
            if not minimumFrequency or not currentFrequency or (minimumFrequency != currentFrequency):
                print "Error: Could not verify that cpu frequency is set to the minimum value of %s" % minimumFrequency
                success = False

            # Run Load Test
            self.minimumFrequencyTestTime = self.runLoadTest()
            if not self.minimumFrequencyTestTime:
                print "Error: Could not retrieve the minimum frequency test's execution time."
                success = False
            else:
                print "Minimum frequency load test time: %.2f" % self.minimumFrequencyTestTime

            # Set the CPU speed to it's highest value as above.
            frequency = self.frequencies[0]
            print "Setting CPU frequency to %u MHz" % (string.atoi(frequency)/1000)
            if not self.setFrequency(frequency):
                success = False

            maximumFrequency = self.getParameter("scaling_max_freq")
            currentFrequency = self.getParameter("scaling_cur_freq")
            if not maximumFrequency or not currentFrequency or (maximumFrequency != currentFrequency):
                print "Error: Could not verify that cpu frequency is set to the maximum value of %s" % maximumFrequency
                success = False

            # Repeat workload test
            self.maximumFrequencyTestTime = self.runLoadTest()
            if not self.maximumFrequencyTestTime:
                print "Error: Could not retrieve the maximum frequency test's execution time."
                success = False
            else:
                print "Maximum frequency load test time: %.2f" % self.maximumFrequencyTestTime

            # Verify MHz increase is comparable to time % decrease
            predictedSpeedup = string.atof(maximumFrequency)/string.atof(minimumFrequency)

            # If "ida" turbo thing, increase the expectation by 8%
            if self.cpuFlags and self.idaFlag in self.cpuFlags:
                print "Note: found %s flag, increasing expected speedup by %.1f%%" % (self.idaFlag, self.idaSpeedupFactor)
                predictedSpeedup = predictedSpeedup*(1.0/(1.0-(self.idaSpeedupFactor/100.0)))

            if self.minimumFrequencyTestTime and self.maximumFrequencyTestTime:
                measuredSpeedup = self.minimumFrequencyTestTime/self.maximumFrequencyTestTime
                print ""
                print "CPU Frequency Speed Up: %.2f" % predictedSpeedup
                print "Measured Speed Up: %.2f" % measuredSpeedup
                differenceSpeedUp =  (abs(measuredSpeedup-predictedSpeedup)/predictedSpeedup)*100
                print "Percentage Difference %.1f%%" % differenceSpeedUp
                if differenceSpeedUp > self.speedUpTolerance:
                    print "Error: measured speedup vs expected speedup is %.1f%% and is not within %.1f%% margin. " % (differenceSpeedUp, self.speedUpTolerance)
                    success = False
            else:
                print "Error: Not enough timing data to calculate speed differences."

        return success

    def runOnDemandTests(self):
        print ""
        print "On Demand Governor Test:"
        print "-------------------------------------------------"
        differenceOnDemandVsMaximum = None
        onDemandTestTime = None
        governor = "ondemand"
        success = True
        if governor not in self.governors:
            print "Note: %s governor not supported" % governor
        else:
            # Set the governor to "ondemand"
            print "Setting governor to %s" % governor
            if not self.setGovernor(governor):
                success = False

            # Wait a fixed period of time, then verify current speed is the slowest in as before
            if not self.verifyMinimumFrequency():
                print "Error: Could not verify that cpu frequency has settled to the minimum value"
                success = False

            # Repeat workload test
            onDemandTestTime = self.runLoadTest()
            if not onDemandTestTime:
		print "Error: No On Demand load test time available."
                success = False
            else:
                print "On Demand load test time: %.2f" % onDemandTestTime

            if onDemandTestTime and self.maximumFrequencyTestTime:
                # Compare the timing to the max results from earlier, again time should be within self.speedUpTolerance
                differenceOnDemandVsMaximum = (abs(onDemandTestTime-self.maximumFrequencyTestTime)/self.maximumFrequencyTestTime)*100
                print "Percentage Difference vs. maximum frequency: %.1f%%" % differenceOnDemandVsMaximum
                if differenceOnDemandVsMaximum > self.speedUpTolerance:
                    print "Error: on demand performance vs maximum of %.1f%% is not within %.1f%% margin" % (differenceOnDemandVsMaximum, self.speedUpTolerance)
                    success = False
            else:
                print "Error: Not enough timing data to calculate speed differences."

            # Verify the current speed has returned to the lowest speed again
            if not self.verifyMinimumFrequency():
                print "Error: Could not verify that cpu frequency has settled to the minimum value"
                success = False

        return success

    def runPerformanceTests(self):
        print ""
        print "Performance Governor Test:"
        print "-------------------------------------------------"
        differencePerformanceVsMaximum = None
        governor = "performance"
        success = True
        if governor not in self.governors:
            print "Note: %s governor not supported" % governor
        else:
            # Set the governor to "performance"
            print "Setting governor to %s" % governor
            if not self.setGovernor(governor):
                success = False

            # Verify the current speed is the same as scaling_max_freq
            maximumFrequency = self.getParameter("scaling_max_freq")
            currentFrequency = self.getParameter("scaling_cur_freq")
            if not maximumFrequency or not currentFrequency or (maximumFrequency != currentFrequency):
                print "Error: Current cpu frequency of %s is not set to the maximum value of %s" % (currentFrequency, maximumFrequency)
                success = False

            # Repeat work load test
            performanceTestTime = self.runLoadTest()
            if not performanceTestTime:
		print "Error: No Performance load test time available."
                success = False
            else:
                print "Performance load test time: %.2f" % performanceTestTime

            if performanceTestTime and self.maximumFrequencyTestTime:
                # Compare the timing to the max results
                differencePerformanceVsMaximum = (abs(performanceTestTime-self.maximumFrequencyTestTime)/self.maximumFrequencyTestTime)*100
                print "Percentage Difference vs. maximum frequency: %.1f%%" % differencePerformanceVsMaximum
                if differencePerformanceVsMaximum > self.speedUpTolerance:
                    print "Error: performance setting vs maximum of %.1f%% is not within %.1f%% margin" % (differencePerformanceVsMaximum, self.speedUpTolerance)
                    success = False
            else:
                print "Error: Not enough timing data to calculate speed differences."

        return success

    def runConservativeTests(self):
        print ""
        print "Conservative Governor Test:"
        print "-------------------------------------------------"
        differenceConservativeVsMinimum = None
        governor = "conservative"
        success = True
        if governor not in self.governors:
            print "Note: %s governor not supported" % governor
        else:
            # Set the governor to "conservative"
            print "Setting governor to %s" % governor
            if not self.setGovernor(governor):
                success = False

            # Set the frequency step to 20, so that it jumps to minimum frequency
            path = os.path.join("conservative", "freq_step")
            if not self.setParameter(path, path, 20, automatch=True):
                success = False

            # Wait a fixed period of time, then verify current speed is the slowest in as before
            if not self.verifyMinimumFrequency(10):
                print "Error: Could not verify that cpu frequency has settled to the minimum value"
                success = False

            # Set the frequency step to 0, so that it doesn't gradually increase
            if not self.setParameter(path, path, 0, automatch=True):
                success = False

            # Repeat work load test
            conservativeTestTime = self.runLoadTest()
            if not conservativeTestTime:
		print "Error: No Conservative load test time available."
                success = False
            else:
                print "Conservative load test time: %.2f" % conservativeTestTime

            if conservativeTestTime and self.minimumFrequencyTestTime:
                # Compare the timing to the max results
                differenceConservativeVsMinimum = (abs(conservativeTestTime-self.minimumFrequencyTestTime)/self.minimumFrequencyTestTime)*100
                print "Percentage Difference vs. minimum frequency: %.1f%%" % differenceConservativeVsMinimum
                if differenceConservativeVsMinimum > self.speedUpTolerance:
                    print "Error: performance setting vs minimum of %.1f%% is not within %.1f%% margin" % (differenceConservativeVsMinimum, self.speedUpTolerance)
                    success = False
            else:
                print "Error: Not enough timing data to calculate speed differences."

        return success

    def restoreGovernors(self):
        print "Restoring original governor to %s" % (self.originalGovernors[0])
        self.setGovernor(self.originalGovernors[0])


def main(args):
    usage = "Usage: %prog [OPTIONS]"
    parser = OptionParser(usage=usage)
    parser.add_option("-q", "--quiet",
                      action="store_true",
                      help="Suppress output.")
    parser.add_option("-c", "--capabilities",
                      action="store_true",
                      help="Only output CPU capabilities.")
    (options, args) = parser.parse_args(args)

    if options.quiet:
        sys.stdout = open(os.devnull, 'a')
        sys.stderr = open(os.devnull, 'a')

    test = CPUScalingTest()
    if not os.path.exists(test.cpufreqDirectory):
        print "CPU Frequency Scaling not supported"
        return 0

    if not test.getSystemCapabilities():
        parser.error("Failed to get system capabilities")

    returnValues = []
    if not options.capabilities:
        returnValues.append(test.runUserSpaceTests())
        returnValues.append(test.runOnDemandTests())
        returnValues.append(test.runPerformanceTests())
        returnValues.append(test.runConservativeTests())
        test.restoreGovernors()

    return 1 if False in returnValues else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
