'''
Created on 30 May 2017

@author: alejomc
'''
'''
coordinates.py

2012 Brandon Mechtley
Arizona State University

Various tools for working with barycentric coordinates of arbitrary dimension
including routines for display. All routines automatically normalize
barycentric coordinates.
'''

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as pp

def project_pointline(p, a, b):
    '''
    Euclidean projection of a point onto a line segment.

    Args:
        p (np.ndarray): point of arbitrary dimensionality.
        a (np.ndarray): first endpoint of the line segment.
        b (np.ndarray): second endpoint of the line segment.

    Returns:
        (1, 2) np.ndarray of the projected cartesian coordinates.
    '''

    return a + (np.dot(p - a, b - a) / np.dot(b - a, b - a)) * (b - a)

def bary2cart(bary, corners=None):
    '''
    Convert barycentric coordinates to cartesian coordinates given the
    cartesian coordinates of the corners.

    Args:
        bary (np.ndarray): barycentric coordinates to convert. If this matrix
            has multiple rows, each row is interpreted as an individual
            coordinate to convert.
        corners (np.ndarray): cartesian coordinates of the corners.

    Returns:
        2-column np.ndarray of cartesian coordinates for each barycentric
        coordinate provided.
    '''

    if corners == None:
        corners = polycorners(bary.shape[-1])

    cart = None

    if len(bary.shape) > 1 and bary.shape[1] > 1:
        cart = np.array([np.sum(b / np.sum(b) * corners.T, axis=1) for b in bary])
    else:
        cart = np.sum(bary / np.sum(bary) * corners.T, axis=1)

    return cart

def polycorners(ncorners=3):
    '''
    Return 2D cartesian coordinates of a regular convex polygon of a specified
    number of corners.

    Args:
        ncorners (int, optional) number of corners for the polygon (default 3).

    Returns:
        (ncorners, 2) np.ndarray of cartesian coordinates of the polygon.
    '''

    center = np.array([0.5, 0.5])
    points = []

    for i in range(ncorners):
        angle = (float(i) / ncorners) * (np.pi * 2) + (np.pi / 2)
        x = center[0] + np.cos(angle) * 0.5
        y = center[1] + np.sin(angle) * 0.5
        points.append(np.array([x, y]))

    return np.array(points)
# import numpy
# corners = polycorners(3)
# print(corners)
# X = numpy.eye(corners.shape[0])
# print(X)
# cart = bary2cart(X, corners=corners)
# print(cart)

def cart2bary(X, P):
    import numpy
    #numpy.set_printoptions(precision=10)
    #P is points
    #X is the simplex
    M = P.shape[0]
    N = P.shape[1]

    assert X.shape[1] == N,  "Simplex X must have same number of columns as point matrix P"

    assert X.shape[0] == (N+1), "Simplex X must have N columns and N+1 rows"
    X1 = X[:-1,] - (numpy.multiply(numpy.ones((N,1)), X[-1,]))

    assert 1.0/numpy.linalg.cond(X1) > 0.00000001, "Degenerate simplex"

    Beta = numpy.matmul((P - numpy.tile(X[-1,], (M,1))), numpy.linalg.inv(X1))

    Beta = numpy.hstack((Beta, (1-numpy.sum(Beta, axis=1).reshape(-1,1))))
    return numpy.round(Beta,10) + 0.0

# print(cart2bary(corners, cart))

def lattice(ncorners=3, sides=False):
    '''
    Create a lattice of linear combinations of barycentric coordinates with
    ncorners corners. This lattice is constructed from the corners, the center
    point between them, points between the corners and the center, and pairwise
    combinations of the corners and the center.

    Args:
        ncorners (int, optional): number of corners of the boundary polygon
            (default 3).
        sides (bool, optional): whether or not to include pairwise combinations
            of the corners (i.e. sides) in the lattice (default False).

    Returns:
        np.ndarray of barycentric coordinates for each point in the lattice.
    '''

    # 1. Corners.
    coords = list(np.identity(ncorners))

    # 2. Center.
    center = np.array([1. / ncorners] * ncorners)
    coords.append(center)

    # 3. Corner - Center.
    for i in range(ncorners):
        for j in range(i + 1, ncorners):
            coords.append((coords[i] + coords[j] + center) / 3)

    # 4. Corner - Corner - Center.
    for i in range(ncorners):
        coords.append((coords[i] + center) / 2)

    # 5. Corner - Corner (Sides)
    if sides:
        for i in range(ncorners):
            for j in range(i + 1, ncorners):
                coords.append((coords[i] + coords[j]) / 2)

    # 6. Return unique coordinates (some duplicates using this method with
    # e.g. ncorners=2)
    return np.array(list(set(tuple(c) for c in coords)))

def verttext(pt, txt, center=[.5,.5], dist=1./15, color='red'):
    '''
    Display a text label for a vertex with respect to the center. The text
    will be a certain distance from the specified vertex in the direction of
    the vector extending from the center point.

    Args:
        pt (np.ndarray): two-dimensional array of cartesian coordinates of the
            point to label.
        txt (str): text to display. Text will be horizontally and vertically
            centered around pt.
        center (np.ndarray, optional): reference center point used for
            arranging the text around pt (default [.5, .5]).
        dist (float, optional): distance between point and the text
            (default 1./15).
        color (str, optional): matplotlib color of the text (default 'red').
    '''

    vert = pt - center
    s = np.sum(np.abs(vert))

    if s == 0:
        vert = np.array([0., 1.])
    else:
        vert /= s

    vert *= dist

    pp.text(
        pt[0] + vert[0],
        pt[1] + vert[1],
        txt,
        horizontalalignment='center',
        verticalalignment='center',
        color=color
    )

def polyshow(
    coords,
    color=None,
    label=None,
    labelvertices=False,
    polycolor=None,
    lines=[]
):
    '''
    Plot a regular convex polygon surrounding one or more barycentric
    coordinates within the it. Vertices and corners will be labeled
    sequentially starting at 0.

    Args:
        coords (np.ndarray or list): one or more barycentric coordinates of
            equal arbitrary dimension. The dimensionality of the coordinates
            will correspond to the number of vertices of the polygon that is
            drawn.
        color (str or list, optional): color in which to draw the coords. If
            color is a list of the same length as coords, each entry will
            correspond to the respective coordinates.
    '''

    # Defaults.
    coords = np.array(coords)
    if len(coords.shape) < 2: coords = [coords]
    for coord in coords:
        if np.sum(coord) > 0:
            coord /= np.sum(coord)

    if color == None: color = 'blue'
    if type(color) == str: color = [color] * len(coords)

    if label == None: label = ''
    if type(label) == str: label = [label] * len(coords)

    # Number of sides.
    d = len(coords[0])

    # Cartesian coordinates of the vertices of the polygon and each point.
    corners = polycorners(d)
    cart = np.array([
        np.sum([
            c * cnr
            for c, cnr in zip(coord, corners)
        ], axis=0)
        for coord in coords
    ])

    # Figure/axes setup.
    f = pp.figure(frameon=False)
    ax = pp.axes(frameon=False)
    ax.axis('equal')
    pp.xticks([])
    pp.yticks([])

    # Add the polygon and its vertices to the figure.
    ax.add_patch(pp.Polygon(corners, closed=True, fill=False, alpha=0.5))
    ax.scatter(corners[:,0], corners[:,1], color='red', s=50)
    if labelvertices:
        map(lambda i: verttext(corners[i], '$v_{%d}$' % i), range(len(corners)))

    # Add any extra lines to the figure.
    map(ax.add_line, lines)

    # Add the interior points and their labels.
    ax.scatter(cart[:,0], cart[:,1], color=color, s=100)

    for c, txt, clr in zip(cart, label, color):
        verttext(c, txt, color=clr)

    return f

def baryedges(coords, sidecoords=None):
    '''
    Return an array of barycentric coordinates corresponding to the closest
    point on each edge of the respective polygon.

    Args:
        coords (np.ndarray): input coordinates.

    Returns:
        np.ndarray of barycentric coordinates.
    '''

    # There ought to be an easier way to do this without switching out of
    # barycentric coordinates, but after obsessing over it for an entire day,
    # I'll have to work on that later.

    d = len(coords)                     # number of edges.
    corners = polycorners(d)            # cart. corners of the polygon.
    cart = bary2cart(coords, corners)   # cart. coordinates of the input point.

    e = np.zeros((d, d))                # bary. coords. for proj. to each side.

    # Pairs of corners that form sides (01, 12, 20 for a triange).
    pairs = zip(np.arange(d), np.roll(np.arange(d), -1))

    # Proj. the pt onto each edge in cart. space, put it back into bary. coords.
    for i1, i2 in pairs:
        proj = project_pointline(cart, corners[i1], corners[i2])

        distances = np.array([
            np.linalg.norm(corners[a] - proj, 2)
            for a in [i1, i2]
        ])

        sidebary = 1.0 - distances / np.sum(distances)
        e[i1, (i1, i2)] = sidebary

    if sidecoords:
        return np.array([e[i1, (i1, i2)] for i1, i2 in pairs])

    return e

def circumcircle(a, b, c):
    '''
    Return the center coordinates of a circle that circumscribes an input
    triangle defined by 2D vertices a, b, and c.

    Args:
        a ((number, number) or np.ndarray): first vertex of the input triangle.
        b ((number, number) or np.ndarray): second vertex of the input triangle.
        c ((number, number) or np.ndarray): third vertex of the input triangle.

    Returns:
        Center coordinates of circumscribing circle in form (x, y).
    '''

    ax, ay = a
    bx, by = b
    cx, cy = c
    ax2, ay2, bx2, by2, cx2, cy2 = [d ** 2 for d in [ax, ay, bx, by, cx, cy]]

    d = 2 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by))

    ux = (
        (ax2 + ay2) * (by - cy) +
        (bx2 + by2) * (cy - ay) +
        (cx2 + cy2) * (ay - by)
    ) / d

    uy = (
        (ax2 + ay2) * (cx - bx) +
        (bx2 + by2) * (ax - cx) +
        (cx2 + cy2) * (bx - ax)
    ) / d

    return ux, uy

def voronoi(x, y):
    '''
    Return a list of voronoi cells for a collection of points.

    Args:
        x (list or np.ndarray): list of coordinates' x components.
        y (list or np.ndarray): list of coordinates' y components.

    Returns:
        (cells, triangles), where cells is a list of voronoi cells, each once
        containing a list of two-dimensional points; and triangles is a list of
        the triangles from a Delaunay triangulation.
    '''

    p = np.array([a for a in zip(x, y)])
    d = mpl.tri.Triangulation(x, y)
    t = d.triangles
    n = t.shape[0]

    # Get circle for each triangle, center will be a voronoi cell point.
    cells = [[] for i in range(x.size)]

    for i in range(n):
        v = [p[t[i,j]] for j in range(3)]
        pt = circumcircle(v[0], v[1], v[2])

        cells[t[i,0]].append(pt)
        cells[t[i,1]].append(pt)
        cells[t[i,2]].append(pt)

    # Reordering cell p in trigonometric way
    for i, cell in enumerate(cells):
        xy = np.array(cell)
        order = np.argsort(np.arctan2(xy[:,1] - y[i], xy[:,0] - x[i]))

        cell = xy[order].tolist()
        cell.append(cell[0])

        cells[i] = cell

    return cells

def baryplot(
    values,
    points=None,
    labels='abc',
    cmap=mpl.cm.BrBG,
    clabel='',
    vmin=None,
    vmax=None
):
    '''
    Create a triangular voronoi cell pseudocolor plot. Create a voronoi diagram
    for each coordinate (points) within the triangle and color each cell
    according to its value (values).

    Args:
        values (list or np.ndarray): list of scalar values for each barycentric
            point.
        points (list or np.ndarray): list of three-parameter barycentric
            points.
        labels (list): list of three label strings, one for each side of the
            triangle.
        cmap (matplotlib.colors.Colormap): colormap for pseudocolor plot.
        clabel (str): colorbar label for values
        vmin (number): minimum value for coloring. If None, use minimum of the
            input values.
        vmax (number): maximum value for coloring. If None, use maximum of the
            input values.
    '''

    if points is None: points = []

    vmin = vmin if vmin is not None else np.amin(values)
    vmax = vmax if vmax is not None else np.amax(values)

    p = bary2cart(points) if len(points) else bary2cart(lattice(3))
    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
    mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    colors = mappable.to_rgba(values)

    #values = (values - np.amin(values)) / (np.amax(values) - np.amin(values))
    cells = voronoi(p[:,0], p[:,1])

    xmin, xmax, xavg = np.amin(p[:,0]), np.amax(p[:,0]), np.mean(p[:,0])
    ymin, ymax, yavg = np.amin(p[:,1]), np.amax(p[:,1]), np.mean(p[:,1])

    s60, c60 = np.sin(np.pi / 3.), np.cos(np.pi / 3.)
    s30, c30 = np.sin(np.pi / 6.), np.cos(np.pi / 6.)

    # Start drawing.
    ax = pp.gca()

    # Clipping triangle for the voronoi patches.
    clip = mpl.patches.Polygon([
        (xmin, ymin), (xmax, ymin), (xavg, ymax),
    ], transform=ax.transData)

    # Draw voronoi patches.
    for i, cell in enumerate(cells):
        codes = [mpl.path.Path.MOVETO] \
            + [mpl.path.Path.LINETO] * (len(cell) - 2) \
            + [mpl.path.Path.CLOSEPOLY]

        pth = mpl.path.Path(cell, codes)

        patch = mpl.patches.PathPatch(
            pth,
            zorder=-1,
            facecolor=colors[i],
            clip_path=clip,
            edgecolor='none'
        )

        ax.add_patch(patch)

    # Add barycentric labels for vertices.
    ax.text(xmin - .0125, ymin - .02, '$(0,1,0)$', ha='right', va='center')
    ax.text(xmax + .0125, ymin - .02, '$(0,0,1)$', ha='left', va='center')
    ax.text(xavg, ymax + .035, '$(1,0,0)$', ha='center', va='bottom')

    # Labels.
    ax.text(
        xavg + c30 * .35, yavg + s30 * .35,
        labels[2], ha='center', va='center', rotation=-60
    )

    ax.text(
        xavg, ymin - .05,
        labels[1], ha='center', va='top'
    )

    ax.text(
        xavg - c30 * .35, yavg + s30 * .35,
        labels[0], ha='center', va='center', rotation=60
    )

    arrowopts = dict(
        width=.00125,
        frac=.0125,
        headwidth=.01,
        transform=ax.transData
    )

    fig = pp.gcf()

    # Arrows along edges.
    ax.add_patch(mpl.patches.YAArrow(
        fig,
        (xmin - c60 * .025, ymin + s60 * .025),
        (xavg - c60 * .025, ymax + s60 * .025),
        **arrowopts
    ))

    ax.add_patch(mpl.patches.YAArrow(
        fig,
        (xmax, ymin - .025),
        (xmin, ymin - .025),
        **arrowopts
    ))

    ax.add_patch(mpl.patches.YAArrow(
        fig,
        (xavg + c60 * .025, ymax + s60 * .025),
        (xmax + c60 * .025, ymin + s60 * .025),
        **arrowopts
    ))

    # Make axes equal, get rid of border.
    pp.axis('equal')
    ax.axis([
        xmin - c60 * .2, xmax + c60 * .2,
        ymin - s60 * .2, ymax + s60 * .2
    ])
    pp.axis('off')

    cax, kw = mpl.colorbar.make_axes(ax, orientation='vertical', shrink=0.7)
    cb = mpl.colorbar.ColorbarBase(
        cax,
        cmap=cmap,
        norm=norm,
        orientation='vertical',
        ticks=np.linspace(vmin, vmax, 5)
    )

    cb.set_label(clabel)


import numpy as np

class Array:

    def stack_it(self):
        to_stack  = np.array(self.values)
        if not np.array_equal(to_stack, self.pile[-1]):
            self.pile = np.vstack((self.pile, np.array(self.values)))

    def __init__(self, values, lower_index=0):
        self.lower_index = lower_index
        self.values = list(values)

        self.pile = np.array(self.values)
        Array.full_array = list(values)

    def swap(self, index1, index2):
        self.values[index2], self.values[index1] = self.values[index1], self.values[index2]
        Array.full_array[self.lower_index + index2], Array.full_array[self.lower_index + index1] = Array.full_array[self.lower_index + index1], Array.full_array[self.lower_index + index2]
        self.stack_it()

    def get_len(self):
        return len(self.values)

def bubble_sort(nums):  # n^2
    # We set swapped to True so the loop looks runs at least once
    swapped = True
    while swapped:
        swapped = False
        for i in range(nums.get_len() - 1):
            if nums.values[i] > nums.values[i + 1]:
                # Swap the elements
                nums.swap(i, i + 1)
                # Set the flag to True so we'll loop again
                swapped = True

# A function to sort the given list using Gnome sort
def gnome_sort(nums):
    idx = 0
    while idx < nums.get_len():
        if idx == 0 or nums.values[idx] > nums.values[idx-1]:
            idx += 1
        else:
            nums.swap(idx, idx-1)
            idx -= 1


# Read and split by =
def read_config(file_location):
	file_to_read = open(file_location, 'r')
	lines = file_to_read.readlines()
	no_blank_lines = [l for l in (line.strip() for line in lines) if l]

	config_dict = {}
	for line in no_blank_lines:
		if line:
			formatted_line = line.split("|")
			config_dict[formatted_line[0]] = formatted_line[1]

	return config_dict

# Write to Config
def update_config(file_location, updated_dict):
	dict_string = ""
	for key, value in updated_dict.items():
		dict_string += '{0}|{1}\n'.format(key, value)

	file_to_write = open(file_location, 'w')
	file_to_write.write(dict_string)
	file_to_write.close()

# Read line by line
def read_line_data(file_location):
	file_to_read = open(file_location, 'r')
	lines = file_to_read.readlines()
	no_blank_lines = [l for l in (line.strip() for line in lines) if l]

	line_list = []
	for line in no_blank_lines:
		if line:
			formatted_line = line.strip()
			line_list.append(formatted_line)

	return line_list


import os
import shutil
import tasks.replace_contents_task

class TestReplaceContentsTask(object):
	task = tasks.replace_contents_task.ReplaceContentsTask()

	def test_execute_task(self):
		shutil.copy('test/res/CodeSample.m', 'test/res/CodeSampleCp.m')
		task_parameters = {'path': 'test/res/CodeSampleCp.m', 'to_replace': 'application(\\s+)', 'replace_with': 'frustration\\1'}
		self.task.execute_task(task_parameters)

		code_sample = ''

		with open('test/res/CodeSampleCp.m', 'r') as code_sample_file:
			code_sample = code_sample_file.read()

		os.remove('test/res/CodeSampleCp.m')

		expected = '#import "AppDelegate.h"\n@implementation AppDelegate\n- (BOOL)application:(UIApplication *)frustration didFinishLaunchingWithOptions:(NSDictionary *)launchOptions {\nreturn YES;\n}\n@end\n'
		assert code_sample == expected


'''
5.Write a Python program to check whether a given number is a narcissistic number
or not.
For example, 371 is a narcissistic number; it has three digits,
and if we cube each digits  33 + 73 + 13 the sum is 371.
Other 3-digit narcissistic numbers are
153 = 13 + 53 + 33
370 = 33 + 73 + 03
407 = 43 + 03 + 73.
There are also 4-digit narcissistic numbers,
some of which are 1634, 8208, 9474 since
1634 = 14+64+34+44
8208 = 84+24+04+84
9474 = 94+44+74+44
'''
n=(input())
s=0
q=len(n)
t=int(n)
for i in range(q):
    a=int(n[i])**len((n))
    s=s+a

if t==s:
     print("number is narcissistic",s)
        #break
else:
    print("number is not narcissistic",s)


from mesa.visualization.modules import CanvasGrid
from mesa.visualization.ModularVisualization import ModularServer
from mesa.visualization.modules import ChartModule

from models import HonestyModel
from helpers import SETTINGS

def agent_portrayal(agent):
    portrayal = {"Shape": "circle", "Filled": "true", "r": 0.5}

    if agent.is_thief:
        portrayal["Color"] = "red"
        portrayal["Layer"] = 0
    elif agent.is_jailed:
        portrayal["Color"] = "gray"
        portrayal["Layer"] = 2
    else:
        portrayal["Color"] = "green"
        portrayal["Layer"] = 1

    return portrayal

grid = CanvasGrid(
    agent_portrayal, SETTINGS["GRID"]["WIDTH"], SETTINGS["GRID"]["HEIGHT"], 500, 500
)

chart = ChartModule(
    [
        {"Label": "thieves", "Color": "Red"},
        {"Label": "normal", "Color": "Green"},
        {"Label": "is_jailed", "Color": "Gray"},
    ],
    data_collector_name="data_collector",
)

server = ModularServer(
    HonestyModel,
    [grid, chart],
    "Honesty Model",
    {
        "initial_normal_agents": SETTINGS["POPULATION"]["INITIAL_HONEST"],
        "initial_thieves": SETTINGS["POPULATION"]["INITIAL_THIEVES"],
        "width": SETTINGS["GRID"]["WIDTH"],
        "height": SETTINGS["GRID"]["HEIGHT"],
    },
)

server.port = 8002
server.launch()


import numpy as np
from scipy.linalg import eigh

class Diffusion:
    """
    Regression weights of kernel Laplacian kernel regularization

    Notes
    -----
    This implementation has not been optimized in term of memory.
    Memory scales in O(p^2 nd) while it could scale in O(nd + p^2).

    Examples
    --------
    >>> import numpy as np
    >>> x_support = np.random.randn(50, 10)
    >>> lap = Diffusion()
    >>> lap.set_support(x_support)
    >>> lap.update_sigma(1, p=20)
    >>> lap.update_mu(1 / len(x_support))
    >>> lap.update_psi(lambd=1e-3)
    >>> x = np.random.randn(30, 10)
    >>> alpha = lap(x)
    """
    def __init__(self, sigma=None, lambd=None, psi=None, mu=None,
                 p=None, full=False, nl=None):
        """

        Parameters
        ----------
        sigma: float
            Bandwidth parameter for the Gaussian kernel
        lambd: float, optional
            Tikhonov regularization parameter
        psi: lambda function, optional
            Filter function
        mu: float, optional
            Regularization parameter for GSVD
        p: int, optional
            Subsampling parameter for low-rank approximation
        full: bool, optional
            Specifies if we use full representation or low-rank representation
            to compute the empirical risk minimizer (default is False).
        nl: int, optional
            Specifies if covariance should be computed only on the first `nl` training points
        """
        self.kernel = GaussianKernel(sigma)

        if lambd is not None:
            self.Tikhonov = True
            if psi is not None:
                raise ValueError('Filter and Tikhonov regularization ' \
                                 + 'can not be specified simultaneously')
            self.lambd = lambd
        else:
            self.Tikhonov = False
            self.psi = psi

        self.mu = mu
        if (p is not None) and full:
            raise NotImplementedError('`p` and `full` can not be specified simultaneously.')
        self.p = p
        self.full = full
        self.nl = nl

    def set_support(self, x_train):
        """Specified input training data"""
        self.n_train = len(x_train)
        self.kernel.set_support(x_train)

    def update_sigma(self, sigma=None, full=None, p=None, nl=None):
        """Setting bandwith parameter

        This should be call after setting the support `x_train`.
        There should be a call to update lambda after setting the bandwith
        """
        if sigma is not None:
            self.kernel.sigma = sigma
            self.kernel.reset()
        self.kernel.set_K()
        self.kernel.set_SZ()

        # Catching arguments
        if full is not None:
            self.full = full
            if self.full:
                self.p = None
        if p is not None:
            self.p = p
        if nl is not None:
            self.nl = nl
        if (self.p is not None) and self.full:
            raise NotImplementedError('`p` and `full` can not be specified simultaneously.')

        # Full representation
        if self.full:
            self.kernel.set_ST()
            self.kernel.set_ZZ()
            self.kernel.set_TZ()
            self.kernel.set_TT()

            ST = self.kernel.ST
            TZ = self.kernel.TZ

            if self.nl is None:
                self.A = ST.T @ ST
                self.A /= self.n_train
            else:
                self.A = ST.T[..., :self.nl] @ ST[:self.nl]
                self.A /= self.nl

            self.B_0 = TZ @ TZ.T
            self.B_0 /= self.n_train

        # Small representation
        else:
            SS = self.kernel.K
            SZ = self.kernel.SZ

            if self.nl is None:
                self.A = SS[:self.p, ...] @ SS[..., :self.p]
                self.A /= self.n_train
            else:
                self.A = SS[:self.p, :self.nl] @ SS[:self.nl, :self.p]
                self.A /= self.nl

            self.B_0 = SZ[:self.p] @ SZ[:self.p].T
            self.B_0 /= self.n_train

    def update_mu(self, mu=None, mu_numerical=10e-7, Tikhonov=False):
        """Setting GSVD regularization parameter"""

        if not hasattr(self, 'B_0'):
            self.update_sigma()

        if mu is not None:
            self.mu = mu
        if self.mu is None:
            raise ValueError('GSVD regularization has not been specified.')

        if self.full:
            self.B = self.B_0 + self.mu * self.kernel.TT
            self.B += mu_numerical * np.eye(self.B.shape[0])
        else:
            self.B = self.B_0 + self.mu * self.kernel.K[:self.p,:self.p]
            self.B += mu_numerical * np.eye(self.B.shape[0])

        if Tikhonov:
            self.Tikhonov = True

        # compute GSVD once and try many filter functions after
        if not self.Tikhonov:
            self.v, self.e = eigh(self.A, self.B)

    def update_psi(self, nl=None, psi=None, lambd=None):
        """Setting Filter function

        Parameters
        ----------
        nl: int, optional
            Number of labelled data among `x_train`, default is `len(x_train)`
        """

        if not hasattr(self, 'B'):
            self.update_mu()

        if nl is None:
            nl = self.nl
        if nl is None:
            nl = self.n_train

        if (psi is not None) and (lambd is not None):
            raise ValueError('Filter and Tikhonov regularization ' \
                             + 'can not be specified simultaneously')

        if psi is not None:
            self.Tikhonov = False
            self.psi = psi

        if lambd is not None:
            self.Tikhonov = True
            self.lambd = lambd

        # full representation
        if self.full:
            b = self.kernel.ST.T[..., :nl] / nl
        # small representation
        else:
            b = self.kernel.K[:self.p, :nl] / nl

        if self.Tikhonov:
            # much faster implementation than GSVD
            self.c = np.linalg.solve(self.A + self.lambd * self.B, b)
        else:
            self.c = (self.e * self.psi(self.v)) @ (self.e.T @ b)

    def __call__(self, x_test):
        """Diffusion scheme computation.

        Parameters
        ----------
        x_test : ndarray
            Points to compute kernel ridge regression weights, of shape (nb_points, input_dim).

        Returns
        -------
        out : ndarray
            Similarity matrix of size (nb_points, n_train) given by kernel Laplacian regularization.
        """
        if not hasattr(self, 'c'):
            self.train()

        if self.full:
            T_x = self.kernel.get_ST(x_test)
            return T_x @ self.c
        else:
            K_x = self.kernel.get_SS(x_test)[:self.p]
            return K_x.T @ self.c

    def train(self, sigma=None, full=None, p=None, n_cov=None, nl=None,
              mu=None, psi=None, lambd=None):
        self.update_sigma(sigma=sigma, full=full, p=p, nl=n_cov)
        if lambd is not None:
            self.Tikhonov = True
        self.update_mu(mu=mu)
        self.update_psi(nl=nl, psi=psi, lambd=lambd)

    def set_phi(self, phi):
        self.c_beta = self.c @ phi

    def call_with_phi(self, x):
        if self.full:
            T_x = self.kernel.get_ST(x)
            return T_x @ self.c_beta
        else:
            K_x = self.kernel.get_SS(x)[:self.p]
            return K_x.T @ self.c_beta

class GaussianKernel:
    """
    Computation of Gaussian kernel and its derivatives
    """

    def __init__(self, sigma=None):
        self.sigma = sigma

    def set_support(self, x):
        """Set train support for kernel method.

        Parameters
        ----------
        x : ndarray
            Training set given as a design matrix, of shape (nb_points, input_dim).
        """
        self.reset()
        self.x = x
        self.n_train = len(x)

    def get_SS(self, x=None):
        """Gaussian kernel.

        .. math:: k(x, y) = exp(-norm{x - y}^2 / (2 * sigma2)).

        Parameters
        ----------
        x: ndarray
            Design matrix of shape (nb_points, input_dim).

        Returns
        -------
        K: ndarray
            Gram matrix of size (n_train, nb_points), similarity between
            training points and testing points.
        """
        if x is None:
            x = self.x
        K = self.x @ x.T
        K *= 2
        if not hasattr(self, "_attr_1"):
            self._attr1 = np.sum(self.x ** 2, axis=1)[:, np.newaxis]
        K -= self._attr1
        K -= np.sum(x ** 2, axis=1)
        K /= 2 * (self.sigma ** 2)
        np.exp(K, out=K)
        return K

    def set_K(self):
        if not hasattr(self, 'K'):
            self.K = self.get_SS()

    def get_SZ(self, x=None, SS=None, reshape=True):
        """First derivative of the Gaussian kernel

        Returns
        -------
        SZ: ndarray
            Array of size (nb_points, d * n_train) or (nb_points, n_train, d)
            SZ[i, j*m] = :math:`\partial_{1,m}` k(x_train[j], x[i])
        """
        if x is None:
            x = self.x
        if SS is None:
            SS = self.get_SS(x).T

        SZ = np.tile(SS[...,np.newaxis], (1, 1, self.x.shape[1]))
        # diff[i,j,k] = x[i,k] - self.x[j,k]
        diff = x[:, np.newaxis, :] - self.x[np.newaxis, ...]
        diff /= self.sigma**2
        # SZ[i, j, k] = (x[i, k] - self.x[j, k]) * k(x[i], self.x[j])
        SZ *= diff

        if reshape:
            n, d = x.shape
            # return SZ.reshape(n, -1, order='F') # slower than the following
            SZ_reshape = np.empty((n, self.n_train*d), SZ.dtype)
            for i in range(d):
                SZ_reshape[:, i*self.n_train:(i+1)*self.n_train] = SZ[..., i]
            return SZ_reshape
        return SZ

    def set_SZ(self):
        if not hasattr(self, 'SZ'):
            self.set_K()
            self.SZ = self.get_SZ(SS=self.K)

    # Matrices for exact ERM computations

    def get_ST(self, x=None, SS=None, SZ=None):
        """Matrix based on derivatives of the Gaussian kernel

        Based on T = [S, Z].
        """
        if x is None:
            x = self.x
        if SS is None:
            SS = self.get_SS(x).T
        if SZ is None:
            SZ = self.get_SZ(x, SS=SS)
        n, d = x.shape
        ST = np.zeros((n, self.n_train*(d+1)), dtype=np.float)
        ST[:, :self.n_train] = SS
        ST[:, self.n_train:] = SZ
        return ST

    def set_ST(self):
        if not hasattr(self, 'ST'):
            self.set_SZ()
            self.ST = self.get_ST(SS=self.K, SZ=self.SZ)

    def get_ZZ(self, x=None, SS=None, reshape=True):
        """Double derivative of the Gaussian kernel

        Returns
        -------
        ZZ: ndarray
            Array of size (n_train * d, nb_points * d) or (n_train, nb_points, d, d)
            ZZ[i*k, j*m] = :math:`\partial_{1,k}\partial_{2,m}` k(x[i], x_train[j])
        """
        if x is None:
            x = self.x
        if SS is None:
            SS = self.get_SS(x).T
        n, d = x.shape
        ZZ = np.tile(SS[...,np.newaxis, np.newaxis], (1, 1, d, d,))
        # diff[i,j,k] = x[i,k] - self.x[j,k]
        diff = x[:, np.newaxis, :] - self.x[np.newaxis, ...]
        # prod_diff[i,j,k,l] = diff[i,j,l]*diff[i,j,k] = (x[i,l] - self.x[j,l]) * (x[i,k] - self.x[j,k])
        prod_diff = diff[:,:, np.newaxis, :]*diff[:,:,:,np.newaxis]
        prod_diff /= self.sigma**4
        prod_diff *= -1
        for i in range(d):
            prod_diff[:, :, i, i] += 1 / (self.sigma**2)
        ZZ *= prod_diff
        if reshape:
            # return ZZ.transpose((0, 2, 1, 3)).reshape(n * d, self.n_train * d, order='F') # slower
            ZZ_reshape = np.empty((n*d, self.n_train*d), ZZ.dtype)
            for i in range(d):
                for j in range(i):
                    ZZ_reshape[n*i:n*(i+1), self.n_train*j:self.n_train*(j+1)] = ZZ[..., i, j]
                    ZZ_reshape[n*j:n*(j+1), self.n_train*i:self.n_train*(i+1)] = ZZ[..., j, i]
                ZZ_reshape[n*i:n*(i+1), self.n_train*i:self.n_train*(i+1)] = ZZ[..., i, i]
            return ZZ_reshape
        return ZZ

    def set_ZZ(self):
        if not hasattr(self, 'ZZ'):
            self.set_K()
            self.ZZ = self.get_ZZ(SS=self.K)

    def get_TZ(self, x=None, SS=None, SZ=None, ZZ=None):
        if x is None:
            x = self.x
        if SS is None:
            SS = self.get_SS(x).T
        if SZ is None:
            SZ = self.get_SZ(x, SS)
        if ZZ is None:
            ZZ = self.get_ZZ(x, SS)
        n, d = x.shape
        TZ = np.zeros((n*(d+1), self.n_train*d), dtype=np.float)
        TZ[:n,:] = SZ
        TZ[n:,:] = ZZ
        return TZ

    def set_TZ(self):
        if not hasattr(self, 'TZ'):
            self.set_K()
            self.set_SZ()
            self.set_ZZ()
            self.TZ = self.get_TZ(SS=self.K, SZ=self.SZ, ZZ=self.ZZ)

    def get_TT(self, x=None, SS=None, SZ=None, ZZ=None):
        if x is None:
            x = self.x
        else:
            raise NotImplementedError('Implementation was not finished for TT')
        if SS is None:
            SS = self.get_SS(x).T
        if SZ is None:
            SZ = self.get_SZ(x)
        if ZZ is None:
            ZZ = self.get_ZZ(x)
        n, d = x.shape
        TT = np.zeros((n*(d+1), self.n_train*(d+1)), dtype=np.float)
        TT[:n, :self.n_train] = SS
        TT[:n, self.n_train:] = SZ
        TT[n:, :self.n_train] = SZ.T
        TT[n:, self.n_train:] = ZZ
        return TT

    def set_TT(self):
        if not hasattr(self, 'TT'):
            self.set_K()
            self.set_SZ()
            self.set_ZZ()
            self.TT = self.get_TT(SS=self.K, SZ=self.SZ, ZZ=self.ZZ)

    def reset(self):
        """Resetting attributes."""
        atts = ['_attr_1', 'K', 'SZ', 'ST', 'ZZ', 'TZ', 'TT']
        for att in atts:
            if hasattr(self, att):
                delattr(self, att)

if __name__=="__main__":
    x_support = np.random.randn(50, 10)
    lap = Diffusion()
    lap.set_support(x_support)
    lap.update_sigma(1, p=20)
    lap.update_sigma(1)
    lap.update_mu(1 / len(x_support))
    lap.update_psi(lambd=1e-3)
    x = np.random.randn(30, 10)
    alpha = lap(x)
    assert(alpha.shape==(30,50))


import random

from zorro.filter import collect_unique_pairs
from zorro.words import get_legal_words
from zorro.counterbalance import find_counterbalanced_subset
from zorro import configs

template1 = {
    'b': 'even {nn1} ever {vbz} {det} {nn2} .',
    'g': 'only {nn1} ever {vbz} {det} {nn2} .',
}

template2 = {
    'b': 'even {nn1} {aux} ever {vb} {det} {nn2} .',
    'g': 'only {nn1} {aux} ever {vb} {det} {nn2} .',
}

def main():
    """
    example:
    "only sarah could ever talk." vs. "even sarah could ever talk"
    """

    vbzs = get_legal_words(tag='VBZ', exclude=('happens', 'says', ))
    vbs = get_legal_words(tag='VB')

    nouns_s = get_legal_words(tag='NN')

    animates_ = (configs.Dirs.legal_words / 'animates.txt').open().read().split()
    animates = find_counterbalanced_subset(animates_, min_size=8, max_size=len(animates_))

    names_ = (configs.Dirs.legal_words / 'names.txt').open().read().split()
    names = find_counterbalanced_subset(names_, min_size=10, max_size=len(names_))

    determiners = ['the', 'this', 'some', 'that'] + ['your', 'his', 'her']

    auxiliaries = ['could', 'can', 'would', 'will']

    def add_argument_after_vb(v: str,
                              argument1: str,
                              ) -> str:
        if v in {'thinks', 'reads'}:
            return f'{v} about'
        elif v in {'lives', 'falls', 'is', 'be'}:
            return f'{v} in'
        elif v in {'stands', 'turns'}:
            return f'{v} on'
        elif v in {'acts', 'looks'}:
            return f'{v} like'
        elif v in {'goes', 'comes'}:
            return f'{v} to'
        elif v in {'gives', 'gives'}:
            return f'{v} {argument1}'
        elif v in {'plays', 'play', 'shows', 'show', 'tells', 'tell'}:
            return f'{v} {argument1}'
        else:
            return v

    while True:

        arg1 = random.choice(['him', 'her'])
        vbz = random.choice(vbzs)
        vb = random.choice(vbs)

        # random choices
        slot2filler = {
            'nn1': random.choice(names + animates),
            'nn2': random.choice(nouns_s),
            'vbz': add_argument_after_vb(vbz, arg1),
            'vb': add_argument_after_vb(vb, arg1),
            'det': random.choice(determiners),
            'aux': random.choice(auxiliaries)
        }

        # add determiner to animate noun
        if slot2filler['nn1'] in animates:
            slot2filler['nn1'] = random.choice(determiners) + ' ' + slot2filler['nn1']

        yield template1['b'].format(**slot2filler)  # bad
        yield template1['g'].format(**slot2filler)  # good

        yield template2['b'].format(**slot2filler)  # bad
        yield template2['g'].format(**slot2filler)  # good

if __name__ == '__main__':
    for n, s in enumerate(collect_unique_pairs(main)):
        print(f'{n//2+1:>12,}', s)


# coding=utf-8

from wikilife_data.dao.base_dao import BaseDAO

READER_NAME_FIELD = "reader_name"
SINCE_ID_FIELD = "since_id"

class TwitterReaderDAO(BaseDAO):

    _collection = None

    def _initialize(self):
        self._collection = self._db.twitter_readers
        self._collection.ensure_index(READER_NAME_FIELD, unique=True)

    def find_reader(self, reader_name):
        filter = {}
        filter[READER_NAME_FIELD] = reader_name
        reader = self._collection.find_one(filter)
        return reader

    def create_reader(self, reader_name, initial_since_id=0):
        reader = {}
        reader[READER_NAME_FIELD] = reader_name
        reader[SINCE_ID_FIELD] = initial_since_id
        created_reader = self._collection.save(reader)
        return created_reader

    def get_reader_since_id(self, reader_name):
        reader = self.find_reader(reader_name)
        return reader[SINCE_ID_FIELD]

    def set_reader_since_id(self, reader_name, since_id):
        reader = self.find_reader(reader_name)
        reader[SINCE_ID_FIELD] = since_id
        updated_reader = self._collection.save(reader)
        return updated_reader


import os, sys
sys.path.append("/home/ip/IPNeuronaleNetze")
sys.path.append("/home/ip/IPNeuronaleNetze/trainers")
from data_loader.datasets import Dataset
import tensorflow as tf
from models.OurModel import OurModel
from trainers.Trainer import Trainer

# Please enter the filepath to your dataset.
filepath = "../data/classification/age/"

def main():
    # This main will first load the datasets then start the training for both, age as well as gender estimation
    # with the default amount of epochs which is the value 312 and an early stopping callback with the patience of 4.
    # You can change the amount of epochs by typing "epochs = X" into the Trainer constructor.
    # If you want to train a model by loading the weights of an allready trained model,
    # use the load_model method from the OurModel class.

    # After the training is done, the models will be saved and evulation files (.csv) will be created.
    # Use the parsers to interpret the .csv files

    dataset = Dataset()

    dataset.createFolder("../data/classification")
    dataset.createClassificationFolders("../data/classification/age/Train")
    dataset.createClassificationFolders("../data/classification/age/Valid")
    dataset.createClassificationFolders("../data/classification/age/Test")
    #LAP
    dataset.downloadAndUnzip("http://158.109.8.102/AppaRealAge/appa-real-release.zip", "../data/appa-real-release.zip",     "../data/appa-real-release")
    dataset.readAndPrintDataImagesLAP("../data/appa-real-release", "train", "../data/classification/age/Train")
    dataset.readAndPrintDataImagesLAP("../data/appa-real-release", "valid", "../data/classification/age/Valid")
    dataset.readAndPrintDataImagesLAP("../data/appa-real-release", "test", "../data/classification/age/Test")

    # Start training for age estimation
    AgeModel = OurModel(0)
    AgeModel = Trainer(AgeModel.model, filepath + "Train",
                        filepath + "Valid", filepath + "Test",
                        identifier = 0)
    AgeModel.train()

if __name__ == "__main__":
    main()


import sys
sys.path.append('..')
from app.client import Client
import unittest

class ClientTest(unittest.TestCase):
    client = {
        "name": "My entreprise",
        "address": {
            "number": "42",
            "street": "Rue du",
            "city": "Paris",
            "postal_code": "34",
            "additional_details": None
        }
    }

    def test_to_dict(self):
        c = Client(**self.client)
        self.assertEqual(c.to_dict(), self.client)

if __name__ == '__main__':
    unittest.main()


import matplotlib.pyplot as plt
import numpy as np
from sklearn.cluster import KMeans
from sklearn import datasets
import pandas as pd

inputfile = 'F:/machineLearning/DateSet/Kmeans/data111.csv'
np.set_printoptions(suppress=True)
data = pd.read_csv(inputfile)
#print(data)
dataset = data.values
X = dataset[:,:3]    #取第一列第二列
print(X)

# 绘制数据分布图
plt.scatter(X[:, 1], X[:, 2], c="red", marker='o', label='see')
plt.xlabel('good number')
plt.ylabel('good scord')
plt.legend(loc=2)
plt.show()

estimator = KMeans(n_clusters=3)   # 构造聚类器
estimator.fit(X) # 聚类
label_pred = estimator.labels_    # 获取聚类标签
print(label_pred)
print(len(label_pred))
# 绘制k-means结果
x0 = X[label_pred == 0]
x1 = X[label_pred == 1]
x2 = X[label_pred == 2]
plt.scatter(x0[:, 1], x0[:, 2], c="red", marker='o', label='label0')
plt.scatter(x1[:, 1], x1[:, 2], c="green", marker='*', label='label1')
plt.scatter(x2[:, 1], x2[:, 2], c="blue", marker='+', label='label2')
plt.xlabel('sepal length')
plt.ylabel('sepal width')
plt.legend(loc=2)
plt.show()


import os
from shutil import rmtree
import xml.etree.ElementTree as ET

from tg import tmpl_context as c
from testfixtures import TempDirectory

from alluratest.controller import setup_basic_test
from allura import model as M
from allura.lib import helpers as h
from allura.scripts.create_sitemap_files import CreateSitemapFiles

class TestCreateSitemapFiles:

    def setup_method(self, method):
        setup_basic_test()

    def run_script(self, options):
        cls = CreateSitemapFiles
        opts = cls.parser().parse_args(options)
        with h.push_config(c, user=M.User.anonymous()):  # tasks & scripts have c.user set
            cls.execute(opts)

    def test_create(self):
        with TempDirectory() as tmpdir:
            rmtree(tmpdir.path)  # needs to be non-existent for the script
            self.run_script(['-o', tmpdir.path])

            tmpdir.compare(['sitemap-0.xml', 'sitemap.xml'],recursive=False, files_only=True)

            xml_index = ET.parse(os.path.join(tmpdir.path, 'sitemap.xml'))
            ns = {'ns0': 'http://www.sitemaps.org/schemas/sitemap/0.9'}
            locs = [loc.text for loc in xml_index.findall('ns0:sitemap/ns0:loc', ns)]
            assert 'http://localhost/allura_sitemap/sitemap-0.xml' in locs

            xml_0 = ET.parse(os.path.join(tmpdir.path, 'sitemap-0.xml'))
            urls = [loc.text for loc in xml_0.findall('ns0:url/ns0:loc', ns)]
            assert 'http://localhost/p/wiki/' not in urls  # blank wiki pages omitted from sitemap
            assert 'http://localhost/p/test/sub1/' in urls


from simplegmail import Gmail
from simplegmail.message import Message

class Watcher:

    gmail: Gmail

    unread: list[Message]

    def __init__(self):
        self.gmail = Gmail()
        self.unread = []
        self.update()

    """
    Returns the number of read emails since the last update
    """
    def update(self) -> int:
        old_unread = self.unread
        new_unread = self.gmail.get_unread_inbox()
        new_unread_ids = [m.id for m in new_unread]

        read_messages: list[Message] = []

        for m in old_unread:
            if m.id not in new_unread_ids:
                read_messages.append(m)

        self.unread = new_unread

        return len(read_messages)

    """
    Returns the number of unread emails
    """
    def get_unread_count(self):
        return len(self.unread)


if __name__ == '__main__':
    f = open('test.txt',encoding='utf-8')
    line = f.readline()
    mydict = eval(line)
    f = open('分词.txt', 'w')
    for key in mydict:
        f.write(key)
        f.write('  :  ')
        f.write(str(mydict[key]))
        f.write('\n')


import sys
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5 import uic, QtWidgets

qtCreatorFile = r"E:\pyWorkspace\newStevenUi\qt\com\tstB.ui"

Ui_MainWindow, QtBaseClass = uic.loadUiType(qtCreatorFile)

class MyDraw(QtWidgets.QWidget, Ui_MainWindow):
    def __init__(self):
        QtWidgets.QWidget.__init__(self)
        Ui_MainWindow.__init__(self)
        self.setupUi(self)

        self.pix = QPixmap(300, 300)
        self.pix.fill(QColor(255, 255, 255))
        self.endPoint = QPoint()
        self.lastPoint = QPoint()
        self.painter = QPainter()

    def paintEvent(self, event):
        self.painter.begin(self)
        self.painter.drawPixmap(0, 0, self.pix)
        self.painter.end()

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.lastPoint = event.pos()
            self.endPoint = self.lastPoint

    def mouseMoveEvent(self, event):
            if event.buttons() == Qt.LeftButton:       # 这里只能用buttons(), 因为button()在mouseMoveEvent()中无论
                self.endPoint = event.pos()            # 按下什么键，返回的都是Qt::NoButton
                self.painter.begin(self.pix)        # 注意这里的参数必须是self.pix，涂鸦只能在这个300*300的白板上进行
                self.painter.setPen(QColor(0, 255, 0))
                self.painter.drawLine(self.lastPoint, self.endPoint)
                self.painter.end()
                self.update()
                self.lastPoint = self.endPoint

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.endPoint = event.pos()
            self.update()

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    window = MyDraw()
    window.show()
    sys.exit(app.exec_())


import json
import codecs
import random
import os, sys

from PIL import Image

import numpy as np

def char_to_unicode(char):
    text = char.encode('unicode-escape')
    ucode = text.decode("utf-8")
    return ucode

def char_to_json(ucode_list, path):
    json_list = [codecs.decode(ucode, 'unicode-escape') for ucode in ucode_list]
    json_entry = {"CN":json_list} # First entry is just a placeholder
    print(json_entry)
    with open(path, 'w') as json_file:
        json.dump(json_entry, json_file)
    return json_entry

def read_json(path):
    with open(path, 'r') as json_file:
        text = json.load(json_file)
        print(text['CN'][1])
    return

def random_color():
    r1 = random.randint(0, 255)
    r2 = random.randint(0, 200)
    r3 = random.randint(0, 200)
    return (r1, r2, r3)

def change_color(image, random_colors):
    image = image.convert("RGBA")
    im_data = np.array(image)
    red, green, blue, alpha = im_data.T
    max_val = 200
    black = (red < max_val) & (blue < max_val) & (green < max_val)
    im_data[..., :-1][black.T] = random_colors
    image = Image.fromarray(im_data)
    return image

if __name__ == "__main__":
    ucode = char_to_unicode("我我")
    print(str(ucode))
    char_to_json(ucode, "charset/cjk.json")
    read_json("charset/cjk.json")


import contextlib
import errno
import glob
import hashlib
import logging
import os
import pipes
import re
import shutil
import stat
import tempfile
import urlparse

from autotest_lib.client.bin import utils as client_utils
from autotest_lib.client.common_lib import error
from autotest_lib.client.common_lib.cros import dev_server
from autotest_lib.server import test
from autotest_lib.server import utils
from autotest_lib.server.cros import cts_expected_failure_parser
from autotest_lib.server.cros import tradefed_chromelogin as login
from autotest_lib.server.cros import tradefed_constants as constants
from autotest_lib.server.cros import tradefed_utils

# For convenience, add to our scope.
parse_tradefed_result = tradefed_utils.parse_tradefed_result
adb_keepalive = tradefed_utils.adb_keepalive

class TradefedTest(test.test):
    """Base class to prepare DUT to run tests via tradefed."""
    version = 1

    # Default and upperbounds of max_retry, based on board and revision
    # after branching (that is, 'y' of R74-12345.y.z).
    #
    # By default, 0<=y<1 does 5 retries and 1<=y does 10. The |max_retry|
    # parameter in control files can override the count, within the
    # _BRANCH_MAX_RETRY limit below.
    _BRANCH_DEFAULT_RETRY = [(0, 5), (1, 10)]  # dev=5, beta=stable=10
    _BRANCH_MAX_RETRY = [(0, 5), (1, 10),      # dev=5, beta=10, stable=99
        (constants.APPROXIMATE_STABLE_BRANCH_NUMBER, 99)]
    # TODO(kinaba): betty-arcnext
    _BOARD_MAX_RETRY = {'betty': 0}

    _SHARD_CMD = None
    _board_arch = None
    _board_name = None
    _release_branch_number = None  # The 'y' of OS version Rxx-xxxxx.y.z
    _android_version = None
    _num_media_bundles = 0
    _perf_results = []

    def _log_java_version(self):
        """Quick sanity and spew of java version installed on the server."""
        utils.run(
            'java',
            args=('-version',),
            ignore_status=False,
            verbose=True,
            stdout_tee=utils.TEE_TO_LOGS,
            stderr_tee=utils.TEE_TO_LOGS)

    def initialize(self,
                   bundle=None,
                   uri=None,
                   host=None,
                   hosts=None,
                   max_retry=None,
                   load_waivers=True,
                   retry_manual_tests=False,
                   warn_on_test_retry=True,
                   hard_reboot_on_failure=False):
        """Sets up the tools and binary bundles for the test."""
        self._install_paths = []
        # TODO(pwang): Remove host if we enable multiple hosts everywhere.
        self._hosts = [host] if host else hosts
        for host in self._hosts:
            logging.info('Hostname: %s', host.host_port)
        self._verify_hosts()

        self._max_retry = self._get_max_retry(max_retry)
        self._warn_on_test_retry = warn_on_test_retry
        # Tests in the lab run within individual lxc container instances.
        if utils.is_in_container():
            cache_root = constants.TRADEFED_CACHE_CONTAINER
        else:
            cache_root = constants.TRADEFED_CACHE_LOCAL

        # TODO(ihf): reevaluate this again when we run out of memory. We could
        # for example use 32 bit java on the first run but not during retries.
        # b/62895114. If select_32bit_java gets deleted for good also remove it
        # from the base image.
        # Try to save server memory (crbug.com/717413).
        # select_32bit_java()

        # The content of the cache survives across jobs.
        self._safe_makedirs(cache_root)
        self._tradefed_cache = os.path.join(cache_root, 'cache')
        self._tradefed_cache_lock = os.path.join(cache_root, 'lock')
        self._tradefed_cache_dirty = os.path.join(cache_root, 'dirty')
        # The content of the install location does not survive across jobs and
        # is isolated (by using a unique path)_against other autotest instances.
        # This is not needed for the lab, but if somebody wants to run multiple
        # TradedefTest instance.
        self._tradefed_install = tempfile.mkdtemp(
            prefix=constants.TRADEFED_PREFIX)
        # Under lxc the cache is shared between multiple autotest/tradefed
        # instances. We need to synchronize access to it. All binaries are
        # installed through the (shared) cache into the local (unshared)
        # lxc/autotest instance storage.
        # If clearing the cache it must happen before all downloads.
        self._clean_download_cache_if_needed()
        # Set permissions (rwxr-xr-x) to the executable binaries.
        permission = (
            stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH
            | stat.S_IXOTH)
        self._install_files(constants.ADB_DIR, constants.ADB_FILES,
                            permission)
        self._install_files(constants.SDK_TOOLS_DIR,
                            constants.SDK_TOOLS_FILES, permission)

        # Install the tradefed bundle.
        bundle_install_path = self._install_bundle(
            uri or self._get_default_bundle_url(bundle))
        self._repository = os.path.join(bundle_install_path,
                                        self._get_tradefed_base_dir())

        # Load expected test failures to exclude them from re-runs.
        self._waivers = set()
        if load_waivers:
            self._waivers.update(
                    self._get_expected_failures('expectations', bundle))
        if not retry_manual_tests:
            self._waivers.update(
                    self._get_expected_failures('manual_tests', bundle))

        # Load modules with no tests.
        self._notest_modules = self._get_expected_failures('notest_modules',
                bundle)
        self._hard_reboot_on_failure = hard_reboot_on_failure

    def cleanup(self):
        """Cleans up any dirtied state."""
        # Kill any lingering adb servers.
        for host in self._hosts:
            try:
                self._run_adb_cmd(host, verbose=True, args=('kill-server',))
            except (error.CmdError, AttributeError):
                pass
        logging.info('Cleaning up %s.', self._tradefed_install)
        try:
            shutil.rmtree(self._tradefed_install)
        except IOError:
            pass

        # Create perf data for Chromeperf.
        for perf in self._perf_results:
            data = dict(
                units='count',
                higher_is_better=False,
                replace_existing_values=True,
            )
            data.update(perf)
            self.output_perf_value(**data)

    def _verify_hosts(self):
        """Verify all hosts' ChromeOS consistency."""
        # Check release builder path. E.g. cave-release/R66-10435.0.0
        release_builder_path = set(host.get_release_builder_path()
                                   for host in self._hosts)
        if len(release_builder_path) > 1:
            raise error.TestFail('Hosts\' CHROMEOS_RELEASE_BUILDER_PATH is '
                                 'different: %s', release_builder_path)

        # Check ChromeOS ARC VERSION. E.g.
        arc_version = set(host.get_arc_version() for host in self._hosts)
        if len(arc_version) > 1:
            raise error.TestFail('Hosts\' CHROMEOS_ARC_VERSION is different: '
                                 '%s', arc_version)

        # Check ChromeOS model for unibuild.
        # TODO(pwang): Adding a check if we found how to detect host's model.

    def _verify_arc_hosts(self):
        """Verify all hosts' Android configuration consistency.

        This method should only be called after all hosts' Android has been
        successfully booted up."""
        # Check all hosts have same Android fingerprint.
        fingerprint = set(self._run_adb_cmd(
            host,
            args=('shell', 'getprop', 'ro.build.fingerprint')).stdout
            for host in self._hosts)
        if len(fingerprint) > 1:
            raise error.TestFail('Hosts\' supported fingerprint is different: '
                                 '%s', fingerprint)

        # Check all hosts support same abilist.
        abilist = set(self._run_adb_cmd(
            host,
            args=('shell', 'getprop', 'ro.product.cpu.abilist')).stdout
            for host in self._hosts)
        if len(abilist) > 1:
            raise error.TestFail('Hosts\' supported abilist is different: %s',
                                 abilist)
        self._abilist = str(list(abilist)[0]).split(',')

    def _calculate_timeout_factor(self, bundle):
        """ Calculate the multiplicative factor for timeout.

        The value equals to the times each test case is run, which is determined
        by the intersection of the supported ABIs of the CTS/GTS bundle and that
        of the tested device."""
        arm_abis = set(('armeabi-v7a', 'arm64-v8a'))
        x86_abis = set(('x86', 'x86_64'))
        if bundle == 'arm':
            tradefed_abis = arm_abis
        elif bundle == 'x86':
            tradefed_abis = x86_abis
        else:
            tradefed_abis = arm_abis | x86_abis
        self._timeout_factor = len(set(self._get_abilist()) & tradefed_abis)

    @contextlib.contextmanager
    def _login_chrome(self, **cts_helper_kwargs):
        """Returns Chrome log-in context manager.

        Please see also cheets_StartAndroid for details about how this works.
        """
        # TODO(pwang): Chromelogin takes 10+ seconds for it to successfully
        #              enter. Parallelize if this becomes a bottleneck.
        instances = []
        for host in self._hosts:
            instances.append(login.ChromeLogin(host, cts_helper_kwargs))

        for instance in instances:
            instance.enter()
        yield instances
        for instance in instances:
            instance.exit()

    def _get_adb_targets(self):
        """Get a list of adb targets."""
        return [self._get_adb_target(host) for host in self._hosts]

    def _get_adb_target(self, host):
        """Get the adb target format.

        This method is slightly different from host.host_port as we need to
        explicitly specify the port so the serial name of adb target would
        match."""
        return '{}:{}'.format(host.hostname, host.port)

    def _run_adb_cmd(self, host=None, **kwargs):
        """Running adb command.

        @param host: DUT that want to connect to. (None if the adb command is
                     intended to run in the server. eg. keygen)
        """
        # As of N, tradefed could not specify which adb socket to use, which use
        # tcp:localhost:5037 by default.
        adb_global_option = ('-H', 'localhost', '-P', '5037')
        if host:
            host_port = self._get_adb_target(host)
            adb_global_option = ('-s', host_port)
        kwargs['args'] = adb_global_option + kwargs.get('args', ())
        result = self._run('adb', **kwargs)
        logging.info('adb %s:\n%s', ' '.join(kwargs.get('args')),
                     result.stdout + result.stderr)
        return result

    def _try_adb_connect(self, host):
        """Attempts to connect to adb on the DUT.

        @param host: DUT that need to be connected.
        @return boolean indicating if adb connected successfully.
        """
        # This may fail return failure due to a race condition in adb connect
        # (b/29370989). If adb is already connected, this command will
        # immediately return success.
        host_port = self._get_adb_target(host)
        result = self._run_adb_cmd(
            host, args=('connect', host_port), verbose=True, ignore_status=True)
        if result.exit_status != 0:
            return False

        result = self._run_adb_cmd(host, args=('devices',))
        if not re.search(r'{}\s+(device|unauthorized)'.format(
                re.escape(host_port)), result.stdout):
            logging.info('No result found in with pattern: %s',
                         r'{}\s+(device|unauthorized)'.format(
                             re.escape(host_port)))
            return False

        # Actually test the connection with an adb command as there can be
        # a race between detecting the connected device and actually being
        # able to run a commmand with authenticated adb.
        result = self._run_adb_cmd(
            host, args=('shell', 'exit'), ignore_status=True)
        return result.exit_status == 0

    def _android_shell(self, host, command):
        """Run a command remotely on the device in an android shell

        This function is strictly for internal use only, as commands do not run
        in a fully consistent Android environment. Prefer adb shell instead.
        """
        host.run('android-sh -c ' + pipes.quote(command))

    def _write_android_file(self, host, filename, data):
        """Writes a file to a location relative to the android container.

        This is an internal function used to bootstrap adb.
        Tests should use adb push to write files.
        """
        android_cmd = 'echo %s > %s' % (pipes.quote(data),
                                        pipes.quote(filename))
        self._android_shell(host, android_cmd)

    def _connect_adb(self, host, pubkey_path):
        """Sets up ADB connection to the ARC container.

        @param host: DUT that should be connected to.
        @param pubkey_path: public key that adb keygen generated.
        """
        logging.info('Setting up adb connection.')
        # Generate and push keys for adb.
        # TODO(elijahtaylor): Extract this code to arc_common and de-duplicate
        # code in arc.py on the client side tests.
        with open(pubkey_path, 'r') as f:
            self._write_android_file(host, constants.ANDROID_ADB_KEYS_PATH,
                                     f.read())
        self._android_shell(
            host, 'restorecon ' + pipes.quote(constants.ANDROID_ADB_KEYS_PATH))

        # This starts adbd.
        self._android_shell(host, 'setprop sys.usb.config mtp,adb')

        # Also let it be automatically started upon reboot.
        self._android_shell(host, 'setprop persist.sys.usb.config mtp,adb')

        # adbd may take some time to come up. Repeatedly try to connect to adb.
        utils.poll_for_condition(
            lambda: self._try_adb_connect(host),
            exception=error.TestFail('Error: Failed to set up adb connection'),
            timeout=constants.ADB_READY_TIMEOUT_SECONDS,
            sleep_interval=constants.ADB_POLLING_INTERVAL_SECONDS)

        logging.info('Successfully setup adb connection.')

    def _wait_for_arc_boot(self, host):
        """Wait until ARC is fully booted.

        Tests for the presence of the intent helper app to determine whether ARC
        has finished booting.
        @param host: DUT that need to be connected to.
        """

        def _intent_helper_running():
            result = self._run_adb_cmd(
                host,
                args=('shell', 'pgrep', '-f', 'org.chromium.arc.intent_helper'),
                ignore_status=True)
            return bool(result.stdout)

        utils.poll_for_condition(
            _intent_helper_running,
            exception=error.TestFail(
                'Error: Timed out waiting for intent helper.'),
            timeout=constants.ARC_READY_TIMEOUT_SECONDS,
            sleep_interval=constants.ARC_POLLING_INTERVAL_SECONDS)

    def _disable_adb_install_dialog(self, host):
        """Disables a dialog shown on adb install execution.

        By default, on adb install execution, "Allow Google to regularly check
        device activity ... " dialog is shown. It requires manual user action
        so that tests are blocked at the point.
        This method disables it.
        """
        logging.info('Disabling the adb install dialog.')
        result = self._run_adb_cmd(
            host,
            verbose=True,
            args=('shell', 'settings', 'put', 'global',
                  'verifier_verify_adb_installs', '0'))
        logging.info('Disable adb dialog: %s', result.stdout)

    def _ready_arc(self):
        """Ready ARC and adb in parallel for running tests via tradefed."""
        # Generate the adb keys on server.
        key_path = os.path.join(self.tmpdir, 'test_key')
        pubkey_path = key_path + '.pub'
        self._run_adb_cmd(verbose=True, args=('keygen', pipes.quote(key_path)))
        os.environ['ADB_VENDOR_KEYS'] = key_path
        # Kill existing adb server to ensure that the env var is picked up.
        self._run_adb_cmd(verbose=True, args=('kill-server',))

        # TODO(pwang): connect_adb takes 10+ seconds on a single DUT.
        #              Parallelize it if it becomes a bottleneck.
        for host in self._hosts:
            self._connect_adb(host, pubkey_path)
            self._disable_adb_install_dialog(host)
            self._wait_for_arc_boot(host)
        self._verify_arc_hosts()

    def _safe_makedirs(self, path):
        """Creates a directory at |path| and its ancestors.

        Unlike os.makedirs(), ignore errors even if directories exist.
        """
        try:
            os.makedirs(path)
        except OSError as e:
            if not (e.errno == errno.EEXIST and os.path.isdir(path)):
                raise

    def _unzip(self, filename):
        """Unzip the file.

        The destination directory name will be the stem of filename.
        E.g., _unzip('foo/bar/baz.zip') will create directory at
        'foo/bar/baz', and then will inflate zip's content under the directory.
        If here is already a directory at the stem, that directory will be used.

        @param filename: Path to the zip archive.
        @return Path to the inflated directory.
        """
        destination = os.path.splitext(filename)[0]
        if os.path.isdir(destination):
            logging.info('Skipping unzip %s, reusing content of %s', filename,
                         destination)
            return destination
        tmp = tempfile.mkdtemp(dir=os.path.dirname(filename))
        logging.info('Begin unzip %s', filename)
        try:
            utils.run('unzip', args=('-d', tmp, filename))
        except:
            logging.error('Failed unzip, cleaning up.')
            # Clean up just created files.
            shutil.rmtree(tmp, ignore_errors=True)
            raise
        logging.info('End unzip %s', filename)
        try:
            os.renames(tmp, destination)
        except:
            logging.error('Failed rename, cleaning up.')
            shutil.rmtree(destination, ignore_errors=True)
            shutil.rmtree(tmp, ignore_errors=True)
            raise
        return destination

    def _dir_size(self, directory):
        """Compute recursive size in bytes of directory."""
        size = 0
        for root, _, files in os.walk(directory):
            for name in files:
                try:
                    size += os.path.getsize(os.path.join(root, name))
                except OSError:
                    logging.error('Inaccessible path (crbug/793696): %s/%s',
                                  root, name)
        return size

    def _invalidate_download_cache(self):
        """Marks the download cache for deferred deletion.

        Used to make cache file operations atomic across failures and reboots.
        The caller is responsible to hold the lock to the cache.
        """
        if not os.path.exists(self._tradefed_cache_dirty):
            os.mkdir(self._tradefed_cache_dirty)

    def _validate_download_cache(self):
        """Validates and unmarks the download cache from deletion.

        Used to make cache file operations atomic across failures and reboots.
        The caller is responsible to hold the lock to the cache.
        """
        shutil.rmtree(self._tradefed_cache_dirty, ignore_errors=True)

    def _clean_download_cache_if_needed(self, force=False):
        """Invalidates cache to prevent it from growing too large."""
        # If the cache is large enough to hold a working set, we can simply
        # delete everything without thrashing.
        # TODO(ihf): Investigate strategies like LRU.
        clean = force
        with tradefed_utils.lock(self._tradefed_cache_lock):
            size = self._dir_size(self._tradefed_cache)
            if size > constants.TRADEFED_CACHE_MAX_SIZE:
                logging.info(
                    'Current cache size=%d got too large. Clearing %s.', size,
                    self._tradefed_cache)
                clean = True
            else:
                logging.info('Current cache size=%d of %s.', size,
                             self._tradefed_cache)
            if os.path.exists(self._tradefed_cache_dirty):
                logging.info('Found dirty cache.')
                clean = True
            if clean:
                logging.warning('Cleaning download cache.')
                shutil.rmtree(self._tradefed_cache, ignore_errors=True)
                self._safe_makedirs(self._tradefed_cache)
                shutil.rmtree(self._tradefed_cache_dirty, ignore_errors=True)

    def _download_to_cache(self, uri):
        """Downloads the uri from the storage server.

        It always checks the cache for available binaries first and skips
        download if binaries are already in cache.

        The caller of this function is responsible for holding the cache lock.

        @param uri: The Google Storage or dl.google.com uri.
        @return Path to the downloaded object, name.
        """
        # We are hashing the uri instead of the binary. This is acceptable, as
        # the uris are supposed to contain version information and an object is
        # not supposed to be changed once created.
        output_dir = os.path.join(self._tradefed_cache,
                                  hashlib.md5(uri).hexdigest())
        # Check for existence of cache entry. We check for directory existence
        # instead of file existence, so that _install_bundle can delete original
        # zip files to save disk space.
        if os.path.exists(output_dir):
            # TODO(crbug.com/800657): Mitigation for the invalid state. Normally
            # this should not happen, but when a lock is force borken due to
            # high IO load, multiple processes may enter the critical section
            # and leave a bad state permanently.
            if os.listdir(output_dir):
                logging.info('Skipping download of %s, reusing content of %s.',
                             uri, output_dir)
                return os.path.join(output_dir,
                    os.path.basename(urlparse.urlparse(uri).path))
            logging.error('Empty cache entry detected %s', output_dir)
        return self._download_to_dir(uri, output_dir)

    def _download_to_dir(self, uri, output_dir):
        """Downloads the gs|http|https uri from the storage server.

        @param uri: The Google Storage or dl.google.com uri.
        @output_dir: The directory where the downloaded file should be placed.
        @return Path to the downloaded object, name.
        """
        # Split uri into 3 pieces for use by gsutil and also by wget.
        parsed = urlparse.urlparse(uri)
        filename = os.path.basename(parsed.path)
        output = os.path.join(output_dir, filename)

        self._safe_makedirs(output_dir)
        if parsed.scheme not in ['gs', 'http', 'https']:
            raise error.TestFail(
                'Error: Unknown download scheme %s' % parsed.scheme)
        if parsed.scheme in ['http', 'https']:
            logging.info('Using wget to download %s to %s.', uri, output_dir)
            # We are downloading 1 file at a time, hence using -O over -P.
            utils.run(
                'wget',
                args=('--report-speed=bits', '-O', output, uri),
                verbose=True)
            return output

        if not client_utils.is_moblab():
            # If the machine can access to the storage server directly,
            # defer to "gsutil" for downloading.
            logging.info('Not in lab. Downloading %s directly to %s.',
                         uri, output)
            # b/17445576: gsutil rsync of individual files is not implemented.
            utils.run('gsutil', args=('cp', uri, output), verbose=True)
            return output

        # We are in the moblab. Because the machine cannot access the storage
        # server directly, use dev server to proxy.
        logging.info('In lab. Downloading %s by staging to %s.',
                     uri, output)

        dirname = os.path.dirname(parsed.path)
        archive_url = '%s://%s%s' % (parsed.scheme, parsed.netloc, dirname)

        # First, request the devserver to download files into the lab network.
        # TODO(ihf): Switch stage_artifacts to honor rsync. Then we don't have
        # to shuffle files inside of tarballs.
        info = self._hosts[0].host_info_store.get()
        ds = dev_server.ImageServer.resolve(info.build)
        ds.stage_artifacts(
            info.build, files=[filename], archive_url=archive_url)

        # Then download files from the dev server.
        # TODO(ihf): use rsync instead of wget. Are there 3 machines involved?
        # Itself, dev_server plus DUT? Or is there just no rsync in moblab?
        ds_src = '/'.join([ds.url(), 'static', dirname, filename])
        logging.info('dev_server URL: %s', ds_src)
        # Calls into DUT to pull uri from dev_server.
        utils.run(
            'wget',
            args=('--report-speed=bits', '-O', output, ds_src),
            verbose=True)
        return output

    def _instance_copyfile(self, cache_path):
        """Makes a copy of a file from the (shared) cache to a wholy owned
        local instance. Also copies one level of cache directoy (MD5 named).
        """
        filename = os.path.basename(cache_path)
        dirname = os.path.basename(os.path.dirname(cache_path))
        instance_dir = os.path.join(self._tradefed_install, dirname)
        # Make sure destination directory is named the same.
        self._safe_makedirs(instance_dir)
        instance_path = os.path.join(instance_dir, filename)
        shutil.copyfile(cache_path, instance_path)
        return instance_path

    def _instance_copytree(self, cache_path):
        """Makes a copy of a directory from the (shared and writable) cache to
        a wholy owned local instance.

        TODO(ihf): Consider using cp -al to only copy links. Not sure if this
        is really a benefit across the container boundary, but it is risky due
        to the possibility of corrupting the original files by an lxc instance.
        """
        # We keep the top 2 names from the cache_path = .../dir1/dir2.
        dir2 = os.path.basename(cache_path)
        dir1 = os.path.basename(os.path.dirname(cache_path))
        instance_path = os.path.join(self._tradefed_install, dir1, dir2)
        logging.info('Copying %s to instance %s', cache_path, instance_path)
        shutil.copytree(cache_path, instance_path)
        return instance_path

    def _install_bundle(self, gs_uri):
        """Downloads a zip file, installs it and returns the local path.

        @param gs_uri: GS bucket that contains the necessary files.
        """
        if not gs_uri.endswith('.zip'):
            raise error.TestFail('Error: Not a .zip file %s.', gs_uri)
        # Atomic write through of file.
        with tradefed_utils.lock(self._tradefed_cache_lock):
            # Atomic operations.
            self._invalidate_download_cache()
            # Download is lazy (cache_path may not actually exist if
            # cache_unzipped does).
            cache_path = self._download_to_cache(gs_uri)
            # Unzip is lazy as well (but cache_unzipped guaranteed to
            # exist).
            cache_unzipped = self._unzip(cache_path)
            # To save space we delete the original zip file. This works as
            # _download only checks existence of the cache directory for
            # lazily skipping download, and unzip itself will bail if the
            # unzipped destination exists. Hence we don't need the original
            # anymore.
            if os.path.exists(cache_path):
                logging.info('Deleting original %s', cache_path)
                os.remove(cache_path)
            # Erase dirty marker from disk.
            self._validate_download_cache()
            # We always copy files to give tradefed a clean copy of the
            # bundle.
            unzipped_local = self._instance_copytree(cache_unzipped)
        return unzipped_local

    def _install_files(self, gs_dir, files, permission):
        """Installs binary tools."""
        for filename in files:
            gs_uri = os.path.join(gs_dir, filename)
            # Atomic write through of file.
            with tradefed_utils.lock(self._tradefed_cache_lock):
                # We don't want to leave a corrupt cache for other jobs.
                self._invalidate_download_cache()
                cache_path = self._download_to_cache(gs_uri)
                # Mark cache as clean again.
                self._validate_download_cache()
                # This only affects the current job, so not part of cache
                # validation.
                local = self._instance_copyfile(cache_path)
            os.chmod(local, permission)
            # Keep track of PATH.
            self._install_paths.append(os.path.dirname(local))

    def _prepare_media(self, cts_uri, needs_push_media):
        """Downloads and offers the cached media files to tradefed."""
        if needs_push_media:
            media = self._install_bundle(cts_uri['media'])
            if os.path.islink(constants.TRADEFED_MEDIA_PATH):
                os.unlink(constants.TRADEFED_MEDIA_PATH)
            if os.path.isdir(constants.TRADEFED_MEDIA_PATH):
                shutil.rmtree(constants.TRADEFED_MEDIA_PATH)
            os.symlink(media, constants.TRADEFED_MEDIA_PATH)

            logging.info('Offered %s as a media directory in %s',
                    media, constants.TRADEFED_MEDIA_PATH)

        # Records the number of existing media bundles, to check later.
        if os.path.isdir(constants.TRADEFED_MEDIA_PATH):
            self._num_media_bundles = len(
                    os.listdir(constants.TRADEFED_MEDIA_PATH))

    def _fail_on_unexpected_media_download(self):
        if os.path.isdir(constants.TRADEFED_MEDIA_PATH):
            contents = os.listdir(constants.TRADEFED_MEDIA_PATH)
            if len(contents) > self._num_media_bundles:
                raise error.TestFail(
                    'Failed: Unexpected media bundle was added %s' % contents)

    def _run(self, *args, **kwargs):
        """Executes the given command line.

        To support SDK tools, such as adb or aapt, this adds _install_paths
        to the extra_paths. Before invoking this, ensure _install_files() has
        been called.
        """
        kwargs['extra_paths'] = (
            kwargs.get('extra_paths', []) + self._install_paths)
        return utils.run(*args, **kwargs)

    def _collect_tradefed_global_log(self, result, destination):
        """Collects the tradefed global log.

        @param result: The result object from utils.run.
        @param destination: Autotest result directory (destination of logs).
        """
        match = re.search(r'Saved log to /tmp/(tradefed_global_log_.*\.txt)',
                          result.stdout)
        if not match:
            logging.error('no tradefed_global_log file is found')
            return

        name = match.group(1)
        dest = os.path.join(destination, 'logs', 'tmp')
        self._safe_makedirs(dest)
        shutil.copy(os.path.join('/tmp', name), os.path.join(dest, name))

    def _parse_result(self, result, waivers=None):
        """Check the result from the tradefed output.

        This extracts the test pass/fail/executed list from the output of
        tradefed. It is up to the caller to handle inconsistencies.

        @param result: The result object from utils.run.
        @param waivers: a set[] of tests which are permitted to fail.
        """
        return parse_tradefed_result(result.stdout, waivers)

    def _get_expected_failures(self, directory, bundle_abi):
        """Return a list of expected failures or no test module.

        @param directory: A directory with expected no tests or failures files.
        @param bundle_abi: 'arm' or 'x86' if the test is for the particular ABI.
                           None otherwise (like GTS, built for multi-ABI.)
        @return: A list of expected failures or no test modules for the current
                 testing device.
        """
        # Load waivers and manual tests so TF doesn't re-run them.
        expected_fail_files = []
        test_board = self._get_board_name()
        test_arch = self._get_board_arch()
        sdk_ver = self._get_android_version()
        expected_fail_dir = os.path.join(self.bindir, directory)
        if os.path.exists(expected_fail_dir):
            expected_fail_files += glob.glob(expected_fail_dir + '/*.yaml')

        waivers = cts_expected_failure_parser.ParseKnownCTSFailures(
            expected_fail_files)
        return waivers.find_waivers(test_arch, test_board, bundle_abi, sdk_ver)

    def _get_abilist(self):
        """Return the abilist supported by calling adb command.

        This method should only be called after the android environment is
        successfully initialized."""
        if not self._abilist:
            self._abilist = self._run_adb_cmd(
                self._hosts[0],
                args=('shell', 'getprop',
                      'ro.product.cpu.abilist')).stdout.split(',')
        return self._abilist

    def _get_release_branch_number(self):
        """Returns the DUT branch number (z of Rxx-yyyyy.z.w) or 0 on error."""
        if not self._release_branch_number:
            ver = (self._hosts[0].get_release_version() or '').split('.')
            self._release_branch_number = (int(ver[1]) if len(ver) >= 3 else 0)
        return self._release_branch_number

    def _get_board_arch(self):
        """Return target DUT arch name."""
        if not self._board_arch:
            self._board_arch = ('arm' if self._hosts[0].get_cpu_arch() == 'arm'
                else 'x86')
        return self._board_arch

    def _get_board_name(self):
        """Return target DUT board name."""
        if not self._board_name:
            self._board_name = self._hosts[0].get_board().split(':')[1]
        return self._board_name

    def _get_android_version(self):
        """Return target DUT Android SDK version"""
        # TODO(kinaba): factor this out to server/hosts/cros_host.py
        if not self._android_version:
            self._android_version = self._hosts[0].run(
                'grep ANDROID_SDK /etc/lsb-release',
                ignore_status=True).stdout.rstrip().split('=')[1]
        return self._android_version

    def _get_max_retry(self, max_retry):
        """Return the maximum number of retries.

        @param max_retry: max_retry specified in the control file.
        @return: number of retries for this specific host.
        """
        if max_retry is None:
            max_retry = self._get_branch_retry(self._BRANCH_DEFAULT_RETRY)
        candidate = [max_retry]
        candidate.append(self._get_board_retry())
        candidate.append(self._get_branch_retry(self._BRANCH_MAX_RETRY))
        return min(x for x in candidate if x is not None)

    def _get_board_retry(self):
        """Return the maximum number of retries for DUT board name.

        @return: number of max_retry or None.
        """
        board = self._get_board_name()
        if board in self._BOARD_MAX_RETRY:
            return self._BOARD_MAX_RETRY[board]
        logging.info('No board retry specified for board: %s', board)
        return None

    def _get_branch_retry(self, table):
        """Returns the retry count for DUT branch number defined in |table|."""
        number = self._get_release_branch_number()
        for lowerbound, retry in reversed(table):
            if lowerbound <= number:
                return retry
        logging.warning('Could not establish channel. Using retry=0.')
        return 0

    def _run_precondition_scripts(self, commands, steps):
        """Run precondition scripts on all the hosts."""
        for host in self._hosts:
            for command in commands:
                # Replace {0} (if any) with the retry count.
                formatted_command = command.format(steps)
                logging.info('RUN: %s\n', formatted_command)
                output = host.run(formatted_command, ignore_status=True)
                logging.info('END: %s\n', output)

    def _run_and_parse_tradefed(self, commands):
        """Kick off the tradefed command.

        Assumes that only last entry of |commands| actually runs tests and has
        interesting output (results, logs) for collection. Ignores all other
        commands for this purpose.

        @param commands: List of lists of command tokens.
        @raise TestFail: when a test failure is detected.
        @return: tuple of (tests, pass, fail, notexecuted) counts.
        """
        target_argument = []
        for host in self._hosts:
            target_argument += ['-s', self._get_adb_target(host)]
        shard_argument = []
        if len(self._hosts) > 1:
            if self._SHARD_CMD:
                shard_argument = [self._SHARD_CMD, str(len(self._hosts))]
            else:
                logging.warning('cts-tradefed shard command isn\'t defined, '
                                'falling back to use single device.')
        commands = [command + target_argument + shard_argument
                    for command in commands]

        try:
            output = self._run_tradefed(commands)
        except Exception as e:
            self._log_java_version()
            if not isinstance(e, error.CmdTimeoutError):
                # In case this happened due to file corruptions, try to
                # force to recreate the cache.
                logging.error('Failed to run tradefed! Cleaning up now.')
                self._clean_download_cache_if_needed(force=True)
            raise

        result_destination = os.path.join(self.resultsdir,
                                          self._get_tradefed_base_dir())
        # Gather the global log first. Datetime parsing below can abort the test
        # if tradefed startup had failed. Even then the global log is useful.
        self._collect_tradefed_global_log(output, result_destination)
        # Result parsing must come after all other essential operations as test
        # warnings, errors and failures can be raised here.
        return self._parse_result(output, waivers=self._waivers)

    def _setup_result_directories(self):
        """Sets up the results and logs directories for tradefed.

        Tradefed saves the logs and results at:
          self._repository/results/$datetime/
          self._repository/results/$datetime.zip
          self._repository/logs/$datetime/
        Because other tools rely on the currently chosen Google storage paths
        we need to keep destination_results in:
          self.resultdir/android-cts/results/$datetime/
          self.resultdir/android-cts/results/$datetime.zip
          self.resultdir/android-cts/results/logs/$datetime/
        To bridge between them, create symlinks from the former to the latter.
        """
        logging.info('Setting up tradefed results and logs directories.')

        results_destination = os.path.join(self.resultsdir,
                                           self._get_tradefed_base_dir())
        logs_destination = os.path.join(results_destination, 'logs')
        directory_mapping = [
            (os.path.join(self._repository, 'results'), results_destination),
            (os.path.join(self._repository, 'logs'), logs_destination),
        ]

        for (tradefed_path, final_path) in directory_mapping:
            if os.path.exists(tradefed_path):
                shutil.rmtree(tradefed_path)
            self._safe_makedirs(final_path)
            os.symlink(final_path, tradefed_path)

    def _install_plan(self, subplan):
        """Copy test subplan to CTS-TF.

        @param subplan: CTS subplan to be copied into TF.
        """
        logging.info('Install subplan: %s', subplan)
        subplans_tf_dir = os.path.join(self._repository, 'subplans')
        if not os.path.exists(subplans_tf_dir):
            os.makedirs(subplans_tf_dir)
        test_subplan_file = os.path.join(self.bindir, 'subplans',
                                         '%s.xml' % subplan)
        try:
            shutil.copy(test_subplan_file, subplans_tf_dir)
        except (shutil.Error, OSError, IOError) as e:
            raise error.TestFail(
                'Error: failed to copy test subplan %s to CTS bundle. %s' %
                (test_subplan_file, e))

    def _should_skip_test(self, _bundle):
        """Some tests are expected to fail and are skipped.

        Subclasses should override with specific details.
        """
        return False

    def _should_reboot(self, steps):
        """Oracle to decide if DUT should reboot or just restart Chrome.

        For now we will not reboot after the first two iterations, but on all
        iterations afterward as before. In particular this means that most CTS
        tests will now not get a "clean" machine, but one on which tests ran
        before. But we will still reboot after persistent failures, hopefully
        not causing too many flakes down the line.
        """
        if steps < 3:
            return False
        return True

    def _run_tradefed_list_results(self):
        """Run the `tradefed list results` command.

        @return: tuple of the last (session_id, pass, fail, all_done?).
        """
        output = self._run_tradefed([['list', 'results']])

        # Parses the last session from the output that looks like:
        #
        # Session  Pass  Fail  Modules Complete ...
        # 0        90    10    1 of 2
        # 1        199   1     2 of 2
        # ...
        lastmatch = None
        for m in re.finditer(r'^(\d+)\s+(\d+)\s+(\d+)\s+(\d+) of (\d+)',
                             output.stdout, re.MULTILINE):
            session, passed, failed, done, total = map(int,
                                                       m.group(1, 2, 3, 4, 5))
            lastmatch = (session, passed, failed, done == total)
        return lastmatch

    def _tradefed_retry_command(self, template, session_id):
        raise NotImplementedError('Subclass should override this function')

    def _tradefed_run_command(self, template):
        raise NotImplementedError('Subclass should override this function')

    def _run_tradefed_with_retries(self,
                                   test_name,
                                   run_template,
                                   retry_template,
                                   timeout,
                                   needs_push_media=False,
                                   target_module=None,
                                   target_plan=None,
                                   bundle=None,
                                   cts_uri=None,
                                   login_precondition_commands=[],
                                   precondition_commands=[],
                                   perf_description=None):
        """Run CTS/GTS with retry logic.

        We first kick off the specified module. Then rerun just the failures
        on the next MAX_RETRY iterations.
        """
        # On dev and beta channels timeouts are sharp, lenient on stable.
        self._timeout = timeout
        if (self._get_release_branch_number() >=
                constants.APPROXIMATE_STABLE_BRANCH_NUMBER):
            self._timeout += 3600

        if self._should_skip_test(bundle):
            logging.warning('Skipped test %s', ' '.join(test_name))
            return

        steps = -1  # For historic reasons the first iteration is not counted.
        self.summary = ''
        accurate = []
        board = self._get_board_name()
        session_id = None

        self._setup_result_directories()
        self._prepare_media(cts_uri, needs_push_media)

        # This loop retries failures. For this reason please do not raise
        # TestFail in this loop if you suspect the failure might be fixed
        # in the next loop iteration.
        while steps < self._max_retry:
            steps += 1
            keep_media = needs_push_media and steps >= 1
            self._run_precondition_scripts(login_precondition_commands, steps)
            with self._login_chrome(
                    board=board,
                    reboot=self._should_reboot(steps),
                    # TODO(rohitbm): Evaluate if power cycle really helps with
                    # Bluetooth test failures, and then make the implementation
                    # more strict by first running complete restart and reboot
                    # retries and then perform power cycle.
                    hard_reboot_on_failure=(self._hard_reboot_on_failure
                                     and steps == self._max_retry),
                    dont_override_profile=keep_media) as current_logins:
                self._ready_arc()
                self._calculate_timeout_factor(bundle)
                self._run_precondition_scripts(precondition_commands, steps)

                # Run tradefed.
                if session_id == None:
                    if target_plan is not None:
                        self._install_plan(target_plan)

                    logging.info('Running %s:', test_name)
                    commands = [self._tradefed_run_command(run_template)]
                else:
                    logging.info('Retrying failures of %s with session_id %d:',
                                 test_name, session_id)
                    commands = [self._tradefed_retry_command(retry_template,
                                                             session_id)]

                # TODO(pwang): Evaluate if it is worth it to get the number of
                #              not-excecuted, for instance, by collecting all
                #              tests on startup (very expensive, may take 30
                #              minutes).
                waived_tests, acc = self._run_and_parse_tradefed(
                    commands)
                self._fail_on_unexpected_media_download()
                result = self._run_tradefed_list_results()
                if not result:
                    logging.error('Did not find any test results. Retry.')
                    for current_login in current_logins:
                        current_login.need_reboot()
                    continue

                waived = len(waived_tests)
                last_session_id, passed, failed, all_done = result
                # If the result is |acc|urate according to the log, or the
                # inaccuracy is recognized by tradefed (not all_done), then
                # it is fine.
                accurate.append(acc or not all_done)
                if failed < waived:
                    logging.error(
                        'Error: Internal waiver bookkeeping has become '
                        'inconsistent (f=%d, w=%d)', failed, waived)

                msg = 'run' if session_id == None else ' retry'
                msg += '(p=%s, f=%s, w=%s)' % (passed, failed, waived)
                self.summary += msg
                logging.info('RESULT: %s %s', msg, result)

                # Check for no-test modules. We use the "all_done" indicator
                # provided by list_results to decide if there are outstanding
                # modules to iterate over (similar to missing tests just on a
                # per-module basis).
                notest = (passed + failed == 0 and all_done)
                if target_module in self._notest_modules:
                    if notest:
                        logging.info('Package has no tests as expected.')
                        return
                    else:
                        # We expected no tests, but the new bundle drop must
                        # have added some for us. Alert us to the situation.
                        raise error.TestFail(
                            'Failed: Remove module %s from '
                            'notest_modules directory!' % target_module)
                elif notest:
                    logging.error('Did not find any tests in module. Hoping '
                                  'this is transient. Retry after reboot.')
                    for current_login in current_logins:
                        current_login.need_reboot()
                    continue

                session_id = last_session_id

                # Check if all the tests passed.
                if failed <= waived and all_done:
                    break

                # TODO(b/127908450) Tradefed loses track of not-executed tests
                # when the commandline pattern included '*', and retry run for
                # them wrongly declares all tests passed. This is misleading.
                # Rather, we give up the retry and report the result as FAIL.
                if not all_done and '*' in ''.join(run_template):
                    break

        # Tradefed finished normally. Record the failures to perf.
        if target_module:
            # Only record the failure by module, which exclude 'all', 'collects-tests-only', etc.
            self._perf_results.append(dict(
                description=perf_description if perf_description else target_module,
                value=failed,
                graph=bundle
            ))

        if session_id == None:
            raise error.TestFail('Error: Could not find any tests in module.')

        if failed <= waived and all_done:
            if not all(accurate):
                # Tests count inaccurate, remove perf to avoid false alarm.
                self._perf_results.pop()
                raise error.TestFail(
                    'Failed: Not all tests were executed. After %d '
                    'retries passing %d tests, waived=%d. %s' % (
                        steps, passed, waived, self.summary))
            # TODO(ihf): Make this error.TestPass('...') once
            # available.
            if steps > 0 and self._warn_on_test_retry:
                raise error.TestWarn(
                    'Passed: after %d retries passing %d tests, '
                    'waived=%d. %s' % (steps, passed, waived,
                                       self.summary))
            return

        raise error.TestFail(
            'Failed: after %d retries giving up. '
            'passed=%d, failed=%d, waived=%d%s%s. %s' %
            (steps, passed, failed, waived, '' if all_done else ', notexec>=1',
             '' if all(accurate) else ', Tests may not be accurate.',
             self.summary))