#!/usr/bin/python

################################################################################
# Copyright 2019 ModalAI Inc.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
#    may be used to endorse or promote products derived from this software
#    without specific prior written permission.
#
# 4. The Software is licensed to be used solely in conjunction with devices
#    provided by ModalAI Inc.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
################################################################################


# Here are some links to the Qualcomm documentation that explains some of
# the thermal management of the Snapdragon 820
# https://source.codeaurora.org/quic/la/kernel/msm-3.14/tree/Documentation/devicetree/bindings/thermal/tsens.txt
# https://developer.qualcomm.com/forum/qdn-forums/software/snapdragon-profiler/29435
# https://source.codeaurora.org/quic/la/kernel/msm-3.18/tree/arch/arm/boot/dts/qcom/msm8996.dtsi?h=LA.HB.1.3.1_rb1.8#n2496

# Other helpful links
# https://developer.qualcomm.com/forum/qdevnet-forums/performance-tools/trepn-profiler/26812
# https://developer.qualcomm.com/forum/qdn-forums/increase-app-performance/trepn-profiler/27801


import os
import sys
import time
import subprocess
import getopt

# ANSI escape sequences
ESCAPE_KEY = chr(27) + '['
ERASE_DISPLAY = '2J'
CURSOR_HOME = '0;0H'
ONE_LINE_UP = 'A'
CLEAR_LINE = 'K'

# CPU names
cpu_names = ['cpu0', 'cpu1', 'cpu2', 'cpu3']

# Here is where you find the CPU frequency information
SYSTEM_READ_COMMAND = "awk '{print $1}' "
SYSTEM_CPU_DIRECTORY = '/sys/devices/system/cpu/'
SYSTEM_SPECIFIC_CPU_DIRECTORY_PREFIX = '/sys/devices/system/cpu/cpu'
SYSTEM_CPU_CURRENT_FREQUENCY = '/cpufreq/scaling_cur_freq'
SYSTEM_CPU_MAXIMUM_FREQUENCY = '/cpufreq/cpuinfo_max_freq'
SYSTEM_CPU_FREQUENCY_COUNTERS = '/cpufreq/stats/time_in_state'

# Here is where you find the GPU frequency information
SYSTEM_GPU_DIRECTORY = '/sys/class/kgsl/kgsl-3d0/'
SYSTEM_GPU_FREQUENCY_COUNTERS = SYSTEM_GPU_DIRECTORY + 'gpu_clock_stats'
SYSTEM_GPU_AVAILABLE_FREQUENCIES = SYSTEM_GPU_DIRECTORY + \
                                   'gpu_available_frequencies'
SYSTEM_GPU_CLOCK = SYSTEM_GPU_DIRECTORY + 'gpuclk'
SYSTEM_GPU_BUSY_COUNTER = SYSTEM_GPU_DIRECTORY + 'gpubusy'


# Here is where you find the system temperatures
# /sys/devices/virtual/thermal/thermal_zone<x>/temp
thermal_zone = [4, 6, 9, 11]
maximum_temperature = 0.0
SYSTEM_THERMAL_DIRECTORY = '/sys/devices/virtual/thermal/thermal_zone'
SYSTEM_CURRENT_TEMPERATURE = '/temp'
POP_MEM_THERMAL_ZONE = 1
GPU_THERMAL_ZONE = 15
THERMAL_SCALING_FACTOR = 0.1

# System constants
NUM_CPU = 4

# Display constants
HEADER_LINE    = 'Name   Freq (MHz) Temp (C) Util (%)'
SEPARATOR_LINE = '-----------------------------------'
CSV_HEADER     = 'Max_Freq_MHz Max_CPU_Temp_C POP_Mem_Temp_C GPU_Temp_C'

# Constants for CPU tick values in /proc/stats
NUM_CPU_TICK_TYPES = 10
CPU_IDLE_COUNT_FIELD = 3 # Based on 0 - 9 fields

# Variables for CPU tick storage and manipulation
cpu_ticks_initialized = False
cpu_ticks = [[0, 0], [0, 0],[0, 0],[0, 0]]
previous_cpu_ticks = [[0, 0], [0, 0],[0, 0],[0, 0]]
IDLE_TICKS = 0
TOTAL_TICKS = 1

# Variables for CPU usage normalization
cpu_max_frequency = [0.0, 0.0, 0.0, 0.0]
frequencies = [[], [], [], []]
previous_frequency_counts = [[], [], [], []]
weighted_frequency = [0.0, 0.0, 0.0, 0.0]
previous_gpu_utilizations = []

# Variables for GPU usage normalization
gpu_max_frequency = 0.0
gpu_frequencies = []
previous_gpu_frequency_counts = []

# Variables and constants for memory use
MEMORY_VALUES = 1
TOTAL_MEMORY = 0
USED_MEMORY = 1
FREE_MEMORY = 2

# voxl-perfmon update rate
UPDATE_RATE_IN_SECONDS = 2.0

# Multiplier for long term GPU utilization
EXTENDED_RATE_MULTIPLIER = 5

# Functions
def get_cpu_current_frequency_filename(cpu):
    return SYSTEM_SPECIFIC_CPU_DIRECTORY_PREFIX + str(cpu) + SYSTEM_CPU_CURRENT_FREQUENCY

def get_cpu_maximum_frequency_filename(cpu):
    return SYSTEM_SPECIFIC_CPU_DIRECTORY_PREFIX + str(cpu) + SYSTEM_CPU_MAXIMUM_FREQUENCY

def get_cpu_current_frequency_command(cpu):
    return SYSTEM_READ_COMMAND + get_cpu_current_frequency_filename(cpu)

def get_cpu_maximum_frequency_command(cpu):
    return SYSTEM_READ_COMMAND + get_cpu_maximum_frequency_filename(cpu)

def get_cpu_current_frequency(cpu):
    cmd = get_cpu_current_frequency_command(cpu)
    freq = 0.0
    if os.path.isfile(get_cpu_current_frequency_filename(cpu)):
        try:
            # Convert the kHz value to a MHz value
            freq = float(os.popen(cmd).readline().rstrip()) / 1000
        except:
            pass
    return freq

def get_cpu_maximum_frequency(cpu):
    cmd = get_cpu_maximum_frequency_command(cpu)
    freq = 0.0
    if os.path.isfile(get_cpu_current_frequency_filename(cpu)):
        try:
            freq = float(os.popen(cmd).readline().rstrip())
        except:
            pass
    return freq

def store_cpu_maximum_frequencies():
    for cpu_id in range(NUM_CPU):
        cpu_max_frequency[cpu_id] = get_cpu_maximum_frequency(cpu_id)

def initialize_frequency_information():
    global frequencies
    global previous_frequency_counts
    global gpu_frequencies
    global gpu_max_frequency
    global previous_gpu_frequency_counts
    global previous_gpu_utilizations

    # First deal with the CPU
    store_cpu_maximum_frequencies()
    for cpu_id in range(NUM_CPU):
        frequency_counter_filename = SYSTEM_SPECIFIC_CPU_DIRECTORY_PREFIX + str(cpu_id) + \
                                     SYSTEM_CPU_FREQUENCY_COUNTERS
        freqs = subprocess.check_output(['cat',
                                         frequency_counter_filename])
        freqs = freqs.split('\n')
        for freq_line in freqs:
            tmp = freq_line.rstrip().split()
            if len(tmp) != 2:
                continue
            frequencies[cpu_id].append(int(tmp[0]))
            previous_frequency_counts[cpu_id].append(int(tmp[1]))

    # Now move on to the GPU
    gpu_frequencies = subprocess.check_output(['cat',
                                              SYSTEM_GPU_AVAILABLE_FREQUENCIES])
    gpu_frequencies = gpu_frequencies.split()
    for i in range(len(gpu_frequencies)):
        gpu_frequencies[i] = float(gpu_frequencies[i])
        gpu_frequencies[i] /= 1000000.0
    gpu_max_frequency = gpu_frequencies[0]
    previous_gpu_frequency_counts = subprocess.check_output(['cat',
                                               SYSTEM_GPU_FREQUENCY_COUNTERS])
    previous_gpu_frequency_counts = previous_gpu_frequency_counts.split()
    if len(previous_gpu_frequency_counts) != len(gpu_frequencies):
        print 'Incorrect number of GPU frequency counts in initialization'
        sys.exit(-1)
    for i in range(len(previous_gpu_frequency_counts)):
        previous_gpu_frequency_counts[i] = int(previous_gpu_frequency_counts[i])
    for i in range(EXTENDED_RATE_MULTIPLIER):
        previous_gpu_utilizations.append(0.0)

def get_current_temperature_filename(zone):
    return SYSTEM_THERMAL_DIRECTORY + \
           str(zone) + SYSTEM_CURRENT_TEMPERATURE

def get_current_temperature_command(zone):
    return SYSTEM_READ_COMMAND + get_current_temperature_filename(zone)

def get_current_temperature(zone):
    cmd = get_current_temperature_command(zone)
    temp = 0.0
    if os.path.isfile(get_current_temperature_filename(zone)):
        try:
            # Documentation states that temperatures are in tenth of
            # a degree Celcius. Convert them to degrees.
            temp = float(os.popen(cmd).readline().rstrip()) * THERMAL_SCALING_FACTOR
        except:
            pass
    return temp

def print_memory_stats():
    print(SEPARATOR_LINE)
    stats = subprocess.check_output(['free', '-m']).split('\n')
    tmp = stats[MEMORY_VALUES].rstrip().split()[1:]
    print('Total memory (MB): ' + tmp[TOTAL_MEMORY])
    print('Used memory  (MB): ' + tmp[USED_MEMORY])
    print('Free memory  (MB): ' + tmp[FREE_MEMORY])
    return 4

def reset_screen():
    print(ESCAPE_KEY + ERASE_DISPLAY)
    print(ESCAPE_KEY + CURSOR_HOME)

def prepare_screen(num_data_rows):
    for i in range(num_data_rows):
        print(ESCAPE_KEY + ONE_LINE_UP)
        print(ESCAPE_KEY + CURSOR_HOME)

def collect_cpu_stats():
    # The CPU time tick statistics are kept in /proc/stat. These are the same
    # numbers that the 'top' program uses
    stats = subprocess.check_output(['cat', '/proc/stat']).split('\n')
    for cpu_id in range(NUM_CPU):
        output_line = cpu_id + 1;
        # Isolate all of the strings that represent clock ticks in
        # a list
        tmp = stats[output_line].rstrip().split()[1:]
# Convert all of the strings into integers
        tmp = [int(x) for x in tmp]
        # Calculate ticks for current period
        idle_ticks = tmp[CPU_IDLE_COUNT_FIELD]
        total_ticks = sum(tmp)
        cpu_ticks[cpu_id][IDLE_TICKS] = \
            idle_ticks - previous_cpu_ticks[cpu_id][IDLE_TICKS]
        cpu_ticks[cpu_id][TOTAL_TICKS] = \
            total_ticks - previous_cpu_ticks[cpu_id][TOTAL_TICKS]
        # Preserve current ticks for next set of calculations
        previous_cpu_ticks[cpu_id][IDLE_TICKS] = idle_ticks
        previous_cpu_ticks[cpu_id][TOTAL_TICKS] = total_ticks

def calculate_weighted_frequency():
    for cpu_id in range(NUM_CPU):
        frequency_counter_filename = SYSTEM_SPECIFIC_CPU_DIRECTORY_PREFIX + str(cpu_id) + \
                                     SYSTEM_CPU_FREQUENCY_COUNTERS
        freqs = subprocess.check_output(['cat',
                                         frequency_counter_filename])
        freqs = freqs.split('\n')
        current_freqs = []
        # Collect all of the current frequency counts
        for freq_line in freqs:
            tmp = freq_line.rstrip().split()
            if len(tmp) != 2:
                continue
            current_freqs.append(int(tmp[1]))
        # Sanity check to make sure the list of possible frequencies is
        # always the same
        if len(current_freqs) != len(previous_frequency_counts[cpu_id]):
            print('Frequency list sizes do not match for CPU ' + str(cpu_id))
            sys.exit(-1)
        # Now collect all of the differences between the current counts and
        # the previous counts to get the counts for this measurement period
        delta_counts = []
        for i in range(len(current_freqs)):
            delta_counts.append(current_freqs[i] - \
                                previous_frequency_counts[cpu_id][i])
        # Add up all of the counts to get a total
        total_frequency_count = float(sum(delta_counts))
        if total_frequency_count > 0:
            # Now calculate weights for each frequency
            weighted_frequencies = []
            for i in range(len(delta_counts)):
                tmp = delta_counts[i] / total_frequency_count
                tmp *= frequencies[cpu_id][i]
                weighted_frequencies.append(tmp)
            # Now calculate the final weighted frequency value
            weighted_frequency[cpu_id] = sum(weighted_frequencies)
        else:
            # CPU is offline during this period
            weighted_frequency[cpu_id] = 0

def print_non_cpu_data():
    lines_printed = 0
    print(SEPARATOR_LINE)
    lines_printed += 1
    print('POP memory temperature: {:6.1f}'.format(
           get_current_temperature(POP_MEM_THERMAL_ZONE)))
    lines_printed += 1
    return lines_printed

def get_cpu_current_utilization(cpu_id):
    cpu_utilization = 0.0
    if cpu_ticks[cpu_id][TOTAL_TICKS] != 0:
        # Utilization is non-idle ticks / total ticks
        cpu_utilization = float(cpu_ticks[cpu_id][TOTAL_TICKS] - \
                                cpu_ticks[cpu_id][IDLE_TICKS]) / \
                          float(cpu_ticks[cpu_id][TOTAL_TICKS])
    cpu_utilization *= 100.0
    # Normalize the CPU utilization based on latest frequency counts
    cpu_utilization *= (weighted_frequency[cpu_id] / cpu_max_frequency[cpu_id])
    return cpu_utilization


def get_gpu_busy():
    gpu_busy = subprocess.check_output(['cat',
                                         SYSTEM_GPU_BUSY_COUNTER])
    gpu_busy = [float(f) for f in gpu_busy.rstrip().split()]

    if sum(gpu_busy) == 0:
        return 0

    return gpu_busy[0] / sum(gpu_busy);


def print_gpu_stats():
    global previous_gpu_frequency_counts

    print(SEPARATOR_LINE)

    # Read instantaneous clock setting if GPU is active
    gpu_clock = subprocess.check_output(['cat',
                                         SYSTEM_GPU_CLOCK])
    gpu_clock = float(gpu_clock) / 1000000.0

    # Read the current GPU temperature
    gpu_temperature = get_current_temperature(GPU_THERMAL_ZONE)

    # Read the busy / (idle + busy) time
    gpu_utilization = get_gpu_busy()

    # Get current frequency counts
    gpu_frequency_counts = subprocess.check_output(['cat',
                                               SYSTEM_GPU_FREQUENCY_COUNTERS])
    gpu_frequency_counts = gpu_frequency_counts.split()
    if len(gpu_frequency_counts) != len(gpu_frequencies):
        print('Incorrect number of GPU frequency counts')
        sys.exit(-1)
    for i in range(len(gpu_frequency_counts)):
        gpu_frequency_counts[i] = int(gpu_frequency_counts[i])

    # Calculate deltas from previous counts
    delta_counts = []
    for i in range(len(gpu_frequency_counts)):
        delta_counts.append(0)
    for i in range(len(delta_counts)):
        delta_counts[i] = gpu_frequency_counts[i] - \
                          previous_gpu_frequency_counts[i]
        # Save current counts as previous counts for next calculations
        previous_gpu_frequency_counts[i] = gpu_frequency_counts[i]

    # The total counts
    total_frequency_count = float(sum(delta_counts))

    if (sum(delta_counts) != 0) or (gpu_utilization == 0) or (total_frequency_count < 1):

        # Compute the weighted frequency
        gpu_weighted_frequency = 0;

        for i in range(len(delta_counts)):
            tmp = delta_counts[i] / total_frequency_count
            tmp *= gpu_frequencies[i]
            gpu_weighted_frequency += tmp
        gpu_weighted_frequency /= max(gpu_frequencies)

        # Get the true utilization
        gpu_utilization = get_gpu_busy() * gpu_weighted_frequency

        # Convert to a percentage
        gpu_utilization *= 100

        ######################################################################## 
        # Old method, Kept as a record in case someone needs to compare old data
        # DO NOT USE!
        ######################################################################## 
        # for i in range(len(delta_counts)):
        #     # Rescale to match stored frequency values
        #     delta_counts[i] = (float(delta_counts[i]) / 1000.0)

        #     # Calculate percentage of time based on frequency
        #     delta_counts[i] = (float(delta_counts[i]) / gpu_frequencies[i]) * 100.0
        #     delta_counts[i] = (float(delta_counts[i]) / UPDATE_RATE_IN_SECONDS)
        #     gpu_utilization += delta_counts[i]
    else:
        gpu_clock = 0.0
        gpu_utilization = 0.0

    previous_gpu_utilizations.append(gpu_utilization)
    previous_gpu_utilizations.remove(previous_gpu_utilizations[0])
    long_term_gpu_utilization = sum(previous_gpu_utilizations) / \
                                float(EXTENDED_RATE_MULTIPLIER)

    print('{:6s} {:10.1f} {:8.1f} {:8.2f} {:d}s {:8.2f} {:d}s'.format(
          'GPU',
          gpu_clock,
          gpu_temperature,
          gpu_utilization,
          int(UPDATE_RATE_IN_SECONDS),
          long_term_gpu_utilization,
          int(UPDATE_RATE_IN_SECONDS) * EXTENDED_RATE_MULTIPLIER))

    return 2

def print_cpu_stats():
    global maximum_temperature
    total_utilization = 0.0
    for cpu_id in range(NUM_CPU):
        frequency = get_cpu_current_frequency(cpu_id)
        temperature = get_current_temperature(thermal_zone[cpu_id])
        if temperature > maximum_temperature:
            maximum_temperature = temperature
        utilization = get_cpu_current_utilization(cpu_id)
        total_utilization += utilization
        print('{:6s} {:10.1f} {:8.1f} {:8.2f}'.format(cpu_names[cpu_id],
                                                      frequency,
                                                      temperature,
                                                      utilization))
    print(SEPARATOR_LINE)
    print('Total {:29.2f}'.format(total_utilization / float(NUM_CPU)))
    return NUM_CPU + 2

def print_csv_line():
    max_temp = 0.0
    max_freq = 0
    for cpu_id in range(NUM_CPU):
        frequency = get_cpu_current_frequency(cpu_id)
        temperature = get_current_temperature(thermal_zone[cpu_id])
        if temperature > max_temp:
            max_temp = temperature
        if frequency > max_freq:
            max_freq = frequency
    print('{:0.1f} {:0.1f} {:0.1f} {:0.1f}'.format(   max_freq,
                                                max_temp,
                                                get_current_temperature(POP_MEM_THERMAL_ZONE),
                                                get_current_temperature(GPU_THERMAL_ZONE)))

def get_number_of_cpus():
    global NUM_CPU
    global cpu_names

    # Get all the directories
    all_files_and_dirs = os.listdir(SYSTEM_CPU_DIRECTORY)

    # Only keep the ones that are "cpuX"
    only_cpus = [f for f in all_files_and_dirs if "cpu" in f]
    only_cpus = [f for f in only_cpus if len(f) == 4]

    if(len(only_cpus) == 0):
        print("Could not find any CPUs in {}".format(SYSTEM_CPU_DIRECTORY))
        sys.exit(-1)

    # Sort them for alphabetical order
    only_cpus.sort()

    # Set the global Variables
    cpu_names = only_cpus
    NUM_CPU = len(cpu_names)

def init_data_structs():
    global cpu_ticks
    global previous_cpu_ticks
    global cpu_max_frequency
    global frequencies
    global previous_frequency_counts
    global weighted_frequency
    global thermal_zone
    global POP_MEM_THERMAL_ZONE
    global THERMAL_SCALING_FACTOR

    cpu_ticks = [[0,0] for _ in range(NUM_CPU)]
    previous_cpu_ticks = [[0,0] for _ in range(NUM_CPU)]
    cpu_max_frequency = [0.0 for _ in range(NUM_CPU)]
    frequencies = [[] for _ in range(NUM_CPU)]
    previous_frequency_counts = [[] for _ in range(NUM_CPU)]
    weighted_frequency = [0.0 for _ in range(NUM_CPU)]

    if(NUM_CPU == 8):
        # 865
        thermal_zone = [1,2,3,4,7,8,9,10]
        POP_MEM_THERMAL_ZONE = 28
        THERMAL_SCALING_FACTOR = 0.001
    else:
        # 820 
        thermal_zone = [4, 6, 9, 11]
        POP_MEM_THERMAL_ZONE = 1
        THERMAL_SCALING_FACTOR = 0.1

# Main
def main(argv):

    # Get the number of cpus
    get_number_of_cpus()

    init_data_structs()

    number_of_screen_data_lines = 0
    csv_mode = 0
    try:
        opts, args = getopt.getopt(argv,"c")
    except getopt.GetoptError:
        print('normal mode: profile.py')
        print('csv mode:    profile.py -c')
        sys.exit(2)
    for opt, arg in opts:
        if opt == '-c':
            csv_mode=1

    reset_screen()
    initialize_frequency_information()


    if csv_mode:
        print(CSV_HEADER)
        while (True):
            try:
                collect_cpu_stats()
                calculate_weighted_frequency()
                print_csv_line()
                time.sleep(1)
            # Prevent the ugly stack trace when Ctrl-C is hit
            except KeyboardInterrupt:
                break
    else:
        while (True):
            try:
                number_of_screen_data_lines = 0
                print(HEADER_LINE)
                print(SEPARATOR_LINE)
                number_of_screen_data_lines += 2
                collect_cpu_stats()
                calculate_weighted_frequency()
                number_of_screen_data_lines += print_cpu_stats()
                number_of_screen_data_lines += print_gpu_stats()
                number_of_screen_data_lines += print_non_cpu_data()
                print('Maximum cpu temperature: {:5.1f}'.format(maximum_temperature))
                number_of_screen_data_lines += 1
                number_of_screen_data_lines += print_memory_stats()
                time.sleep(UPDATE_RATE_IN_SECONDS)
                prepare_screen(number_of_screen_data_lines)
            # Prevent the ugly stack trace when Ctrl-C is hit
            except KeyboardInterrupt:
                break

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