from math import pi
from functools import wraps

from einops import rearrange, repeat
from einops.layers.torch import Reduce
import torch
from torch import nn, einsum
import torch.nn.functional as F

# helpers

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def cache_fn(f):
    cache = dict()
    @wraps(f)
    def cached_fn(*args, _cache = True, key = None, **kwargs):
        if not _cache:
            return f(*args, **kwargs)
        nonlocal cache
        if key in cache:
            return cache[key]
        result = f(*args, **kwargs)
        cache[key] = result
        return result
    return cached_fn

def fourier_encode(x, max_freq, num_bands = 4):
    x = x.unsqueeze(-1)
    device, dtype, orig_x = x.device, x.dtype, x

    scales = torch.linspace(1., max_freq / 2, num_bands, device = device, dtype = dtype)
    scales = scales[(*((None,) * (len(x.shape) - 1)), Ellipsis)]

    x = x * scales * pi
    x = torch.cat([x.sin(), x.cos()], dim = -1)
    x = torch.cat((x, orig_x), dim = -1)
    return x

# helper classes

class PreNorm(nn.Module):
    def __init__(self, dim, fn, context_dim = None):
        super().__init__()
        self.fn = fn
        self.norm = nn.LayerNorm(dim)
        self.norm_context = nn.LayerNorm(context_dim) if exists(context_dim) else None

    def forward(self, x, **kwargs):
        x = self.norm(x)

        if exists(self.norm_context):
            context = kwargs['context']
            normed_context = self.norm_context(context)
            kwargs.update(context = normed_context)

        return self.fn(x, **kwargs)

class GEGLU(nn.Module):
    def forward(self, x):
        x, gates = x.chunk(2, dim = -1)
        return x * F.gelu(gates)

class FeedForward(nn.Module):
    def __init__(self, dim, mult = 4, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, dim * mult * 2),
            GEGLU(),
            nn.Linear(dim * mult, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    def __init__(self, query_dim, context_dim = None, heads = 8, dim_head = 64, dropout = 0., scale=None):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)
        if scale:
            self.scale = scale #**-1
        else:
            self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias = False)
        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias = False)

        self.dropout = nn.Dropout(dropout)
        self.to_out = nn.Linear(inner_dim, query_dim)

    def forward(self, x, context=None, mask=None):
        h = self.heads

        q = self.to_q(x)
        context = default(context, x)
        k, v = self.to_kv(context).chunk(2, dim = -1)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))

        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale

        if exists(mask):
            mask = rearrange(mask, 'b ... -> b (...)')
            max_neg_value = -torch.finfo(sim.dtype).max
            mask = repeat(mask, 'b j -> (b h) () j', h = h)
            sim.masked_fill_(~mask, max_neg_value)

        # attention, what we cannot get enough of
        A = sim.softmax(dim = -1)
        attn = self.dropout(A)

        out = einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h = h)

        if context.shape != x.shape:
            return self.to_out(out), A
        else:
            return self.to_out(out)

class DualQueryCrossAttention(nn.Module):
    def __init__(self, query_dim, context_dim = None, heads = 8, dim_head = 64, dropout = 0., scale=None):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)
        if scale:
            self.scale = nn.Parameter(torch.tensor([scale])) #**-1
        else:
            self.scale = dim_head ** -0.5
        self.heads = heads

        self.to_q = nn.Linear(query_dim, inner_dim, bias = False)
        self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias = False)

        self.dropout = nn.Dropout(dropout)
        self.to_out = nn.Linear(inner_dim, query_dim)

        # Attention ranking
        self.to_score_q = nn.Linear(query_dim, inner_dim, bias = False)
        self.to_score_out = nn.Linear(inner_dim, query_dim)

    def forward(self, x, score_x, context=None, mask=None):
        h = self.heads

        q = self.to_q(x)
        score_q = self.to_score_q(score_x)
        k, v = self.to_kv(context).chunk(2, dim = -1)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q, k, v))

        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale

        score_sim = einsum('b i d, b j d -> b i j', score_q, k) * self.scale

        if exists(mask):
            mask = rearrange(mask, 'b ... -> b (...)')
            max_neg_value = -torch.finfo(sim.dtype).max
            mask = repeat(mask, 'b j -> (b h) () j', h = h)
            sim.masked_fill_(~mask, max_neg_value)

        # attention, what we cannot get enough of
        A = sim.softmax(dim = -1)
        attn = self.dropout(A)

        score_attn = score_sim.softmax(dim = -1)

        out = einsum('b i j, b j d -> b i d', attn, v)
        out = rearrange(out, '(b h) n d -> b n (h d)', h = h)

        score_out = einsum('b i j, b j d -> b i d', score_attn, v)
        score_out = rearrange(score_out, '(b h) n d -> b n (h d)', h = h)

        return self.to_out(out), A, self.to_score_out(score_out), score_attn

# Based on the merging approach from Truong et al. "How Transferable are Self-supervised Features in Medical Image Classification Tasks?"
class Merger(nn.Module):
    def __init__(self, proj_dim):
        super(Merger, self).__init__()
        self.swav_head = nn.Linear(2048, proj_dim)
        self.dino_head = nn.Linear(2048, proj_dim)
        self.dino_v2_head = nn.Linear(1024, proj_dim)

    def forward(self, data):
        dino_out = self.dino_head(data['dino'])
        dino_v2_out = self.dino_v2_head(data['dino_v2'])
        swav_out = self.swav_head(data['swav'])
        joint = torch.cat([dino_out, dino_v2_out, swav_out], dim=-1)
        return joint

class Perceiver(nn.Module):
    def __init__(
        self,
        *,
        num_freq_bands,
        depth,
        max_freq,
        input_channels = 3,
        input_axis = 2,
        num_latents = 1024,
        latent_dim = 512,
        cross_heads = 1,
        latent_heads = 8,
        cross_dim_head = 64,
        latent_dim_head = 64,
        n_classes = 1000,
        attn_dropout = 0.,
        ff_dropout = 0.,
        weight_tie_layers = False,
        fourier_encode_data = True,
        self_per_cross_attn = 1,
        latent_bounds = 2,
        scale = None,
    ):
        """The shape of the final attention mechanism will be:
        depth * (cross attention -> self_per_cross_attn * self attention)

        Args:
          num_freq_bands: Number of freq bands, with original value (2 * K + 1)
          depth: Depth of net.
          max_freq: Maximum frequency, hyperparameter depending on how
              fine the data is.
          freq_base: Base for the frequency
          input_channels: Number of channels for each token of the input.
          input_axis: Number of axes for input data (2 for images, 3 for video)
          num_latents: Number of latents, or induced set points, or centroids.
              Different papers giving it different names.
          latent_dim: Latent dimension.
          cross_heads: Number of heads for cross attention. Paper said 1.
          latent_heads: Number of heads for latent self attention, 8.
          cross_dim_head: Number of dimensions per cross attention head.
          latent_dim_head: Number of dimensions per latent self attention head.
          num_classes: Output number of classes.
          attn_dropout: Attention dropout
          ff_dropout: Feedforward dropout
          weight_tie_layers: Whether to weight tie layers (optional).
          fourier_encode_data: Whether to auto-fourier encode the data, using
              the input_axis given. defaults to True, but can be turned off
              if you are fourier encoding the data yourself.
          self_per_cross_attn: Number of self attention blocks per cross attn.
          final_classifier_head: mean pool and project embeddings to number of classes (num_classes) at the end
        """
        super().__init__()
        self.input_axis = input_axis
        self.max_freq = max_freq
        self.num_freq_bands = num_freq_bands
        self.n_classes = n_classes

        self.fourier_encode_data = fourier_encode_data
        fourier_channels = (input_axis * ((num_freq_bands * 2) + 1)) if fourier_encode_data else 0
        self.merge_embeddings = Merger(input_channels//3)
        input_dim = fourier_channels + input_channels

        self.latents = nn.Parameter(
            torch.nn.init.trunc_normal_(
                torch.zeros((num_latents, latent_dim)),
                mean=0,
                std=0.02,
                a=-latent_bounds,
                b=latent_bounds))

        self.score_latents = nn.Parameter(
            torch.nn.init.trunc_normal_(
                torch.zeros((1, latent_dim)),
                mean=0,
                std=0.02,
                a=-latent_bounds,
                b=latent_bounds))

        # Cross-Attention Layer
        get_cross_attn = lambda: PreNorm(latent_dim, DualQueryCrossAttention(latent_dim, input_dim, heads = cross_heads, dim_head = cross_dim_head, dropout = attn_dropout, scale=scale), context_dim = input_dim) #new
        get_cross_ff = lambda: PreNorm(latent_dim, FeedForward(latent_dim, dropout = ff_dropout))
        get_mil_ff = lambda: PreNorm(latent_dim, FeedForward(latent_dim, dropout = ff_dropout))

        get_latent_attn = lambda: PreNorm(latent_dim, Attention(latent_dim, heads = latent_heads, dim_head = latent_dim_head, dropout = attn_dropout))
        get_latent_ff = lambda: PreNorm(latent_dim, FeedForward(latent_dim, dropout = ff_dropout))

        get_cross_attn, get_cross_ff, get_latent_attn, get_latent_ff, get_mil_ff = map(cache_fn, (get_cross_attn, get_cross_ff, get_latent_attn, get_latent_ff, get_mil_ff))

        self.layers = nn.ModuleList([])
        for i in range(depth):
            should_cache = i > 0 and weight_tie_layers
            cache_args = {'_cache': should_cache}

            self_attns = nn.ModuleList([])

            for block_ind in range(self_per_cross_attn):
                self_attns.append(nn.ModuleList([
                    get_latent_attn(**cache_args, key = block_ind),
                    get_latent_ff(**cache_args, key = block_ind)
                ]))
            self.layers.append(nn.ModuleList([
                get_cross_attn(**cache_args),
                get_cross_ff(**cache_args),
                get_mil_ff(**cache_args),
                self_attns
            ]))

        self.to_logits = nn.Sequential(
            Reduce('b n d -> b d', 'mean'),
            nn.LayerNorm(latent_dim),
            nn.Linear(latent_dim, n_classes)
        )

        self.to_score_logits = nn.Sequential(
            Reduce('b n d -> b d', 'mean'),
            nn.LayerNorm(latent_dim),
            nn.Linear(latent_dim, n_classes)
        )

    def forward(
        self,
        data,
        mask = None,
        return_embeddings = False,
    ):
        data = self.merge_embeddings(data)

        if len(data.shape)==2: # flops
            data= data.unsqueeze(0)  # flops
        b, *axis, _, device, dtype = *data.shape, data.device, data.dtype
        assert len(axis) == self.input_axis, 'input data must have the right number of axis'

        if self.fourier_encode_data:
            # calculate fourier encoded positions in the range of [-1, 1], for all axis

            axis_pos = list(map(lambda size: torch.linspace(-1., 1., steps=size, device=device, dtype=dtype), axis))
            pos = torch.stack(torch.meshgrid(*axis_pos, indexing = 'ij'), dim = -1)
            enc_pos = fourier_encode(pos, self.max_freq, self.num_freq_bands)
            enc_pos = rearrange(enc_pos, '... n d -> ... (n d)')
            enc_pos = repeat(enc_pos, '... -> b ...', b = b)

            data = torch.cat((data, enc_pos), dim = -1)

        # concat to channels of data and flatten axis
        data = rearrange(data, 'b ... d -> b (...) d')

        x = repeat(self.latents, 'n d -> b n d', b = b)

        score_x = repeat(self.score_latents, 'n d -> b n d', b = b)

        # layers
        for cross_attn, cross_ff, mil_ff, self_attns in self.layers:
            x_attn, A_raw, score_x_attn, score_A  = cross_attn(x=x, score_x=score_x, context=data, mask=mask)
            x = x_attn + x
            x = cross_ff(x) + x
            score_x = score_x_attn + score_x
            score_x = mil_ff(score_x) + score_x

            for self_attn, self_ff in self_attns:
                x = self_attn(x) + x
                x = self_ff(x) + x

        # to logits
        logits = self.to_logits(x)
        results_dict={'student_logits':self.to_score_logits(score_x), 'features_teacher':x, 'features_student':score_x}
        Y_hat = torch.argmax(logits, dim=1)
        Y_prob = F.softmax(logits, dim = 1)
        return logits, Y_prob, Y_hat, score_A, results_dict

if __name__ == "__main__":
    data = torch.randn((1, 6000, 512)).cuda()
    model = Perceiver(input_channels=512,
                      input_axis=1,
                      num_freq_bands=6,
                      max_freq=10.,
                      depth=6,
                      num_latents=1024,
                      latent_dim=512,
                      cross_heads=1,
                      latent_heads=8,
                      cross_dim_head=64,
                      latent_dim_head=64,
                      n_classes=2,
                      attn_dropout=0.2,
                      ff_dropout=0.2,
                      weight_tie_layers=True,
                      fourier_encode_data=False,
                      self_per_cross_attn=1,
                      final_classifier_head=True,
                      latent_bounds=2,
                      scale=0.125).cuda()
    print(model.eval())
    results = model(data=data)
    print(results)


from PIL import Image
import numpy as np
import pickle
import os

def write_csv(df, filename):
    """
    Args:
        df: pandas dataframe to write
        filename (str): path to store the dataframe
    """
    if not os.path.exists(os.path.dirname(filename)):
        print('\n\nDirectory path %s does not exist. Creating it...' % filename)
        os.makedirs(os.path.dirname(filename))

    df.to_csv(filename, index=False)

def save_obj(obj, name):
    """
    Store the object in a pkl file
    :param obj: python object to be stored
    :param name: file name (Not insert .pkl)
    :return:
    """
    with open(name + '.pkl', 'wb') as f:
        pickle.dump(obj, f)

def save_np(npy, filename):
    """
    Store numpy to memory.
    Args:
        npy: numpy to save
        filename (str): filename
    """
    if not os.path.exists(os.path.dirname(filename)):
        print('\n\nDirectory path %s does not exist. Creating it...' % filename)
        os.makedirs(os.path.dirname(filename))

    np.save(filename, npy)

def save_image(image, filename):
    """
    Store an image to hard disk
    Args:
        image (pytorch tensor): image to save
        filename (str): filename
    """
    if not os.path.exists(os.path.dirname(filename)):
        os.makedirs(os.path.dirname(filename))

    im = Image.fromarray(image)
    im.save(filename)


from components import Direction
from events import GameAction
import input
import logging
import utils

class BehaviourManager:
    """The Bahavoiur manager stores component behevoiurs."""
    logger = logging.getLogger('behaviour.BehaviourManager')

    def __init__(self, event_manager, entitiy_manager, config=None):
        self.event_manager = event_manager
        self.entity_manager = entitiy_manager
        self.behaviours = {}
        if config is not None:
            self.configure(config)

    def configure(self, config):
        """Configure behaviourse base on a configuration.

        Ars:
            config: {'component_behaviours': [{component_type, name: ''}]}

        """
        if 'component_behaviours' in config:
            for comp_behaviour_config in config['component_behaviours']:
                self.add_comp_behaviour_from_config(comp_behaviour_config)
        if 'min_turn_time' in config:
            self.min_turn_time = config['min_turn_time']

    def add_comp_behaviour_from_config(self, config):
        """Associates the behaviour specified be the name with the specified
        compoenent_type.

        Args:
            component_type (str): The component type associated with the
                behaviour
            behaviour_impl (str): The implementation of the behaviour.

        """
        impl = config['impl']
        behaviour_class = utils.class_for_name(impl[0], impl[1])
        behaviour = behaviour_class(self.event_manager, self.entity_manager)
        self.add_component_behaviour(config['component_type'], behaviour)

    def add_component_behaviour(self, component_type, behaviour):
        """Register the given behaviour with the specified component type."""
        self.logger.debug("Register %s for component type %s",
                          component_type, behaviour)
        self.behaviours.update({component_type: behaviour})

    def update(self, round, delta_time):
        """Updates all behaviours for all entitites with a associated
        components."""
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)
            for entity, component in components.items():
                behaviour.entity = entity
                behaviour.component = component
                behaviour.update(round, delta_time)

class TurnBehaviourManager(BehaviourManager):

    logger = logging.getLogger('behaviour.TurnBehaviourManager')

    def __init__(self, event_manager, entitiy_manager, config=None):
        BehaviourManager.__init__(self, event_manager,
                                  entitiy_manager, config)
        self.locked_entities = []

    def update(self, round, delta_time):
        """Updates all behaviours for all entitites with a associated
        component, but only if there are no locked entities."""

        # TODO: Make the method prettier ^^
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)

            for entity, component in components.items():
                s = self.locked_entities
                turn = self.entity_manager.get(entity, "Turn")
                if (turn.ticks == 0 and turn.delta >= turn.min_turn_time):
                    # It's the entity's Turn!
                    if (turn.locking and entity not in s):
                        # Lock it
                        self.locked_entities.append(entity)
                    if (entity in s or len(s) == 0):
                        behaviour.entity = entity
                        behaviour.component = component
                        ticks = behaviour.update(round, delta_time)
                        if (ticks is not None):
                            # The behaviour has made an action
                            if (entity in s):
                                self.locked_entities.remove(entity)
                            self._add_ticks(entity, ticks)
                else:
                    turn.delta += delta_time
        self._normalize()

    def _add_ticks(self, entity, ticks):
        """ Count up the Tick Counter of entity, and reset its delta time """
        turn_comp = self.entity_manager.get(entity, "Turn")
        turn_comp.ticks += ticks
        turn_comp.delta = 0

    def _normalize(self):
        """ Reduce the Tick counter of all turn components by the minimum
        Tick count, to prevent tick counters from increasing indefinitely.
        This assures an only relative tick count """
        min_ticks = None
        for component_type, behaviour in self.behaviours.items():
            components = self.entity_manager.get_all(component_type)
            ticks = [self.entity_manager.get(entity, "Turn").ticks
                     for entity in components.keys()]
            if (len(ticks) > 0):
                minimum_of_type = min(ticks)
                if (min_ticks is None or min_ticks > minimum_of_type):
                    min_ticks = minimum_of_type
        if min_ticks is not None:
            for component_type, behaviour in self.behaviours.items():
                components = self.entity_manager.get_all(component_type)
                for turn_comp in [self.entity_manager.get(entity, "Turn")
                                  for entity in components.keys()]:
                    turn_comp.ticks -= min_ticks

class EntityComponentBehaviour(object):
    """Implements logic for entities with specific components."""

    def __init__(self, event_manager, entitiy_manager):
        self.entity_manager = entitiy_manager
        self.event_manager = event_manager
        self.entity = None
        self.component = None

    def update(self, round, delta_time):
        pass

class InputBehaviour(EntityComponentBehaviour):
    """Implements User Input. Controls all entites with an InputComponent."""

    logger = logging.getLogger('behaviours.InputBehaviour')

    def __init__(self, event_manager, entitiy_manager):
        super(InputBehaviour, self).__init__(event_manager, entitiy_manager)
        # TODO make configurable
        self.isometric = True

    def update(self, round, delta_time):
        """Converts input to an appropriate InputAction."""
        # TODO Implement gamestatus aware behaviour manager to keep the turn
        # based logic as central as possible and possibly enable switching
        # between realtime and turn based. (A realtime behaviour would not
        # check for a state
        speed = self.move()

        if input.is_pressed(input.K_ENTER):
            # TODO: Implement Targeting Inputs (combined input or
            # sequential?)
            # TODISCUSS: Context actions, autotargeting
            speed = self.use(0, 0)
        return speed

    def set_input_direction(self):
        direction = self.component.direction
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
                input.K_KP7):
            direction.set(Direction.D_LEFT)
        else:
            direction.set(Direction.D_LEFT, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_KP3) or input.is_pressed(
            input.K_KP6) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_RIGHT)
        else:
            direction.set(Direction.D_RIGHT, 0)
        if input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
            input.K_KP2) or input.is_pressed(
                input.K_KP3):
            direction.set(Direction.D_DOWN)
        else:
            direction.set(Direction.D_DOWN, 0)
        if input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP7) or input.is_pressed(
            input.K_KP8) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_UP)
        else:
            direction.set(Direction.D_UP, 0)

    def set_iso_input_direction(self):
        direction = self.component.direction
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
            input.K_KP7) or input.is_pressed(
                input.K_KP8):
            direction.set(Direction.D_LEFT)
        else:
            direction.set(Direction.D_LEFT, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP2) or input.is_pressed(
            input.K_KP3) or input.is_pressed(
                input.K_KP6):
            direction.set(Direction.D_RIGHT)
        else:
            direction.set(Direction.D_RIGHT, 0)
        if input.is_pressed(
            input.K_LEFT) or input.is_pressed(
            input.K_DOWN) or input.is_pressed(
            input.K_KP4) or input.is_pressed(
            input.K_KP1) or input.is_pressed(
                input.K_KP2):
            direction.set(Direction.D_DOWN)
        else:
            direction.set(Direction.D_DOWN, 0)
        if input.is_pressed(
            input.K_RIGHT) or input.is_pressed(
            input.K_UP) or input.is_pressed(
            input.K_KP6) or input.is_pressed(
            input.K_KP8) or input.is_pressed(
                input.K_KP9):
            direction.set(Direction.D_UP)
        else:
            direction.set(Direction.D_UP, 0)

    def move(self):
        """Throws a MoveAction."""
        if self.isometric:
            self.set_iso_input_direction()
        else:
            self.set_input_direction()
        return 1 if self.component.direction.direction > 0 else None

    def use(self, dx, dy):
        """ Throws a UseEntity Event """
        event = self.event_manager.create(GameAction.UseEntityAction)
        event.user = self.entity
        self.event_manager.throw(event)
        return 2


import numpy as np
import matplotlib.pyplot as plt
import numpy.linalg as linalg

from Models.cross_validation import cross_validation
from Models.error_score import *
from Models.import_data import import_data

def main():
    rbf_scales()
    # rbf_centres()

def rbf_centres():
    wine_data, wine_features = import_data('winequality-red.csv')
    inputmtx = wine_data[:, 0:11]
    targets = wine_data[:, 11]
    train_inputs, train_targets, test_inputs, test_targets = cross_validation(inputmtx, targets, 0.25)
    scale = 2831 #Fix Scale
    scales = np.logspace(3, -2)
    centres, sd = classify_inputs(train_inputs, train_targets)
    # centres = [8.319637273,	0.527820513,	0.27097561,	2.538805503,	0.087466542,	15.87492183,	46.46779237,	0.996746679,	3.311113196,	0.658148843,	10.42298311,	5.636022514]
    centres = np.array(centres)
    feature_mapping = construct_rbf_feature_mapping(centres, scale)
    designmtx = feature_mapping(train_inputs)
    weights = ml_weights(designmtx, train_targets)
    rbf_approx = construct_feature_mapping_approx(feature_mapping, weights)
    prediction_values = rbf_approx(test_inputs)
    rmse, mae, medae, mape, variance = error_score(test_targets, prediction_values)

    print("Best RSME: ", rmse)

def classify_inputs(inputs, targets):
    mintarg = int(min(targets))
    maxtarg = int(max(targets))
    t = {}
    f = {}
    m = {}
    v = {}
    centres = np.zeros((inputs.shape[0], inputs.shape[1]))
    sd = np.zeros((inputs.shape[0], inputs.shape[1]))

    for n in range(mintarg, maxtarg + 1):
        t["group{0}".format(n)] = []
        f["group{0}".format(n)] = []
        m["group{0}".format(n)] = []
        v["group{0}".format(n)] = []

    for i in enumerate(targets):
        for n in range(mintarg, maxtarg + 1):
            if int(i[1]) == n:
                t["group{0}".format(n)].append(i)
                f["group{0}".format(n)].append(inputs[i[0], :])
                centres[i[0], :] = int(n)
                sd[i[0], :] = int(n)

    for group in f:
        array = np.array(f[group])
        for j in array.T:
            m[group].append(sum(j) / len(j))
            v[group].append(np.var(j))

    for i in range(0, centres.shape[0]):
        for j in range(0, centres.shape[1]):
            ind = int(centres[i, j])
            val = m["group{0}".format(ind)][j]
            var = v["group{0}".format(ind)][j]
            centres[i, j] = val
            sd[i, j] = var

    return centres, sd

def rbf_scales():
    wine_data, wine_features = import_data('winequality-red.csv')
    inputmtx = wine_data[:, 0:11]
    targets = wine_data[:, 11]
    inputs, targets, test_inputs, test_targets = cross_validation(inputmtx, targets, 0.25)

    # specify the centres of the rbf basis functions
    # centres = np.linspace(0, 1, 10)
    # centres, sd = classify_inputs(inputs, targets)
    centres = np.array([8.319637273,	0.527820513,	0.27097561,	2.538805503,	0.087466542,	15.87492183,	46.46779237,	0.996746679,	3.311113196,	0.658148843,	10.42298311,	5.636022514])
    # scales = np.logspace(4, 0)
    scales = np.linspace(100, 300, 200)
    errors = []
    for scale in scales:
        # print(scale)
    # define the feature mapping for the data
        feature_mapping = construct_rbf_feature_mapping(centres, scale)
        designmtx = feature_mapping(inputs)
        weights = ml_weights(designmtx, targets)
        rbf_approx = construct_feature_mapping_approx(feature_mapping, weights)
        prediction_values = rbf_approx(test_inputs)
        rmse, mae, medae, mape, variance = error_score(test_targets, prediction_values)
        errors.append(rmse)

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(scales, errors, 'r')
    ax.set_xlabel("Scale")
    ax.set_ylabel("rmse")
    plt.show()

    min_rmse, min_scale = minimise_arr_error(errors, scales)
    print("Best RSME: ", min_rmse)
    print("Best Scale: ", min_scale)

def construct_rbf_feature_mapping(centres, scale):
    """
    parameters
    ----------
    centres - a DxM matrix (numpy array) where D is the dimension of the space
        and each row is the central position of an rbf basis function.
        For D=1 can pass an M-vector (numpy array).
    scale - a float determining the width of the distribution. Equivalent role
        to the standard deviation in the Gaussian distribution.

    returns
    -------
    feature_mapping - a function which takes an NxD data matrix and returns
        the design matrix (NxM matrix of features)
    """
    #  to enable python's broadcasting capability we need the centres
    # array as a 1xDxM array
    if len(centres.shape) == 1:
        centres = centres.reshape((1, 1, centres.size))
    else:
        centres = centres.reshape((1, centres.shape[1], centres.shape[0]))
    # the denominator
    denom = 2*scale**2
    # now create a function based on these basis functions
    def feature_mapping(datamtx):
        #  to enable python's broadcasting capability we need the datamtx array
        # as a NxDx1 array
        if len(datamtx.shape) == 1:
            # if the datamtx is just an array of scalars, turn this into
            # a Nx1x1 array
            datamtx = datamtx.reshape((datamtx.size, 1, 1))
        else:
            # if datamtx is NxD array, then we reshape matrix as a
            # NxDx1 array
            datamtx = datamtx.reshape((datamtx.shape[0], datamtx.shape[1], 1))
        sub = datamtx - centres
        summ = -np.sum((sub) ** 2, 2)
        calc = np.exp(summ/denom)
        return calc
    # return the created function
    return feature_mapping

def construct_feature_mapping_approx(feature_mapping, weights):
    """
    This function creates and returns a prediction function based on a
    feature mapping and some weights.

    The returned prediction function takes a set of input values and returns
    the predicted output for each.
    """
    # here is a function that is created on the fly from the input feature
    # mapping and weights
    def prediction_function(xs):
        designmtx = np.matrix(feature_mapping(xs))
        return linear_model_predict(designmtx, weights)
    # we return the function reference (handle) itself. This can be used like
    # any other function
    return prediction_function

def linear_model_predict(designmtx, weights):
    ys = np.matrix(designmtx)*np.matrix(weights).reshape((len(weights),1))
    return np.array(ys).flatten()

def ml_weights(inputmtx, targets):
    """
    This method returns the weights that give the best linear fit between
    the processed inputs and the targets.
    """
    Phi = np.matrix(inputmtx)
    targets = np.matrix(targets).reshape((len(targets),1))
    weights = linalg.pinv(Phi.transpose()*Phi)*Phi.transpose()*targets
    return np.array(weights).flatten()

if __name__ == '__main__':
    # this bit only runs when this script is called from the command line
    main()


#!/usr/bin/env python3

# -*- coding: utf-8 -*-
"""
Created on Thu Mar  28 2019

python script for plasmids

@author: ded
"""

import argparse
import os
import re
from collections import  defaultdict
from Bio import SeqIO
from itertools import combinations
from io import StringIO
from Bio.Blast.Applications import NcbiblastnCommandline
from Bio.Blast import NCBIXML
import datetime

import sys

parser = argparse.ArgumentParser(description="")
parser.add_argument("-a", "--adapter_dir", help="directory containing trimmomatic logs", default="Adapter_trim_logs")
parser.add_argument("-s", "--spade_completeion_log", help="completeion log checking for Spades completeing sucessfully", default="Spade_completeion_log_summary.txt")
parser.add_argument("-cd", "--contig_dir", help="directory of contigs in fasta format.", default="assemblies/All_contigs/")
parser.add_argument("-r", "--reference_dir", help="directory containing reference files", default="02_Ref")
parser.add_argument("-ss", "--sign_up_sheet", help="generate command files based on signup-sheet read in", default="SignUpSheet.tsv")
parser.add_argument("-l", "--local", help="flag to run on local machine and not run every aspect. used primarily for testing", action='store_true')
parser.add_argument("--commands_only", help="only generate commands files for adapter, spades, breseq based on input file", action='store_true')
parser.add_argument("--stats_only", help="only generate stats files for adapter, and spades output", action='store_true')

# show help if no arguments passed
# if len(sys.argv) < 2:
#     parser.print_help()
#     sys.exit(1)

args = parser.parse_args()

print("\nScript executed at: " + str(datetime.datetime.now()) + "\nWith the following command line options:")
print("\t" + "\n\t".join(map(str, [str(arg) + "\t" + str(getattr(args, arg)) for arg in vars(args)])))  # log all input options

#Test for script failure conditions
fail_to_print = []
if not os.path.exists("Raw_Reads") or len([x for x in os.listdir("Raw_Reads") if x.endswith(".gz")]) == 0:
    fail_to_print.extend(["Reads should be in `Raw_Reads` directory. Try the following suggested commands from inside the directory containing the raw reads and rerun this script", "\tmkdir Raw_Reads", "\tmv *.gz Raw_Reads/"])

if not os.path.exists(args.sign_up_sheet):
    fail_to_print.append("Sign up sheet not located. Specified by default or at command line to be: `%s`. Please make sure this file exists or specify an alternative at the command line using `-ss` or `--sign_up_sheet` flags") % args.sign_up_sheet

if len(fail_to_print) > 0:
    print("\n".join(map(str, ["\n\n\n!!!!!!\nSCRIPT FAIL!!\n!!!!!!\n\n\n"] + fail_to_print + [os.getcwd()])))
    sys.exit()

#generate lists for warnings and completed tasks and failed tasks  #TODO redo as single dictionary see final print area for more details
warnings_to_print = []
completed_tasks_to_print = []
failed_tasks_to_print = []
warned_tasks_to_print = []

adapter_dict = {}
Contig_dict = {}

def adapter_trimming_stats():
    print("adapter trimming stats")
    for f in os.listdir(args.adapter_dir):

        if f.endswith(".log"):
            name = re.sub(".log$", "", f)
            assert name not in adapter_dict and name not in Contig_dict, "duplicated file name:%s" % f
            Contig_dict[name] = {'R1R2': [], 'R1': [], 'R2': [], 'P1P2U1U2': [], 'P1P2': [], 'P1': [], 'P2': [], 'P1P2U1': [], 'P1P2U2': [],  'P1U1': [], 'P2U2': []}

            with open(args.adapter_dir + "/" + f, "r") as f_in:
                for line in f_in:
                    line=line.rstrip()
                    if re.match("Input Read Pairs", line):
                        line = line.split(" ")
                        # [input, both, forward, reverse, drop]
                        adapter_dict[name] = [line[3], line[6], line[11], line[16], line[19]]

    print("\t".join(map(str, ["Sample", "Raw read pairs", "Both", "Forward only", "Reverse only", "Drop"])))
    for name in sorted(adapter_dict):
        print("\t".join(map(str, [name] + adapter_dict[name])))

def spades_stats():
    # identify failed plasmid assembiles. NOTE this is failed to complete assembly pipeline, NOT failed to assemble any contigs
    print("calculating spade stats")
    with open(args.spade_completeion_log, "r") as sl_in:

        for line in sl_in:
            line=line.rstrip()
            if re.search(".log:1$", line):
                continue
            assert re.search(".log:0$", line), "Spade completion log lists sample that somehow completed multiple times. This does not make sense. \n%s" % line
            assert line.count("/") == 1 and line.count(".log:0") == 1, "Unexpected assembly log format log files expected to be stored in a single directory, and end with '.log'. this file should also have ended in :0 as currently looking at/for failed assemblies.\n%s" % line
            name = re.sub(".log:0","", re.sub(".*\/", "", line))  # strip directory info from line
            name = name.split(".")
            Contig_dict[name[0]]["".join(name[1:])] = "NA"

def contig_read_in():
    #Read contigs into dictionary
    print("reading contigs in")
    for f in os.listdir(args.contig_dir):

        if f.endswith(".contigs.fasta"):
            base = re.sub(".contigs.fasta$", "", f)
            assert len(base.split(".")) == 2, "unable to identify name from contig file: %s\nName expected to be between '.' and file name to end with .contigs.fasta" % f
            name = base.split(".")[1]

            assert base.split(".")[0].split("_")[1] == "reads", "unable to idtnify assembly conditions from contig file: %s\nFormat expected to be 'Raw|Trim_reads_' followed by R|P|U 1|2 each separated by '_' marks" % f
            assembly = "".join(base.split(".")[0].split("_")[2:])

            assert Contig_dict[name][assembly] is not "NA", "Contig file found for assembly previously identified as a failed assembly. This is different than an assembly which produces no contig but successfully completes assembly pipeline. %s" % f

            with open(args.contig_dir + "/" + f, "rU") as handle:
                for record in SeqIO.parse(handle, "fasta"):
                    Contig_dict[name][assembly].append((record.name, record.seq))

    all_single_contig_dict = {}
    print("\t".join(map(str, ["Sample", 'R1R2', 'R1', 'R2', 'P1P2U1U2', 'P1P2', 'P1', 'P2', 'P1P2U1', 'P1P2U2', 'P1U1', 'P2U2', "tot Fail", "tot None", "tot >1", "tot 1"])))
    for sample in sorted(Contig_dict):
        if args.local:
            # select one of the following to run or not run this block if running/testing locally
            continue
            #pass
        to_print = [sample]
        all_single_contig_dict[sample] = []
        for assembly_type in ['R1R2', 'R1', 'R2', 'P1P2U1U2', 'P1P2', 'P1', 'P2', 'P1P2U1', 'P1P2U2', 'P1U1', 'P2U2']:
            if Contig_dict[sample][assembly_type] == "NA":
                to_print.append("NA")
            else:
                to_print.append(len(Contig_dict[sample][assembly_type]))
                if len(Contig_dict[sample][assembly_type]) == 1:
                    all_single_contig_dict[sample].append(Contig_dict[sample][assembly_type][0][1])
        to_print += [to_print[1:].count("NA"), to_print[1:].count(0), sum(i>1 for i in to_print[1:] if i is not "NA"), to_print[1:].count(1)]
        print("\t".join(map(str, to_print)))

    #check for identical
    testing = 0
    for sample in sorted(Contig_dict):
        if args.local:
            # select one of the following to run or not run this block if running/testing locally
            continue
            #pass

        if sample != "JCM-10" : continue
        print(sample)

        subject_ref_file = "Reference_sequences/JCM-10.fa"

        sample_score_matrix = []
        if len(all_single_contig_dict[sample]) > 1:
            if len(set(all_single_contig_dict[sample])) > 1:
                num_contigs = 0
                fasta_files = []
                for seq in set(all_single_contig_dict[sample]):
                    num_contigs +=1
                    fasta_files.append("contig_check_temp_files/" + sample + "_" + str(num_contigs) + ".fa")
                    if not os.path.exists("contig_check_temp_files/" + sample + "_" + str(num_contigs) + ".fa"):
                        with open("contig_check_temp_files/" + sample + "_" + str(num_contigs) + ".fa", "w") as output:
                            if num_contigs == 2:
                                print(">" + sample + "\n" + seq +seq, file=output)
                            else:
                                print(">" + sample + "\n" + seq, file=output)
                    print(sample + "_" + str(num_contigs) + ".fa" + "\t" + str(len(seq)))

                for c in combinations(fasta_files, 2):
                    print("\t".join(map(str, [sample, c[0], c[1]])))

                    blast_output = NcbiblastnCommandline(query=c[0], subject=c[1], outfmt=5)()[0]
                    blast_result = NCBIXML.read(StringIO(blast_output))

                    perfects = 0
                    for aln in blast_result.alignments:
                        print(len(aln.hsps), "alignments")

                        for hsp in aln.hsps:
                            if (hsp.align_length - hsp.identities == 0) and (hsp.gaps == 0):
                                #print hsp.identities
                                print("following is perfect")
                                perfects += 1
                            print('****Alignment****')
                            print('sequence:', aln.title)
                            print('length:', hsp.align_length)
                            print('e value:', hsp.expect)
                            print(hsp.query)
                            print(hsp.match)
                            print(hsp.sbjct)
                            print(hsp.identities, "matches")
                    print(perfects, "perfects")
                testing +=1
                if testing >= 1:
                    sys.exit()

# fail pairwise2 would have been inside contig_read_in
# skip = True  #temp for testing. this block takes the longest by far
# for sample in sorted(Contig_dict):
#     if skip:
#         continue
#     sample_score_matrix = []
#     if len(all_single_contig_dict[sample]) > 1:
#         if len(set(all_single_contig_dict[sample])) > 1:
#             for c in combinations(set(all_single_contig_dict[sample]), 2):
#                 if len(c[0]) >= len(c[1]):
#                     sample_score_matrix.append(max(([pairwise2.align.globalxx(c[0] + c[0], c[1], score_only=True),
#                                                      pairwise2.align.globalxx(c[0] + c[0], c[1].reverse_complement(), score_only=True)])) / len(c[0]))
#                 else:
#                     sample_score_matrix.append(max(([pairwise2.align.globalxx(c[1] + c[1], c[0], score_only=True),
#                                                      pairwise2.align.globalxx(c[1] + c[1], c[0].reverse_complement(), score_only=True)])) / len(c[1]))
#             print "\t".join(map(str, [sample, sum(sample_score_matrix) / float(len(sample_score_matrix)), sample_score_matrix]))
#         else:
#             print "\t".join(map(str, [sample, 1.0, "ident string"]))
#     else:
#         if len(all_single_contig_dict[sample]) == 1:
#             print "\t".join(map(str, [sample, 1.0, "single positive"]))
#         else:
#             print "\t".join(map(str, [sample, "NA", "No single contig"]))

#print("starting ref checks")
# #check contigs against references
# for ref in os.listdir(args.reference_dir):
#     if args.local:
#         # select one of the following to run or not run this block if running/testing locally
#         continue
#         #pass
#     if ref.endswith('.fa') or ref.endswith('.fasta'):
#         name = re.sub(".fasta$", "", ref)
#         name = re.sub(".fa$", "", name)
#         assert name in Contig_dict, "unmatched reference file name %s" % ref
#         ref_compare_score = []
#         if len(set(all_single_contig_dict[name])) >= 1:
#             best_contig = "NA"
#             with open(args.reference_dir + "/" + ref, "rU") as handle:
#                 fasta_count = 0
#                 for record in SeqIO.parse(handle, "fasta"):
#                     fasta_count += 1
#                     for contig in set(all_single_contig_dict[name]):
#                         if len(contig) >= len(record.seq):
#                             ref_compare_score.append(max(([pairwise2.align.globalxx(contig + contig, record.seq, score_only=True), pairwise2.align.globalxx(contig +contig, record.seq.reverse_complement(), score_only=True)]))/ len(contig))
#                         else:
#                             ref_compare_score.append(max(([pairwise2.align.globalxx(record.seq + record.seq, contig, score_only=True),
#                                                            pairwise2.align.globalxx(record.seq + record.seq, contig.reverse_complement(), score_only=True)])) / len(record.seq))
#                         if ref_compare_score[-1] == max(ref_compare_score):
#                             best_contig = contig
#                 assert best_contig != "NA", "No contig scored best? %s %s" % (name, ref_compare_score)
#                 print max(ref_compare_score)
#                 if len(best_contig) >= len(record.seq):
#                     if pairwise2.align.globalxx(best_contig + best_contig, record.seq, score_only=True) / len(best_contig) > pairwise2.align.globalxx(best_contig + best_contig, record.seq.reverse_complement(), score_only=True) / len(best_contig):
#                         print pairwise2.format_alignment(*pairwise2.align.globalxx(best_contig + best_contig, record.seq, one_alignment_only=True)[0])
#                     else:
#                         print pairwise2.format_alignment(*pairwise2.align.globalxx(best_contig + best_contig, record.seq.reverse_complement(), one_alignment_only=True)[0])
#                 else:
#                     if pairwise2.align.globalxx(record.seq + record.seq, best_contig, score_only=True) / len(record.seq) > pairwise2.align.globalxx(record.seq + record.seq, best_contig.reverse_complement(), score_only=True) / len(record.seq):
#                         print pairwise2.format_alignment(*pairwise2.align.globalxx(record.seq + record.seq, best_contig, one_alignment_only=True)[0])
#                     else:
#                         print pairwise2.format_alignment(*pairwise2.align.globalxx(record.seq + record.seq, best_contig.reverse_complement(), one_alignment_only=True)[0])
#
#                 if fasta_count > 1:
#                     print "\n\n\n!!!!!!\n\n\n!!!!!!!\nWARNING multiple fasta sequences found in file: %s\n\n\n" % ref
#
#                 print "\t".join(map(str, [name, sum(ref_compare_score)/ float(len(ref_compare_score)), ref_compare_score]))
#         else:
#             print "\t".join(map(str, [name, "No single contig to compare"]))
#
#                 #for assembly_type in Contig_dict[]

def spades_commands_gen():
    #check if this will run:
    if len(os.listdir("Trim_Reads/")) == 0:
        print("Spades commands can not be run yet, but commands generated. commands can be run upon completion of adapter trimming commands")

    print("Spades command gen started")

    assembly_dirs_check = ["assemblies", "assemblies/Trim_reads_P1_P2_U1_U2", "assemblies/Trim_reads_P1_P2", "assemblies/Trim_reads_P1", "assemblies/Trim_reads_P2", "assemblies/Trim_reads_P1_P2_U1", "assemblies/Trim_reads_P1_P2_U2", "assemblies/Trim_reads_P1_U1", "assemblies/Trim_reads_P2_U2", "Spades_Logs"]
    for dir_check in assembly_dirs_check:
        if not os.path.exists(dir_check):
            os.makedirs(dir_check)

    with open(args.sign_up_sheet, "r") as SUS, open("spades_commands", "w") as spades_command_out:
        Header_search = True
        for line in SUS:
            line = line.rstrip()
            if Header_search:
                assert re.match("^Researcher\tSample_ID\tReference_file", line), "header line missing in %s" % line
                Header_search = False
                continue
            if re.match("Your name", line):
                continue

            line = line.split("\t")

            Name = line[0].replace(" ", "_") + "." + line[1]
            P1 = Name + "_1P.fastq.gz"
            P2 = Name + "_2P.fastq.gz"
            U1 = Name + "_1U.fastq.gz"
            U2 = Name + "_2U.fastq.gz"

            print("plasmidspades.py -o assemblies/Trim_reads_P1_P2_U1_U2/%s -1 %s -2 %s -s %s -s %s >& Spades_Logs/%s.P1.P2.U1.U2.log" %(Name, P1, P2, U1, U2, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P1_P2/%s -1 %s -2 %s >& Spades_Logs/%s.P1.P2.log" % (Name, P1, P2, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P1/%s -s %s >& Spades_Logs/%s.P1.log" % (Name, P1, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P2/%s -s %s >& Spades_Logs/%s.P2.log" % (Name, P2, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P1_P2_U1/%s -1 %s -2 %s -s %s >& Spades_Logs/%s.P1.P2.U1.log" % (Name, P1, P2, U1, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P1_P2_U2/%s -1 %s -2 %s -s %s >& Spades_Logs/%s.P1.P2.U2.log" % (Name, P1, P2, U2, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P1_U1/%s -s %s -s %s >& Spades_Logs/%s.P1.U1.log" % (Name, P1, U1, Name), file=spades_command_out)
            print("plasmidspades.py -o assemblies/Trim_reads_P2_U2/%s -s %s -s %s >& Spades_Logs/%s.P2.U2.log" % (Name, P2, U2, Name), file=spades_command_out)

def SignUpSheet_read_in():
    with open(args.sign_up_sheet, "r") as SUS:
        all_raw_reads = [x for x in os.listdir("Raw_Reads") if x.endswith(".gz") and not re.match("^" + "Undetermined_S0_L001_", x)]
        used_read_file = []
        # for raw_read_file in os.listdir("Raw_Reads"):
        #     if raw_read_file.endswith(".gz"):
        #         all_raw_reads.append(raw_read_file)
        assert len(all_raw_reads) > 0, "No reads were found in `Raw_Reads` move gzipped raw read files to that directory. Possible suggested command:\n\tmv *.gz Raw_Reads/\nPossible script failure. This should have been caught in failure checks"
        all_ids = []

        Header_search = True
        for line in SUS:
            line = line.rstrip()
            if Header_search:
                assert re.match("^Researcher\tSample_ID\tReference_file",line), "header line missing in %s" %line
                Header_search = False
                continue
            if re.match("Your name", line):
                continue
            line = line.split("\t")
            researcher = line[0].replace(" ", "_")
            sample_ID = line[1]
            reference_file = line[2]
            assert sample_ID not in master_dict[researcher], "duplicated sampleID for single researcher. %s\nmaster_dict:\n%s" % (line, list(master_dict[researcher].keys()))
            assert sample_ID not in all_ids, "duplicated sampleID for different researchers %s\nall ids:\n%s" % (line, all_ids)

            master_dict[researcher][sample_ID] = {"Ref": reference_file, "Raw_reads": []}
            all_ids.append(sample_ID)

            for raw_read in all_raw_reads:
                if re.match("^" + sample_ID +"_", raw_read):
                    master_dict[researcher][sample_ID]["Raw_reads"].append(raw_read)
                    used_read_file.append(raw_read)

        unused_reads = [x for x in all_raw_reads if x not in used_read_file]
        if len(unused_reads) > 0:
            warnings_to_print.append("The following read file exists without a sample ID and will not be trimmed.\n\t" + "\n\t".join(sorted(unused_reads)))

        samples_missing_reads = [" from ".join([s_ID, res]) for res in master_dict for s_ID in master_dict[res] if len(master_dict[res][s_ID]["Raw_reads"]) == 0]
        if len(samples_missing_reads) > 0:
            warnings_to_print.append("No read files were identified for the following samples, and samples will not be analyzed.\n\t" + "\n\t".join(sorted(samples_missing_reads)))

        all_refs = list(set([master_dict[res][s_ID]["Ref"] for res in master_dict for s_ID in master_dict[res]]))
        for ref in all_refs:
            if os.path.exists("02_Ref/" + ref):
                pass
            else:
                if "signupsheet read in" not in warned_tasks_to_print:
                    warned_tasks_to_print.append("signupsheet read in")
                warnings_to_print.append("%s reference file missing. move file to 02_Ref/ directory" % ref)

            if ref.endswith(".gbk") or ref.endswith(".gb") or ref.endswith(".fa") or ref.endswith(".fasta"):
                pass
            else:
                if "signupsheet read in" not in warned_tasks_to_print:
                    warned_tasks_to_print.append("signupsheet read in")
                warnings_to_print.append("%s reference file is of unknown type. rename or verify file is correctly formated genbank or fasta file" % ref)

        if "signupsheet read in" not in warned_tasks_to_print:
            completed_tasks_to_print.append("signupsheet read in")

def adapter_trim_commands_gen():
    if not os.path.exists("Trim_Reads"): os.makedirs("Trim_Reads")
    if not os.path.exists("Adapter_trim_logs"): os.makedirs("Adapter_trim_logs")
    with open("adapter_trim_commands", "w") as adapter_command_out:
        for researcher in master_dict:
            for sample_ID in master_dict[researcher]:
                sample_reads = ["Raw_Reads/" + x for x in master_dict[researcher][sample_ID]["Raw_reads"]]
                sample_reads = " ".join(sample_reads)
                print("trimmomatic PE %s -baseout Trim_Reads/%s.fastq.gz ILLUMINACLIP:illumina_truseq_6bp_dual_barcode_PE.fasta:4:30:10 MINLEN:30 >& Adapter_trim_logs/%s.log" % (
                sample_reads, sample_ID, sample_ID), file=adapter_command_out)

    if not os.path.exists("illumina_truseq_6bp_dual_barcode_PE.fasta"):
        warnings_to_print.append(
            "adapter fasta file not found. run the following command to allow adapter trimming:\n\tcp /corral-repl/utexas/breseq/genomes/adapters/illumina_truseq_6bp_dual_barcode_PE.fasta .")
        warned_tasks_to_print.append("adapter command generation")
    else:
        completed_tasks_to_print.append("adapter command generation")

def breseq_commands_gen():
    if not os.path.exists("03_Output"): os.makedirs("03_Output")
    if not os.path.exists("04_Logs"): os.makedirs("04_Logs")

    with open("breseq_commands", "w") as breseq_command_out, open("gdtools_commands", "w") as gd_command_out:
        gdtools_compare_dict = defaultdict(lambda: defaultdict(list))
        for researcher in master_dict:

            if not os.path.exists("03_Output/" + researcher): os.makedirs("03_Output/" + researcher)

            for sample_ID in master_dict[researcher]:

                gdtools_compare_dict[researcher][master_dict[researcher][sample_ID]["Ref"]].append(sample_ID + ".gd")

                trim_read_files = " ".join(["Trim_Reads/" + sample_ID + x for x in ["_1P.fastq.gz", "_2P.fastq.gz", "_1U.fastq.gz"]])  # , "_2U.fastq.gz"  # eliminated as several are 0 and dont have quick fix yet

                print("breseq -o 03_Output/%s/%s -r 02_Ref/%s %s >& 04_Logs/%s.%s.log.txt"  % (researcher, sample_ID, master_dict[researcher][sample_ID]["Ref"], trim_read_files, researcher, sample_ID), file=breseq_command_out)

        #TODO work this into run on tacc
        print("The following commands can not be run as a script:", file=gd_command_out)
        for researcher in gdtools_compare_dict:
            print(researcher, file=gd_command_out)
            for ref in gdtools_compare_dict[researcher]:
                print("gdtools compare -r ../../../02_Ref/%s -o %s.%s.html %s" % (ref, researcher, ref, " ".join(sorted(gdtools_compare_dict[researcher][ref]))), file=gd_command_out)

    #TODO need stuff here about sucess/failure

    #Check if commands can run
    # if len(os.listdir("Trim_Reads/")) == 0:  # TODO change this to warn commands can't be run until after adapters complete rather than looking for reads
    #     print "Breseq commands can not be generated at this time as Trimmed reads folder is currently empty. Likely cause of this is adapter commands not having ben run yet. rerun script after adapters are trimmed"
    #     return
    #
    # print "breseq command gen started"
    # gdtools_compare_dict = {}
    # if not os.path.exists("03_Output_Trim"): os.makedirs("03_Output_Trim")
    # if not os.path.exists("04_Logs"): os.makedirs("04_Logs")
    #
    # with open(args.sign_up_sheet, "r") as f_in, open("breseq_commands", "w") as command_out, open("gdtools_commands", "w") as gd_command_out:
    #
    #     Header_search = True
    #     Samples_lacking_reads = []
    #     for line in f_in:
    #         line = line.rstrip()
    #
    #         if Header_search:
    #             assert re.match("^Researcher\tSample_ID\tReference_file",line), "header line missing in %s" % line
    #             Header_search = False
    #             continue
    #         if re.match("Your name", line):
    #             continue

            # line = line.split("\t")
            # Name = line[0].replace(" ", "_")
            # Base_sample_ID = line[1]
            # ref_file = line[2]
            #
            # if Name not in gdtools_compare_dict:
            #     gdtools_compare_dict[Name] = {}
            #
            # if ref_file not in gdtools_compare_dict[Name]:
            #     gdtools_compare_dict[Name][ref_file] = []
            #
            # gdtools_compare_dict[Name][ref_file].append(Base_sample_ID+".gd")
            #
            # trim_reads = ["Trim_Reads/" + x for x in os.listdir("Trim_Reads") if re.match("^" + Base_sample_ID + "_", x)]
            # if len(trim_reads) == 0:
            #     Samples_lacking_reads.append(Base_sample_ID)
            #     continue
            # trim_reads = " ".join(trim_reads)
            #
            # Output_dir = Name + Base_sample_ID
            # if not os.path.exists("03_Output_Trim/" + Output_dir):os.makedirs("03_Output_Trim/" + Output_dir)
            # print>>command_out, "breseq -o 03_Output_Trim/%s -r 02_Ref/%s %s >& 04_Logs/%s.%s.Trim.log.txt" % (Output_dir, ref_file, trim_reads, Name, Base_sample_ID)

        # print>>gd_command_out, "The following commands can not be run as a script:"
        # for researcher in gdtools_compare_dict:
        #     print>>gd_command_out, researcher
        #     for ref_entry in gdtools_compare_dict[researcher]:
        #         print>>gd_command_out, "gdtools compare -r ../../../02_Ref/%s -o %s.%s.html %s" % (ref_entry, researcher, ref_entry, " ".join(gdtools_compare_dict[researcher][ref_entry]))

        # if len(Samples_lacking_reads) > 0:
        #     print "The following samples did not have breseq commands generated as reads were not identified corresponding to the sample:"
        #     print "\t" + "\n\t".join(map(str, Samples_lacking_reads))

master_dict = defaultdict(lambda: defaultdict())

SignUpSheet_read_in()  # must be called to allow anything else to run
if args.commands_only:
    adapter_trim_commands_gen()
    breseq_commands_gen() # won't work until trimming done
    spades_commands_gen() # commands won't work until trimming done

if args.stats_only:
    pass

if args.local:
    #comment out parts you dont want to run
    #adapter_trim_commands_gen()
    #breseq_commands_gen() # won't work until trimming done
    #spades_commands_gen() # commands won't execute until trimming done
    #adapter_trimming_stats()
    #spades_stats()
    #contig_read_in() #requires adapter trimming stats to run
    pass

# else: # run all elements
#     adapter_trimming_stats()
#     breseq_commands_gen()
#     spades_stats()
#     contig_read_in()

#TODO rework as dictionary of tasks: d[task] == {} -- sucess ... d[fail/warn] = [failures/warnings] use asserts to make sure task isn't in d, and warnings last thing generated. fail should exit function always

if len(completed_tasks_to_print) > 0:
    print("The following tasks were successfully completed:")
    print("\t"+"\n\t".join(map(str, completed_tasks_to_print)))

if len(failed_tasks_to_print) > 0:
    print("The following tasks failed:")
    print("\t"+"\n\t".join(map(str, failed_tasks_to_print)))

if len(warned_tasks_to_print) > 0:
    assert len(warnings_to_print) > 0, "Task warned without warning being generated. Script problem. \n%s" % warned_tasks_to_print
    print("The following tasks generated warnings:")
    print("\t"+"\n\t".join(map(str, warned_tasks_to_print)))

if len(warnings_to_print) != 0:
    print("\n".join(map(str, ["\n\n\n!!!!!!\nSCRIPT WARNINGS!!\n!!!!!!\n\n\n"] + warnings_to_print)))
else:
    print("Script completed successfully without generating any warnings.")


from lib.mode import Mode
# from cmd.prediction import Prediction

def main():
    inputCSV = "data/mode/e28b969e-53a0-4b01-8711-a9c18e61ebab/hog2.csv"
    myMode = Mode(inputCSV)
    winnerClass, dictinaryData = myMode.calculate()
    print("Winner is : ", winnerClass)
    print("dictionary: ", dictinaryData)

if __name__ == '__main__':
    main()


from pyspark.sql import SparkSession
from pyspark import SparkConf

if __name__ == '__main__':
    # CREATING SPARKSESSION
    spark = SparkSession \
            .builder \
            .appName("join-employees-and-offices-py") \
            .getOrCreate()

    # SHOW CONFIGURED
    print(SparkConf().getAll())

    # SET LOG LEVEL
    spark.sparkContext.setLogLevel("INFO")

    # PATH OF OBJECTS IN CLOUD STORAGE
    get_offices = "gs://landing_z/input/offices/offices"
    get_employees = "gs://landing_z/input/employees/employees"

    # CREATING DATAFRAME
    df_offices = spark.read \
            .format("csv") \
            .option("inferSchema", "true") \
            .option("header", "true") \
            .load(get_offices)

    # CREATING DATAFRAME
    df_employees = spark.read \
        .format("csv") \
        .option("inferSchema", "true") \
        .option("header", "true") \
        .load(get_employees)

    print("NUMBERS PARTITIONS -> ",df_offices.rdd.getNumPartitions())
    print("NUMBERS PARTITIONS -> ",df_employees.rdd.getNumPartitions())

    print("SCHEMA DATAFRAME")
    df_offices.printSchema()

    print("SCHEMA DATAFRAME")
    df_employees.printSchema()

    print("DISPLAYING DATAFRAME")
    df_offices.show()

    print("DISPLAYING DATAFRAME")
    df_employees.show()

    print("COUNTING DATAFRAME ->",df_offices.count())
    print("COUNTING DATAFRAME -> ",df_employees.count())

    # CREATING VIEW TABLES
    df_offices.createOrReplaceTempView("offices")
    df_employees.createOrReplaceTempView("employees")

    # EXECUTINH JOIN ENTER BOTH TABLES
    result_join = spark.sql("""
    SELECT
    employees.employee_number,
    employees.last_name,
    employees.first_name,
    employees.extension,
    employees.email,
    employees.office_code AS emp_office_code,
    employees.reports_to,
    employees.job_Title,
    offices.office_code AS ofc_office_code,
    offices.city,
    offices.phone,
    offices.address_line1,
    offices.address_line2,
    offices.state,
    offices.country,
    offices.postal_code,
    offices.territory
    FROM employees INNER JOIN offices ON employees.office_code = offices.office_code;
    """)

    print("EXPLAIN")
    result_join.explain()

    print("COUNTING DATAFRAME -> ",result_join.count())
    result_join.show()

    # CREATING PARQUET FILE WITH THE RESULT
    result_join.write.format("parquet").mode("overwrite").save("gs://gold-zone/vendendores")

    #
    spark.stop()


# Import libraries
import time
import math
import numpy as np
from tqdm import tqdm

def axissummary(chunk):
    mean = np.mean(chunk)
    median = np.median(chunk)
    min = np.mean(chunk)
    max = np.median(chunk)
    std = np.std(chunk)
    return [mean, median, min, max, std]

def chunksummary(chunk):
    details = []
    for i in range(21):
        details.extend(axissummary(chunk[:,i]))
    return details

def replaysummary(frames, notes):
    summary = []
    for note in notes:
        if (note[0] != 0 and note[1] > 0.05 and note[1] < 29.95):
            id, time = note[0], note[1]
            frame = math.floor(time*30)
            result = [id]
            result.extend(chunksummary(frames[max(0,frame-30):frame]))
            result.extend(chunksummary(frames[frame:min(900,frame+30)]))
            summary.append(result)
    return summary

def usersummary(user):
    summary = []
    replays, notes = user

    for i in range(len(replays)):
        summary.extend(replaysummary(replays[i], notes[i]))

    return np.vstack([np.array(summary)[:5000], np.array(summary)[-5000:]])

def nairfeatures(replays, notes):
    users = []
    for user in tqdm(list(zip(replays, notes))):
        users.append(usersummary(user))
    return np.array(users)

# Start measuring performance
start_time = time.time()

# Load the input data
print("Loading input data...")
replays = np.load('./data/user-replays.npy')
notes = np.load('./data/user-notes.npy')

# Control group
print("Processing control feautres...")
nair = nairfeatures(replays, notes)
np.save('./data/nair', nair)

# Log performance results
end_time = time.time()
print("Finished in %s Minutes" % ((end_time - start_time) / 60))


# https://github.com/rgeirhos/texture-vs-shape/blob/master/code/probabilities_to_decision.py
import numpy as np
from abc import ABC, abstractmethod
print('test')

import helper.human_categories as hc

class ProbabilitiesToDecisionMapping(ABC):
    @abstractmethod
    def probabilities_to_decision(self, probabilities):
        pass

    def check_input(self, probabilities):
        assert type(probabilities) is np.ndarray
        assert (probabilities >= 0.0).all() and (probabilities <= 1.0).all()

class ImageNetProbabilitiesTo16ClassesMapping(ProbabilitiesToDecisionMapping):
    def __init__(self, aggregation_function=np.mean):

        self.aggregation_function = aggregation_function

    def probabilities_to_decision(self, probabilities):
        self.check_input(probabilities)
        assert len(probabilities) == 1000

        max_value = -float("inf")
        category_decision = None
        c = hc.HumanCategories()
        for category in hc.get_human_object_recognition_categories():
            indices = c.get_imagenet_indices_for_category(category)
            values = np.take(probabilities, indices)
            aggregated_value = self.aggregation_function(values)
            if aggregated_value > max_value:
                max_value = aggregated_value
                category_decision = category

        return category_decision


import matplotlib

matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

class SystematicStrats():

    def __init__(self, df, stocks_list, prices):
        self.df = df
        self.stock_list = stocks_list
        self.ranked_df = self.contruct_ranking()
        self.prices = prices

    def contruct_ranking(self):
        col_names = self.df.columns[1:]
        mask = self.df[col_names].applymap(lambda x: isinstance(x, (int, float)))
        self.df[col_names] = self.df[col_names].where(mask)
        self.df = self.df.ffill()
        self.df.set_index('Ticker', inplace=True)
        self.df[col_names] = self.df[col_names].apply(lambda x: x.astype('float'))

        # return only filtered df
        self.df = self.df[self.df.index.str.contains('|'.join(self.stock_list))]
        dfr = pd.DataFrame()
        for z in col_names:
            # df[str(z)+str('_R')] = 0
            dfr[str(z)] = self.df[z].rank()

        return dfr

    def generate_signals(self, nlong):
        self.ranked_dfT = self.ranked_df.transpose()
        self.ranked_dfT.index = pd.to_datetime(self.ranked_dfT.index)
        colnames = list(self.ranked_dfT.columns)
        colnames.sort()
        self.ranked_dfT = self.ranked_dfT[colnames]

        self.ranked_dfT_s = self.ranked_dfT.shift()
        #self.ranked_dfT_s.replace(np.nan, 0, inplace=True)
        df_array = self.ranked_dfT_s.values

        #long portfolio
        ones_array_l = (df_array < nlong + 1).astype(int)
        signals_l = pd.DataFrame(ones_array_l, index=self.ranked_dfT.index)
        signals_l.columns = colnames

        #short portfolio
        ones_array_s = (df_array > 17 - nlong).astype(int)
        signals_s = pd.DataFrame(ones_array_s, index=self.ranked_dfT.index)
        signals_s.columns = colnames
        return signals_l,signals_s

    def generate_systematic_strat_returns(self, sig_l, sig_s, nlong):

        rtn_l = self.generate_rtns_long_short(sig_l,nlong)
        rtn_s = self.generate_rtns_long_short(sig_s,nlong)

        return rtn_l,rtn_s

    def generate_rtns_long_short(self, sig, nlong):
        ones_array = sig.values
        ones_array = (ones_array * (1 / nlong)).round(2)
        prices_r = self.prices.reindex(self.ranked_dfT.index)
        prices_chg_r = prices_r.pct_change()
        prices_chg_r.replace(np.nan, 0, inplace=True)
        prices_chg_r_array = prices_chg_r.values

        rtns = np.multiply(prices_chg_r_array, ones_array).round(4)
        rtns_df_ls = pd.DataFrame(rtns, index=prices_chg_r.index)

        return rtns_df_ls

if __name__ == '__main__':

    n_long = 3  # number of long positions
    filename = 'Russell3000.xlsx'
    Lfile = pd.ExcelFile(filename)

    prices = pd.read_csv('Russell_price.csv', index_col='date', parse_dates=True)
    stocks = list(prices.columns)
    stocks.sort()
    prices = prices[stocks]

    ebitda = pd.read_excel(Lfile, 'EBITDA', skiprows=1)
    eps = pd.read_excel(Lfile, 'EPS', skiprows=1)
    fcf = pd.read_excel(Lfile, 'FCF', skiprows=1)
    rev = pd.read_excel(Lfile, 'REVENUE', skiprows=1)
    opt = pd.read_excel(Lfile, 'OPERATING PROFIT AFTER TAX', skiprows=1)
    # qoq = pd.read_excel(Lfile, 'QoQ Return', skiprows=1)

    # run strategy

    combo = pd.DataFrame()
    features = [eps, ebitda, fcf, rev, opt] #change as required

    for f in features:
        sys_strats = SystematicStrats(f, stocks, prices)
        signals_l,signals_s = sys_strats.generate_signals(n_long)
        sys_strat_returns_l,sys_strat_returns_s = sys_strats.generate_systematic_strat_returns(signals_l,signals_s, n_long)

        rtn_ts_l = sys_strat_returns_l.sum(axis=1)
        rtn_ts_s = sys_strat_returns_s.sum(axis=1)

        cum_rtns_l = np.cumproduct(rtn_ts_l + 1)
        cum_rtns_s = np.cumproduct(rtn_ts_s + 1)
        cum_rtn_ls = cum_rtns_l + cum_rtns_s
        rtn_ls_chg = cum_rtn_ls.pct_change()
        cum_rtn_ls_chg = np.cumproduct(rtn_ls_chg+1)-1
        cum_rtn_ls_chg.replace(np.nan, 0, inplace=True)

        combo = pd.concat([combo, cum_rtn_ls_chg], axis=1)

    # extract S&P data
    spx = pd.read_csv('spx_long.csv', index_col='Date', parse_dates=True)['Adj Close']
    spx = spx.reindex(combo.index)
    spx = spx.pct_change()
    cum_spx = np.cumproduct(spx + 1) - 1
    cum_spx.replace(np.nan, 0, inplace=True)
    combo = pd.concat([combo, cum_spx], axis=1)

    combo.columns = ['EPS', 'EBITDA', 'FCF', 'REVENUE', 'OPT','SPX'] #change as required
    combo.plot()
    plt.legend(combo.columns)
    plt.title('Systematic Long Only Strategy Based On Ranking Russell3000')
    plt.show()

    print(combo)


import requests
import numpy as np
from config import Sensors
from mysql_connect import Connect
from datetime import date, timedelta
import json
from sd_send_email import send_mail
import sys, getopt

def obtain_monitor_event(id=None):
    if id:
        query_sql = "select id, event_name, where_condition, expectation, " \
                "standard_deviation, event_product, description, aggregator from xsl_sensors_data.track_event_monitor " \
                "where id in(%s) order by event_product" %id
    else:
        query_sql = "select id, event_name, where_condition, expectation, " \
                "standard_deviation, event_product, description, aggregator from xsl_sensors_data.track_event_monitor " \
                "where status = 1 order by event_product"
    conn = Connect()
    events = conn.sql_query(query_sql)
    conn.close_connect()
    return events

def update_E_and_SD(E_and_SD):
    update_sql = "update xsl_sensors_data.track_event_monitor " \
                 "set expectation = %s, standard_deviation = %s " \
                 "where id = %s "
    conn = Connect()
    try:
        conn.cursor.executemany(update_sql, E_and_SD)
        conn.connection.commit()
        conn.close_connect()
    except BaseException as e:
        print(e)
        conn.close_connect()

def calculate_E_and_SD(update=True, monitor_id=None):
    # 当传update值,进行均值和标准差更新
    print("进入更新均值和标准差更新过程")
    if update:
        from_date = (date.today() - timedelta(days=date.today().isoweekday()+7)).strftime('%Y-%m-%d')
        to_date = (date.today() - timedelta(days=date.today().isoweekday())).strftime('%Y-%m-%d')
        events = obtain_monitor_event(monitor_id)
        # 使用配置表中的event_name和conditions,注意event的index随obtain_monitor_event中的sql字段变化
        # 之后几行使用较多List下标, 注意变动影响
        results = [(id, obtain_event_from_sensors(event_name, aggregator, condition, from_date, to_date))
                   for (id, event_name, condition, e, sd, e_prod, desc, aggregator) in events]
        E_and_SD = []
        for result in results:
            if result[1]['rows']:
                result_values = [value[0] for value in result[1]['rows'][0]['values']]
            else:
                result_values = [0]
            E_and_SD.append((round(np.average(result_values),2), round(np.std(result_values),2), result[0]))
        update_E_and_SD(E_and_SD)
        print("全部事件每日期望值和标准差已更新!")
    else:
        pass

def obtain_event_from_sensors(event_name, aggregator, condition, from_date, to_date):
    url = Sensors.url + '/api/events/report?token=' + Sensors.token
    data = {
        "measures": [
            {
                "event_name": event_name,
                "aggregator": aggregator
            }
        ],
        "unit": "day",
        "from_date": from_date,
        "to_date": to_date
    }
    if condition:
        condition = json.loads(condition)
        data["filter"] = condition

    event_result = requests.post(url, json=data).json()
    return event_result

def generate_event_url(event_name, aggregator, condition):
    # 由于不懂SD的url拼接规则,只好按常规访问的url模仿
    from_date = (date.today() - timedelta(days=30)).strftime('%Y-%m-%d')
    to_date = (date.today() - timedelta(days=1)).strftime('%Y-%m-%d')
    url = Sensors.url + "/segmentation/#"
    tmp = "measures[0][event_name]={event_name}&measures[0][aggregator]={aggregator}&" \
          "unit=day&from_date={from_date}&to_date={to_date}"\
            .format(event_name=event_name, aggregator=aggregator,from_date=from_date,to_date=to_date)
    condition_str = ""
    if condition:
        condition = json.loads(condition)
        if 'relation' in condition:
            condition_str += "&filter[relation]={}".format(condition['relation'])
        for i in range(len(condition['conditions'])):
            # 这里当param同时出现几个时未做处理
            condition_str += "&filter[conditions][{}][field]={}" \
                             "&filter[conditions][{}][function]={}" \
                             .format(i, condition['conditions'][i]['field'],
                                     i, condition['conditions'][i]['function'])
            if 'params' in condition['conditions'][i]:
                for param in condition['conditions'][i]['params']:
                    condition_str += "&filter[conditions][{}][params][]={}".format(i,param)
    url = url + tmp + condition_str
    return url

def obtain_event_monitor_data(monitor_id=None):
    print("开始从SensorsData所需监控的数据;")
    events = obtain_monitor_event(monitor_id)
    from_date = (date.today() - timedelta(days=1)).strftime('%Y-%m-%d')
    to_date = (date.today() - timedelta(days=1)).strftime('%Y-%m-%d')
    # 注意从MySQL配置中取回的event顺序
    results = [(id, obtain_event_from_sensors(event_name, aggregator, condition, from_date, to_date))
               for (id, event_name, condition, e, sd, e_prod, desc, aggregator) in events]
    event_result_list= []
    for result in results:
        for event in events:
            if event[0] == result[0] and result[1]['num_rows'] == 1:
                # 拼接跳转的链接url
                link_url = generate_event_url(event[1],event[7], event[2])
                tmp_list = list(event)
                tmp_list[2] = result[1]['rows'][0]['values'][0][0]
                tmp_list.append(link_url)
                event_result_list.append(tuple(tmp_list))
            elif event[0] == result[0]:
                link_url = generate_event_url(event[1], event[7], event[2])
                tmp_list = list(event)
                tmp_list[2] = 0
                tmp_list.append(link_url)
                event_result_list.append(tuple(tmp_list))
    print("SensorsData所需监控的数据完成;")
    return event_result_list

if __name__ == '__main__':
    # 通过输入监控id来进行测试,同时传入测试的邮箱,发送给测试人数据, 如果测试则直接调用calculate_E_and_SD
    try:
        opts, args = getopt.getopt(sys.argv[1:], "i:e:", ["monitor_id=", "email="])
    except getopt.GetoptError:
        print('sd_track_monitor.py -i <monitor_id list> -e <email address>')
        sys.exit(2)

    if opts:
        args = []
        for opt, arg in opts:
            args.append(arg)
    if args:
        if len(args) != 2:
            print("请输入测试监控id和发送测试监控的邮箱。格式sd_track_monitor.py -i <monitor_id list> -e <email address>")
            sys.exit(2)
        print("进入埋点监控测试,测试id为{args[0]},发送测试邮箱为{args[1]}".format(args=args))
        calculate_E_and_SD(update=True, monitor_id=args[0])
        event_list = obtain_event_monitor_data(monitor_id=args[0])
        send_mail(event_list=event_list, email_list=args[1])

    else:
        event_list = obtain_event_monitor_data()
        if date.today().isoweekday() == 1:
            calculate_E_and_SD()
        send_mail(event_list=event_list, email_list="xx@xxx.com,xxx@xxx.com")


import socket

import pyautogui

# Set up a server socket
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind(("0.0.0.0", 12345))
server_socket.listen(1)

print("Server listening...")

# Accept incoming connections
client_socket, client_address = server_socket.accept()
print(f"Connection from {client_address}")

while True:
    command = client_socket.recv(1024).decode()
    if not command:
        break
    if command == "move_cursor":
        # Receive coordinates (x, y)
        x, y = map(int, client_socket.recv(1024).decode().split(","))
        # Move the cursor
        pyautogui.moveTo(x, y)

client_socket.close()
server_socket.close()


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.autograd import Variable
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt

from dataset import BatchPyntCloudToTensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class MLP(nn.Sequential):
    def __init__(self, layer_sizes, dropout=0.7):
        layers = []
        for i in range(len(layer_sizes)-1):
            layers.append(nn.Linear(layer_sizes[i], layer_sizes[i+1]))
            layers.append(nn.BatchNorm1d(layer_sizes[i+1]))
            layers.append(nn.Dropout(dropout))
            layers.append(nn.ReLU())
        super(MLP, self).__init__(*layers)

class SharedMLP(nn.Sequential):
    def __init__(self, layer_sizes):
        layers = []
        for i in range(len(layer_sizes)-1):
            layers.append(nn.Conv1d(layer_sizes[i], layer_sizes[i+1], 1))
            layers.append(nn.BatchNorm1d(layer_sizes[i+1]))
            layers.append(nn.ReLU())
        super(SharedMLP, self).__init__(*layers)

class T_net(nn.Module):
    def __init__(self, size, dropout=0.7, bn_momentum=None):
        super(T_net, self).__init__()
        self.size = size

        self.shared_mlp = SharedMLP([size, 64, 128, 1024])

        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, size*size, bias=False)
        self.fc3.requires_grad_(False)

        self.bn1 = nn.BatchNorm1d(512, momentum=bn_momentum)
        self.bn2 = nn.BatchNorm1d(256, momentum=bn_momentum)

    def forward(self, x):
        '''
            Input: B x size x N
        '''
        out = self.shared_mlp(x)
        out = F.max_pool1d(out, kernel_size=x.size(-1))
        out = out.view(-1, 1024)
        out = F.relu(self.bn1(self.fc1(out)))
        out = F.relu(self.bn2(self.fc2(out)))

        out = self.fc3(out)
        out = out.view(-1, self.size, self.size)
        bias = torch.eye(self.size, requires_grad=True).expand(x.size(0), -1, -1).to(device)
        return out + bias

class InputTransform(nn.Module):
    def __init__(self):
        super(InputTransform, self).__init__()
        self.T_net = T_net(3)

    def forward(self, x):
        out = self.T_net(x)
        return torch.bmm(x.transpose(1, 2), out).transpose(1, 2)

class FeatureTransform(nn.Module):
    def __init__(self):
        super(FeatureTransform, self).__init__()
        self.T_net = T_net(64)

    def forward(self, x):
        out = self.T_net(x)
        self.A = out
        return torch.bmm(x.transpose(1, 2), out).transpose(1, 2)

class ClassificationNN(nn.Module):
    def __init__(self, num_classes):
        super(ClassificationNN, self).__init__()
        self.input_transform = InputTransform()
        self.feature_transform = FeatureTransform()

        self.shared_mlp_1 = SharedMLP([3, 64, 64])
        self.shared_mlp_2 = SharedMLP([64, 64, 128, 1024])
        self.mlp = MLP([1024, 512, 256, num_classes])

    def forward(self, x):
        out = self.input_transform(x)
        out = self.shared_mlp_1(out)
        feat_out = self.feature_transform(out)
        out = self.shared_mlp_2(feat_out)
        out = F.max_pool1d(out, x.size(-1)).view(x.size(0), -1)
        out = self.mlp(out)
        return out, feat_out

    def predict(self, x):
        return self.forward(x)[0]

class PartSegmentationNN(nn.Module):
    def __init__(self, num_features: int):
        super(PartSegmentationNN, self).__init__()
        self.input_transform = InputTransform()
        self.feature_transform = FeatureTransform()

        self.mlp_1 = SharedMLP([3, 64, 64])
        self.mlp_2 = SharedMLP([64, 64, 128, 1024])
        self.mlp_3 = SharedMLP([1088, 512, 256, 128])
        self.mlp_4 = SharedMLP([128, 128, num_features])

    def forward(self, x):
        out = self.input_transform(x)
        out = self.mlp_1(out)
        feat_out = self.feature_transform(feat_out)
        global_feature = self.mlp_2(out)
        global_feature = F.max_pool1d(global_feature, x.size(2))
        global_feature = global_feature.expand(-1, -1, x.size(-1))
        out = torch.cat([out, global_feature], 1)
        out = self.mlp_3(out)
        out = self.mlp_4(out)
        return out, feat_out

    def predict(self, x):
        return self.forward(x)[0]

class SegmentationNN(nn.Module):
    def __init__(self, num_features: int):
        super(SegmentationNN, self).__init__()
        self.input_transform = InputTransform()
        self.feature_transform = FeatureTransform()

        self.mlp_1 = SharedMLP([3, 64, 64])
        self.mlp_2 = SharedMLP([64, 64, 128, 1024])
        self.mlp_3 = SharedMLP([1088, 512, 256, 128])
        self.mlp_4 = SharedMLP([128, 128, num_features])

    def forward(self, x):
        out = self.input_transform(x)
        out = self.mlp_1(out)
        feat_out = self.feature_transform(feat_out)
        global_feature = self.mlp_2(out)
        global_feature = F.max_pool1d(global_feature, x.size(2))
        global_feature = global_feature.expand(-1, -1, x.size(-1))
        out = torch.cat([out, global_feature], 1)
        out = self.mlp_3(out)
        out = self.mlp_4(out)
        return out, feat_out

    def predict(self, x):
        return self.forward(x)[0]

def loss_fn(preds, labels, feature_transform, reg=0.0001):
    loss = torch.nn.NLLLoss()
    def feat_loss(A):
        I = torch.eye(64, requires_grad=True).expand(A.size(0), -1, -1).to(device)
        AA_T = torch.bmm(A, A.transpose(1, 2))
        return torch.linalg.norm(I - AA_T, ord='fro', dim=(1,2))
    return loss(preds, labels) + reg * torch.mean(feat_loss(feature_transform))

def train(model, trainset, validset, optimizer, epochs=10, batch_size=32, device=torch.device('cpu')):
    train_dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
    valid_dataloader = DataLoader(validset, batch_size=batch_size, shuffle=False)

    train_losses, train_accs = [], []
    valid_losses, valid_accs = [], []

    model = model.to(device)

    for epoch in range(epochs):
        model.train()
        # Train and get training loss and accuracy
        train_loss, train_num_correct = [], []
        for x, y in tqdm(train_dataloader, unit='batch'):
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            pred, feat = model(x)
            loss = loss_fn(pred, y, feat)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
            train_num_correct.append(torch.sum(pred.argmax(1) == y).item())
        train_losses.append(np.mean(train_loss))
        train_accs = np.asarray(train_num_correct) / len(trainset)

        model.eval()
        # Get validation loss and accuracy
        with torch.no_grad():
            valid_loss, valid_num_correct = [], []
            for x, y in tqdm(valid_dataloader, unit='batch'):
                x, y = x.to(device), y.to(device)
                pred, feat = model(x)
                loss = loss_fn(pred, y, feat)
                valid_loss.append(loss.item())
                valid_num_correct.append(torch.sum(pred.argmax(1) == y).item())

        print('Finished Epoch {}\n training loss: {}, validation loss: {} \n training accuracy: {}, validation accuracy: {}'
                .format(epoch+1, train_losses[-1], valid_losses[-1], train_accs[-1], valid_accs[-1]))

    return train_losses, valid_losses, train_accs, valid_accs

def plot_stats(epochs, train_losses, valid_losses, train_accs, valid_accs):
    epochs = range(0, self.params['epochs'] + 1)
    _, axs = plt.subplots(1, 2, layout='constrained', sharex=True, sharey=True)
    axs[0].plot(epochs, [0] + train_losses, label='training')
    axs[0].plot(epochs, [0] + valid_losses, label='validation')
    axs[0].title('Epochs vs Loss')
    axs[0].xlabel('Epochs')
    axs[0].ylabel('Loss')
    axs[0].legend()

    axs[1].plot(epochs, [0] + train_accs, label='training')
    axs[1].plot(epochs, [0] + valid_accs, label='validation')
    axs[1].title('Epochs vs Accuracy')
    axs[1].xlabel('Epochs')
    axs[1].ylabel('Accuracy')
    axs[1].legend()

if __name__ == '__main__':
    from pytorch_model_summary import summary
    rand_data = Variable(torch.rand(32, 3, 2000)) # B X 3 X N
    shared_mlp = SharedMLP([3, 64, 64])
    out = shared_mlp(rand_data)
    print(summary(shared_mlp, rand_data), '\n', out.size()) # B X 64 X N

    rand_data = Variable(torch.rand(32, 3, 2000)) # B X 3 X N
    input_transform = InputTransform()
    out = input_transform(rand_data)
    print(summary(input_transform, rand_data), '\n', out.size()) # B X 3 X 3

    rand_data = Variable(torch.rand(32, 64, 2000)) # B X 64 X N
    input_transform = FeatureTransform()
    out = input_transform(rand_data)
    print(summary(input_transform, rand_data), '\n', out.size()) # B X 64 X 64

    rand_data = Variable(torch.rand(32, 3, 2000)) # B X 3 X N
    cls_net = ClassificationNN(num_classes=10)
    out = cls_net(rand_data)
    print(summary(cls_net, rand_data),' \n', out.size())

    rand_data = Variable(torch.rand(32, 3, 2000)) # B X 3 X N
    seg_net = SegmentationNN(num_features=10)
    out = seg_net(rand_data)
    print(summary(seg_net, rand_data), '\n', out.size())


from pathlib import Path
from queue import PriorityQueue
from typing import Set, Tuple, List

import numpy as np
import numpy.typing as npt

from hw1.utils import neighbors, plot_GVD, PathPlanMode, distance

def cell_to_GVD_gradient_ascent(
    grid: npt.ArrayLike, GVD: Set[Tuple[int, int]], cell: Tuple[int, int]
) -> List[Tuple[int, int]]:
    """Find the shortest path from any cell in the enviroment to a cell on the
    GVD using gradient ascent.
    Args:
        grid (numpy): NxN numpy array representing the world, with obstacles,
        walls, and the distance from each cell to the obstacles.
        GVD (set[tuple]): A set of tuples containing the cells on the GVD.
        cell (tuple): The starting/ending cell of this path.
    Returns:
        list<tuple>: list of tuples of the path.
    """

    path = [cell]

# initially set the highest value to the starting/ending cell of path
    high_value = cell

# loop will run until we hit a value in GVD
    while high_value not in GVD:
        # setting the first and second values of the tuple to an
        # individual variable so that i can use them in neighbors function
        cell_i, cell_j = high_value
        # iterate through each neighbor that neighbor function gives me and check if they are
        # than current highest value. If so set to high_value variable. Append each iteration of
        # high_value to path list and return path once we hot GVD and while loop is broken
        for neighbor in neighbors(grid, cell_i, cell_j):
            if neighbor > high_value:
                high_value=neighbor

            path.append(high_value)

    return path

def cell_to_GVD_a_star(
    grid: npt.ArrayLike, GVD: Set[Tuple[int, int]], cell: Tuple[int, int],
    goal: Tuple[int, int]
) -> List[Tuple[int, int]]:
    """Find the shortest path from any cell in the enviroment to the GVD using
    A* with L2 distance heurstic.
    Args:
        grid (numpy): NxN numpy array representing the world, with obstacles,
        walls, and the distance from each cell to the obstacles.
        GVD (set<tuple>): A set of tuples containing the cells on the GVD.
        cell (tuple): The starting/ending cell of this path.
    Returns:
        list[tuple], dict, list[int]: list of tuples of the path, and the reached
        dictionary, and the list of frontier sizes.
    """

    # define a priority queue
    frontier = PriorityQueue()
    frontier.put((0, cell))
    frontier_size = [0]

    # construct a reached table using python dictionary. The key is the (x, y)
    # tuple of the cell position, the value is dictiionary with the cell's cost,
    # and cell parent.
    reached = {cell: {"cost": 0, "parent": None}}

    # created a while loop that lasts until frontier is empty or a path to GVD is found

    while not frontier.empty():
        # intialized current value from frontier and i and j coordinates from current tuple
        low_f, current = frontier.get()
        i_cord, j_cord = current
        # check if we have hit GVD if current tuple is in GVD set. If so create a path to cell
        #  by going back through parent history in reached table tell we hit given start/end point.
        if current in GVD:
            path=[]
            while current !=cell:
                path.append(current)
                current = reached[current]["parent"]
            path.append(cell)
            return path[::-1], reached, frontier_size

        # for loop to go through each neighbor around current cell. Created cost function that increases
        # based on distance from current cell to neighbor.
        for neighbor in neighbors(grid, i_cord, j_cord):
            cost = reached[current]["cost"]+distance(current, neighbor)
            # check to see if neighbor is in reached table or if current cost is less than neighbor cost
            # if check is passed, we add neighbor to reached table with cost and parent.
            # Then calculated f function where g = cost and h = neighbor distance to goal and set
            # it to the priority variable. This was then put into the frontier
            if neighbor not in reached or cost < reached[neighbor]["cost"]:
                reached[neighbor] = {"cost": cost, "parent":current}
                priority = cost + distance(neighbor, goal)
                frontier.put((priority, neighbor))
                frontier_size.append(frontier.qsize())

    return None, reached, frontier_size

def GVD_path(
    grid: npt.ArrayLike,
    GVD: Set[Tuple[int, int]],
    A: Tuple[int, int],
    B: Tuple[int, int],
    mode: PathPlanMode
) -> List[Tuple[int, int]]:
    """Find the shortest path between two points on the GVD using
    Breadth-First-Search and DFS
    Args:
        grid (numpy): NxN numpy array representing the world, with obstacles,
        walls, and the distance from each cell to the obstacles.
        A (tuple): The starting cell of the path.
        B (tuple): The ending cell of the path.
    Returns:
        list[tuple], dict, list[int]: return the path, pointers, and frontier
        size array.
    """

    # the set of cells on the GVD
    GVD = set(GVD)

    # the set of visited cells
    closed = set([])

    # the set of cells on the current frontier
    frontier = [A]

    # back pointers to find the path once reached the goal B. The keys
    # should both be tuples of cell positions (x, y)
    pointers = {}

    # the length of the frontier array, update this variable at each step.
    frontier_size = [0]

    # while loop that runs as long as frontier is not empty. Check if mode is DFS or BFS.
    # Intiate while loop that goes through frontier and pops a node which is then set to current variable
    # with the i and j coordinates pulled out from it to use in neighbors function. While loop breaks if current
    # node is ever equal to B/goal tuple.
    while len(frontier) > 0:
        if mode == PathPlanMode.DFS:
            pointers[A] = None
            while frontier:
                current=frontier.pop()
                i_cord, j_cord = current
                if current == B:
                    break
                # For each neighbor of current node check if they are not in pointers and in GVD.
                # If so add each neighbor to frontier and add current frontier size to frontier_size list
                # Add each neighbor to closed list, so we know what cells we visited. Break for loop once we find B
                for neighbor in neighbors(grid, i_cord, j_cord):
                    if neighbor not in pointers and neighbor in GVD:
                        frontier.append(neighbor)
                        pointers[neighbor] = current
                        frontier_size.append(len(frontier))
                        closed.add(neighbor)
                        if neighbor == B:
                            break
            # Once we find B, initialize path that starts at B and goes back through each pointer and
            # adds it to path list. Once we reach point A, reverse the list and we have our path from A to B
            path = []
            if B in pointers:
                node=B
                while node !=A:
                    path.append(node)
                    node= pointers[node]
                path.append(A)
                path.reverse()
            else:
                path = None

        # BFS was implemented the same as above but we always pop the first element in our frontier
        elif mode == PathPlanMode.BFS:
            pointers[A] = None
            while frontier:
                current = frontier.pop(0)
                i_cord, j_cord = current
                if current == B:
                    break
                for neighbor in neighbors(grid, i_cord, j_cord):
                    if neighbor not in pointers and neighbor in GVD:
                        frontier.append(neighbor)
                        pointers[neighbor] = current
                        frontier_size.append(len(frontier))
                        closed.add(neighbor)
                        if neighbor == B:
                            break
            path = []
            if B in pointers:
                node= B
                while node != A:
                    path.append(node)
                    node = pointers[node]
                path.append(A)
            else:
                path = None

    return path, pointers, frontier_size

def compute_path(
    grid,
    GVD: set[tuple],
    start: tuple,
    goal: tuple,
    outmode: PathPlanMode = PathPlanMode.GRAD,
    inmode: PathPlanMode = PathPlanMode.DFS):

    """ Compute the path on the grid from start to goal using the methods
    implemented in this file.
    Returns:
        list: a list of tuples represent the planned path.
    """

    if outmode == PathPlanMode.GRAD:
        start_path = cell_to_GVD_gradient_ascent(grid, GVD, start)
        end_path = list(reversed(cell_to_GVD_gradient_ascent(grid, GVD, goal)))
    else:
        start_path = cell_to_GVD_a_star(grid, GVD, start, goal)[0]
        end_path = list(reversed(cell_to_GVD_a_star(grid, GVD, goal, start)[0]))
    mid_path, reached, frontier_size = GVD_path(
        grid, GVD, start_path[-1], end_path[0], inmode)
    return start_path + mid_path[1:-1] + end_path

def test_world(
    world_id,
    start,
    goal,
    outmode: PathPlanMode = PathPlanMode.GRAD,
    inmode: PathPlanMode = PathPlanMode.DFS,
    world_dir="worlds"):

    print(f"Testing world {world_id} with modes {inmode} and {outmode}")
    grid = np.load(f"{world_dir}/world_{world_id}.npy")
    GVD = set([tuple(cell) for cell in np.load(
        f"{world_dir}/world_{world_id}_gvd.npy")])
    path = compute_path(grid, GVD, start, goal, outmode=outmode, inmode=inmode)
    print(f"Path length: {len(path)} steps")
    plot_GVD(grid, world_id, GVD, path)