"""raytracing helper functions for deferred rendering using pyembree as a backend

Currently, meshes are represented as simple dictionaries with vertices, faces, uvs and optionally additional per-vertex
attributes (e.g., tangents, bitangents and normals). This is due to a limitation in the trimesh.Trimesh class, namely
that texture coordinates can only be specified per vertex and not per face, preventing proper unwrapping. Inputs of type
trimesh.Trimesh or trimesh.Scene are therefore converted to dictionaries or lists of dictionaries with the above keys.

The most important function is embree_render_deferred():

embree_render_deferred() produces inputs for deferred shading. It takes a list of meshes and a camera dict, as well as a
point light position as input and returns np.ndarray buffers with all relevant geometric quantities per intersected
pixel: 3D intersection points, interpolated vertex normals and tangents, normalized local and unnormalized global light
and view directions.

For additional light sources, get_local_dirs() can be used with the buffers returned from embree_render_deferred() and
a light position.

embree_render_deferred() calls the following functions: embree_create_scene(), embree_intersect_scene(),
interpolate_vertex_attributes() and get_local_dirs().
- embree_create_scene() constructs an EmbreeScene object, given the list of meshes. Camera parameters are either user-
specified, or will be set automatically (random camera position facing the scene center). Unless specified, the camera
focal length is automatically set so that the scene tightly fits onto the camera sensor. Unless explicitly disabled, it
computes per-vertex tangent frames, using each mesh's texture coordinates.
- embree_intersect_scene() performs the actual ray tracing and returns a pixel mask with the camera's resolution
indicating which pixels are hit, as well as buffers for all intersected pixels with the following attributes:
geometry and triangle IDs for each intersection, intersection depth (distance from camera), 3D intersection point,
barycentric coordinates within each triangle
- interpolate_vertex_attributes() computes per pixel texture coordinates and tangent frames by interpolating with the
barycentric coordinates returned from the embree_intersect_scene()

missing features:
- tracing shadow rays
- camera distortion model
- fallback tangent computing for meshes without texture coordinates
"""
from copy import deepcopy
import numpy as np
import trimesh.visual.texture
from warnings import warn
from pyembree.rtcore_scene import EmbreeScene
from pyembree.mesh_construction import TriangleMesh

from pysmtb.image import assign_masked
from pysmtb.utils import Dct, find_dim, dims_execpt, safe_divide

try:
    # we only optionally depend on trimesh, a lot of functionality also works without
    from trimesh import Scene, Trimesh
except ModuleNotFoundError:
    from types import NoneType
    Scene = NoneType
    Trimesh = NoneType

def normalize(vec, axis=0):
    """normalize input array along specified dimension"""
    return safe_divide(vec, np.linalg.norm(vec, axis=axis, keepdims=True))

def get_vertices(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the vertices as NV x 3 array"""
    if isinstance(mesh, Trimesh):
        vertices = mesh.vertices
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        vertices = np.concatenate([g.vertices for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'vertices'):
            raise Exception('mesh dict should have field vertices')
        vertices = mesh['vertices']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return vertices

def get_normals(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the vertex normals as NV x 3 array"""
    if isinstance(mesh, Trimesh):
        normals = mesh.normals
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        normals = np.concatenate([g.normals for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'normals'):
            raise Exception('mesh dict should have field normals')
        normals = mesh['normals']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return normals

def get_vertex_normals(mesh):
    """return per-vertex normal vectors, which are usually obtained by weighted averaging of each vertex's surrounding
    face normals"""
    if isinstance(mesh, Trimesh):
        vertex_normals = mesh.vertex_normals
    elif isinstance(mesh, Scene):
        vertex_normals = np.concatenate([g.vertex_normals for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'vertex_normals'):
            raise Exception('mesh dict should have field vertex_normals')
        vertex_normals = mesh['vertex_normals']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return vertex_normals

def get_uv_coords(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the texture coordinates as NV x 2 array"""
    if isinstance(mesh, Trimesh) and hasattr(mesh, 'visual') and hasattr(mesh.visual, 'uv'):
        uvs = mesh.visual.uv
    elif isinstance(mesh, Trimesh) and 'ply_raw' in mesh.metadata.keys():
        props = mesh.metadata['ply_raw']['vertex']['properties'].keys()
        u_key = 'u' if 'u' in props else 's'
        v_key = 'v' if 'v' in props else 't'
        u = mesh.metadata['ply_raw']['vertex']['data'][u_key]
        v = mesh.metadata['ply_raw']['vertex']['data'][v_key]
        uvs = np.stack((u, v), axis=1)
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        uvs = np.concatenate([g.visual.uv for g in mesh.geometry.values()], axis=0)
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'uvs'):
            raise Exception('mesh dict should have field uvs')
        uvs = mesh['uvs']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))
    return uvs

def get_faces(mesh, join_geometries=False):
    """given an input mesh (trimesh.Trimesh, trimesh.Scene, dict), return the faces as NF x 3 array"""
    if isinstance(mesh, Trimesh):
        faces = mesh.faces
    elif isinstance(mesh, Scene):
        if not join_geometries and len(mesh.geometry) > 1:
            raise Exception('mesh has multiple geometries, set join_geometries=True')
        raise NotImplementedError('join_geometries is not supported for faces')
    elif isinstance(mesh, dict):
        if not hasattr(mesh, 'faces'):
            raise Exception('mesh dict should have field faces')
        faces = mesh['faces']
    else:
        raise Exception('unsupported input type: ' + str(type(mesh)))

    if faces.ndim == 1:
        faces = np.stack([face[1] for face in faces], axis=0)
    if faces.shape[1] == 4:
        faces = np.concatenate((faces[:, :3], faces[:, [0, 2, 3]]), axis=0)
    return faces

def get_tangent_frames(mesh):
    """compute tangent frames per vertex, returns the mesh as dict, even if input was a Trimesh"""
    # algorithm:
    # source: Eric Lengyel, Foundations of Game Engine Development - Volume 2: Rendering
    # originally under http://www.terathon.com/code/tangent.html
    # https://web.archive.org/web/20190211192552/http://www.terathon.com/code/tangent.html
    # given triangle vertices p_0, p_1, p_2 (counterclockwise) with corresponding texture coordinates (u_i, v_i)
    # it holds: p_i − p_j = (u_i − u_j) * t + (v_i − v_j) * b
    #
    # define:
    # e_1 = p_1 - p_0,   (x_1, y_1) = (u_1 - u_0, v_1 - v_0)
    # e_2 = p_2 - p_0,   (x_2, y_2) = (u_2 - u_0, v_2 - v_0)
    #
    # we solve for tangent and bitangent vectors t and b using a system of equations:
    # e_1 = x_1 * t + y_1 * b
    # e_2 = x_2 * t + y_2 * b
    #
    # [t, b] = 1 / (x_1 * y_2 - x_2 * y_1) * [e_1, e_2] * [[y_2, -x_2], [-y_1, x_1]]
    #
    # once solved, compute per-vertex tangents and bitangents by summing up the vectors from each triangle and
    # re-orthonormalizing, given per-vertex normal vectors:
    # t' = normalize(t - dot(t, n) * n)
    # b' = normalize(b - dot(b, n) * n - dot(b, t') * t')

    # convert trimesh inputs to dicts
    if isinstance(mesh, Trimesh):
        mesh = trimesh_to_dict(mesh)

    if isinstance(mesh, Scene):
        meshes = trimesh_scene_to_dicts(mesh)
        for mesh_id in range(len(meshes)):
            meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])
        return meshes

    if hasattr(mesh, 'vertex_tangent_frames'):
        # tangent space already computed
        return mesh

    vertices = get_vertices(mesh)
    faces = get_faces(mesh)
    uvs = get_uv_coords(mesh)
    vertex_tangents = np.zeros_like(vertices)
    vertex_bitangents = np.zeros_like(vertices)
    vertex_normals = mesh.vertex_normals

    face_verts = vertices[faces.T, :]
    face_uvs = uvs[faces.T, :]

    e1 = face_verts[1] - face_verts[0]
    e2 = face_verts[2] - face_verts[0]

    x1 = face_uvs[1, :, 0:1] - face_uvs[0, :, 0:1]
    y1 = face_uvs[1, :, 1:2] - face_uvs[0, :, 1:2]
    x2 = face_uvs[2, :, 0:1] - face_uvs[0, :, 0:1]
    y2 = face_uvs[2, :, 1:2] - face_uvs[0, :, 1:2]

    r = safe_divide(1, (x1 * y2 - x2 * y1))
    t = (e1 * y2 - e2 * y1) * r
    b = (e2 * x1 - e1 * x2) * r

    # sum (bi)tangents that are associated to a vertex from its neighboring triangles
    vertex_tangents[faces[:, 0], :] += t
    vertex_tangents[faces[:, 1], :] += t
    vertex_tangents[faces[:, 2], :] += t
    vertex_bitangents[faces[:, 0], :] += b
    vertex_bitangents[faces[:, 1], :] += b
    vertex_bitangents[faces[:, 2], :] += b

    # normalize summed vectors
    vertex_tangents = normalize(vertex_tangents, axis=1)
    vertex_bitangents = normalize(vertex_bitangents, axis=1)

    # re-orthonormalization
    vertex_tangents = normalize(vertex_tangents - np.sum(vertex_normals * vertex_tangents, axis=1, keepdims=True) * vertex_normals, axis=1)
    vertex_bitangents = normalize(vertex_bitangents - np.sum(vertex_bitangents * vertex_normals, axis=1, keepdims=True) * vertex_normals
                                  - np.sum(vertex_bitangents * vertex_tangents, axis=1, keepdims=True) * vertex_tangents, axis=1)

    # determine handedness of tangent frame
    sgn_mask = np.sum(np.cross(vertex_tangents, vertex_bitangents, axis=1) * vertex_normals, axis=1, keepdims=True) > 0
    det_sgn = np.ones_like(sgn_mask, dtype=np.float32)
    det_sgn[~sgn_mask] = -1.

    # NV x 3 x 3 array (tangent, bitangent & normal stacked in rows)
    matrix_to_tangent_space = np.stack((vertex_tangents, det_sgn * vertex_bitangents, vertex_normals), axis=1)

    mesh.vertex_tangent_frames = matrix_to_tangent_space
    mesh.det_sgn = det_sgn
    return mesh

def create_unit_sphere(res=100):
    """create sphere centered on origin with radius 1 and res many segments in azimuth axis"""
    thetas = np.linspace(0, np.pi, res // 2)
    phis = np.linspace(0, 2 * np.pi, res)
    phis, thetas = np.meshgrid(phis, thetas)
    cp = np.cos(phis)
    sp = np.sin(phis)
    ct = np.cos(thetas)
    st = np.sin(thetas)
    xyz = np.stack((cp * st, sp * st, ct), axis=2).reshape((-1, 3))
    normals = xyz
    uvs = np.stack((phis / (2 * np.pi), thetas / np.pi), axis=2).reshape((-1, 2))

    nv = res * (res // 2)
    faces = np.r_[:nv].reshape(res // 2, res, order='C')

    # # faces:
    # # 0, 1, 2,
    # # 3, 4, 5,
    # # 6, 7, 8
    # faces = np.stack((
    #     faces[:-1],  # 0, 1, 2,
    #     np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5,
    #     np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
    #     np.roll(np.roll(faces, -1, axis=0)[:-1], -1, axis=1),  # 4, 5, 3
    #     np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
    #     np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5
    # ), axis=2).reshape((-1, 3))

    # faces:
    # 0, 1, 2,
    # 3, 4, 5,
    # 6, 7, 8
    faces = np.stack((
        faces[:-1],  # 0, 1, 2,
        np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5,
        np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
        np.roll(np.roll(faces, -1, axis=0)[:-1], -1, axis=1),  # 4, 5, 3
        np.roll(faces, -1, axis=1)[:-1],  # 1, 2, 0
        np.roll(faces, -1, axis=0)[:-1],  # 3, 4, 5
    ), axis=2).reshape((-1, 3))

    # mesh = Trimesh(vertices=xyz.reshape((-1, 3)), faces=faces.reshape((-1, 3)))
    # mesh.show()

    return Dct(vertices=xyz, uvs=uvs, faces=faces, vertex_normals=normals)

def create_unit_cube():
    """create cube centered on origin with side length 1"""
    # x: side, y: front-back, z: up-down
    xyz = np.array([[0, 0, 0],  # 0, bottom
                    [1, 0, 0],  # 1
                    [1, 1, 0],  # 2
                    [0, 1, 0],  # 3
                    [0, 0, 0],  # 4 front
                    [1, 0, 0],  # 5
                    [1, 0, 1],  # 6
                    [0, 0, 1],  # 7
                    [1, 0, 0],  # 8 right
                    [1, 1, 0],  # 9
                    [1, 1, 1],  # 10
                    [1, 0, 1],  # 11
                    [1, 1, 0],  # 12 back
                    [0, 1, 0],  # 13
                    [0, 1, 1],  # 14
                    [1, 1, 1],  # 15
                    [0, 1, 0],  # 16 left
                    [0, 0, 0],  # 17
                    [0, 0, 1],  # 18
                    [0, 1, 1],  # 19
                    [0, 0, 1],  # 20 top
                    [1, 0, 1],  # 21
                    [1, 1, 1],  # 22
                    [0, 1, 1],  # 23
                    ], dtype=np.float32)
    normals = np.array(([[ 0,  0, -1]] * 4) +  # bottom
                       ([[ 0, -1,  0]] * 4) +  # front
                       ([[ 1,  0,  0]] * 4) +  # right
                       ([[ 0,  1,  0]] * 4) +  # back
                       ([[-1,  0,  0]] * 4) +  # left
                       ([[ 0,  0,  1]] * 4)    # top
                       )
    uvs = np.array([[0.25, 0.375],  # 0, bottom
                    [0.50, 0.375],  # 1
                    [0.50, 0.125],  # 2
                    [0.25, 0.125],  # 3
                    [0.25, 0.375],  # 4 front
                    [0.50, 0.375],  # 5
                    [0.50, 0.625],  # 6
                    [0.25, 0.625],  # 7
                    [0.50, 0.375],  # 8 right
                    [0.75, 0.375],  # 9
                    [0.75, 0.625],  # 10
                    [0.50, 0.625],  # 11
                    [0.75, 0.375],  # 12 back
                    [1.00, 0.375],  # 13
                    [1.00, 0.625],  # 14
                    [0.75, 0.625],  # 15
                    [0.00, 0.375],  # 16 left
                    [0.25, 0.375],  # 17
                    [0.25, 0.625],  # 18
                    [0.00, 0.625],  # 19
                    [0.25, 0.625],  # 20 top
                    [0.50, 0.625],  # 21
                    [0.50, 0.875],  # 22
                    [0.25, 0.875],  # 23
                    ])

    # from pysmtb.plotting import text3
    # th = text3(xyz)
    xyz = xyz - np.r_[0.5, 0.5, 0.5][None]

    faces = np.array([
        [0, 2, 1],  # bottom
        [2, 0, 3],
        [4, 5, 6],  # front
        [6, 7, 4],
        [8, 9, 10],  # right
        [10, 11, 8],
        [12, 13, 14],  # back
        [14, 15, 12],
        [16, 17, 18],  # left
        [18, 19, 16],
        [20, 21, 22],  # top
        [22, 23, 20],
                      ])
    # mesh = Trimesh(vertices=xyz.reshape((-1, 3)), faces=faces.reshape((-1, 3)))
    # mesh.show()

    return Dct(vertices=xyz, uvs=uvs, faces=faces, vertex_normals=normals)

def get_bbox(inp, dims: tuple = None):
    """get bounding box of vertices"""
    bbox = None
    if isinstance(inp, Scene):
        return inp.bounds

    if isinstance(inp, list):
        for mesh in inp:
            bbox_ = get_bbox(mesh)
            bbox = bbox_ if bbox is None else  np.array([np.minimum(bbox[0, :], bbox_[0, :]),
                                                         np.maximum(bbox[1, :], bbox_[1, :])])
        return bbox

    if isinstance(inp, Trimesh):
        inp = inp.vertices
    elif isinstance(inp, dict):
        inp = inp['vertices']
    elif not isinstance(inp, np.ndarray):
        raise Exception('input must be trimesh.Trimesh, trimesh.Scene or list of dicts with vertices key')

    if dims is None:
        dims = dims_execpt(inp, find_dim(inp, size=3))
    lower = np.min(inp, axis=dims).squeeze()
    upper = np.max(inp, axis=dims).squeeze()
    bbox = np.r_[lower[None], upper[None]]
    return bbox

def get_bbox_corners(bbox):
    """construct 8 corner points given extreme diagonal points of AABB"""
    xl, yl, zl = bbox[0]
    xu, yu, zu = bbox[1]
    corners = np.r_[np.r_[xl, yl, zl][None],
                    np.r_[xu, yl, zl][None],
                    np.r_[xu, yu, zl][None],
                    np.r_[xl, yu, zl][None],
                    np.r_[xl, yl, zu][None],
                    np.r_[xu, yl, zu][None],
                    np.r_[xu, yu, zu][None],
                    np.r_[xl, yu, zu][None]
                    ]
    return corners

def trimesh_scene_to_dicts(scene: Scene):
    """convert trimesh.Scene into list of dicts with vertices, faces and uv coordinates"""
    output = []
    for mesh in list(scene.geometry.values()):
        output.append(trimesh_to_dict(mesh))
    return output

def trimesh_to_dict(mesh: Trimesh):
    """extract vertices, faces and uv coordinates from trimesh.Trimesh and return them in dict"""
    return Dct(vertices=get_vertices(mesh),
               faces=get_faces(mesh),
               uvs=get_uv_coords(mesh),
               vertex_normals=get_vertex_normals(mesh))

def cam_extrinsics(cam):
    """compute camera's extrinsics matrix given position, lookat and up vectors in the input dict"""
    if not hasattr(cam, 'forward'):
        cam.forward = normalize(cam.lookat - cam.position)

    if not hasattr(cam, 'up'):
        cam.up = np.r_[0., 1., 0.]

    # re-orthonormalization
    cam_side = normalize(np.cross(cam.up, cam.forward))
    cam.up = normalize(np.cross(cam.forward, cam_side))

    # set up extrinsics matrix
    cam.extrinsics = np.r_[np.r_[np.r_[cam_side[None], cam.up[None], cam.forward[None]].T, np.r_[0., 0., 0.][None]].T, np.r_[cam.position, 1.][None]].T

    return cam

def set_tight_cam(vertices, cam: dict = None, use_bbox: bool = True, visualize: bool = False):
    # backward compatibility
    warn('set_tight_cam() is deprecated and should be replaced by cam_auto_zoom()')
    return cam_auto_zoom(vertices=vertices, cam=cam, use_bbox=use_bbox, visualize=visualize)

def cam_auto_zoom(vertices, cam: dict = None, use_bbox: bool = True, visualize: bool = False, debug: bool = False):
    """automatically set camera focal length to enclose a scene (either all vertices, or their axis aligned bounding
    box) in the view frustum"""
    if not isinstance(vertices, np.ndarray) and not isinstance(vertices, list):
        raise Exception('vertices must be nv x 3 np.ndarray or list of such arrays')

    bbox = get_bbox(vertices)
    bbox_diam = np.linalg.norm(np.diff(bbox, axis=0))
    scene_center = np.mean(bbox, axis=0)

    if use_bbox:
        # project only scene bounding box corners, not the vertices themselves to estimate the camera's FOV
        # (less accurate but faster)
        vertices = get_bbox_corners(bbox)

    if isinstance(vertices, list):
        vertices = np.concatenate(vertices, axis=0)

    # set up camera if not provided: lookat: scene center, position: random direction & distance from lookat
    if cam is None:
        cam = Dct(res_x=512, res_y=256)
    cam = Dct(cam)

    if not hasattr(cam, 'cx'):
        cam.cx = cam.res_x / 2.
        cam.cy = cam.res_y / 2.

    if not hasattr(cam, 'extrinsics'):
        if not hasattr(cam, 'position'):
            random_dir = normalize(np.random.rand(3) - 0.5)
            random_dist = (bbox_diam / 2 + 10. * np.random.rand())
            cam.position = scene_center + random_dist * random_dir

        if not hasattr(cam, 'lookat'):
            cam.lookat = scene_center

        if not hasattr(cam, 'up'):
            cam.up = np.r_[0., 1., 0.]

        # compute extrinsics matrix from position, lookat and up vectors
        cam = cam_extrinsics(cam)

    # bring bbox corners into camera space
    corners_cam = (np.linalg.inv(cam.extrinsics) @ np.r_[vertices.T, np.ones((vertices.shape[0], 1)).T]).T[:, :3]
    if np.abs(corners_cam[:, 2]).min() == 0.0:
        raise Exception('camera inside object, cannot automatically set fov / focal length')

    # perspective projection
    corners_projected = corners_cam / corners_cam[:, 2:3]
    # compute bbox of projected points in image plane
    bbox_projected = np.r_[np.min(corners_projected, axis=0)[None], np.max(corners_projected, axis=0)[None]]

    # at unit focal length (image plane) compute sensor size necessary to enclose projected scene bbox
    min_sensor_width = 2 * np.maximum(np.abs(bbox_projected[0, 0]), bbox_projected[1, 0])
    min_sensor_height = 2 * np.maximum(np.abs(bbox_projected[0, 1]), bbox_projected[1, 1])
    xy_ratio = min_sensor_width / min_sensor_height

    # account for physical aspect ratio
    film_xy_ratio = cam.res_x / cam.res_y

    # adjust sensor dimension with its actual aspect ratio
    if xy_ratio > film_xy_ratio:
        # x too wide --> increase y
        min_sensor_height = min_sensor_width / film_xy_ratio
    else:
        # y too high --> increase x
        min_sensor_width = film_xy_ratio * min_sensor_height
    cam.sensor_width = min_sensor_width
    cam.sensor_height = min_sensor_height

    cam.fov_x_deg = np.rad2deg(2 * np.arctan(cam.sensor_width / 2))
    cam.fov_y_deg = np.rad2deg(2 * np.arctan(cam.sensor_height / 2))

    # now compute focal length for a 36mm sensor (along its larger axis)
    if cam.res_x > cam.res_y:
        film_width = 1  # 36
        film_height = film_width * cam.res_y / cam.res_x
    else:
        film_height = 1  # 36
        film_width = film_height * cam.res_x / cam.res_y
    cam.pix_per_mm_x = cam.res_x / film_width
    cam.pix_per_mm_y = cam.res_y / film_height

    cam.fx_mm = film_width / (2 * np.tan(np.deg2rad(cam.fov_x_deg) / 2))
    cam.fy_mm = film_height / (2 * np.tan(np.deg2rad(cam.fov_y_deg) / 2))
    cam.fx = cam.pix_per_mm_x * cam.fx_mm
    cam.fy = cam.pix_per_mm_y * cam.fy_mm

    cam.intrinsics = np.array([[cam.fx, 0., cam.cx], [0., cam.fy, cam.cy], [0., 0., 1.]])

    if debug:
        output = Dct(cam=cam, corners_projected=corners_projected, bbox_projected=bbox_projected)
    else:
        output = cam

    if visualize:
        corners_bbox_projected_world = (cam.extrinsics @ np.r_[get_bbox_corners(bbox_projected).T, np.ones((8, 1)).T]).T[:, :3]
        ray_origins, ray_dirs = create_rays(cam)
        corners_projected_world = (cam.extrinsics @ np.r_[corners_projected.T, np.ones((corners_projected.shape[0], 1)).T])[:3].T
        from pysmtb.plotting import plot3, quiver3, scatter3
        sh = scatter3(cam.position, marker='s', color=np.r_[0, 0, 0])
        ax = sh.axes
        quiver3(cam.position[None], cam.up[None], color=np.r_[0.6, 0.6, 0.6], axes=ax)
        quiver3(cam.position[None], scene_center[None] - cam.position[None], color=np.r_[0.2, 0.2, 0.2], axes=ax)
        plot3(corners_bbox_projected_world[np.r_[np.r_[:4], 0]].T, axes=ax, color=np.r_[0.5, 0.4, 0.0])
        if use_bbox:
            for ci in range(8):
                plot3(np.r_[vertices[ci:ci+1, :], corners_projected_world[ci:ci+1, :], cam.position[None]].T, axes=ax, color=0.85 * np.r_[1, 1, 1])
            scatter3(vertices, axes=ax, color=np.r_[0.1, 0.4, 0.85])
            plot3(vertices[np.r_[np.r_[:4], 0, np.r_[4:8], 4, 5, 1, 2, 6, 7, 3]].T, axes=ax, color=np.r_[0.1, 0.4, 0.85])
            scatter3(corners_projected_world, axes=ax, color=np.r_[0.85, 0.6, 0.1])
            plot3(corners_projected_world[np.r_[np.r_[:4], 0, np.r_[4:8], 4, 5, 1, 2, 6, 7, 3]].T, axes=ax, color=np.r_[0.85, 0.6, 0.1])
        else:
            scatter3(vertices[::123, :], marker='.', axes=ax, color=np.r_[0.1, 0.4, 0.85])
            scatter3(corners_projected_world[::123, :], marker='.', axes=ax, color=np.r_[0.85, 0.6, 0.1])

        scatter3(ray_origins[::3*123] + ray_dirs[::3*123], marker='.', axes=ax)
        # quiver3(ray_origins[::123], ray_dirs[::123], axes=ax)
    return output

def cam_auto_zoom_trajectory(vertices, cam_positions, cam: dict = None, use_bbox: bool = True, visualize: bool = False):
    """given scene vertices and an N x 3 array of camera positions, as well as a camera dict with initial intrisics
    (resolution, principal point, ...), set the camera's focal length to ensure the projected scene is always contained
    in the view frustum; returns cam dict with fx and fy set to the mininum over all focal lengths, as well as a list of
    extrinsics matrices corresponding to the camera positions"""
    cam_orig = deepcopy(cam)
    lookat = np.mean(get_bbox(vertices), axis=0)
    fx = []
    fy = []
    extrinsics = []
    centers = []
    corners_projected = []
    bbox_projected = []
    for cam_pos in cam_positions:
        cam = deepcopy(cam_orig)
        cam.position = cam_pos
        cam.lookat = lookat
        cam = cam_extrinsics(cam)

        debug = cam_auto_zoom(vertices=vertices, cam=cam, use_bbox=use_bbox, debug=True)
        corners_projected.append(debug.corners_projected)
        bbox_projected.append(debug.bbox_projected)
        cam = debug.cam
        centers.append(cam.lookat)
        fx.append(cam.fx)
        fy.append(cam.fy)
        extrinsics.append(cam.extrinsics)
    # final focal length is the minimum (largest field of view) over each position's focal length
    cam.fx = np.min(fx)
    cam.fy = np.min(fy)

    if visualize:
        from pysmtb.plotting import scatter3
        sh = scatter3(vertices, '.')
        ax = sh.axes
        scatter3(cam_positions, 's', axes=ax)
        scatter3(np.array([(ex @ np.concatenate((bbp, np.ones((bbp.shape[0], 1))), axis=1).T).T[:, :3] for ex, bbp in
                           zip(extrinsics, bbox_projected)]).reshape((-1, 3)), '.', axes=ax)

    return cam, extrinsics

def create_rays(cam, repeat_origin=True, xs: np.ndarray = None, ys: np.ndarray = None):
    """given camera specifications (resolution in x and y, extrinsics and intrinsics), compute ray origins and
    directions for the full pixel raster (or optionally a patch specified by x and y ranges)"""
    if xs is None:
        xs = np.r_[0.5: float(cam.res_x) + 0.5]
    if ys is None:
        ys = np.r_[0.5: float(cam.res_y) + 0.5]
    xs, ys = np.meshgrid(xs, ys)
    pts_pixels = np.stack((xs.ravel(), ys.ravel(), np.ones(xs.size)), axis=0).astype(np.float32)

    # construct ray origins and directions
    ray_origins = np.array(cam.extrinsics @ np.r_[0, 0, 0, 1])[None, :3].astype(np.float32)
    if repeat_origin:
        ray_origins = ray_origins.repeat(xs.size, axis=0)
    ray_dirs = (np.linalg.inv(cam.intrinsics) @ pts_pixels).T
    ray_dirs = normalize(ray_dirs, axis=1)
    ray_dirs = (cam.extrinsics[:3, :3] @ ray_dirs.T).astype(np.float32).T
    return ray_origins, ray_dirs

def embree_create_scene(meshes,
                        with_tangent_frames: bool = True,
                        scene: EmbreeScene = None,
                        cam: dict = None,
                        auto_cam: bool = True,
                        auto_cam_bbox: bool = True,
                        auto_cam_visualize: bool = False):
    """given a scene as a list of meshes, construct a pyembree scene, optionally adjust the camera to automatically
    focus the scene, returns embree scene and camera"""

    if scene is None:
        # create new scene
        scene = EmbreeScene()

    if isinstance(meshes, Trimesh):
        meshes = meshes.scene()

    if isinstance(meshes, Scene):
        meshes = trimesh_scene_to_dicts(meshes)

    if isinstance(meshes, dict):
        meshes = [meshes]

    if not isinstance(meshes, list):
        raise Exception('input of type ' + str(type(meshes)) + ' is not supported')

    # precompute per-vertex tangent frames
    if with_tangent_frames:
        for mesh_id in range(len(meshes)):
            meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])

    if auto_cam and auto_cam_bbox:
        bbox = get_bbox(meshes)
    all_vertices = []
    for mesh in meshes:
        TriangleMesh(scene, mesh['vertices'], mesh['faces'])
        if auto_cam and not auto_cam_bbox:
            all_vertices.append(mesh.vertices)

    if auto_cam:
        cam = cam_auto_zoom(bbox if auto_cam_bbox else all_vertices, cam=cam, use_bbox=auto_cam_bbox, visualize=auto_cam_visualize)
    return scene, cam

def embree_intersect_scene(scene: EmbreeScene,
                           cam: dict,
                           return_as_rasters: bool = False,
                           xs: np.ndarray = None,
                           ys: np.ndarray = None,
                           ):
    """
    Given a pyembree scene and a camera dict, shoot rays and intersect the scene, returning a dict with either
    N x C or H x W x C arrays if return_as_rasters == True. Optionally, a range of pixels can be specified as separate
    x- and y-coordinate arrays to only trace a region of interest. Such pixel coordinates should be in the range from 0
    to cam.res_x - 1 and 0 to cam.res_y - 1 respectively.

    The returned dict has keys:
    hit: binary mask indicating hit / miss (always H x W x 1)
    tfar: depth in camera space, 1e37 for missed pixels
    Ng: (unnormalized!) geometric (per-face) normal vectors
    u, v: barycentric coordinates for each intersected triangle
    primID: triangle ID
    geomID: sub-mesh ID
    """
    cam = Dct(cam)
    res_x = cam.res_x
    res_y = cam.res_y
    intrinsics = np.array([[cam.fx, 0, cam.cx], [0, cam.fy, cam.cy], [0, 0, 1]])

    if hasattr(cam, 'extrinsics'):
        extrinsics = cam.extrinsics
    else:
        extrinsics = np.eye(4)

    # generate pixel array
    if xs is None:
        xs = np.r_[0.5: float(res_x) + 0.5]
    elif np.remainder(xs[0], 1) == 0.:
        xs += 0.5
    if ys is None:
        ys = np.r_[float(res_y) + 0.5: 0.5: -1]
    elif np.remainder(ys[0], 1) == 0.:
        ys += 0.5
    xs, ys = np.meshgrid(xs, ys)
    pts_pixels = np.stack((xs.ravel(), ys.ravel(), np.ones(xs.size)), axis=0).astype(np.float32)
    num_pix = res_y * res_x

    # construct ray origins and directions
    ray_origins = np.array(extrinsics @ np.r_[0, 0, 0, 1])[None, :3].astype(np.float32).repeat(xs.size, axis=0)
    ray_dirs = (np.linalg.inv(intrinsics) @ pts_pixels).T
    ray_dirs = normalize(ray_dirs, axis=1)
    ray_dirs = (extrinsics[:3, :3] @ ray_dirs.T).astype(np.float32).T

    # trace rays
    its = scene.run(ray_origins, ray_dirs, output=1)
    its = Dct(its)

    # compute hit mask
    hit = its.tfar < 1e37
    # always reshape hit mask to bitmap resolution
    its.hit = hit.reshape((res_y, res_x, 1))

    # normalize normals
    its.Ng[hit, :] = normalize(its.Ng[hit, :], axis=1)

    # reshape all fields to image size, optionally multiply with hit mask to set missed pixels to 0
    # geomIDs are 0-based so we cannot overwrite their background with 0
    for key in its.keys():
        if return_as_rasters:
            its[key] = its[key].reshape((res_y, res_x, -1))
        elif key not in ['hit']:
            its[key] = its[key].reshape((num_pix, -1))
            its[key] = its[key][hit, :]

    # compute intersection points
    its['points'] = ray_origins[hit, :] + its['tfar'] * ray_dirs[hit, :]
    return its

def interpolate_vertex_attributes(its, meshes, return_as_rasters: bool = False):
    """given intersection dict with barycentric coordinates ('u', 'v'), geometry IDs and triangle IDs, as well as the
    corresponding list of meshes, compute texture coordinates coordinates for each intersected pixel"""

    if isinstance(meshes, Scene):
        meshes = list(meshes.geometry.values())
    elif not isinstance(meshes, list):
        raise Exception('meshes should be list of trimesh.Trimesh objects or one trimesh.Scene object')

    # compute uv coordinates
    hit = its['hit']
    num_hit = np.count_nonzero(hit)
    its['uvs'] = np.zeros((num_hit, 2), dtype=np.float32)
    its['tangent_frames'] = np.zeros((num_hit, 3, 3), dtype=np.float32)
    for mesh_id in range(len(meshes)):
        # select only those intersections corresponding to the current mesh
        mesh_mask = its['geomID'] == mesh_id

        # barycentric coordinates are used to interpolate vertex attributes
        u = its['u'][mesh_mask]
        v = its['v'][mesh_mask]
        uvw = np.stack([1 - u - v, u, v], axis=0)

        # get per-triangle-vertex attributes (currently uv coordinates and tangent frames)
        triangle_ids = its['primID'][mesh_mask]
        its_vertex_ids = get_faces(meshes[mesh_id])[triangle_ids, :].T
        face_vertex_uvs = get_uv_coords(meshes[mesh_id])[its_vertex_ids, :]
        meshes[mesh_id] = get_tangent_frames(meshes[mesh_id])
        face_vertex_tangent_frames = meshes[mesh_id].vertex_tangent_frames[its_vertex_ids, :, :]  # 3 x (3 * NF) x 3 x 3

        # interpolate per pixel texture coordinates
        its['uvs'][mesh_mask.flatten(), :] = np.sum(face_vertex_uvs * uvw[..., None], axis=0)
        # interpolate per pixel tangent frames
        its['tangent_frames'][mesh_mask.flatten(), :, :] = np.sum(face_vertex_tangent_frames * uvw[:, :, None, None], axis=0)

    if return_as_rasters:
        # reshape to image size
        its['uvs'] = its['uvs'].reshape(its['u'].shape[:2] + (2, ))
        its['tangent_frames'] = its['tangent_frames'].reshape(its['u'].shape[:2] + (3, 3))
    return its

def get_local_dirs(its, global_position, normalized=True):
    """given intersection dict with per-pixel interpolated tangent frames and intersection points, as well as a global
    positions (e.g. light or view), compute per pixel local directions and unnormalized global directions (e.g., for
    computing light attenuation)"""
    dirs_global = global_position[None, :] - its['points']
    frames = its['tangent_frames']
    dirs_local = np.einsum('nij,nj->ni', frames, dirs_global)
    if normalized:
        dirs_local = normalize(dirs_local, axis=-1)
    return Dct(dirs_local=dirs_local, dirs_global=dirs_global)

def embree_render_deferred(meshes = None,
                           with_tangent_frames: bool = True,
                           cam: dict = None,
                           auto_cam: bool = True,
                           auto_cam_bbox: bool = False,
                           light_position: tuple = (0., 10., 0.),
                           return_as_rasters: bool = False):
    """given a list of meshes, optionally a camera and a point light position, render the scene via ray tracing;
    returns pixel buffers with:
    - 3D intersection points
    -
    - view and light directions in tangent space
    - """

    if meshes is None:
        meshes = []
        meshes.append(create_unit_sphere())
        meshes.append(create_unit_cube())

    bbox = get_bbox(meshes)
    bbox_diam = np.linalg.norm(np.diff(bbox, axis=0))

    if cam is None:
        cam = Dct(res_x=384, res_y=256)
    if not hasattr(cam, 'cx'):
        cam.cx = cam.res_x / 2.
        cam.cy = cam.res_y / 2.

    if not hasattr(cam, 'position'):
        cam.position = np.mean(bbox, axis=0) + (bbox_diam / 2 + 10. * np.random.rand()) * normalize(np.random.rand(3) - 0.5)

    # scene & camera setup
    scene, cam = embree_create_scene(meshes=meshes, with_tangent_frames=with_tangent_frames, cam=cam, auto_cam=auto_cam,
                                     auto_cam_bbox=auto_cam_bbox, auto_cam_visualize=False)

    # ray tracing
    buffers = embree_intersect_scene(scene=scene, cam=cam, return_as_rasters=return_as_rasters)

    # per pixel tangent frames
    buffers = interpolate_vertex_attributes(buffers, meshes, return_as_rasters=return_as_rasters)
    h, w = cam.res_y, cam.res_x

    # local light & view directions
    view_dirs = get_local_dirs(buffers, cam.position, normalized=True)
    light_dirs = get_local_dirs(buffers, light_position, normalized=True)

    if return_as_rasters:
        mask = buffers['hit']
        buffers['view_dirs_local'] = assign_masked(mask, view_dirs.pop('dirs_local'))
        buffers['view_dirs_global'] = assign_masked(mask, view_dirs.pop('dirs_global'))
        buffers['light_dirs_local'] = assign_masked(mask, light_dirs.pop('dirs_local'))
        buffers['light_dirs_global'] = assign_masked(mask, light_dirs.pop('dirs_global'))
    else:
        buffers['view_dirs_local'] = view_dirs.pop('dirs_local')
        buffers['view_dirs_global'] = view_dirs.pop('dirs_global')
        buffers['light_dirs_local'] = light_dirs.pop('dirs_local')
        buffers['light_dirs_global'] = light_dirs.pop('dirs_global')

    return buffers


import numpy as np
import cvxpy as cp
import const
import pickle
import general
import simulate

def iterative(D: np.array, M: np.array):
    '''
    Estimates state according to iterative MLE.

    :param D   : N array of measured data
        datatype: D[i] = [index of POVM]
    :param M   : Nxdxd array of POVM set
    :param iter: number of iterations
    :return: dxd array of iterative MLE estimator
    '''
    dim = M.shape[-1]
    n   = general.count(D, np.zeros(len(M), dtype=np.int))

    iter_max = 500
    dist     = float(1)

    rho_1 = np.eye(dim)/dim
    rho_2 = np.eye(dim)/dim

    j = 0
    while j<iter_max and dist>1e-14:
        p      = np.einsum('ik,nki->n', rho_1, M)
        R      = np.einsum('n,n,nij->ij', n, 1/p, M)
        update = R@rho_1@R
        rho_1  = update/np.trace(update)

        if j>=40 and j%20==0:
            dist  = general.infidelity(rho_1, rho_2)
        rho_2 = rho_1

        j += 1

    return rho_1

def gradient(D: np.array, M: np.array):
    '''
    Estimates state according to iterative MLE using gradient descent.

    :param D   : N array of measured data
        datatype: D[i] = [index of POVM]
    :param M   : Nxdxd array of POVM set
    :param iter: number of iterations
    :return: dxd array of iterative MLE estimator
    '''
    N = len(D)
    dim = M.shape[-1]

    # build likelihood
    n = general.count(D, np.zeros(N, dtype=np.int))
    M_D = M[D]

    # initialise gradient descent
    rho = cp.Variable((dim, dim), hermitian=True)
    constraints = [rho>>0, cp.real(cp.trace(rho))==1]

    # loglike = cp.sum( [cp.prod([n[i], cp.log(cp.real(cp.trace(cp.matmul(M[i], rho))))]) for i in range(N)] )
    # loglike = cp.real(cp.trace(rho))
    loglike = cp.sum( [cp.log(cp.real(cp.trace(cp.matmul(M_D[i], rho)))) for i in range(6)] )
    prob    = cp.Problem(cp.Maximize(loglike), constraints)

    # solve
    prob.solve()
    return rho.value

def gradient_scipy(D: np.array, M: np.array):
    '''
    Estimates state according to iterative MLE using gradient descent.

    :param D   : N array of measured data
        datatype: D[i] = [index of POVM]
    :param M   : Nxdxd array of POVM set
    :param iter: number of iterations
    :return: dxd array of iterative MLE estimator
    '''
    N = len(D)
    dim = M.shape[-1]

    # build likelihood
    n = general.count(D, np.zeros(N, dtype=np.int))
    M_D = M[D]

    L = lambda rho: np.prod([np.trace(rho.reshape(dim, dim)@M[i])**n[i] for i in range(len(M))])

    rho_0 = np.eye(dim)/dim
    cons = {'type': 'eq', 'fun': lambda rho: np.trace(rho.reshape(dim, dim)) - 1,
            'type': 'eq', 'fun': lambda rho: np.sum(np.abs( rho.reshape(dim, dim) - general.H(rho.reshape(dim, dim)) ))}

    minimize(L, rho_0, method='SLQP', constraints=cons)

def two_step(rho: np.array, M0: np.array, N: int, N0: int, f_align, cup=True, prec=1e-14):
    '''
    Estimates with one intermediate step of POVM realignment.

    :param rho    : dxd array of density matrix
    :param M0     : Nxdxd array of POVM set
    :param N      : total number of measurements
    :param N0     : number of measurements in the first step
    :param f_align: function about how to align after firs step
    :param cup    : all data or only data after first estimate
    :return: dxd array of adaptive mle estimator
    '''
    # initial constants
    dim = rho.shape[-1]

    # initial estimate
    D0    = simulate.measure(rho, np.minimum(N, N0), M0)
    rho_0 = iterative(D0, M0)

    if N<=N0:
        return rho_0
    else:
        # rallignment according to initial estimate
        M1 = f_align(rho_0, M0)

        # true state
        N1 = int(N-N0)
        D1 = simulate.measure(rho, N1, M1)
        if cup:
            D = np.concatenate([D0, D1+len(M1)], axis=0)
            M = np.concatenate([M0, M1], axis=0)
        else:
            D = D1
            M = M1

        return iterative(D, M)


# -*- coding: utf-8 -*-
"""
Created on Sat Nov 30 09:53:01 2019

@author: ferna
"""

from sys import argv

script, filename = argv

print(f"We're going to erase {filename}." )
print("If you don't want that hit CTRL-C (^C).")
print("If you do want that, hit RETURN.")

input("?") # any input works

print("Opening the file...")
target = open(filename, 'w') # allows program to write on file
# r is read, wr is read and write, a is append,wont truncate just add
print("Truncating the file. Goodbye!")
target.truncate() # not necesesary because already truncates the file

print("Now I'm going to ask you for three lines.")

line1 = input("line 1: ")
line2 = input("line 2: ")
line3 = input("line 3: ")

print("I'm going to write these to file.")

target.write(line1)
target.write("\n")
target.write(line2)
target.write("\n")
target.write(line3)

print("And finally, we close it.")
target.close()


#服务端
#导入socket模块
import socket

#创建socket
s = socket.socket()
#绑定端口
s.bind(("127.0.0.1",6666))

#等待客户端连接
s.listen(5)

while True:
    #建立客户端连接
    c,addr = s.accept()
    print("连接地址：",addr)
    c.send(b'Welcome')
    #关闭连接
    c.close()


#!/usr/bin/env python

import os
import sys

path = '/usr/share/archivematica/dashboard'
if path not in sys.path:
    sys.path.append(path)
os.environ['DJANGO_SETTINGS_MODULE'] = 'settings.common'
from fpr.models import FPRule
from main.models import FileFormatVersion, Transfer

def is_extractable(f):
    try:
        format = f.fileformatversion_set.get().format_version
    except FileFormatVersion.DoesNotExist:
        return False

    extract_rules = FPRule.active.filter(purpose="extract", format=format)
    if extract_rules:
        return True
    else:
        return False

def main(sip_uuid):
    transfer = Transfer.objects.get(uuid=sip_uuid)
    for f in transfer.file_set.iterator():
        if is_extractable(f):
            return 0

    return 1

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


# -*- coding: utf-8 -*-
# This file is part of Viper - https://github.com/viper-framework/viper
# See the file 'LICENSE' for copying permission.
from __future__ import unicode_literals

import os
import sys
from tests.conftest import FIXTURE_DIR
from viper.common.objects import File, MispEvent
import pytest

class TestMispEvent:

    @pytest.mark.parametrize("filename", ["58e902cd-dae8-49b9-882b-186c02de0b81.json"])
    def test_mispevent(self, capsys, filename):
        mispevent = MispEvent(os.path.join(FIXTURE_DIR, filename))
        mispevent.online()
        mispevent.offline()
        ips = mispevent.get_all_ips()
        domains = mispevent.get_all_domains()
        urls = mispevent.get_all_urls()
        hashes = mispevent.get_all_hashes()
        assert '191.101.230.149' in ips
        assert not domains
        assert not urls
        assert '722050c1b3f110c0ac9f80bc80723407' in hashes[0]
        assert not hashes[1]

class TestFile:
    @pytest.mark.parametrize("filename, name", [
        ("string_handling/ascii.txt", "ascii.txt"),
        ("string_handling/with blank.txt", "with blank.txt")
        ])
    def test_init(self, capsys, filename, name):
        instance = File(os.path.join(FIXTURE_DIR, filename))

        assert isinstance(instance, File)
        assert instance.path == os.path.join(FIXTURE_DIR, filename)
        assert instance.name == name

        out, err = capsys.readouterr()
        assert out == ""

    @pytest.mark.skipif(sys.version_info < (3, 3), reason="requires at least python3.3")
    @pytest.mark.parametrize("filename, name", [
        ("string_handling/dümmy.txt", "dümmy.txt"),
        ])
    def test_init_unicode(self, capsys, filename, name):
        instance = File(os.path.join(FIXTURE_DIR, filename))

        assert isinstance(instance, File)
        assert instance.path == os.path.join(FIXTURE_DIR, filename)
        assert instance.name == name

        out, err = capsys.readouterr()
        assert out == ""


import networkx as nx
import numpy as np

class Sniper:

    def __init__(self, prop):
        self.y = 0.5
        self.r = 0.05
        self.graph = nx.Graph()
        self.proportion = prop

    def score(self, grad):

        nodes = [int(k) for k in list(grad.keys())]
        self.graph.add_nodes_from(nodes)

        length = len(grad)
        distance_matrix = np.zeros([length, length])
        for i in range(length):
            for j in range(i + 1, length):
                distance_matrix[i][j] = np.sqrt(np.sum(np.square(np.asarray(grad[str(i)]) - np.asarray(grad[str(j)]))))

        max_val = distance_matrix.max()
        min_val = distance_matrix.min()
        for i in range(length):
            for j in range(i + 1, length):
                distance_matrix[i][j] = (distance_matrix[i][j] - min_val) / (max_val - min_val)

        tr = self.y
        max_len = 0
        while tr < 1 and max_len < length * self.proportion:
            self.add_edges(tr, distance_matrix)
            cliques = nx.find_cliques(self.graph)
            max_len = 0
            for item in cliques:
                if len(item) > max_len:
                    max_len = len(item)
                    max_clique = item
            tr += self.r

        score = {}
        for key in grad:
            if int(key) in max_clique:
                score[key] = 0
            else:
                score[key] = 1
        return score

    def add_edges(self, threshold, matrix):

        self.graph.remove_edges_from(self.graph.edges())

        length = len(matrix)
        for i in range(length):
            for j in range(i + 1, length):
                if matrix[i][j] < threshold:
                    self.graph.add_edge(i, j)
        return


import base64
import json

from aws_lambda_powertools.utilities.data_classes import (
    KinesisFirehoseEvent,
    event_source,
)
from aws_lambda_powertools.utilities.typing import LambdaContext

@event_source(data_class=KinesisFirehoseEvent)
def lambda_handler(event: KinesisFirehoseEvent, context: LambdaContext):
    result = []

    for record in event.records:
        # if data was delivered as json; caches loaded value
        data = record.data_as_json

        processed_record = {
            "recordId": record.record_id,
            "data": base64.b64encode(json.dumps(data).encode("utf-8")),
            "result": "Ok",
        }

        result.append(processed_record)

    # return transformed records
    return {"records": result}


"""
上半部的程式碼在跑跑測試資料會有部分出問題
"""

import re
letterA, letterB = input().split(" ")
sentence_input = input()
def count_word_pairs_in_paragraph(word1, word2, paragraph):
    paragraph = paragraph.strip()
    #逗點、分號、句點、驚嘆號和問號切割句子
    sentences = re.split(r'[.!?]',paragraph)

    #sentences = re.split(r'[.!?]',paragraph)
    print(sentences)
    #del sentences[-1]
    #print(sentences)

    total_sentences = len(sentences)
    word_pair_count = 0

    for sentence in sentences:
        sentence = sentence.strip()
        if sentence:
            words = sentence.split()
            if (word1 in words) and (word2 in words):
                word_pair_count += 1

    return total_sentences, word_pair_count, sentences

total_sentences, word_pair_count ,sentences= count_word_pairs_in_paragraph(letterA, letterB, sentence_input )
print(str(total_sentences-1), str(word_pair_count))

"""
下半部是公布的正確解答程式碼
"""
"""
請複製其中一段即可
"""
# 讀入英文單詞及英文段落
word1, word2 = input().split()
paragraph = input()

# 給定標點符號
sentence_sep = '.!?'
delimeters = ',;'

# 計算句子數量
sentence_cnt = 0
for sep in sentence_sep:
    sentence_cnt += paragraph.count(sep)

# 移除非句子分隔符號的標點符號
for deli in delimeters:
    while deli in paragraph:
        paragraph = paragraph.replace(deli, ' ')

# 先把句子切開
sentences = paragraph.split('.')
for sep in '!?':
    for part in sentences:
        if sep in part:
            sentences.remove(part)
            for item in part.split(sep):
                sentences.append(item)

# 計算兩個英文單詞同時出現的句子數
target_cnt = 0
for sentence in sentences:
    sentence = ' ' + sentence + ' '
    if ' ' + word1 + ' ' in sentence and ' ' + word2 + ' ' in sentence:
        target_cnt += 1

# 印出結果
print(sentence_cnt, target_cnt)


start=[1,2,3,-1,4,5,6,7,8]
desination=[1,2,3,4,5,-1,6,7,8]
def func(start,desination,depth):
    for limit in range(0,depth+1):
        visited_states=[]
        if dfs(start,desination,limit,visited_states):
            return True
    return False

def gen(state,m,b):
    temp=state[:]
    if m=='l': temp[b],temp[b-1]=temp[b-1],temp[b]
    if m=='r':temp[b],temp[b+1]=temp[b+1],temp[b]
    if m=='u':temp[b],temp[b-3]=temp[b-3],temp[b]
    if m=='d':temp[b],temp[b+3]=temp[b+3],temp[b]
    return temp

def possible_moves(state,visited_states):
    b=state.index(-1)
    d=[]
    pos_moves=[]
    if b<=5: d.append('d')
    if b>=3 : d.append('u')
    if b%3 > 0: d.append('l')
    if b%3 < 2: d.append('r')
    for i in d:
        temp=gen(state,i,b)
        if not temp in visited_states: pos_moves.append(temp)
    print(pos_moves)
    return pos_moves

def dfs(start,desination,limit,visited_states):
    if start==desination: return True
    if limit <= 0 :return False
    visited_states.append(start)
    adj=possible_moves(start,visited_states)
    for move in adj:
        if dfs(move,desination,limit-1,visited_states): return True
    return False

print(func(start,desination,1))


import sys
import numpy as np
import pandas as pd
from optparse import OptionParser

# Construct the list of predictors
PREDICTORS = [
    'yyyymmdd',
    'game_no',
    'weekday',
    'visiting_team',
    'visiting_team_league',
    'home_team',
    'home_team_league',
    'daytime',
    'park_id',
]

# Add the starting line-up for both teams:
# id, name, and defensive position
for prefix in ['home_', 'visiting_']:
    for j in range(1, 10):
        PREDICTORS.append(prefix+f'starting_player{j}_id')
        PREDICTORS.append(prefix+f'starting_player{j}_name')
        PREDICTORS.append(prefix+f'starting_player{j}_def_position')

def evaluate(alg, train_data, test_data, batchsize = 10):

    # Allow the algorithm to train itself
    alg.train(train_data)

    # Predictions container
    predicted = pd.Series(index = test_data.index)

    # Loop through test set and predict
    num_test = test_data.shape[0]
    j = 0
    while j < num_test:

        # Create batch
        batchlen = min(batchsize, num_test - j)
        batch = test_data.iloc[j:j+batchlen]

        # Create identifying information
        batched_test_info = batch[PREDICTORS]

        # Predict
        predicted[batch.index] = alg.predict(batched_test_info)

        # Update
        alg.update(batch)

        # Increment j
        j += batchlen

    return predicted

# Helper functions to load the algorithms
def load_modules(alg_classes):
    """Each agent class must be in module class_name.lower().
    Returns a dictionary class_name->class"""

    def load(class_name):
        module_name = str(class_name).lower()  # by convention / fiat
        module = __import__(module_name)
        agent_class = module.__dict__[class_name]
        return (class_name, agent_class)

    return dict(map(load, alg_classes))

def main(args):

    # Fetch classnames
    usage_msg = 'python3 evaluation.py [algorithms]'
    parser = OptionParser(usage=usage_msg)
    (options, args) = parser.parse_args()

    if len(args) == 0:
        args = ['Baseline']

    # Load classes
    alg_classes = load_modules(args)
    alg_instances = {key:alg_classes[key]() for key in alg_classes}

    # Load training/test data
    nrows = 1000
    train_data = pd.read_csv('data/processed/train.csv', nrows=nrows)
    test_data = pd.read_csv('data/processed/test.csv', nrows=nrows)

    # Tiny bit of formatting for test set
    test_data = test_data.sort_values('yyyymmdd')

    # Evaluate
    all_predictions = pd.DataFrame(index=test_data.index, columns=args)
    for alg_name in alg_instances:
        all_predictions[alg_name] = evaluate(
            alg_instances[alg_name], train_data, test_data
        )
    print(all_predictions)

if __name__ == '__main__':

    main(sys.argv)


#!/usr/bin/python
# -*- coding: UTF-8 -*-

import unittest

import dn_kw
import ew_kw
import oh_kw
import od_kw
import es_kw

MojeSlupki = es_kw.MojeSlupki

class FakturoweRoczneSlupki(MojeSlupki):
    def __init__(self, tgk, aqr, dnw, rn_after):
        '''
        FakturoweRoczneSlupki:
        '''
        MojeSlupki.__init__(self, tgk, aqr, dnw, rn_after)

    def ustaw_skalowanie_obrazu(self):
        '''
        FakturoweRoczneSlupki:
        '''
        self.ustaw_skalowany_obraz()

    def wyznacz_etykiete(self, pocz):
        '''
        FakturoweRoczneSlupki:
        '''
        return str(dn_kw.RokDnia(pocz))

class TestFakturowychRocznychSlupkow(unittest.TestCase):
    def test_fakturowych_rocznych_slupkow(self):
        '''
        TestFakturowychRocznychSlupkow:
        '''
        tgk = od_kw.PseudoTGK()
        tgk.wyznacz_unikalny_moment_dla_grafiki()
        aqr = ew_kw.SzkieletDatDlaFakturMiesLat()
        dzien_pocz = 13149
        dzien_kon = 13879
        szkielet_lat = dn_kw.daty_lat(dzien_pocz, dzien_kon)
        aqr.przypisz_szkielet(szkielet_lat)
        lp_wykresu = 0
        dnw = oh_kw.SimpleDNW(lp_wykresu)
        rn_after = 2
        obk = FakturoweRoczneSlupki(tgk, aqr, dnw, rn_after)
        self.assertEqual(obk.szerokosc_dx_skali, 0)
        self.assertEqual(obk.szerokosc_slupka, 30)
        self.assertEqual(obk.wysokosc_obrazu, 150)
        self.assertEqual(obk.margines_dy_powyzej_slupka, 20)
        self.assertEqual(obk.wsp_y_na_dole_slupka, 130)
        self.assertEqual(obk.gorna_mniejsza, 140)
        self.assertEqual(obk.linii_na_dole(), 1)
        self.assertEqual(obk.chce_bez_tresci, 0)


from __future__ import (absolute_import, division,
                        print_function, unicode_literals, )
from future.builtins import range

import numpy as np
import tt
from tt.optimize import tt_min
import copy
import scipy

import ttrecipes as tr

def search_subspace(t, Xs, Ys, tsq=None, tmin=None, tmax=None, hinv=True, vinv=True, dims=None, eps=1e-3):
    """
    Find a region that bests match a given template (both have to be axis-aligned).

    :param Xs: a list of M ndarrays (each of dimension M) encoding the positions to be searched. Note: non-rectangular regions will be extended to their axis-aligned bounding box
    :param Ys: an ndarray of dimension M containing the values to be searched. Must have the same size as the elements of Xs
    :param hinv: whether to search with invariance to horizontal translation (default: True). If True, the values in Xs must lie between 0 and 1
    :param vinv: whether to search with invariance to vertical translation (default: True). If True, the values in Ys must lie between 0 and 1
    :param dims: the dimensions along which the subspace should be searched. Must be a list with M elements. If None (default), the best tuple of M dimensions will be chosen
    :return: match: the center of the best-matching window; dims: the dimensions along which the best subspace was found; borders: a list of M pairs
     (start, end) with the coordinates for the region

    TODO: support non-invariances
    TODO: support search in 2D and more
    TODO: use masks

    """

    N = t.d
    M = Ys.ndim
    if M != 1:
        raise NotImplementedError("Only 1D search is supported for now")
    if hinv == False or vinv == False:
        raise NotImplementedError
    assert len(Xs) == M
    assert np.all([X.ndim == M for X in Xs])
    if hinv:
        assert np.all([np.all(X >= 0) for X in Xs])
        assert np.all([np.all(X <= 1) for X in Xs])
    if vinv:
        assert np.all(Ys >= 0)
        assert np.all(Ys <= 1)
    if dims is not None:
        assert len(dims) == M

    if dims == None:
        dims = np.arange(N)
    inds = np.argsort(Xs[0]) # Sort the values, in case they are reversed
    Xs[0] = Xs[0][inds]
    Ys = Ys[inds]
    record_value = float('inf')

    for dim in dims:
        Xs_adj = (Xs[0]-np.min(Xs[0])-1e-9)*t.n[dim]  # Xs adjusted to this dimension's grid ticks
        Ys_adj = scipy.interpolate.interp1d(Xs_adj, Ys)(np.arange(np.floor(Xs_adj[-1])+1))  # Adjusted Ys
        Ys_adj = Ys_adj*(tmax - tmin) + tmin  # Ys are now mapped to the tensor output's range
        Ys_adj -= np.mean(Ys_adj)  # Vertical invariance

        # Compute the loss tensor: for each position, the squared norm (dot product of itself) of a windowed fiber
        # consisting of the difference between the original fiber (x minus its mean t, for vertical invariance) and
        # the vertically-centered template fiber y (with zero mean). We use the expansion:
        # (x - t - y)*(x - t - y) = x*x - 2*x*t - 2*x*y + 2*t*y + t*t + y*y = x*x - t*t - 2*x*y + y*y
        # where we have used that mean(y) = 0

        if tsq is None:  # This part may be reused by subsequent queries, so we memoize it
            tsq = tt.multifuncrs2([t], lambda x: x**2, eps=eps, verb=False)

        box_filter = np.ones([1, Ys_adj.size, 1]) / Ys_adj.size  # To average TT cores along their 2nd dimension

        cores = tt.vector.to_list(tsq)
        cores[dim] = scipy.signal.convolve(cores[dim], box_filter, mode='valid')
        tsquarebox = tt.vector.from_list(cores)

        cores = tt.vector.to_list(t)
        cores[dim] = scipy.signal.convolve(cores[dim], box_filter, mode='valid')
        tbox = tt.vector.from_list(cores)
        tboxsquare = tt.multifuncrs2([tbox], lambda x: x**2, eps=eps, verb=False)

        cores = tt.vector.to_list(t)
        cores[dim] = scipy.signal.convolve(cores[dim], np.reshape(Ys_adj[::-1], [1, Ys_adj.size, 1]), mode='valid')
        tconv = tt.vector.from_list(cores)

        tfinal = Ys_adj.size*tsquarebox - Ys_adj.size*tboxsquare - 2*tconv + np.dot(Ys_adj,Ys_adj)*tt.ones(tsquarebox.n)

        # Find an approximate minimum
        value, match = tt_min.min_tens(tfinal, rmax=5, nswp=5, verb=False)

        if value < record_value:
            record_value = value
            record_match = [int(m) for m in match]
            record_match[dim] += len(Ys_adj)/2  # Sum half the window size to get the final focus point
            record_dim = [dim]
            record_borders = [[record_match[dim] - len(Ys_adj)/2, record_match[dim]+len(Ys_adj)/2-1]]

    return record_match, record_dim, record_borders

def best_subspace(t, ndim=1, target='max', mode='cross', eps=1e-6, verbose=False, **kwargs):
    """
    Find an axis-aligned subspace of a certain dimensionality that has the highest/lowest variance. Example applications:
    - In visualization, to find interesting subspaces
    - In factor fixing (sensitivity analysis), to find the set of parameters (and their values) that will minimize the uncertainty of a model

    TODO: only uniform independently distributed inputs are supported now

    :param t: a TT
    :param ndim: dimensionality of the subspace sought (default is 2)
    :param target: if 'max' (default), the highest variance will be sought; if 'min', the lowest one
    :param mode: 'cross' (default) or 'kronecker'
    :param verbose:
    :param kwargs: arguments for the cross-approximation
    :return: (a) a list of indices, with slice(None) in the free subspace's dimensions, and (b) the variance of that subspace

    """

    assert mode in ('cross', 'kronecker')
    assert target in ('max', 'min')

    # Build up a tensor that contains variances of all possible subspaces of any dimensionality, using the formula E(X^2) - E(X)^2
    if mode == 'cross':
        cores = tt.vector.to_list(tt.multifuncrs2([t], lambda x: x ** 2, eps=eps, verb=verbose, **kwargs))
    else:
        cores = tt.vector.to_list((t*t).round(0))
    cores = [np.concatenate([np.mean(core, axis=1, keepdims=True), core], axis=1) for core in cores]
    part1 = tt.vector.from_list(cores)  # E(X^2)

    cores = tt.vector.to_list(t)
    cores = [np.concatenate([np.mean(core, axis=1, keepdims=True), core], axis=1) for core in cores]
    part2 = tt.vector.from_list(cores)
    if mode == 'cross':
        part2 = tt.multifuncrs2([part2], lambda x: x ** 2, eps=eps, verb=verbose, **kwargs)
    else:
        part2 = (part2*part2).round(0)  # E(X)^2

    variances = (part1 - part2).round(0)

    # Filter out encoded subspaces that do not have the target dimensionality
    mask = tt.vector.to_list(tr.core.hamming_eq_mask(t.d, t.d-ndim))
    mask = [np.concatenate([core[:, 0:1, :], np.repeat(core[:, 1:, :], sh, axis=1)], axis=1) for core, sh in zip(mask, t.n)]
    mask = tt.vector.from_list(mask)

    # Find and return the best candidate
    if target == 'max':
        prod = tt.vector.round(variances*mask, eps=eps)
        val, point = tt_min.min_tens(-prod, verb=verbose)
        val = -val
    else:
        shift = -1e3*tt_min.min_tens(-variances, verb=False, rmax=1)[0]
        variances_shifted = variances - tt.vector.from_list([np.ones([1, sh+1, 1]) for sh in t.n])*shift
        val, point = tt_min.min_tens(variances_shifted*mask, verb=verbose)
        val += shift
    nones = np.where(np.array(point) == 0)[0]
    point = [p - 1 for p in point]
    for i in nones:
        point[i] = slice(None)
    return point, val

def moments(t, modes, order, centered=False, normalized=False, keepdims=False, eps=1e-3, verbose=False, **kwargs):
    """
    Given an N-dimensional TT and a list of M modes, returns a TT of dimension N - M that contains the k-th order moments along these modes

    :param t: a TT
    :param modes: a list of M integers
    :param order: an integer
    :param centered: if True the moments will be computed about their mean. Default is False
    :param normalized: if True the moments will be divided by sigma^order. Default is False
    :param eps: accuracy for cross-approximation (default is 1e-3)
    :return: a TT of dimension N - M

    """

    N = t.d
    assert np.all(0 <= np.array(modes))
    assert np.all(np.array(modes) < N)
    if not hasattr(modes, '__len__'):
        modes = [modes]
    assert len(modes) == len(set(modes))  # Modes may not be repeated
    assert 1 <= len(modes) <= N

    if centered or normalized:
        central_cores = []
        cores = tt.vector.to_list(t)
        for n in range(N):
            if n in modes:
                central_cores.append(np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1))
            else:
                central_cores.append(cores[n])
        central = t - tt.vector.from_list(central_cores)
    if centered:
        if order == 1:
            moments = copy.deepcopy(central)
        else:
            moments = tt.multifuncrs2([central], lambda x: x**order, eps=eps, verb=verbose, **kwargs)
    else:
        if order == 1:
            moments = copy.deepcopy(t)
        else:
            moments = tt.multifuncrs2([t], lambda x: x**order, eps=eps, verb=verbose, **kwargs)
    cores = tt.vector.to_list(moments)
    for mode in modes:
        cores[mode] = np.mean(cores[mode], axis=1, keepdims=True)
    moments = tt.vector.from_list(cores)
    if normalized:
        central = tt.multifuncrs2([central], lambda x: x**2, eps=eps, verb=verbose, **kwargs)
        cores = tt.vector.to_list(central)
        for mode in modes:
            cores[mode] = np.mean(cores[mode], axis=1, keepdims=True)
        variances = tt.vector.from_list(cores)
        moments = tt.multifuncrs2([moments, variances], lambda x: x[:, 0] / (x[:, 1] ** (order/2.)), eps=eps, verb=verbose, **kwargs)
    if not keepdims:
        moments = tr.core.squeeze(moments, modes=modes)
    return moments

def means(t, modes, **kwargs):
    """
    Convenience function for the first moment (see :func: `moments`)
    """

    return moments(t, modes, order=1, centered=False, normalized=False, **kwargs)

def variances(t, modes, **kwargs):
    """
    Convenience function for the second centered unnormalized moment (see :func: `moments`)
    """

    return moments(t, modes, order=2, centered=True, normalized=False, **kwargs)

# 0, 2
# 0 1 2 3 4
# 2 1 0 3 4

# 0 1
# 1 0

# 0 1 2 3 4 5
# 1 0 2 3 4 5
# 2 0 1 3 4 5
# 3 0 1 2 4 5
# 4 0 1 2 3 5
# 5 0 1 2 3 4

# 0 1 2 3 4 5
# 1 0 2 3 4 5
# 2 1 0 3 4 5
# 3 1 2 0 4 5
# 4 1 2 3 0 5
# 5 1 2 3 4 0

# 1 0 2 3 4 5
# 0 1 2 3 4 5
# 0 2 1 3 4 5
# 0 3 2 1 4 5
# 0 4 2 3 1 5
# 0 5 1 2 3 4

# 0 1 2 :
# 0 : 2 3
# t1 x y z : s
# s2 x y z t :
# x y z t1 s
# x y z t s2

# 0 - 1
# 1 0 -

# 0 1 2
# 1 0 2
# 2 0 1

# f(x1, r, s) = f(x2, r, s) forall r, s
# t(x1, r, s) = f(r, x2, s) forall r, s

# 0 - 1 2
# 1 0 - 2
# 2 0 1 -

# x, r, s, t
# f(x1, r, s, t) = f(x2, r, s, t) forall r, s, t
# f(x1, s, t) = f(x2, r, t)
# t(x1, s, t) = t(r, x2, t) forall r, s, t

# f(x1, x1, x2, t) = f(x2, x1, x2, t)

# x 1 2
# 0 y 2

# 0 1 2 3 4
# 2 1 0 3 4
# 0 1 2 3 4

def pca_ensemble(t, modes, degree=3, orthogonal=False, anchor_mean=True, eps=1e-3, verbose=False):
    """
    Given a tensor and a list of modes, consider the space of subtensors obtained
    by fixing these modes. Find a PCA basis for that space and return the coefficients
    for each subtensor

    :param t: a TT
    :param modes: the ensemble modes. Can also be a list of lists; in that case
    the ensembles will share bases
    :param degree: how large a PCA basis to take (default is 3)
    :param orthogonal:
    :param anchor_mean:
    :param eps: relative tolerance while permuting modes. Default is 1e-3
    :param verbose:
    :return: (ensemble, preserved):
        - ensemble: a TT of size M1 x ... x Mk x `degree`, where M1, etc. are the sizes of `t` along `modes`
        - preserved: the fraction (between 0 and 1; higher is better) of norm preserved by the embedding

    """

    N = t.d
    if verbose:
        print('Computing PCA ensemble...')
    if not hasattr(modes, '__len__'):
        modes = [modes]
    modesarray = np.array(modes)
    single = False
    if modesarray.ndim == 1:
        single = True
        modesarray = np.array([modesarray])
    if not all([len(modesarray[m]) == len(modesarray[0]) for m in range(len(modesarray))]):
        raise ValueError("All mode lists must have the same length")
    # modesarray = np.array([[3, 6], [1, 4]])
    # Group target modes next to each other
    first = int(np.round(np.mean(modesarray.flatten())-modesarray.shape[1]/2.))
    # print(np.rint(0.50000))

    import matplotlib.pyplot as plt
    plt.figure()
    # plt.plot(t.full()[20, :] - np.mean(t.full()[20, :]))
    # plt.plot(t.full()[:, 4] - np.mean(t.full()[:, 4]))
    # plt.imshow(t.full(), origin='lower')
    # print(t.full()[17, :])
    # print(t.full()[:, 4])
    # print(t.full()[17, :] - np.mean(t.full()[17, :]))
    # print(t.full()[:, 4] - np.mean(t.full()[:, 4]))
    # plt.show()
    # assert 0

    print(modesarray, first)
    stack = []
    for modes in modesarray:
        if len(modes) != len(np.unique(modes)):
            raise ValueError('Modes may not be repeated')
        assert 0 <= np.min(modes)
        assert np.max(modes) < N

        if verbose:
            print('\tTransposing tensor... 0', end='')
        # part1 = np.arange(first)
        all1 = list(range(first))
        all2 = list(range(first, N))
        part1 = np.setdiff1d(all1, modes)
        part2 = np.setdiff1d(all2, modes)
        print('We now transpose to', list(part1)+list(modes)+list(part2))
        ttransp = tr.core.transpose(t, list(part1)+list(modes)+list(part2), eps=eps)
        for m in modes:
            if m in all1:
                all1[all1.index(m)] = -1
            if m in all2:
                all2[all2.index(m)] = -1
        all1.append(-2)  # Extra dimension to stack the tensors
        print(all1, all2)
        converted = np.array(all1+list(modes)+all2)
        assert first == np.where(converted == -2)[0][0]
        # shape = t.n[modes]
        idx = np.where(converted < 0)[0]
        shape = np.ones_like(idx)
        shape[np.where(converted[idx] == -1)[0]] = t.n[modes]
        # shape[np.where(converted[idx] == -1)[0]] = 1
        print('idx = {}, shape = {}'.format(idx, shape))
        # idx1 = []
        # idx2 = np.setdiff1d(np.arange(first, N), part2)
        # print('idx1:', idx1)
        # print('idx2:', idx2)
        ttransp = tr.core.insert_dummies(ttransp, idx, shape=shape)
        print('shape:', shape)
        # print(ttransp)
        # print(ttransp[[0, 0, 0, 0]+[0]*6])
        # print(ttransp[[0, 0, 0, 1]+[0]*6])
        # print(ttransp[[0, 0, 0, 2]+[0]*6])
        # print(ttransp)
        #
        # print((1./np.prod(t.n[modes])))
        # assert 0
        # ttransp *= (np.prod(t.n[modes]))
        # cores = tt.vector.to_list(ttransp)
        # cores[2] /= 32
        # ttransp = tt.vector.from_list(cores)
        stack.append(ttransp)
    stack = tr.core.concatenate(stack, axis=first, eps=eps)

    # first = first-1
    last = first+len(modes)+1
    # print(first, last)
    # print(t, first, last)
    # assert 0
    # print(stack)
    # print(stack[0, 15, 0, 25])
    # print(t[15, 25])
    # print(stack[0, 15, 0, 15])
    # print(t[15, 15])
    # print(stack[0, 13, 0, 25])
    # print(t[13, 25])
    # print(stack[1, 25, 12, 25])
    # print(t[12, 25])
    # blah = stack.full()
    # blah = blah - np.mean(blah, axis=(2, 3), keepdims=True)
    # print(np.mean(blah, axis=(2, 3), keepdims=True).shape)
    # print(blah.shape, np.linalg.norm(blah))
    # print(first, last)
    # print(modesarray)
    t = stack
    # assert 0
    # print('*******')
    # print(t)
    # print(tt.vector.round(t, eps=eps))
    # # print(tt.vector.norm(t))
    # assert 0

    if anchor_mean:
        # We now project each vector on the subspace spanned by the constant vector
        # This is achieved by subtracting each vector's within-mean
        cores = copy.deepcopy(tt.vector.to_list(t))
        for n in range(t.d):
            if n < first or n >= last:
                cores[n] = np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1)
        withinmeanrep = tt.vector.from_list(cores)
        print("Let's do means")
        print('t=', t)
        print('modes =', list(range(first))+list(range(last, t.d)))
        withinmean = tr.core.means(t, modes=list(range(first))+list(range(last, t.d)))
        t = tt.vector.round(t-withinmeanrep, eps=eps)
        degree -= 1

    # print(tt.vector.norm(t))
    # assert 0
    # print(withinmean)
    # assert 0
    # withinmeanrep2 = withinmeanrep - tr.core.constant_tt(shape=withinmean.n, fill=tr.core.mean(withinmeanrep))
    withinmean2 = withinmean - tr.core.constant_tt(shape=withinmean.n, fill=tr.core.mean(withinmean))
    # sumsq = np.sum(withinmean2)
    # print(withinmean2)
    # sumsq1 = tr.core.variances(withinmean2, modes=np.arange(len(modes))).full() #* np.prod(t.n)
    # sumsq1 = np.asscalar(sumsq1)
    sumsq1 = np.asscalar(np.sum(withinmean2.full()**2))# * np.prod(t.n[np.delete(np.arange(N), modes)])
    # targetsq1 = np.var(withinmean.full())*np.prod(t.n[modes])

    sumsq1 = np.sqrt(sumsq1)
    # targetsq1 = np.sqrt(targetsq1)
    print('sumsq1 = {}'.format(sumsq1))
    # print('sumsq1 = {}, targetsq1 = {}, sumsq1/targetsq1 = {}'.format(sumsq1, targetsq1, sumsq1/targetsq1))

    # print(np.sqrt(sumsq1/targetsq1))
    # print(tt.vector.norm(withinmean2)**2)
    # sumsq = np.sum(withinmean2.full()**2)
    # print(sumsq)
    # assert 0

    # Subtract the cross-mean (mean of the collection of vectors)
    cores = copy.deepcopy(tt.vector.to_list(t))
    for n in range(first, last):
        cores[n] = np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1)
    crossmean = tt.vector.from_list(cores)
    t = tt.vector.round(t-crossmean, eps=eps)

    # sumsq2 = tt.vector.norm(t) / np.sqrt(np.prod(t.n[np.delete(np.arange(N), modes)]))
    # sumsq2 = tt.vector.norm(t) / np.sqrt(np.prod(t.n[np.delete(np.arange(t.d), [0, 1])]))
    sumsq2 = tt.vector.norm(t) / np.sqrt(np.prod(t.n[np.arange(first)])*np.prod(t.n[np.arange(last, t.d)]))
    # sumsq2 = np.sum(tr.core.variances(t, modes=list(range(final_pos))+list(range(final_pos+len(modes), N))).full())\
    #          * np.prod(t.n[modes])
    # sumsq2 = np.sqrt(sumsq2)
    print('sumsq2 = {}'.format(sumsq2))
    # t = t*(sumsq2 / sumsq1 * targetsq1)
    # print('Now 0:', tt.vector.norm(t))
    # t *= (targetsq1 / sumsq1)
    # print('Now 1:', tt.vector.norm(t))
    # assert 0
    # print(t)

    # Orthogonalize to the left and right of these modes
    cores = tt.vector.to_list(t)
    if verbose:
        print('\n\tOrthogonalization...')
    for i in range(first):
        tr.core.left_orthogonalize(cores, i)
    for i in range(t.d-1, last-1, -1):
        tr.core.right_orthogonalize(cores, i)

    # Cut out the relevant sequence of modes
    cores = cores[first:last]
    # print(first, last, len(cores))
    # assert 0

    # Convert the sequence into a TT (convert ranks to spatial modes)
    cores = [np.eye(cores[0].shape[0])[np.newaxis, :, :]] + cores + [np.eye(cores[-1].shape[2])[:, :, np.newaxis]]
    t = tt.vector.from_list(cores)
    # sumsq2 = np.sum(t.full()**2)/(64**(N-2))
    # print(sumsq2)
    # print(t)
    # assert 0

    # Put the rank modes together at the very right
    if verbose:
        print('\tPutting rank modes together...')
    t = tr.core.shift_mode(t, 0, t.d-2, eps=eps)
    cores = tt.vector.to_list(t)

    # Merge the two rank dimensions into one
    cores[-2] = np.einsum('iaj,jbk->ikab', cores[-2], cores[-1])
    del cores[-1]
    cores[-1] = np.reshape(cores[-1], [-1, cores[-1].shape[2]*cores[-1].shape[3], 1])
    t = tt.vector.from_list(cores)

    # Subtract the cross-mean (mean of the collection of vectors)
    # print('t1:', t)
    # cores = [np.repeat(np.mean(c, axis=1, keepdims=True), c.shape[1], axis=1) for c in cores[:-1]] + [cores[-1]]
    # crossmean = tt.vector.from_list(cores)
    # t = tt.vector.round(t-crossmean, eps=eps)
    # print('t2:', t)
    print('Norma ara:', tt.vector.norm(t))
    norm = tt.vector.norm(t)
    if norm > 0:
        t *= (sumsq2 / norm)
    print('Norma ara:', tt.vector.norm(t))
    # print(t)

    # Compute the truncated PCA of the space spanned by the rank modes
    if verbose:
        print('\tPCA compression...')
    cores = tt.vector.to_list(t)
    tr.core.orthogonalize(cores, t.d-1)

    print('********')
    t = tt.vector.from_list(cores)
    print(t)
    for i in range(32):
        print(i, tt.vector.norm(t[0, i, :] - t[1, 24, :]))
    # assert 0
    # sumsq2 = np.sum(cores[-1]**2)
    # sumsq2 = np.sum(tr.core.variances(tt.vector.from_list(cores), modes=[0, 1]).full())
    # print(tt.vector.from_list(cores).full().shape)
    # sumsq2 = np.sum(np.sum(tt.vector.from_list(cores).full()**2, axis=-1)) / np.prod(tt.vector.from_list(cores).n)
    # print(tr.core.means(tt.vector.from_list(cores), modes=[0, 1]).full())
    # print(sumsq2)
    # assert 0
    cores[-1] = cores[-1][..., 0]
    svd = np.linalg.svd(cores[-1], full_matrices=False)
    cores[-1] = svd[0][:, :degree]
    if not orthogonal:
        cores[-1] = cores[-1].dot(np.diag(svd[1][:degree]))
    cores[-1] = cores[-1][:, :, np.newaxis]
    if cores[-1].shape[1] < degree:  # Degenerate case: add a slice of zeros
        cores[-1] = np.concatenate([cores[-1], np.zeros([cores[-1].shape[0], degree-cores[-1].shape[1],
                                                         cores[-1].shape[2]])], axis=1)
    t = tt.vector.from_list(cores)

    # print('********')
    # t = tt.vector.from_list(cores)
    # print(t)
    # for i in range(32):
    #     print(i, tt.vector.norm(t[0, i, :] - t[1, 24, :]))
    # assert 0
    # print('Norma ara:', tt.vector.norm(t))
    # t *= (1./tt.vector.norm(t))
    # assert 0

    # print('Now:', np.sqrt(np.var(t.full())))
    if anchor_mean:
        withinmean = tt.vector.from_list(tt.vector.to_list(withinmean) + [np.ones([1, 1, 1])])
        # print('**********')
        # print(withinmean, t)
        # print(tt.vector.norm(withinmean))
        # print(tt.vector.norm(t))
        # assert 0
        t = tr.core.concatenate([withinmean, t], axis=withinmean.d-1, eps=eps)

    if verbose:
        print()
    print(t.full())
    # assert 0
    # print(tt.vector.norm(t))
    print()
    # print(np.sum(t.full()**2))
    # print(t)
    tfull = t.full()
    # print(print(np.mean(np.mean(t.full(), axis=0), axis=0)))
    print('Obtained:')
    print(np.linalg.norm(tfull[..., 0]-np.mean(tfull[..., 0])))
    print(np.linalg.norm(withinmean2.full()))
    obtained = np.linalg.norm(tfull[..., 1:]-np.mean(tfull[..., 1:]))
    print(obtained)
    print()
    preserved = (sumsq1**2 + obtained**2) / (sumsq1**2 + sumsq2**2)
    print('This embedding preserves {:g}% of the original norm'.format(preserved*100))
    # print(np.sqrt(np.var(t[..., 0])))
    # print(np.sqrt(np.var(t[..., 1:])))
    # print(t[0, :, 0].full().flatten()[:10])
    # print(t[0, :, 1].full().flatten()[:10])
    # print(t[0, :, 2].full().flatten()[:10])
    # assert 0
    tlist = []
    for n in range(t.n[0]):
        tlist.append(t[[n] + [slice(None)]*(t.d-1)])
    if single:
        tlist = tlist[0]
    print()
    print('***************** tlist:', tlist)
    # assert 0
    return tlist, preserved

def pca_ensemble2(t, modes, degree=3, orthogonal=False, anchor_mean=True, eps=1e-3, verbose=False):
    """
    Given a tensor and a list of modes, consider the space of subtensors obtained
    by fixing these modes. Find a PCA basis for that space and return the coefficients
    for each subtensor

    :param t: a TT
    :param modes: the ensemble modes
    :param degree: how large a PCA basis to take (default is 3)
    :param orthogonal:
    :param anchor_mean:
    :param eps: relative tolerance while permuting modes. Default is 1e-3
    :param verbose:
    :return: (ensemble, preserved):
        - ensemble: a TT of size M1 x ... x Mk x `degree`, where M1, etc. are the sizes of `t` along `modes`
        - preserved: the fraction (between 0 and 1; higher is better) of norm preserved by the embedding
        - residual: the tensor containing the trailing R-degree principal components

    """

    # modes = [3, 6]
    # N = t.d
    # part1 = np.arange(modes[0])
    # part2 = np.setdiff1d(np.arange(modes[0], N), modes)
    # # print('**************', modes, ' **** ', list(part1)+list(modes)+list(part2))
    # order = list(part1)+list(modes)+list(part2)
    # idx = np.empty(len(order))
    # idx[order] = np.arange(len(order))
    # d = t.d
    # print(order)
    # assert 0
    # tr.core.transpose(t, order)
    # assert 0

    N = t.d
    shape = t.n
    if not hasattr(modes, '__len__'):
        modes = [modes]
    modesarray = np.array(modes)
    if modesarray.ndim == 1:
        modesarray = np.array([modesarray])
    if not all([len(modesarray[m]) == len(modesarray[0]) for m in range(len(modesarray))]):
        raise ValueError("All mode lists must have the same length")
    for modes in modesarray:
        if len(modes) != len(np.unique(modes)):
            raise ValueError('Modes may not be repeated')
        assert 0 <= np.min(modes)
        assert np.max(modes) < N

        if verbose:
            print('Computing PCA ensemble...')

        # Group target modes next to each other
        if verbose:
            print('\tTransposing tensor... 0', end='')
        # idx = list(range(N))
        # def list_shift(idx, mode, shift):
        #     assert mode + shift >= 0
        #     assert mode + shift < len(idx)
        #     idx.insert(mode+shift, int(idx[mode]))
        #     del idx[mode]
        # pivot = modes[0]  # We will organize modes from left-to-right after this one
        # for i in range(1, len(modes)):
        #     source = idx.index(modes[i])
        #     target = idx.index(pivot) + i
        #     if source < target:
        #         target -= 1
        #     t = tr.core.shift_mode(t, source, target-source, eps=eps)
        #     list_shift(idx, source, target-source)
        #     if verbose:
        #         print(' {}'.format(i), end='')
        # final_pos = idx.index(pivot)

        # part1 = np.arange(modes[0])
        part1 = np.setdiff1d(np.arange(modes[0]), modes)
        part2 = np.setdiff1d(np.arange(modes[0], N), modes)
        print('**************', modes, ' **** ', list(part1)+list(modes)+list(part2))
        order = list(part1)+list(modes)+list(part2)
        # t = tt.core.tools.permute(t, list(part1)+list(modes)+list(part2), eps=eps)
        t = tr.core.transpose(t, order, eps=eps)
        print(t)
        # assert 0
        # final_pos = modes[0]
        final_pos = order.index(modes[0])
        # final_pos = min(modes)

        if anchor_mean:
            # We now project each vector on the subspace spanned by the constant vector
            # This is achieved by subtracting each vector's within-mean
            cores = copy.deepcopy(tt.vector.to_list(t))
            for n in range(N):
                if n < final_pos or n >= final_pos+len(modes):
                    cores[n] = np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1)
            withinmeanrep = tt.vector.from_list(cores)
            withinmean = tr.core.means(t, modes=list(range(final_pos))+list(range(final_pos+len(modes), N)))
            t = tt.vector.round(t-withinmeanrep, eps=eps)
            degree -= 1

        # withinmeanrep2 = withinmeanrep - tr.core.constant_tt(shape=withinmean.n, fill=tr.core.mean(withinmeanrep))
        withinmean2 = withinmean - tr.core.constant_tt(shape=withinmean.n, fill=tr.core.mean(withinmean))
        # sumsq = np.sum(withinmean2)
        # print(withinmean2)
        # sumsq1 = tr.core.variances(withinmean2, modes=np.arange(len(modes))).full() #* np.prod(t.n)
        # sumsq1 = np.asscalar(sumsq1)
        sumsq1 = np.asscalar(np.sum(withinmean2.full()**2))# * np.prod(t.n[np.delete(np.arange(N), modes)])
        # targetsq1 = np.var(withinmean.full())*np.prod(t.n[modes])

        sumsq1 = np.sqrt(sumsq1)
        # targetsq1 = np.sqrt(targetsq1)
        print('sumsq1 = {}'.format(sumsq1))
        # print('sumsq1 = {}, targetsq1 = {}, sumsq1/targetsq1 = {}'.format(sumsq1, targetsq1, sumsq1/targetsq1))

        # print(np.sqrt(sumsq1/targetsq1))
        # print(tt.vector.norm(withinmean2)**2)
        # sumsq = np.sum(withinmean2.full()**2)
        # print(sumsq)
        # print(modes)
        # assert 0

        # Subtract the cross-mean (mean of the collection of vectors)
        cores = copy.deepcopy(tt.vector.to_list(t))
        print(t, final_pos, final_pos+len(modes))
        for n in range(final_pos, final_pos+len(modes)):
            cores[n] = np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1)
        crossmean = tt.vector.from_list(cores)
        t = tt.vector.round(t-crossmean, eps=eps)

        sumsq2 = tt.vector.norm(t) / np.sqrt(np.prod(shape[np.delete(np.arange(N), modes)]))
        # sumsq2 = np.sum(tr.core.variances(t, modes=list(range(final_pos))+list(range(final_pos+len(modes), N))).full())\
        #          * np.prod(t.n[modes])
        # sumsq2 = np.sqrt(sumsq2)
        print('sumsq2 = {}'.format(sumsq2))
        # t = t*(sumsq2 / sumsq1 * targetsq1)
        # print('Now 0:', tt.vector.norm(t))
        # t *= (targetsq1 / sumsq1)
        # print('Now 1:', tt.vector.norm(t))
        # assert 0
        # print(t)

        # Orthogonalize to the left and right of these modes
        cores = tt.vector.to_list(t)
        if verbose:
            print('\n\tOrthogonalization...')
        for i in range(final_pos):
            tr.core.left_orthogonalize(cores, i)
        for i in range(N-1, final_pos+len(modes)-1, -1):
            tr.core.right_orthogonalize(cores, i)

        # Cut out the relevant sequence of modes
        cores = cores[final_pos:final_pos+len(modes)]

        # Convert the sequence into a TT (convert ranks to spatial modes)
        cores = [np.eye(cores[0].shape[0])[np.newaxis, :, :]] + cores + [np.eye(cores[-1].shape[2])[:, :, np.newaxis]]
        t = tt.vector.from_list(cores)
        # sumsq2 = np.sum(t.full()**2)/(64**(N-2))
        # print(sumsq2)
        # assert 0

        # Put the rank modes together at the very right
        if verbose:
            print('\tPutting rank modes together...')
        t = tr.core.shift_mode(t, 0, t.d-2, eps=eps)
        cores = tt.vector.to_list(t)

        # Merge the two rank dimensions into one
        cores[-2] = np.einsum('iaj,jbk->ikab', cores[-2], cores[-1])
        del cores[-1]
        cores[-1] = np.reshape(cores[-1], [-1, cores[-1].shape[2]*cores[-1].shape[3], 1])
        t = tt.vector.from_list(cores)

        # Subtract the cross-mean (mean of the collection of vectors)
        # print('t1:', t)
        # cores = [np.repeat(np.mean(c, axis=1, keepdims=True), c.shape[1], axis=1) for c in cores[:-1]] + [cores[-1]]
        # crossmean = tt.vector.from_list(cores)
        # t = tt.vector.round(t-crossmean, eps=eps)
        # print('t2:', t)
        print('Norma ara:', tt.vector.norm(t))
        norm = tt.vector.norm(t)
        if norm > 0:
            t *= (sumsq2 / norm)
        print('Norma ara:', tt.vector.norm(t))
        print(t)

        # Compute the truncated PCA of the space spanned by the rank modes
        if verbose:
            print('\tPCA compression...')
        cores = tt.vector.to_list(t)
        tr.core.orthogonalize(cores, t.d-1)
        # sumsq2 = np.sum(cores[-1]**2)
        # sumsq2 = np.sum(tr.core.variances(tt.vector.from_list(cores), modes=[0, 1]).full())
        # print(tt.vector.from_list(cores).full().shape)
        # sumsq2 = np.sum(np.sum(tt.vector.from_list(cores).full()**2, axis=-1)) / np.prod(tt.vector.from_list(cores).n)
        # print(tr.core.means(tt.vector.from_list(cores), modes=[0, 1]).full())
        # print(sumsq2)
        # assert 0
        cores[-1] = cores[-1][..., 0]
        svd = np.linalg.svd(cores[-1], full_matrices=False)

        # Compute error
        cores2 = copy.deepcopy(cores)
        cores2[-1] = svd[0][:, degree:].dot(np.diag(svd[1])[degree:])[:, :, np.newaxis]
        error = tt.vector.from_list(cores2)

        cores[-1] = svd[0][:, :degree]
        if not orthogonal:
            cores[-1] = cores[-1].dot(np.diag(svd[1][:degree]))
        cores[-1] = cores[-1][:, :, np.newaxis]
        if cores[-1].shape[1] < degree:  # Degenerate case: add a slice of zeros
            cores[-1] = np.concatenate([cores[-1], np.zeros([cores[-1].shape[0], degree-cores[-1].shape[1],
                                                             cores[-1].shape[2]])], axis=1)
        t = tt.vector.from_list(cores)

        # print('Norma ara:', tt.vector.norm(t))
        # t *= (1./tt.vector.norm(t))
        # assert 0

        # print('Now:', np.sqrt(np.var(t.full())))
        if anchor_mean:
            withinmean = tt.vector.from_list(tt.vector.to_list(withinmean) + [np.ones([1, 1, 1])])
            t = tr.core.concatenate([withinmean, t], axis=len(modes), eps=eps)

        if verbose:
            print()

    # print(tt.vector.norm(t))
    print()
    # print(np.sum(t.full()**2))
    # print(t)
    tfull = t.full()
    # print(print(np.mean(np.mean(t.full(), axis=0), axis=0)))
    print('Obtained:')
    print(np.linalg.norm(tfull[..., 0]-np.mean(tfull[..., 0])))
    print(np.linalg.norm(withinmean2.full()))
    obtained = np.linalg.norm(tfull[..., 1:]-np.mean(tfull[..., 1:]))
    print(obtained)
    print()
    preserved = (sumsq1**2 + obtained**2) / (sumsq1**2 + sumsq2**2)
    print('This embedding preserves {:g}% of the original norm'.format(preserved*100))
    # print(np.sqrt(np.var(t[..., 0])))
    # print(np.sqrt(np.var(t[..., 1:])))
    # assert 0
    return t, preserved, error

# def pca_ensemble(t, modes, degree=3, orthogonal=False, anchor_mean=True, eps=1e-3, verbose=False):
#     """
#     Given a tensor and a list of modes, consider the space of subtensors obtained
#     by fixing these modes. Find a PCA basis for that space and return the coefficients
#     for each subtensor
#
#     :param t: a TT
#     :param modes: the ensemble modes
#     :param degree: how large a PCA basis to take (default is 3)
#     :param orthogonal:
#     :param anchor_mean:
#     :param eps: relative tolerance while permuting modes. Default is 1e-3
#     :param verbose:
#     :return: M1 x ... x Mk x degree, where M1, etc. are the sizes of `t` along `modes`
#
#     """
#
#     N = t.d
#     if not hasattr(modes, '__len__'):
#         modes = [modes]
#     if len(modes) != len(np.unique(modes)):
#         raise ValueError('Modes may not be repeated')
#     assert 0 <= np.min(modes)
#     assert np.max(modes) < N
#
#     if verbose:
#         print('Computing PCA ensemble...')
#
#     # Group target modes next to each other
#     if verbose:
#         print('\tTransposing tensor... 0', end='')
#     idx = list(range(N))
#     def list_shift(idx, mode, shift):
#         assert mode + shift >= 0
#         assert mode + shift < len(idx)
#         tmp = idx[mode]
#         del idx[mode]
#         if shift < 0:
#             idx.insert(mode+shift, tmp)
#         else:
#             idx.insert(mode+shift, tmp)
#     pivot = modes[0]  # We will organize modes from left-to-right after this one
#     for i in range(1, len(modes)):
#         source = idx.index(modes[i])
#         target = idx.index(pivot) + i
#         if source < target:
#             target -= 1
#         t = tr.core.shift_mode(t, source, target-source, eps=eps)
#         list_shift(idx, source, target-source)
#         if verbose:
#             print(' {}'.format(i), end='')
#     final_pos = idx.index(pivot)
#
#     if anchor_mean:
#         # We now project each vector on the subspace spanned by the constant vector
#         # This is achieved by subtracting each vector's within-mean
#         cores = copy.deepcopy(tt.vector.to_list(t))
#         for n in range(N):
#             if n < final_pos or n >= final_pos+len(modes):
#                 cores[n] = np.repeat(np.mean(cores[n], axis=1, keepdims=True), cores[n].shape[1], axis=1)
#         tmeanrep = tt.vector.from_list(cores)
#         withinmean = tr.core.means(t, modes=list(range(final_pos))+list(range(final_pos+len(modes), N)))
#         t = tt.vector.round(t-tmeanrep, eps=eps)
#         degree -= 1
#
#     # Orthogonalize to the left and right of these modes
#     cores = tt.vector.to_list(t)
#     if verbose:
#         print('\n\tOrthogonalization...')
#     for i in range(final_pos):
#         tr.core.left_orthogonalize(cores, i)
#     for i in range(N-1, final_pos+len(modes)-1, -1):
#         tr.core.right_orthogonalize(cores, i)
#
#     # Cut out the relevant sequence of modes
#     cores = cores[final_pos:final_pos+len(modes)]
#
#     # Convert the sequence into a TT (convert ranks to spatial modes)
#     cores = [np.eye(cores[0].shape[0])[np.newaxis, :, :]] + cores + [np.eye(cores[-1].shape[2])[:, :, np.newaxis]]
#     t = tt.vector.from_list(cores)
#
#     # Put the rank modes together at the very right
#     if verbose:
#         print('\tPutting rank modes together...')
#     t = tr.core.shift_mode(t, 0, t.d-2, eps=eps)
#     cores = tt.vector.to_list(t)
#
#     # Merge the two rank dimensions into one
#     cores[-2] = np.einsum('iaj,jbk->ikab', cores[-2], cores[-1])
#     del cores[-1]
#     cores[-1] = np.reshape(cores[-1], [-1, cores[-1].shape[2]*cores[-1].shape[3], 1])
#
#     # Subtract the cross-mean (mean of the collection of vectors)
#     t = tt.vector.from_list(cores)
#     cores = [np.repeat(np.mean(c, axis=1, keepdims=True), c.shape[1], axis=1) for c in cores[:-1]] + [cores[-1]]
#     crossmean = tt.vector.from_list(cores)
#     t = tt.vector.round(t-crossmean, eps=eps)
#
#     # Compute the truncated PCA of the space spanned by the rank modes
#     if verbose:
#         print('\tPCA compression...')
#     cores = tt.vector.to_list(t)
#     tr.core.orthogonalize(cores, t.d-1)
#     cores[-1] = cores[-1][..., 0]
#     svd = np.linalg.svd(cores[-1], full_matrices=False)
#     cores[-1] = svd[0][:, :degree]
#     if not orthogonal:
#         cores[-1] = cores[-1].dot(np.diag(svd[1][:degree]))
#     cores[-1] = cores[-1][:, :, np.newaxis]
#     if cores[-1].shape[1] < degree:  # Degenerate case: add a slice of zeros
#         cores[-1] = np.concatenate([cores[-1], np.zeros([cores[-1].shape[0], degree-cores[-1].shape[1],
#                                                          cores[-1].shape[2]])], axis=1)
#     t = tt.vector.from_list(cores)
#
#     if anchor_mean:
#         withinmean = tt.vector.from_list(tt.vector.to_list(withinmean) + [np.ones([1, 1, 1])])
#         t = tr.core.concatenate([withinmean, t], axis=len(modes), eps=eps)
#
#     if verbose:
#         print()
#
#     return t

# def pca_ensemble(t, modes, degree=3, eps='same', verbose=False):
#     """
#     Given a tensor and a list of modes, consider the space of subtensors obtained
#     by fixing these modes. Find a PCA basis for that space and return the coefficients
#     for each subtensor
#
#     :param t: a TT
#     :param modes: the ensemble modes
#     :param degree: how large a PCA basis to take (default is 3)
#     :param eps: relative tolerance while permuting modes. Default behavior is to
#     preserve the original number of ranks
#     :param verbose:
#     :return: M1 x ... x Mk x degree, where M1, etc. are the sizes of `t` along `modes`
#
#     """
#
#     N = t.d
#
#     if verbose:
#         print('Computing PCA ensemble...')
#
#     # Group target modes next to each other
#     if verbose:
#         print('\tTransposing tensor... 0', end='')
#     modes = np.sort(modes)
#     for i in range(1, len(modes)):
#         t = tr.core.shift_mode(t, modes[i], modes[0] + i - modes[i], eps=eps)
#         if verbose:
#             print(' {}'.format(i), end='')
#     cores = tt.vector.to_list(t)
#
#     # Orthogonalize to the left and right of these modes
#     if verbose:
#         print('\n\tOrthogonalization...')
#     for i in range(modes[0]):
#         tr.core.left_orthogonalize(cores, i)
#     for i in range(N-1, modes[0]+len(modes)-1, -1):
#         tr.core.right_orthogonalize(cores, i)
#
#     # Decompress (end ranks are not 1 in general)
#     if verbose:
#         print('\tMerging spatial dimensions...')
#     subtensors = tt.vector.from_list(cores[modes[0]:modes[0]+len(modes)])
#     stack = tr.core.full(subtensors, keep_end_ranks=True)
#
#     # Merge all spatial dimensions in one
#     stack = np.transpose(stack, [0, stack.ndim-1] + list(range(1, stack.ndim-1)))
#     stack = np.reshape(stack, [stack.shape[0]*stack.shape[1], -1])
#
#     # Compress (PCA) along the spatial dimension
#     if verbose:
#         print('\tPCA compression...')
#     svd = np.linalg.svd(stack, full_matrices=False)
#     ensemble = np.diag(svd[1][:degree]).dot(svd[2][:degree, :]).T
#
#     # Reorganize the ensemble as expected
#     ensemble = np.reshape(ensemble, list(subtensors.n) + [degree])
#
#     if verbose:
#         print('\tDone')
#
#     return ensemble