import logging
from conf import LisaLogging
LisaLogging.setup()

import getpass
from subprocess import Popen, PIPE

import os
import pandas as pd
import scipy.integrate as integrate
import re
import json

# Support to access the remote target
import devlib
from env import TestEnv

import trappy

import socket

from time import sleep

import netifaces as ni

"""
Explanation: Run Tests on ChromeOS using test_that
This notebook shows how to run tests on a Chromebook using test_that. At the end of the test results are collected as a dictionary as well as trace events if you specify so in the target configuration.
NOTE: if you want to receive a token from the benchmark before starting trace collection and power measurement, you need to modify the benchmark such that a specific UDP packet is sent to the host machine. As a reference, consider the following python script and convert it to the language in which the benchmark is written accordingly:
python
import socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.sendto("POWER", (&lt;TARGET_IP&gt;, 1234))
sock.close()
You should put this right before the instruction that starts the execution of the actual workload in the benchmark.
In case of Acquarium, you should put this in
shell
&lt;CROS_PATH&gt;/trunk/src/third_party/autotest/files/client/site_tests/graphics_WebGLAquarium/graphics_WebGLAcquarium.py:run_fish_test()
Being CROS_PATH the ChromeOS chroot.
End of explanation
"""


# Define your path to the ChromeOS installation folder
CROS_BASE = "/data/chromiumos"

# Setup a target configuration
my_conf = {
    
    # Target platform and board
    "platform"    : 'linux',
    
    # Target board IP/MAC address
    "host"        : '192.168.0.1',
    
    # Login credentials
    "username"    : 'root',
    "password"    : 'test0000',
    
    # Tools required by the experiments
    "tools"   : [ 'trace-cmd' ], 
    
    # FTrace events to collect for all the tests configuration which have
    # the "ftrace" flag enabled
    "ftrace"  : {
        "events" : [
           "cpu_frequency",
           "cpu_idle",
           "sched_switch"
        ],
        "buffsize" : 10 * 1024,
    },
}

# Initialize a test environment using:
# the provided target configuration (my_conf)
te = TestEnv(my_conf)
target = te.target

def get_host_ip():
    """Returns the IP of the local host"""
    ifs = ni.interfaces()
    for interface in ifs:
        if interface == 'lo':
            continue

        addresses = ni.ifaddresses(interface)
        if addresses.has_key(socket.AF_INET):
            return addresses[socket.AF_INET][0]['addr']

HOST_IP = get_host_ip()

"""
Explanation: Target Configuration
End of explanation
"""


CROS_SDK_BIN_PATH = CROS_BASE + "/chromium/tools/depot_tools/cros_sdk"
username = !id -un
CROS_PATH = CROS_BASE + "/chroot/home/" + username[0]

"""
Explanation: Set ChromeOS paths
End of explanation
"""


def parse_graphics_WebGLAquarium(results_dir):
    results_file = os.path.join(
        CROS_PATH,
        os.path.basename(results_dir),
        'results-1-graphics_WebGLAquarium/graphics_WebGLAquarium/results/keyval'
    )
    
    data = {}
    with open(results_file) as data_file:
        for line in data_file:
            if line.strip():
                key, val = line.split('=')
                data[key] = float(val)
        
    return data

parse_results = {
    # Acquarium
    'graphics_WebGLAquarium' : parse_graphics_WebGLAquarium
}

"""
Explanation: Test-specific Parser Functions
End of explanation
"""


def CrosSdkSession(password):
    """
    Create cros_sdk session. The user will be asked to type his password.
    
    :param password: host machine password
    :type password: str
    """
    cros_sdk_session = Popen(['sudo -Sk {}'.format(CROS_SDK_BIN_PATH)],
                             bufsize=1,
                             stdin=PIPE,
                             stdout=PIPE,
                             stderr=PIPE,
                             cwd=CROS_PATH,
                             shell=True)
    cros_sdk_session.stdin.write(password)
    cros_sdk_session.stdin.write('\n')
    return cros_sdk_session

def test_that(password, te, test, pwr_time_s, get_token=False):
    """
    Run a specific test using the test_that command.
    
    :param password: host machine password
    :type password: str
    
    :param te: Test Environment object
    :type te: env.TestEnv
    
    :param test: name of the test to be run
    :type test: str
    
    :param pwr_time_s: power measurement duration in seconds
    :type pwr_time_s: int
    
    :param get_token: if True wait for token before collecting traces
    :type get_token: bool
    """
    
    results_dir = "~/results-dir"
    pwr_file = "~/power.txt"
    
    test_cmd = 'test_that -b oak {} --results_dir {} {}\n'.format(te.ip,
                                                                  results_dir,
                                                                  test)
    pwr_cmd = 'dut-control -t {} -y dvfs1_mw dvfs2_mw > {}\n'.format(pwr_time_s,
                                                                     pwr_file)
    
    # Create cros_sdk session
    cros_sdk_session = CrosSdkSession(password)
        
    logging.info('#### Start %s execution', test)
    cros_sdk_session.stdin.write(test_cmd)
    
    if get_token:
        # Setup socket to get token from target
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind((HOST_IP, 1234))
        logging.debug('     Waiting for token...')
        while True:
            data, addr = sock.recvfrom(32)
            if addr[0] == te.ip and data == "POWER":
                break
        sock.close()    
        logging.debug('     Token received....')
    
    logging.debug('     Start trace collection')
    # Check if trace events need to be collected
    if te.ftrace:
        te.ftrace.start()
        
    sleep(5)
    # Start measuring power
    cros_sdk_session.stdin.write(pwr_cmd)
    
    # communicate will close the session when the command terminates
    cros_sdk_session.communicate()
    logging.info('#### Completed %s execution', test)
    
    if te.ftrace:
        te.ftrace.stop()
        te.ftrace.get_trace(os.path.join(te.res_dir, 'trace.dat'))
    
    # Parse results using test-specific parser
    results = parse_results[test](results_dir)
    
    # Copy results to our Test Environment results directory
    with open(os.path.join(te.res_dir, 'results.json'), 'w') as outfile:
        json.dumps(results, outfile)
    
    return {
        "results"  : results,
        "pwr_file" : pwr_file
    }

# ask user for host password
password = getpass.getpass()

# Run Acquarium and collect results
acquarium_res = test_that(password,
                          te,
                          'graphics_WebGLAquarium',
                          10,
                          get_token=True)

# Print Acquarium results
df = pd.DataFrame.from_dict(acquarium_res['results'], orient="index")
df.columns = ['Values']
df

"""
Explanation: Run a Test
End of explanation
"""


if te.ftrace:
    trappy.plotter.plot_trace(te.res_dir)

"""
Explanation: Plot trace data using TRAPpy
End of explanation
"""


# Collect SERVO trace
servo_trace_file = os.path.join(CROS_PATH,
                                os.path.basename(acquarium_res['pwr_file']))
time = []
dvfs1_wm = []
dvfs2_wm = []
with open(servo_trace_file, 'r') as f:
    for l in f:
        if l.startswith("@@"):
            continue

        info = re.split(r'[ :]+', l)
        if info[1] == "dvfs1_mw":
            time.append(float(info[0]))
            dvfs1_wm.append(float(info[2]))
        else:
            dvfs2_wm.append(float(info[2]))

# Create dataframes for power data from SERVO board
big_pwr = pd.DataFrame(dvfs1_wm, index=time, columns=['Power'])
little_pwr = pd.DataFrame(dvfs2_wm, index=time, columns=['Power'])

x = big_pwr.index.get_values()
y = big_pwr.Power.get_values()
bnrg = integrate.simps(y, x=x)

x = little_pwr.index.get_values()
y = little_pwr.Power.get_values()
lnrg = integrate.simps(y, x=x)
results = {
    'big Cluster' : bnrg,
    'LITTLE Cluster' : lnrg,
    'Total' : bnrg + lnrg,
}

df = pd.DataFrame.from_dict(results, orient="index")
df.columns = ['Energy [mJ]']
df

"""
Explanation: Compute Energy
End of explanation
"""


