#from localflavor.br.models import BRCPFField, BRCNPJField
from django.db import models
from django.core.mail import send_mail
from model_utils.models import TimeStampedModel

class AreopagusBusiness(TimeStampedModel):
    name = models.TextField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Areopagus Business'
        verbose_name_plural = 'Areopagus Business'

    def __str__(self):
        return self.name

class Type(models.Model):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Type'
        verbose_name_plural = 'Typies'

    def __str__(self):
        return self.name

class CommunicationTool(TimeStampedModel):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Communication Tool'
        verbose_name_plural = 'Communication Tools'

    def __str__(self):
        return self.name

class Collaborator(TimeStampedModel):
    name = models.CharField(max_length=255)
    areopagus_business_id = models.ForeignKey(AreopagusBusiness, on_delete=models.CASCADE)
    cpf = models.CharField(max_length=11, default=0)
    cnpj = models.CharField(max_length=14, default=0)
    rg = models.CharField(max_length=255)
    born_date = models.DateTimeField()
    type_id = models.ForeignKey(Type, on_delete=models.CASCADE)

    class Meta:
        ordering = ['id']
        verbose_name = 'Collaborator'
        verbose_name_plural = 'Collaborators'

    def __str__(self):
        return self.name

class ValueCollaborator(TimeStampedModel):
    value = models.FloatField()
    collaborator_id = models.ForeignKey(Collaborator, on_delete=models.CASCADE)
    date_end = models.DateTimeField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Value Collaborator'
        verbose_name_plural = 'Collaborator Values '

    def __str__(self):
        return self.value

class Justify(TimeStampedModel):
    collaborator_id = models.ForeignKey(Collaborator, on_delete=models.CASCADE)
    description = models.TextField()
    date_time = models.DateTimeField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Justify'
        verbose_name_plural = 'Justifies'
    def __str__(self):
        return self.description

class Estate(TimeStampedModel):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Estate'
        verbose_name_plural = 'Estates'

    def __str__(self):
        return self.name

class StatusContract(TimeStampedModel):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Status Contract'
        verbose_name_plural = 'Status Contracts'

    def __str__(self):
        return self.name

class Contract(TimeStampedModel):
    cnpj = models.CharField(max_length=255, default=0)
    playlist_link = models.CharField(max_length=500)
    meet_id = models.IntegerField()
    tech_lead_name = models.CharField(max_length=255)
    communication_tool_id = models.ForeignKey(CommunicationTool, on_delete=models.CASCADE)
    coorporative_email = models.CharField(max_length=255)
    estate_id = models.ForeignKey(Estate, on_delete=models.CASCADE)
    status_contract_id = models.ForeignKey(StatusContract, on_delete=models.CASCADE)
    begin_date_contract = models.DateTimeField()
    date_end_contract = models.DateTimeField()
    hours_control_link = models.TextField()
    is_full_time = models.BooleanField(default=False)
    is_international = models.BooleanField(default=False)
    value = models.FloatField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Contract'
        verbose_name_plural = 'Contracts'

    def __str__(self):
        return self.cnpj

class CollaboratorContract(TimeStampedModel):
    contract_id = models.ForeignKey(Contract, on_delete=models.CASCADE)
    collaborator_id = models.ForeignKey(Collaborator, on_delete=models.CASCADE)
    status = models.IntegerField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Collaborator Contract'
        verbose_name_plural = 'Collaborators Contract'

    def __str__(self):
        return self.status

class TaskTool(TimeStampedModel):
    name = models.CharField(max_length=500)

    class Meta:
        ordering = ['id']
        verbose_name = 'Task Tool'
        verbose_name_plural = 'Task Tools'

    def __str__(self):
        return self.name

class TaskToolContract(TimeStampedModel):
    contract_id = models.ForeignKey(Contract, on_delete=models.CASCADE)
    task_tool_id = models.ForeignKey(TaskTool, on_delete=models.CASCADE)
    link = models.CharField(max_length=500, default="")

    class Meta:
        ordering = ['id']
        verbose_name = 'Task Tool Contract'
        verbose_name_plural = 'Task Tool Contracts'

    def __str__(self):
        return self.link

class SourceCodeTool(TimeStampedModel):
    name = models.CharField(max_length=500)

    class Meta:
        ordering = ['id']
        verbose_name = 'Source Code Tool'
        verbose_name_plural = 'Source Code Tools'

    def __str__(self):
        return self.name

class SourceCodeToolContract(TimeStampedModel):
    contract_id = models.ForeignKey(Contract, on_delete=models.CASCADE)
    source_code_tool_id = models.ForeignKey(SourceCodeTool, on_delete=models.CASCADE)
    link = models.CharField(max_length=500, default="")

    class Meta:
        ordering = ['id']
        verbose_name = 'Source Code Tool Contract'
        verbose_name_plural = 'Source Code Tool Contracts'

    def __str__(self):
        return self.link

class StatusProject(TimeStampedModel):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Status Project'
        verbose_name_plural = 'Status Projects'

    def __str__(self):
        return self.name

class Project(TimeStampedModel):
    name = models.CharField(max_length=255)
    contract_id = models.ForeignKey(Contract, on_delete=models.CASCADE)
    status_project_id = models.ForeignKey(StatusProject, on_delete=models.CASCADE)
    documentaion_id = models.IntegerField()
    doc_link = models.TextField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Project'
        verbose_name_plural = 'Projects'

    def __str__(self):
        return self.name

class StatusTask(TimeStampedModel):
    name = models.CharField(max_length=255)

    class Meta:
        ordering = ['id']
        verbose_name = 'Status Task'
        verbose_name_plural = 'Status Tasks'

    def __str__(self):
        return self.name

class Repository(TimeStampedModel):
    link = models.CharField(max_length=500, default="")
    ssh_link = models.CharField(max_length=500, default="")
    project_id = models.ForeignKey(Project, on_delete=models.CASCADE)

    class Meta:
        ordering = ['id']
        verbose_name = 'Repository'
        verbose_name_plural = 'Repositories'

    def __str__(self):
        return self.link

class Commits(TimeStampedModel):
    task_repository_id = models.ForeignKey(Repository, on_delete=models.CASCADE)
    description = models.TextField()
    date_time = models.DateTimeField()
    collaborator_id = models.ForeignKey(Collaborator, on_delete=models.CASCADE)

    class Meta:
        ordering = ['id']
        verbose_name = 'Commit'
        verbose_name_plural = 'Commits'

    def __str__(self):
        return self.description

class Task(TimeStampedModel):
    title = models.CharField(max_length=500)
    description = models.TextField()
    date_in_todo = models.DateTimeField()
    date_in_progress = models.DateTimeField()
    date_in_review = models.DateTimeField()
    date_done = models.DateTimeField()
    date_blocked = models.DateTimeField()
    date_expected = models.DateTimeField()
    status_task_id = models.ForeignKey(StatusTask, on_delete=models.CASCADE)
    project_id = models.ForeignKey(Project, on_delete=models.CASCADE)
    estimate = models.IntegerField()
    link_branch = models.CharField(max_length=500, default="")
    pr_link = models.CharField(max_length=500, default="")
    total_commits = models.IntegerField()
    link_video = models.CharField(max_length=500, default="")

    class Meta:
        ordering = ['id']
        verbose_name = 'Task'
        verbose_name_plural = 'Tasks'

    def __str__(self):
        return self.title

class TaskRepository(TimeStampedModel):
    task_id = models.ForeignKey(Task, on_delete=models.CASCADE)
    repository_id = models.ForeignKey(Repository, on_delete=models.CASCADE)
    total_commits = models.IntegerField()

    class Meta:
        ordering = ['id']
        verbose_name = 'Task Repository'
        verbose_name_plural = 'Task Repositories'

    def __str__(self):
        return self.total_commits

class Calendar(TimeStampedModel):
    date = models.DateField()
    collaborator_id = models.ForeignKey(Collaborator, on_delete=models.CASCADE)
    task_id = models.ForeignKey(Task, on_delete=models.CASCADE)
    status = models.IntegerField()
    is_english = models.BooleanField(default=False)
    is_levi_signature = models.BooleanField(default=False)
    is_pr = models.BooleanField(default=False)

    class Meta:
        ordering = ['id']
        verbose_name = 'Calendar'
        verbose_name_plural = 'Calendars'

    def __str__(self):
        return self.date


import numpy as np
from scipy.signal import gammatone
from matplotlib import pyplot as plt

from modules.normalization import max_norm

def make_filterbank(n_filters, kernel_size):
    return np.stack([
        gammatone(
            freq=x,
            ftype='fir',
            order=4,
            numtaps=kernel_size,
            fs=22050)[0] for x in np.geomspace(20, 11000, n_filters)
        ], axis=0)

if __name__ == '__main__':
    bank = make_filterbank(64, 512)
    bank = bank / (bank.max(axis=1, keepdims=True) + 1e-8)
    spec = np.abs(np.fft.rfft(bank, axis=-1, norm='ortho'))

    for filt in spec:
        plt.plot(filt)

    plt.show()


import numpy as np
import pygame
from time import time
from utils import calculate_fps
from flappy import FlappyBird, FlappyFlock, FlappyBrain
from constants import WIDTH_BETWEEN_PIPES, HEIGHT_BETWEEN_PIPES, STARTING_RANGE, \
    SCREEN_SIZE, MAP_MOVEMENT, RANGE_INCREMENT, S_HEIGHT, S_WIDTH, SCORE_Y, SCORE_SPACING, \
    GAME_HEIGHT, PIPE_HEIGHT, MAX_RANGE
from position import Position

class Simulator:

    def __init__(self, model, display, window, fps=None):
        self.caption = "AI Flappy Bird!"
        self.fps = fps
        self.clock = pygame.time.Clock()
        self.display = display
        self.model = model
        self.current_timestamp = time()
        self.calc_fps = 0
        self.window = window
        self.game_over = False

    def start(self):
        run = True
        pygame.display.set_caption(self.caption)
        # - - - - - - - - - - - - - - - - - - - - - - - - - -
        while run:
            if self.fps is not None and run:
                self.clock.tick(self.fps)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    run = False
                    if self.model is not None:
                        self.model.handle_close_event()
                        self.quit()
                    break
            t = self.current_timestamp
            self.current_timestamp = time()
            if t is not None:
                self.calc_fps = calculate_fps(max(self.current_timestamp - t, 1))
            keys_pressed = pygame.key.get_pressed()
            run = self.model.update_state(keys_pressed) and run
            self.update_display()
        pygame.quit()

    def update_display(self):
        self.display.render(self.window, self.model, self.game_over)
        pygame.display.update()

    def reset(self):
        pass

    def quit(self):
        pass

class Game:

    def __init__(self, bird):
        """
        Model for the flappy bird game
        """
        # Every 10 levels increase the starting range
        self.pipe_range = STARTING_RANGE
        self.score = 0
        self.height = S_WIDTH
        self.width = S_HEIGHT
        self.bird = bird
        self.pipe_manager = PipeManager(S_WIDTH, S_HEIGHT, self.pipe_range)

    def init(self):
        pass

    def update_state(self, keys_pressed):
        self.pipe_manager.update(self.score)
        self.bird.score(self.handle_score())
        collision = self.handle_collision()
        self.bird.dead = collision
        inputs = self.get_inputs()
        self.bird.update(inputs, keys_pressed)
        if collision:
            self.reset()
        return True

    def get_inputs(self):
        nearest_pipe = self.pipe_manager.get_current_pipe()
        assert nearest_pipe is not None, "Nearest Pipe is None"
        distance = nearest_pipe.upper_rect.x - self.bird.rect.x
        bird_height = self.bird.rect.centery
        upper = nearest_pipe.upper_rect.bottom
        lower = nearest_pipe.lower_rect.top
        return np.array([distance, bird_height, upper, lower])

    def compare_position(self):
        pass

    def handle_score(self):
        pipe = self.pipe_manager.get_current_pipe()
        if self.bird.rect.x >= pipe.x and not pipe.passed:
            pipe.passed = True
            self.score += 1

    def handle_collision(self):
        return self.pipe_manager.handle_collision(self.bird)

    def reset(self):
        self.bird.reset()

    def handle_close_event(self):
        pass

class GeneticGame:

    def __init__(self, size, mini_batch_size):
        """
        Model for the flappy bird game
        """
        # Every 10 levels increase the starting range
        self.pipe_range = STARTING_RANGE
        self.score = 0
        self.height = S_WIDTH
        self.width = S_HEIGHT
        self.flock = FlappyFlock(size, mini_batch_size)
        self.pipe_manager = PipeManager(S_WIDTH, S_HEIGHT, self.pipe_range)

    def init(self):
        pass

    def update_state(self, keys_pressed):
        self.flock.update(keys_pressed, self)
        return True

    def reset(self):
        self.pipe_range = STARTING_RANGE
        self.score = 0
        self.height = S_WIDTH
        self.width = S_HEIGHT
        self.pipe_manager = PipeManager(S_WIDTH, S_HEIGHT, self.pipe_range)

    def handle_close_event(self):
        pass

class PipeManager:

    def __init__(self, width, height, starting_range, green=True):
        self.pipes = []
        self.top_mask = None
        self.bottom_mask = None
        self.num_pipes = 0
        self.width = width
        self.height = height
        self.range = starting_range
        self.green = green
        self.generate_pipes(0)
        self.init()

    def render(self, window):
        for pipe in self.pipes:
            pipe.render(window)

    def init(self):
        s = PipeSet.generate_set(0, 0)
        self.top_mask = pygame.mask.from_surface(s.top_pipe)
        self.bottom_mask = pygame.mask.from_surface(s.bottom_pipe)

    def generate_pipes(self, score):
        if score % 10 == 0 and score != 0:
            self.range += RANGE_INCREMENT
            self.range = min(self.range, MAX_RANGE)
        self.pipes.append(PipeSet.generate_set(self.width, self.range))
        self.num_pipes += 1

    def update(self, score):
        # check the last pipe to see if we need to generate another set
        if self.width - self.pipes[self.num_pipes - 1].x >= WIDTH_BETWEEN_PIPES:
            self.generate_pipes(score)
        delete = False
        for i, pipe in enumerate(self.pipes):
            delete = pipe.update() or delete
        if delete:
            self.num_pipes -= 1
            self.pipes.pop(0)

    def get_current_pipe(self):
        return self.pipes[0]

    def handle_collision(self, bird):
        bird_x = bird.current_image.get_width()
        pipe = self.find_closest_pipe(bird_x)
        bird_mask = pygame.mask.from_surface(bird.current_image)
        # check mask
        # else check if height is above the cutoff or below the cutoff
        offset = (bird.rect.x - pipe.upper_rect.x), (bird.rect.y - pipe.upper_rect.y)
        c_upper = self.top_mask.overlap(bird_mask, offset)
        offset = (bird.rect.x - pipe.lower_rect.x), (bird.rect.y - pipe.lower_rect.y)
        c_lower = self.bottom_mask.overlap(bird_mask, offset)
        if c_upper is not None \
                or c_lower is not None \
                or bird.rect.y >= GAME_HEIGHT \
                or (pipe.in_between(bird_x) and bird.rect.y < pipe.upper_rect.y + PIPE_HEIGHT):
            return True
        return False

    def find_closest_pipe(self, bird_x):
        for pipe in self.pipes:
            if bird_x < pipe.x or pipe.in_between(bird_x):
                return pipe
        return None

class PipeSet:

    def __init__(self, x, y_upper, y_lower, green=True):
        self.green_pipe = pygame.image.load("./assets/sprites/pipe-green.png").convert_alpha()
        self.red_pipe = pygame.image.load("./assets/sprites/pipe-red.png").convert_alpha()
        self.bottom_pipe = self.green_pipe if green else self.red_pipe
        self.top_pipe = pygame.transform.rotate(self.bottom_pipe.copy(), 180)
        self.x = x
        self.upper_rect = pygame.Rect(
            x,
            y_upper - self.bottom_pipe.get_height(),
            self.bottom_pipe.get_width(),
            self.bottom_pipe.get_height()
        )
        self.lower_rect = pygame.Rect(
            x,
            y_lower,
            self.bottom_pipe.get_width(),
            self.bottom_pipe.get_height()
        )
        self.passed = False

    def update(self):
        delete = False
        self.x -= MAP_MOVEMENT
        self.upper_rect.move_ip(-MAP_MOVEMENT, 0)
        self.lower_rect.move_ip(-MAP_MOVEMENT, 0)
        if self.x + self.lower_rect.width <= 0:
            delete = True
        return delete

    def render(self, window):
        window.blit(self.bottom_pipe, self.lower_rect)
        window.blit(self.top_pipe, self.upper_rect)

    def in_between(self, x):
        x2 = self.x + self.bottom_pipe.get_width()
        return x2 > x > self.x

    @staticmethod
    def generate_set(width, pipe_range):
        """
        :return: a top Pipe and a lower pipe
        """
        r = np.random.rand() * pipe_range
        r = (-r if np.random.rand() >= .4 else r) * (.4 * GAME_HEIGHT)
        x = width
        top_pipe = (.5 * GAME_HEIGHT) + r - (.5 * HEIGHT_BETWEEN_PIPES)
        bottom_pipe = (.5 * GAME_HEIGHT) + r + (.5 * HEIGHT_BETWEEN_PIPES)
        return PipeSet(x, top_pipe, bottom_pipe)

class GameBoard:

    def __init__(self, board_width, day=True):
        self.background_day = pygame.image.load("./assets/sprites/background-day.png").convert()
        self.background_night = pygame.image.load("./assets/sprites/background-night.png").convert()
        self.game_over = pygame.image.load("./assets/sprites/gameover.png").convert_alpha()
        self.start_message = pygame.image.load("./assets/sprites/message.png").convert_alpha()
        self.bg = self.background_day if day else self.background_night
        self.ground = Ground(board_width)
        self.game_height = self.bg.get_height() - self.ground.ground.get_width()
        self.score_board = ScoreBoard(0, 0)

    def render(self, window, model, game_over):
        size = window.get_size()
        window.blit(self.bg, (0, 0))
        ground, x = self.ground.step()
        model.pipe_manager.render(window)
        window.blit(ground, (-x, self.bg.get_height() - ground.get_height()))
        self.score_board.render(window, model.score)
        model.bird.render(window)

class GeneticGameBoard(GameBoard):

    def __init__(self, board_width, day=True):
        super().__init__(board_width, day=day)

    def render(self, window, model, game_over):
        window.blit(self.bg, (0, 0))
        ground, x = self.ground.step()
        model.pipe_manager.render(window)
        window.blit(ground, (-x, self.bg.get_height() - ground.get_height()))
        self.score_board.render(window, model.score)
        model.flock.render(window)

class Ground:

    def __init__(self, board_width):
        self.x = 0
        self.board_width = board_width
        self.ground = pygame.image.load("./assets/sprites/base.png").convert()

    def step(self):
        if self.x + MAP_MOVEMENT >= (self.ground.get_width() - self.board_width) - 10:
            self.x = 0
        self.x += MAP_MOVEMENT
        return self.ground, self.x

class ScoreBoard:

    def __init__(self, x, y, scale=1):
        """
        :param x:
        :param y:
        :param scale:
        :param height:
        :param width:
        """
        self.scale = scale
        self.x = x
        self.y = y
        self.number_images = self.init_number_images()

    def init_number_images(self):
        images = []
        for i in range(10):
            image = pygame.image.load("./assets/sprites/{}.png".format(i)).convert_alpha()
            size = image.get_size()
            images.append(pygame.transform.smoothscale(image, (size[0] * self.scale, size[1] * self.scale)))
        return images

    def render(self, window, score):
        images = self.get_images(score)
        t_width = 0
        for image in images:
            t_width += image.get_width()
        t_width += (SCORE_SPACING * len(images))
        offset = (S_WIDTH - t_width) * .5
        for image in images:
            window.blit(image, (offset, SCORE_Y))
            offset += image.get_width() + SCORE_SPACING

    def get_images(self, score):
        score = str(score)
        ret = []
        for char in score:
            ret.append(self.number_images[int(char)])
        return ret


import dataclasses
import gc
import inspect
import logging
import os
from dataclasses import asdict, dataclass
from typing import Callable, List, NamedTuple, Optional, Sequence, Union

import numpy as np
import optuna
import torch
from optuna.pruners import HyperbandPruner
from optuna.samplers import TPESampler
from optuna.visualization import plot_optimization_history, plot_param_importances
from torch.utils.tensorboard.writer import SummaryWriter

import wandb
from rl_algo_impls.a2c.optimize import sample_params as a2c_sample_params
from rl_algo_impls.runner.config import Config, EnvHyperparams, RunArgs
from rl_algo_impls.runner.running_utils import (
    ALGOS,
    base_parser,
    get_device,
    hparam_dict,
    load_hyperparams,
    make_policy,
    set_seeds,
)
from rl_algo_impls.shared.callbacks import Callback
from rl_algo_impls.shared.callbacks.hyperparam_transitions import HyperparamTransitions
from rl_algo_impls.shared.callbacks.optimize_callback import (
    Evaluation,
    OptimizeCallback,
    evaluation,
)
from rl_algo_impls.shared.callbacks.reward_decay_callback import RewardDecayCallback
from rl_algo_impls.shared.callbacks.self_play_callback import SelfPlayCallback
from rl_algo_impls.shared.stats import EpisodesStats
from rl_algo_impls.shared.vec_env import make_env, make_eval_env
from rl_algo_impls.wrappers.self_play_wrapper import SelfPlayWrapper
from rl_algo_impls.wrappers.vectorable_wrapper import find_wrapper

@dataclass
class StudyArgs:
    load_study: bool
    study_name: Optional[str] = None
    storage_path: Optional[str] = None
    n_trials: int = 100
    n_jobs: int = 1
    n_evaluations: int = 4
    n_eval_envs: int = 8
    n_eval_episodes: int = 16
    timeout: Union[int, float, None] = None
    wandb_project_name: Optional[str] = None
    wandb_entity: Optional[str] = None
    wandb_tags: Sequence[str] = dataclasses.field(default_factory=list)
    wandb_group: Optional[str] = None
    virtual_display: bool = False

class Args(NamedTuple):
    train_args: Sequence[RunArgs]
    study_args: StudyArgs

def parse_args() -> Args:
    parser = base_parser()
    parser.add_argument(
        "--load-study",
        action="store_true",
        help="Load a preexisting study, useful for parallelization",
    )
    parser.add_argument("--study-name", type=str, help="Optuna study name")
    parser.add_argument(
        "--storage-path",
        type=str,
        help="Path of database for Optuna to persist to",
    )
    parser.add_argument(
        "--wandb-project-name",
        type=str,
        default="rl-algo-impls-tuning",
        help="WandB project name to upload tuning data to. If none, won't upload",
    )
    parser.add_argument(
        "--wandb-entity",
        type=str,
        help="WandB team. None uses the default entity",
    )
    parser.add_argument(
        "--wandb-tags", type=str, nargs="*", help="WandB tags to add to run"
    )
    parser.add_argument(
        "--wandb-group", type=str, help="WandB group to group trials under"
    )
    parser.add_argument(
        "--n-trials", type=int, default=100, help="Maximum number of trials"
    )
    parser.add_argument(
        "--n-jobs", type=int, default=1, help="Number of jobs to run in parallel"
    )
    parser.add_argument(
        "--n-evaluations",
        type=int,
        default=4,
        help="Number of evaluations during the training",
    )
    parser.add_argument(
        "--n-eval-envs",
        type=int,
        default=8,
        help="Number of envs in vectorized eval environment",
    )
    parser.add_argument(
        "--n-eval-episodes",
        type=int,
        default=16,
        help="Number of episodes to complete for evaluation",
    )
    parser.add_argument("--timeout", type=int, help="Seconds to timeout optimization")
    parser.add_argument(
        "--virtual-display", action="store_true", help="Use headless virtual display"
    )
    # parser.set_defaults(
    #     algo=["a2c"],
    #     env=["CartPole-v1"],
    #     seed=[100, 200, 300],
    #     n_trials=5,
    #     virtual_display=True,
    # )
    train_dict, study_dict = {}, {}
    for k, v in vars(parser.parse_args()).items():
        if k in inspect.signature(StudyArgs).parameters:
            study_dict[k] = v
        else:
            train_dict[k] = v

    study_args = StudyArgs(**study_dict)
    # Hyperparameter tuning across algos and envs not supported
    assert len(train_dict["algo"]) == 1
    assert len(train_dict["env"]) == 1
    train_args = RunArgs.expand_from_dict(train_dict)

    if not all((study_args.study_name, study_args.storage_path)):
        hyperparams = load_hyperparams(train_args[0].algo, train_args[0].env)
        config = Config(train_args[0], hyperparams, os.getcwd())
        if study_args.study_name is None:
            study_args.study_name = config.run_name(include_seed=False)
        if study_args.storage_path is None:
            study_args.storage_path = (
                f"sqlite:///{os.path.join(config.runs_dir, 'tuning.db')}"
            )
    # Default set group name to study name
    study_args.wandb_group = study_args.wandb_group or study_args.study_name

    return Args(train_args, study_args)

def objective_fn(
    args: Sequence[RunArgs], study_args: StudyArgs
) -> Callable[[optuna.Trial], float]:
    def objective(trial: optuna.Trial) -> float:
        if len(args) == 1:
            return simple_optimize(trial, args[0], study_args)
        else:
            return stepwise_optimize(trial, args, study_args)

    return objective

def simple_optimize(trial: optuna.Trial, args: RunArgs, study_args: StudyArgs) -> float:
    base_hyperparams = load_hyperparams(args.algo, args.env)
    base_config = Config(args, base_hyperparams, os.getcwd())
    if args.algo == "a2c":
        hyperparams = a2c_sample_params(trial, base_hyperparams, base_config)
    else:
        raise ValueError(f"Optimizing {args.algo} isn't supported")
    config = Config(args, hyperparams, os.getcwd())

    wandb_enabled = bool(study_args.wandb_project_name)
    if wandb_enabled:
        wandb.init(
            project=study_args.wandb_project_name,
            entity=study_args.wandb_entity,
            config=asdict(hyperparams),
            name=f"{config.model_name()}-{str(trial.number)}",
            tags=study_args.wandb_tags,
            group=study_args.wandb_group,
            sync_tensorboard=True,
            monitor_gym=True,
            save_code=True,
            reinit=True,
        )
        wandb.config.update(args)

    tb_writer = SummaryWriter(config.tensorboard_summary_path)
    set_seeds(args.seed, args.use_deterministic_algorithms)

    env = make_env(
        config, EnvHyperparams(**config.env_hyperparams), tb_writer=tb_writer
    )
    device = get_device(config, env)
    policy = make_policy(config, env, device, **config.policy_hyperparams)
    algo = ALGOS[args.algo](policy, device, tb_writer, **config.algo_hyperparams)

    self_play_wrapper = find_wrapper(env, SelfPlayWrapper)
    eval_env = make_eval_env(
        config,
        EnvHyperparams(**config.env_hyperparams),
        override_hparams={"n_envs": study_args.n_eval_envs},
        self_play_wrapper=self_play_wrapper,
    )
    optimize_callback = OptimizeCallback(
        policy,
        eval_env,
        trial,
        tb_writer,
        step_freq=config.n_timesteps // study_args.n_evaluations,
        n_episodes=study_args.n_eval_episodes,
        deterministic=config.eval_hyperparams.get("deterministic", True),
    )
    callbacks: List[Callback] = [optimize_callback]
    if config.hyperparams.reward_decay_callback:
        callbacks.append(
            RewardDecayCallback(
                config,
                env,
                **(config.hyperparams.reward_decay_callback_kwargs or {}),
            )
        )
    if config.hyperparams.hyperparam_transitions_kwargs:
        callbacks.append(
            HyperparamTransitions(
                config,
                env,
                algo,
                **config.hyperparams.hyperparam_transitions_kwargs,
            )
        )
    if self_play_wrapper:
        callbacks.append(SelfPlayCallback(policy, self_play_wrapper))
    try:
        algo.learn(config.n_timesteps, callbacks=callbacks)

        if not optimize_callback.is_pruned:
            optimize_callback.evaluate()
            if not optimize_callback.is_pruned:
                policy.save(config.model_dir_path(best=False))

        eval_stat: EpisodesStats = callback.last_eval_stat  # type: ignore
        train_stat: EpisodesStats = callback.last_train_stat  # type: ignore

        tb_writer.add_hparams(
            hparam_dict(hyperparams, vars(args)),
            {
                "hparam/last_mean": eval_stat.score.mean,
                "hparam/last_result": eval_stat.score.mean - eval_stat.score.std,
                "hparam/train_mean": train_stat.score.mean,
                "hparam/train_result": train_stat.score.mean - train_stat.score.std,
                "hparam/score": optimize_callback.last_score,
                "hparam/is_pruned": optimize_callback.is_pruned,
            },
            None,
            config.run_name(),
        )
        tb_writer.close()

        if wandb_enabled:
            wandb.run.summary["state"] = (  # type: ignore
                "Pruned" if optimize_callback.is_pruned else "Complete"
            )
            wandb.finish(quiet=True)

        if optimize_callback.is_pruned:
            raise optuna.exceptions.TrialPruned()

        return optimize_callback.last_score
    except AssertionError as e:
        logging.warning(e)
        return np.nan
    finally:
        env.close()
        eval_env.close()
        gc.collect()
        torch.cuda.empty_cache()

def stepwise_optimize(
    trial: optuna.Trial, args: Sequence[RunArgs], study_args: StudyArgs
) -> float:
    algo = args[0].algo
    env_id = args[0].env
    base_hyperparams = load_hyperparams(algo, env_id)
    base_config = Config(args[0], base_hyperparams, os.getcwd())
    if algo == "a2c":
        hyperparams = a2c_sample_params(trial, base_hyperparams, base_config)
    else:
        raise ValueError(f"Optimizing {algo} isn't supported")

    wandb_enabled = bool(study_args.wandb_project_name)
    if wandb_enabled:
        wandb.init(
            project=study_args.wandb_project_name,
            entity=study_args.wandb_entity,
            config=asdict(hyperparams),
            name=f"{str(trial.number)}-S{base_config.seed()}",
            tags=study_args.wandb_tags,
            group=study_args.wandb_group,
            save_code=True,
            reinit=True,
        )

    score = -np.inf

    for i in range(study_args.n_evaluations):
        evaluations: List[Evaluation] = []

        for arg in args:
            config = Config(arg, hyperparams, os.getcwd())

            tb_writer = SummaryWriter(config.tensorboard_summary_path)
            set_seeds(arg.seed, arg.use_deterministic_algorithms)

            env = make_env(
                config,
                EnvHyperparams(**config.env_hyperparams),
                normalize_load_path=config.model_dir_path() if i > 0 else None,
                tb_writer=tb_writer,
            )
            device = get_device(config, env)
            policy = make_policy(config, env, device, **config.policy_hyperparams)
            if i > 0:
                policy.load(config.model_dir_path())
            algo = ALGOS[arg.algo](policy, device, tb_writer, **config.algo_hyperparams)

            self_play_wrapper = find_wrapper(env, SelfPlayWrapper)
            eval_env = make_eval_env(
                config,
                EnvHyperparams(**config.env_hyperparams),
                normalize_load_path=config.model_dir_path() if i > 0 else None,
                override_hparams={"n_envs": study_args.n_eval_envs},
                self_play_wrapper=self_play_wrapper,
            )

            start_timesteps = int(i * config.n_timesteps / study_args.n_evaluations)
            train_timesteps = (
                int((i + 1) * config.n_timesteps / study_args.n_evaluations)
                - start_timesteps
            )

            callbacks = []
            if config.hyperparams.reward_decay_callback:
                callbacks.append(
                    RewardDecayCallback(
                        config,
                        env,
                        start_timesteps=start_timesteps,
                        **(config.hyperparams.reward_decay_callback_kwargs or {}),
                    )
                )
            if config.hyperparams.hyperparam_transitions_kwargs:
                callbacks.append(
                    HyperparamTransitions(
                        config,
                        env,
                        algo,
                        start_timesteps=start_timesteps,
                        **config.hyperparams.hyperparam_transitions_kwargs,
                    )
                )
            if self_play_wrapper:
                callbacks.append(SelfPlayCallback(policy, self_play_wrapper))
            try:
                algo.learn(
                    train_timesteps,
                    callbacks=callbacks,
                    total_timesteps=config.n_timesteps,
                    start_timesteps=start_timesteps,
                )

                evaluations.append(
                    evaluation(
                        policy,
                        eval_env,
                        tb_writer,
                        study_args.n_eval_episodes,
                        config.eval_hyperparams.get("deterministic", True),
                        start_timesteps + train_timesteps,
                    )
                )

                policy.save(config.model_dir_path())

                tb_writer.close()

            except AssertionError as e:
                logging.warning(e)
                if wandb_enabled:
                    wandb_finish("Error")
                return np.nan
            finally:
                env.close()
                eval_env.close()
                gc.collect()
                torch.cuda.empty_cache()

        d = {}
        for idx, e in enumerate(evaluations):
            d[f"{idx}/eval_mean"] = e.eval_stat.score.mean
            d[f"{idx}/train_mean"] = e.train_stat.score.mean
            d[f"{idx}/score"] = e.score
        d["eval"] = np.mean([e.eval_stat.score.mean for e in evaluations]).item()
        d["train"] = np.mean([e.train_stat.score.mean for e in evaluations]).item()
        score = np.mean([e.score for e in evaluations]).item()
        d["score"] = score

        step = i + 1
        wandb.log(d, step=step)

        print(f"Trial #{trial.number} Step {step} Score: {round(score, 2)}")
        trial.report(score, step)
        if trial.should_prune():
            if wandb_enabled:
                wandb_finish("Pruned")
            raise optuna.exceptions.TrialPruned()

    if wandb_enabled:
        wandb_finish("Complete")
    return score

def wandb_finish(state: str) -> None:
    wandb.run.summary["state"] = state  # type: ignore
    wandb.finish(quiet=True)

def optimize() -> None:
    from pyvirtualdisplay.display import Display

    train_args, study_args = parse_args()
    if study_args.virtual_display:
        virtual_display = Display(visible=False, size=(1400, 900))
        virtual_display.start()

    sampler = TPESampler(**TPESampler.hyperopt_parameters())
    pruner = HyperbandPruner()
    if study_args.load_study:
        assert study_args.study_name
        assert study_args.storage_path
        study = optuna.load_study(
            study_name=study_args.study_name,
            storage=study_args.storage_path,
            sampler=sampler,
            pruner=pruner,
        )
    else:
        study = optuna.create_study(
            study_name=study_args.study_name,
            storage=study_args.storage_path,
            sampler=sampler,
            pruner=pruner,
            direction="maximize",
        )

    try:
        study.optimize(
            objective_fn(train_args, study_args),
            n_trials=study_args.n_trials,
            n_jobs=study_args.n_jobs,
            timeout=study_args.timeout,
        )
    except KeyboardInterrupt:
        pass

    best = study.best_trial
    print(f"Best Trial Value: {best.value}")
    print("Attributes:")
    for key, value in list(best.params.items()) + list(best.user_attrs.items()):
        print(f"  {key}: {value}")

    df = study.trials_dataframe()
    df = df[df.state == "COMPLETE"].sort_values(by=["value"], ascending=False)
    print(df.to_markdown(index=False))

    fig1 = plot_optimization_history(study)
    fig1.write_image("opt_history.png")

    fig2 = plot_param_importances(study)
    fig2.write_image("param_importances.png")

if __name__ == "__main__":
    optimize()


import pandas as pd
import datetime

import matplotlib.pyplot as plt
PATH = "C:\\datasets\\"
FILE = 'shampoo.csv'
df   = pd.read_csv(PATH + FILE, index_col=0)
df.info()

# Plot data before differencing.
df.plot()
plt.xticks(rotation=45)
plt.show()

# Perform differencing.
dfDifferenced = df.diff()

# Plot data after differencing.
plt.plot(dfDifferenced)
plt.xticks(rotation=75)
plt.show()


from typing import List
from typing import Dict
from oakvar import BaseConverter

class Converter(BaseConverter):
    def check_format(self, f) -> bool:
        """
        Detect the format of an input file.

        Arguments:
            f: a file handle to an input file
        Returns:
            bool: True if the input file is for this converter,
                  False if not.

        The example below checks if the input file's first line indicates
        VCF file format.
        """
        line = f.readline()
        return line.startswith("##fileformat=VCF")

    # If your converter module needs something else than
    # the standard way of opening a text input file,
    # read line by line, and coverting each line into
    # a list of dictionaries of variants,
    # you may want to start with modifying
    # convert_file method. In that case, uncomment
    # the below convert_file method and add your implementation.
    #
    # def convert_file(
    #     self, file, *__args__, exc_handler=None, **__kwargs__
    # ) -> Iterator[Tuple[int, List[dict]]]:
    #     line_no = 0
    #     for line in file:
    #         line_no += 1
    #         try:
    #             yield line_no, self.convert_line(line)
    #         except Exception as e:
    #             if exc_handler:
    #                 exc_handler(line_no, e)
    #             else:
    #                 raise e
    #     return None

    def convert_line(self, line) -> List[Dict]:
        """
        Converts a line from an input file to OakVar's variant dict.

        Arguments:
            l: a string of a line from an input file
        Returns:
            dict: a list of dicts, each dict for a variant collected
                  from the input line. Each dict should have
                  the following required fields:

                  chrom: chromosome name [str]
                  pos: chromosomal position [int]
                  ref_base: reference bases [str]
                  alt_base: altername bases [str]

                  Optional fields for each dict are:

                  sample_id: the ID or name of a sample having the variant [list[str]]
                  tags: a custom tag given to the variant [list[str]]
        """
        _ = line
        var_dicts = []
        var_dict = {
            "chrom": "chr1",
            "pos": 2878349,
            "ref_base": "A",
            "alt_base": "T",
            "sample_id": "sample1",
        }
        var_dicts.append(var_dict)
        return var_dicts


from PIL import Image
from pylab import *
#import os
#pil_im = Image.open('fruitas.jpg').convert('L')

'''
#converts all the picture in the file as jps
for infile in filelist:
    outfile = os.path.splitext(infile)[0]+".jpg"
    if infle != otule:
        try:
            Image.open(infile).save(outfile)
        except IOError:
            print("cannot convert", infile)
'''
#read imget to array
im = array(Image.open('fruitas.jpg'))

#plot the image
imshow(im)

#some points
x=[100, 100, 400, 400]
y=[200, 500, 200, 500]

#plot the points with red star-makers
plot(x, y, 'g+')

#line plot connecting the first two points
plot(x[:2], y[:2])

#add title and show the plot
title('Plotting: "fruitas.jpg"')
show()


from bs4 import BeautifulSoup
import requests
import re
import pandas as pd
url = 'https://www.marca.com/futbol/sevilla.html?intcmp=MENUESCU&s_kw=sevilla'

page = requests.get(url)
soup = BeautifulSoup(page.content, 'html.parser')

data_noticias = pd.DataFrame()
lista_noticias = []
for eq in soup.find_all('h3', class_="mod-title"):
    eq = str(eq)
    m = re.search('title=(.+?)>', eq)
    noticias = m.group(1)
    lista_noticias.append(noticias)

df = pd.DataFrame(lista_noticias)

project_id = 'news_MARCA'

from google.colab import auth
auth.authenticate_user()

from googleapiclient.discovery import build
gcs_service = build('storage', 'v1')
bucket_name = 'segmento_noticias'

file_name = 'noticias.csv'

from googleapiclient.http import MediaFileUpload

media = MediaFileUpload(df,
                        mimetype='text/plain',
                        resumable=True)

request = gcs_service.objects().insert(bucket=bucket_name,
                                       name=file_name,
                                       media_body=media)

response = None
while response is None:

  _, response = request.next_chunk()

print('Upload completo')


import os
import json

inventory = os.listdir("labs/")
inventory.sort()

flagged_packages = ["tomcat", "tomcat8",]
found_versions = []
host_results = []

for file in inventory:
    with open(os.path.join("labs", file)) as json_inventory:
        json_packages = json.load(json_inventory)
        for package in json_packages:
            if package in flagged_packages:
                print("Hostname: " + file)
                for key in json_packages[package]:
                    version = key["version"]
                    print("Version: " + package + "-" + version + "\n")


# get lora model info from a pile of jsons saved from Civitai
# can be tweaked to do more, like ckpts, TIs, etc, and list more hashes...
# currently is bare minimum to populate list hopefully for Model Keyword to use

import re
from pprint import pprint
import sys
import os
from pathlib import Path
import json

for file in Path('.').iterdir():
    # Check if the file is a regular file
    if file.is_file():
        # if it's a json
        if file.suffix in [".json"]:
            #print(f"opening {file.stem}")
            with open(file) as model_file:
                 model = json.load(model_file)
                 #pprint(model)
                 model_type = model["type"]
                 if model_type == "LORA":
                     for items in model["modelVersions"]:
                        #pprint(items)
                        twords = []
                        for word in items['trainedWords']:
                            if re.match("<lora.*>$",word):
                               # it's not a trigger word
                               pass
                            else:
                               twords.append(word.strip().strip(",").strip())
                        trainedwords = re.sub("\|\|","|", re.sub(", ?","|","|".join(twords)).rstrip("|"))
                        if trainedwords == "":
                           break
                        for fileinfo in items['files']:
                           #if 'AutoV2' in fileinfo['hashes'].keys():
                           #   hash = fileinfo['hashes']['AutoV2']
                           #   print(f"{hash}, {trainedwords}, {file.stem}")
                           if 'AutoV1' in fileinfo['hashes'].keys():
                              hash = fileinfo['hashes']['AutoV1']
                              print(f"{hash}, {trainedwords}, {file.stem}")


__metaclass__ = type

from datetime import (
    datetime,
    timedelta,
    )
import logging
import re
from smtplib import SMTPException
import StringIO
import unittest

import pytz
from storm.store import Store
from testtools.matchers import (
    MatchesRegex,
    Not,
    )
from transaction import commit
from zope.component import (
    getSiteManager,
    getUtility,
    )
from zope.interface import implementer
from zope.sendmail.interfaces import IMailDelivery

from lp.app.enums import InformationType
from lp.bugs.adapters.bugchange import (
    BranchLinkedToBug,
    BranchUnlinkedFromBug,
    BugAttachmentChange,
    BugDuplicateChange,
    BugInformationTypeChange,
    BugTagsChange,
    BugTaskStatusChange,
    BugTitleChange,
    BugWatchAdded,
    BugWatchRemoved,
    CveLinkedToBug,
    CveUnlinkedFromBug,
    )
from lp.bugs.enums import BugNotificationStatus
from lp.bugs.interfaces.bug import (
    CreateBugParams,
    IBug,
    IBugSet,
    )
from lp.bugs.interfaces.bugnotification import IBugNotificationSet
from lp.bugs.interfaces.bugtask import (
    BugTaskImportance,
    BugTaskStatus,
    )
from lp.bugs.mail.bugnotificationrecipients import BugNotificationRecipients
from lp.bugs.model.bugnotification import (
    BugNotification,
    BugNotificationFilter,
    BugNotificationRecipient,
    )
from lp.bugs.model.bugsubscriptionfilter import BugSubscriptionFilterMute
from lp.bugs.model.bugtask import BugTask
from lp.bugs.scripts.bugnotification import (
    construct_email_notifications,
    get_activity_key,
    get_email_notifications,
    notification_batches,
    notification_comment_batches,
    process_deferred_notifications,
    SendBugNotifications,
    )
from lp.registry.interfaces.person import IPersonSet
from lp.registry.interfaces.product import IProductSet
from lp.services.config import config
from lp.services.database.interfaces import IStore
from lp.services.database.sqlbase import (
    flush_database_updates,
    sqlvalues,
    )
from lp.services.mail.helpers import (
    get_contact_email_addresses,
    get_email_template,
    )
from lp.services.mail.sendmail import set_immediate_mail_delivery
from lp.services.mail.stub import TestMailer
from lp.services.messages.interfaces.message import IMessageSet
from lp.services.propertycache import cachedproperty
from lp.testing import (
    login,
    person_logged_in,
    TestCase,
    TestCaseWithFactory,
    )
from lp.testing.dbuser import (
    lp_dbuser,
    switch_dbuser,
    )
from lp.testing.fixture import ZopeUtilityFixture
from lp.testing.layers import LaunchpadZopelessLayer
from lp.testing.matchers import Contains

@implementer(IBug)
class MockBug:
    """A bug which has only the attributes get_email_notifications() needs."""

    duplicateof = None
    information_type = InformationType.PUBLIC
    messages = []

    def __init__(self, id, owner):
        self.id = id
        self.initial_message = getUtility(IMessageSet).fromText(
            'Bug Title', 'Initial message.', owner=owner)
        self.owner = owner
        self.bugtasks = []
        self.tags = []

    @property
    def title(self):
        return "Mock Bug #%s" % self.id

    def getBugNotificationRecipients(self, level=None):
        recipients = BugNotificationRecipients()
        no_priv = getUtility(IPersonSet).getByEmail(
            'no-priv@canonical.com')
        recipients.addDirectSubscriber(no_priv)
        return recipients

    def __eq__(self, other):
        """Compare by id to make different subclasses of MockBug be equal."""
        return self.id == other.id

class ExceptionBug(MockBug):
    """A bug which causes an exception to be raised."""

    def getBugNotificationRecipients(self, level=None):
        raise Exception('FUBAR')

class DBExceptionBug(MockBug):
    """A bug which causes a DB constraint to be triggered."""

    def getBugNotificationRecipients(self, level=None):
        # Trigger a DB constraint, resulting in the transaction being
        # unusable.
        firefox = getUtility(IProductSet).getByName('firefox')
        bug_one = getUtility(IBugSet).get(1)
        BugTask(bug=bug_one, product=firefox, owner=self.owner)

class MockBugNotificationRecipient:
    """A mock BugNotificationRecipient for testing."""

    def __init__(self):
        self.person = getUtility(IPersonSet).getByEmail(
            'no-priv@canonical.com')
        self.reason_header = 'Test Rationale'
        self.reason_body = 'Test Reason'

class MockBugNotification:
    """A mock BugNotification used for testing.

    Using a real BugNotification won't allow us to set the bug to a mock
    object.
    """

    def __init__(self, message, bug, is_comment, date_emailed):
        self.message = message
        self.bug = bug
        self.is_comment = is_comment
        self.date_emailed = date_emailed
        self.recipients = [MockBugNotificationRecipient()]
        self.activity = None

class FakeNotification:
    """An even simpler fake notification.

    Used by TestGetActivityKey, TestNotificationCommentBatches and
    TestNotificationBatches."""

    class Message(object):
        pass

    def __init__(self, is_comment=False, bug=None, owner=None):
        self.is_comment = is_comment
        self.bug = bug
        self.message = self.Message()
        self.message.owner = owner
        self.activity = None

@implementer(IBugNotificationSet)
class FakeBugNotificationSetUtility:
    """A notification utility used for testing."""

    def getRecipientFilterData(self, bug, recipient_to_sources,
                               notifications):
        return dict(
            (recipient, {'sources': sources, 'filter descriptions': []})
            for recipient, sources in recipient_to_sources.items())

class MockBugActivity:
    """A mock BugActivity used for testing."""

    def __init__(self, target=None, attribute=None,
                 oldvalue=None, newvalue=None):
        self.target = target
        self.attribute = attribute
        self.oldvalue = oldvalue
        self.newvalue = newvalue

class TestGetActivityKey(TestCase):
    """Tests for get_activity_key()."""

    def test_no_activity(self):
        self.assertEqual(get_activity_key(FakeNotification()), None)

    def test_normal_bug_attribute_activity(self):
        notification = FakeNotification()
        notification.activity = MockBugActivity(attribute='title')
        self.assertEqual(get_activity_key(notification), 'title')

    def test_collection_bug_attribute_added_activity(self):
        notification = FakeNotification()
        notification.activity = MockBugActivity(
            attribute='cves', newvalue='some cve identifier')
        self.assertEqual(get_activity_key(notification),
                         'cves:some cve identifier')

    def test_collection_bug_attribute_removed_activity(self):
        notification = FakeNotification()
        notification.activity = MockBugActivity(
            attribute='cves', oldvalue='some cve identifier')
        self.assertEqual(get_activity_key(notification),
                         'cves:some cve identifier')

    def test_bugtask_attribute_activity(self):
        notification = FakeNotification()
        notification.activity = MockBugActivity(
            attribute='status', target='some bug task identifier')
        self.assertEqual(get_activity_key(notification),
                         'some bug task identifier:status')

class TestGetEmailNotifications(TestCase):
    """Tests for the exception handling in get_email_notifications()."""
    layer = LaunchpadZopelessLayer

    def setUp(self):
        """Set up some mock bug notifications to use."""
        super(TestGetEmailNotifications, self).setUp()
        switch_dbuser(config.malone.bugnotification_dbuser)
        sample_person = getUtility(IPersonSet).getByEmail(
            'test@canonical.com')
        self.now = datetime.now(pytz.timezone('UTC'))

        # A normal comment notification for bug 1
        msg = getUtility(IMessageSet).fromText(
            'Subject', "Comment on bug 1", owner=sample_person)
        self.bug_one_notification = MockBugNotification(
            message=msg, bug=MockBug(1, sample_person),
            is_comment=True, date_emailed=None)

        # Another normal comment notification for bug one.
        msg = getUtility(IMessageSet).fromText(
            'Subject', "Comment on bug 1", owner=sample_person)
        self.bug_one_another_notification = MockBugNotification(
            message=msg, bug=MockBug(1, sample_person),
            is_comment=True, date_emailed=None)

        # A comment notification for bug one which raises an exception.
        msg = getUtility(IMessageSet).fromText(
            'Subject', "Comment on bug 1", owner=sample_person)
        self.bug_one_exception_notification = MockBugNotification(
            message=msg, bug=ExceptionBug(1, sample_person),
            is_comment=True, date_emailed=None)

        # A comment notification for bug one which raises a DB exception.
        msg = getUtility(IMessageSet).fromText(
            'Subject', "Comment on bug 1", owner=sample_person)
        self.bug_one_dbexception_notification = MockBugNotification(
            message=msg, bug=DBExceptionBug(1, sample_person),
            is_comment=True, date_emailed=None)

        # We need to commit the transaction, since the error handling
        # will abort the current transaction.
        commit()

        sm = getSiteManager()
        self._original_utility = sm.getUtility(IBugNotificationSet)
        sm.unregisterUtility(self._original_utility)
        self._fake_utility = FakeBugNotificationSetUtility()
        sm.registerUtility(self._fake_utility)

    def tearDown(self):
        super(TestGetEmailNotifications, self).tearDown()
        sm = getSiteManager()
        sm.unregisterUtility(self._fake_utility)
        sm.registerUtility(self._original_utility)

    def _getAndCheckSentNotifications(self, notifications_to_send):
        """Return the notifications that were successfully sent.

        It calls get_email_notifications() with the supplied
        notifications and return the ones that were actually sent. It
        also checks that the notifications got sent to the correct
        addresses.
        """
        email_notifications = get_email_notifications(notifications_to_send)
        to_addresses = set()
        sent_notifications = []
        for notifications, omitted, messages in email_notifications:
            for message in messages:
                to_addresses.add(message['to'])
            recipients = {}
            for notification in notifications:
                for recipient in notification.recipients:
                    for address in get_contact_email_addresses(
                        recipient.person):
                        recipients[address] = recipient
            expected_to_addresses = recipients.keys()
            self.assertEqual(
                sorted(expected_to_addresses), sorted(to_addresses))
            sent_notifications += notifications
        return sent_notifications

    def test_catch_simple_exception_last(self):
        # Make sure that the first notification is sent even if the
        # last one causes an exception to be raised.
        notifications_to_send = [
            self.bug_one_notification,
            self.bug_one_exception_notification,
            ]
        sent_notifications = self._getAndCheckSentNotifications(
            notifications_to_send)
        self.assertEqual(sent_notifications, notifications_to_send)

    def test_catch_simple_exception_in_the_middle(self):
        # Make sure that the first and last notifications are sent even
        # if the middle one causes an exception to be raised.
        notifications_to_send = [
            self.bug_one_notification,
            self.bug_one_exception_notification,
            self.bug_one_another_notification,
            ]
        sent_notifications = self._getAndCheckSentNotifications(
            notifications_to_send)
        self.assertEqual(
            sent_notifications,
            notifications_to_send)

    def test_catch_db_exception_last(self):
        # Make sure that the first notification is sent even if the
        # last one causes an exception to be raised. Also make sure that
        # the current transaction is in a usable state.
        notifications_to_send = [
            self.bug_one_notification,
            self.bug_one_dbexception_notification,
            ]
        sent_notifications = self._getAndCheckSentNotifications(
            notifications_to_send)
        self.assertEqual(sent_notifications, notifications_to_send)

        # The transaction should have been rolled back and restarted
        # properly, so getting something from the database shouldn't
        # cause any errors.
        bug_four = getUtility(IBugSet).get(4)
        self.assertEqual(bug_four.id, 4)

    def test_catch_db_exception_in_the_middle(self):
        # Make sure that the first and last notifications are sent even
        # if the middle one causes an exception to be raised. Also make
        # sure that the current transaction is in a usable state.
        notifications_to_send = [
            self.bug_one_notification,
            self.bug_one_dbexception_notification,
            self.bug_one_another_notification,
            ]
        sent_notifications = self._getAndCheckSentNotifications(
            notifications_to_send)
        self.assertEqual(
            sent_notifications, notifications_to_send)

        # The transaction should have been rolled back and restarted
        # properly, so getting something from the database shouldn't
        # cause any errors.
        bug_four = getUtility(IBugSet).get(4)
        self.assertEqual(bug_four.id, 4)

    def test_early_exit(self):
        # When not-yet-exhausted generators need to be deallocated Python
        # raises a GeneratorExit exception at the point of their last yield.
        # The get_email_notifications generator was catching that exception in
        # a try/except and logging it, leading to bug 994694.  This test
        # verifies that the fix for that bug (re-raising the exception) stays
        # in place.

        # Set up logging so we can later assert that no exceptions are logged.
        log_output = StringIO.StringIO()
        logger = logging.getLogger()
        log_handler = logging.StreamHandler(log_output)
        logger.addHandler(logging.StreamHandler(log_output))
        self.addCleanup(logger.removeHandler, log_handler)

        # Make some data to feed to get_email_notifications.
        person = getUtility(IPersonSet).getByEmail('test@canonical.com')
        msg = getUtility(IMessageSet).fromText('', '', owner=person)
        bug = MockBug(1, person)
        # We need more than one notification because we want the generator to
        # stay around after being started.  Consuming the first starts it but
        # since the second exists, the generator stays active.
        notifications = [
            MockBugNotification(
                message=msg, bug=bug, is_comment=True, date_emailed=None),
            MockBugNotification(
                message=msg, bug=bug, is_comment=True, date_emailed=None),
            ]

        # Now we create the generator, start it, and then close it, triggering
        # a GeneratorExit exception inside the generator.
        email_notifications = get_email_notifications(notifications)
        email_notifications.next()
        email_notifications.close()

        # Verify that no "Error while building email notifications." is logged.
        self.assertEqual('', log_output.getvalue())

class TestNotificationCommentBatches(unittest.TestCase):
    """Tests of `notification_comment_batches`."""

    def test_with_nothing(self):
        # Nothing is generated if an empty list is passed in.
        self.assertEqual([], list(notification_comment_batches([])))

    def test_with_one_non_comment_notification(self):
        # Given a single non-comment notification, a single tuple is
        # generated.
        notification = FakeNotification(False)
        self.assertEqual(
            [(1, notification)],
            list(notification_comment_batches([notification])))

    def test_with_one_comment_notification(self):
        # Given a single comment notification, a single tuple is generated.
        notification = FakeNotification(True)
        self.assertEqual(
            [(1, notification)],
            list(notification_comment_batches([notification])))

    def test_with_two_notifications_comment_first(self):
        # Given two notifications, one a comment, one not, and the comment
        # first, two tuples are generated, both in the same group.
        notification1 = FakeNotification(True)
        notification2 = FakeNotification(False)
        notifications = [notification1, notification2]
        self.assertEqual(
            [(1, notification1), (1, notification2)],
            list(notification_comment_batches(notifications)))

    def test_with_two_notifications_comment_last(self):
        # Given two notifications, one a comment, one not, and the comment
        # last, two tuples are generated, both in the same group.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notifications = [notification1, notification2]
        self.assertEqual(
            [(1, notification1), (1, notification2)],
            list(notification_comment_batches(notifications)))

    def test_with_three_notifications_comment_in_middle(self):
        # Given three notifications, one a comment, two not, and the comment
        # in the middle, three tuples are generated, all in the same group.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notification3 = FakeNotification(False)
        notifications = [notification1, notification2, notification3]
        self.assertEqual(
            [(1, notification1), (1, notification2), (1, notification3)],
            list(notification_comment_batches(notifications)))

    def test_with_more_notifications(self):
        # Given four notifications - non-comment, comment, non-comment,
        # comment - four tuples are generated. The first three notifications
        # are in the first group, the last notification is in a group on its
        # own.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notification3 = FakeNotification(False)
        notification4 = FakeNotification(True)
        notifications = [
            notification1, notification2,
            notification3, notification4,
            ]
        self.assertEqual(
            [(1, notification1), (1, notification2),
             (1, notification3), (2, notification4)],
            list(notification_comment_batches(notifications)))

class TestNotificationBatches(unittest.TestCase):
    """Tests of `notification_batches`."""

    def test_with_nothing(self):
        # Nothing is generated if an empty list is passed in.
        self.assertEqual([], list(notification_batches([])))

    def test_with_one_non_comment_notification(self):
        # Given a single non-comment notification, a single batch is
        # generated.
        notification = FakeNotification(False)
        self.assertEqual(
            [[notification]],
            list(notification_batches([notification])))

    def test_with_one_comment_notification(self):
        # Given a single comment notification, a single batch is generated.
        notification = FakeNotification(True)
        self.assertEqual(
            [[notification]],
            list(notification_batches([notification])))

    def test_with_two_notifications_comment_first(self):
        # Given two similar notifications, one a comment, one not, and the
        # comment first, a single batch is generated.
        notification1 = FakeNotification(True)
        notification2 = FakeNotification(False)
        notifications = [notification1, notification2]
        self.assertEqual(
            [[notification1, notification2]],
            list(notification_batches(notifications)))

    def test_with_two_notifications_comment_last(self):
        # Given two similar notifications, one a comment, one not, and the
        # comment last, a single batch is generated.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notifications = [notification1, notification2]
        self.assertEqual(
            [[notification1, notification2]],
            list(notification_batches(notifications)))

    def test_with_three_notifications_comment_in_middle(self):
        # Given three similar notifications, one a comment, two not, and the
        # comment in the middle, one batch is generated.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notification3 = FakeNotification(False)
        notifications = [notification1, notification2, notification3]
        self.assertEqual(
            [[notification1, notification2, notification3]],
            list(notification_batches(notifications)))

    def test_with_more_notifications(self):
        # Given four similar notifications - non-comment, comment,
        # non-comment, comment - two batches are generated. The first three
        # notifications are in the first batch.
        notification1 = FakeNotification(False)
        notification2 = FakeNotification(True)
        notification3 = FakeNotification(False)
        notification4 = FakeNotification(True)
        notifications = [
            notification1, notification2,
            notification3, notification4,
            ]
        self.assertEqual(
            [[notification1, notification2, notification3], [notification4]],
            list(notification_batches(notifications)))

    def test_notifications_for_same_bug(self):
        # Batches are grouped by bug.
        notifications = [FakeNotification(bug=1) for number in range(5)]
        observed = list(notification_batches(notifications))
        self.assertEqual([notifications], observed)

    def test_notifications_for_different_bugs(self):
        # Batches are grouped by bug.
        notifications = [FakeNotification(bug=number) for number in range(5)]
        expected = [[notification] for notification in notifications]
        observed = list(notification_batches(notifications))
        self.assertEqual(expected, observed)

    def test_notifications_for_same_owner(self):
        # Batches are grouped by owner.
        notifications = [FakeNotification(owner=1) for number in range(5)]
        observed = list(notification_batches(notifications))
        self.assertEqual([notifications], observed)

    def test_notifications_for_different_owners(self):
        # Batches are grouped by owner.
        notifications = [
            FakeNotification(owner=number) for number in range(5)]
        expected = [[notification] for notification in notifications]
        observed = list(notification_batches(notifications))
        self.assertEqual(expected, observed)

    def test_notifications_with_mixed_bugs_and_owners(self):
        # Batches are grouped by bug and owner.
        notifications = [
            FakeNotification(bug=1, owner=1),
            FakeNotification(bug=1, owner=2),
            FakeNotification(bug=2, owner=2),
            FakeNotification(bug=2, owner=1),
            ]
        expected = [[notification] for notification in notifications]
        observed = list(notification_batches(notifications))
        self.assertEqual(expected, observed)

    def test_notifications_with_mixed_bugs_and_owners_2(self):
        # Batches are grouped by bug and owner.
        notifications = [
            FakeNotification(bug=1, owner=1),
            FakeNotification(bug=1, owner=1),
            FakeNotification(bug=2, owner=2),
            FakeNotification(bug=2, owner=2),
            ]
        expected = [notifications[0:2], notifications[2:4]]
        observed = list(notification_batches(notifications))
        self.assertEqual(expected, observed)

    def test_notifications_with_mixed_bugs_owners_and_comments(self):
        # Batches are grouped by bug, owner and comments.
        notifications = [
            FakeNotification(is_comment=False, bug=1, owner=1),
            FakeNotification(is_comment=False, bug=1, owner=1),
            FakeNotification(is_comment=True, bug=1, owner=1),
            FakeNotification(is_comment=False, bug=1, owner=2),
            ]
        expected = [notifications[0:3], notifications[3:4]]
        observed = list(notification_batches(notifications))
        self.assertEqual(expected, observed)

class EmailNotificationTestBase(TestCaseWithFactory):

    layer = LaunchpadZopelessLayer

    def setUp(self):
        super(EmailNotificationTestBase, self).setUp()
        login('foo.bar@canonical.com')
        self.product_owner = self.factory.makePerson(name="product-owner")
        self.person = self.factory.makePerson(name="sample-person")
        self.product = self.factory.makeProduct(owner=self.product_owner)
        self.product_subscriber = self.factory.makePerson(
            name="product-subscriber")
        self.product.addBugSubscription(
            self.product_subscriber, self.product_subscriber)
        self.bug_subscriber = self.factory.makePerson(name="bug-subscriber")
        self.bug_owner = self.factory.makePerson(name="bug-owner")
        self.bug = self.factory.makeBug(
            target=self.product, owner=self.bug_owner,
            information_type=InformationType.USERDATA)
        self.reporter = self.bug.owner
        self.bug.subscribe(self.bug_subscriber, self.reporter)
        [self.product_bugtask] = self.bug.bugtasks
        commit()
        login('test@canonical.com')
        switch_dbuser(config.malone.bugnotification_dbuser)
        self.now = datetime.now(pytz.UTC)
        self.ten_minutes_ago = self.now - timedelta(minutes=10)
        self.notification_set = getUtility(IBugNotificationSet)
        for notification in self.notification_set.getNotificationsToSend():
            notification.date_emailed = self.now
        flush_database_updates()

    def tearDown(self):
        for notification in self.notification_set.getNotificationsToSend():
            notification.date_emailed = self.now
        flush_database_updates()
        super(EmailNotificationTestBase, self).tearDown()

    def get_messages(self):
        notifications = self.notification_set.getNotificationsToSend()
        email_notifications = get_email_notifications(notifications)
        for (bug_notifications,
             omitted_notifications,
             messages) in email_notifications:
            for message in messages:
                yield message, message.get_payload(decode=True)

class EmailNotificationsBugMixin:

    change_class = change_name = old = new = alt = unexpected_text = None

    def change(self, old, new):
        self.bug.addChange(
            self.change_class(
                self.ten_minutes_ago, self.person, self.change_name,
                old, new))

    def change_other(self):
        self.bug.addChange(
            BugInformationTypeChange(
                self.ten_minutes_ago, self.person, "information_type",
                InformationType.PUBLIC, InformationType.USERDATA))

    def test_change_seen(self):
        # A smoketest.
        self.change(self.old, self.new)
        message, body = self.get_messages().next()
        self.assertThat(body, Contains(self.unexpected_text))

    def test_undone_change_sends_no_emails(self):
        self.change(self.old, self.new)
        self.change(self.new, self.old)
        self.assertEqual(list(self.get_messages()), [])

    def test_undone_change_is_not_included(self):
        self.change(self.old, self.new)
        self.change(self.new, self.old)
        self.change_other()
        message, body = self.get_messages().next()
        self.assertThat(body, Not(Contains(self.unexpected_text)))

    def test_multiple_undone_changes_sends_no_emails(self):
        self.change(self.old, self.new)
        self.change(self.new, self.alt)
        self.change(self.alt, self.old)
        self.assertEqual(list(self.get_messages()), [])

class EmailNotificationsBugNotRequiredMixin(EmailNotificationsBugMixin):
    # This test collection is for attributes that can be None.
    def test_added_removed_sends_no_emails(self):
        self.change(None, self.old)
        self.change(self.old, None)
        self.assertEqual(list(self.get_messages()), [])

    def test_removed_added_sends_no_emails(self):
        self.change(self.old, None)
        self.change(None, self.old)
        self.assertEqual(list(self.get_messages()), [])

    def test_duplicate_marked_changed_removed_sends_no_emails(self):
        self.change(None, self.old)
        self.change(self.old, self.new)
        self.change(self.new, None)
        self.assertEqual(list(self.get_messages()), [])

class EmailNotificationsBugTaskMixin(EmailNotificationsBugMixin):

    def change(self, old, new, index=0):
        self.bug.addChange(
            self.change_class(
                self.bug.bugtasks[index], self.ten_minutes_ago,
                self.person, self.change_name, old, new))

    def test_changing_on_different_bugtasks_is_not_undoing(self):
        with lp_dbuser():
            product2 = self.factory.makeProduct(owner=self.product_owner)
            self.bug.addTask(self.product_owner, product2)
        self.change(self.old, self.new, index=0)
        self.change(self.new, self.old, index=1)
        message, body = self.get_messages().next()
        self.assertThat(body, Contains(self.unexpected_text))

class EmailNotificationsAddedRemovedMixin:

    old = new = added_message = removed_message = None

    def add(self, item):
        raise NotImplementedError
    remove = add

    def test_added_seen(self):
        self.add(self.old)
        message, body = self.get_messages().next()
        self.assertThat(body, Contains(self.added_message))

    def test_added_removed_sends_no_emails(self):
        self.add(self.old)
        self.remove(self.old)
        self.assertEqual(list(self.get_messages()), [])

    def test_removed_added_sends_no_emails(self):
        self.remove(self.old)
        self.add(self.old)
        self.assertEqual(list(self.get_messages()), [])

    def test_added_another_removed_sends_emails(self):
        self.add(self.old)
        self.remove(self.new)
        message, body = self.get_messages().next()
        self.assertThat(body, Contains(self.added_message))
        self.assertThat(body, Contains(self.removed_message))

class TestEmailNotificationsBugTitle(
    EmailNotificationsBugMixin, EmailNotificationTestBase):

    change_class = BugTitleChange
    change_name = "title"
    old = "Old summary"
    new = "New summary"
    alt = "Another summary"
    unexpected_text = '** Summary changed:'

class TestEmailNotificationsBugTags(
    EmailNotificationsBugMixin, EmailNotificationTestBase):

    change_class = BugTagsChange
    change_name = "tags"
    old = ['foo', 'bar', 'baz']
    new = ['foo', 'bar']
    alt = ['bing', 'shazam']
    unexpected_text = '** Tags'

    def test_undone_ordered_set_sends_no_email(self):
        # Tags use ordered sets to generate change descriptions, which we
        # demonstrate here.
        self.change(['foo', 'bar', 'baz'], ['foo', 'bar'])
        self.change(['foo', 'bar'], ['baz', 'bar', 'foo', 'bar'])
        self.assertEqual(list(self.get_messages()), [])

class TestEmailNotificationsBugDuplicate(
    EmailNotificationsBugNotRequiredMixin, EmailNotificationTestBase):

    change_class = BugDuplicateChange
    change_name = "duplicateof"
    unexpected_text = 'duplicate'

    def _bug(self):
        with lp_dbuser():
            return self.factory.makeBug()

    old = cachedproperty('old')(_bug)
    new = cachedproperty('new')(_bug)
    alt = cachedproperty('alt')(_bug)

class TestEmailNotificationsBugTaskStatus(
    EmailNotificationsBugTaskMixin, EmailNotificationTestBase):

    change_class = BugTaskStatusChange
    change_name = "status"
    old = BugTaskStatus.TRIAGED
    new = BugTaskStatus.INPROGRESS
    alt = BugTaskStatus.INVALID
    unexpected_text = 'Status: '

class TestEmailNotificationsBugWatch(
    EmailNotificationsAddedRemovedMixin, EmailNotificationTestBase):

    # Note that this is for bugwatches added to bugs.  Bugwatches added
    # to bugtasks are separate animals AIUI, and we don't try to combine
    # them here for notifications.  Bugtasks have only zero or one
    # bugwatch, so they can be handled just as a simple bugtask attribute
    # change, like status.

    added_message = '** Bug watch added:'
    removed_message = '** Bug watch removed:'

    @cachedproperty
    def tracker(self):
        with lp_dbuser():
            return self.factory.makeBugTracker()

    def _watch(self, identifier='123'):
        with lp_dbuser():
            # This actually creates a notification all by itself.  However,
            # it won't be sent out for another five minutes.  Therefore,
            # we send out separate change notifications.
            return self.bug.addWatch(
                self.tracker, identifier, self.product_owner)

    old = cachedproperty('old')(_watch)
    new = cachedproperty('new')(lambda self: self._watch('456'))

    def add(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BugWatchAdded(
                    self.ten_minutes_ago, self.product_owner, item))

    def remove(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BugWatchRemoved(
                    self.ten_minutes_ago, self.product_owner, item))

class TestEmailNotificationsBranch(
    EmailNotificationsAddedRemovedMixin, EmailNotificationTestBase):

    added_message = '** Branch linked:'
    removed_message = '** Branch unlinked:'

    def _branch(self):
        with lp_dbuser():
            return self.factory.makeBranch()

    old = cachedproperty('old')(_branch)
    new = cachedproperty('new')(_branch)

    def add(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BranchLinkedToBug(
                    self.ten_minutes_ago, self.person, item, self.bug))

    def remove(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BranchUnlinkedFromBug(
                    self.ten_minutes_ago, self.person, item, self.bug))

class TestEmailNotificationsCVE(
    EmailNotificationsAddedRemovedMixin, EmailNotificationTestBase):

    added_message = '** CVE added:'
    removed_message = '** CVE removed:'

    def _cve(self, sequence):
        with lp_dbuser():
            return self.factory.makeCVE(sequence)

    old = cachedproperty('old')(lambda self: self._cve('2020-1234'))
    new = cachedproperty('new')(lambda self: self._cve('2020-5678'))

    def add(self, item):
        with lp_dbuser():
            self.bug.addChange(
                CveLinkedToBug(
                    self.ten_minutes_ago, self.person, item))

    def remove(self, item):
        with lp_dbuser():
            self.bug.addChange(
                CveUnlinkedFromBug(
                    self.ten_minutes_ago, self.person, item))

class TestEmailNotificationsAttachments(
    EmailNotificationsAddedRemovedMixin, EmailNotificationTestBase):

    added_message = '** Attachment added:'
    removed_message = '** Attachment removed:'

    def _attachment(self):
        with lp_dbuser():
            # This actually creates a notification all by itself, via an
            # event subscriber.  However, it won't be sent out for
            # another five minutes.  Therefore, we send out separate
            # change notifications.
            return self.bug.addAttachment(
                self.person, 'content', 'a comment', 'stuff.txt')

    old = cachedproperty('old')(_attachment)
    new = cachedproperty('new')(_attachment)

    def add(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BugAttachmentChange(
                    self.ten_minutes_ago, self.person, 'attachment',
                    None, item))

    def remove(self, item):
        with lp_dbuser():
            self.bug.addChange(
                BugAttachmentChange(
                    self.ten_minutes_ago, self.person, 'attachment',
                    item, None))

class TestEmailNotificationsWithFilters(TestCaseWithFactory):
    """Ensure outgoing mails have corresponding headers, accounting for mutes.

    Every filter that could have potentially caused a notification to
    go off has a `BugNotificationFilter` record linking a `BugNotification`
    and a `BugSubscriptionFilter`.

    From those records, we include all BugSubscriptionFilter.description
    in X-Subscription-Filter-Description headers in each email.

    Every team filter that caused notifications might be muted for a
    given recipient.  These can cause headers to be omitted, and if all
    filters that caused the notifications are omitted then the
    notification itself will not be sent.
    """

    layer = LaunchpadZopelessLayer

    def setUp(self):
        super(TestEmailNotificationsWithFilters, self).setUp()
        self.bug = self.factory.makeBug()
        subscriber = self.factory.makePerson()
        self.subscription = self.bug.default_bugtask.target.addSubscription(
            subscriber, subscriber)
        self.filter_count = 0
        self.notification = self.addNotification(subscriber)

    def addNotificationRecipient(self, notification, person):
        # Manually insert BugNotificationRecipient for
        # construct_email_notifications to work.
        # Not sure why using SQLObject constructor doesn't work (it
        # tries to insert a row with only the ID which fails).
        Store.of(notification).execute("""
            INSERT INTO BugNotificationRecipient
              (bug_notification, person, reason_header, reason_body)
              VALUES (%s, %s, %s, %s)""" % sqlvalues(
                          notification, person,
                          u'reason header', u'reason body'))

    def addNotification(self, person):
        # Add a notification along with recipient data.
        # This is generally done with BugTaskSet.addNotification()
        # but that requires a more complex set-up.
        message = self.factory.makeMessage()
        notification = BugNotification(
            message=message, activity=None, bug=self.bug,
            is_comment=False, date_emailed=None)
        self.addNotificationRecipient(notification, person)
        return notification

    def addFilter(self, description, subscription=None):
        if subscription is None:
            subscription = self.subscription
            filter_count = self.filter_count
            self.filter_count += 1
        else:
            # For a non-default subscription, always use
            # the initial filter.
            filter_count = 0

        # If no filters have been requested before,
        # use the initial auto-created filter for a subscription.
        if filter_count == 0:
            bug_filter = subscription.bug_filters.one()
        else:
            bug_filter = subscription.newBugFilter()
        bug_filter.description = description
        BugNotificationFilter(
            bug_notification=self.notification,
            bug_subscription_filter=bug_filter)

    def getSubscriptionEmailHeaders(self, by_person=False):
        filtered, omitted, messages = construct_email_notifications(
            [self.notification])
        if by_person:
            headers = {}
        else:
            headers = set()
        for message in messages:
            if by_person:
                headers[message['to']] = message.get_all(
                    "X-Launchpad-Subscription", [])
            else:
                headers = headers.union(
                    set(message.get_all(
                        "X-Launchpad-Subscription", [])))
        return headers

    def getSubscriptionEmailBody(self, by_person=False):
        filtered, omitted, messages = construct_email_notifications(
            [self.notification])
        if by_person:
            filter_texts = {}
        else:
            filter_texts = set()
        for message in messages:
            filters_line = None
            for line in message.get_payload().splitlines():
                if line.startswith("Matching subscriptions: "):
                    filters_line = line
                    break
            if filters_line is not None:
                if by_person:
                    filter_texts[message['to']] = filters_line
                else:
                    filter_texts.add(filters_line)
        return filter_texts

    def test_header_empty(self):
        # An initial filter with no description doesn't cause any
        # headers to be added.
        self.assertContentEqual([],
                                self.getSubscriptionEmailHeaders())

    def test_header_single(self):
        # A single filter with a description makes all emails
        # include that particular filter description in a header.
        self.addFilter(u"Test filter")

        self.assertContentEqual([u"Test filter"],
                                self.getSubscriptionEmailHeaders())

    def test_header_multiple(self):
        # Multiple filters with a description make all emails
        # include all filter descriptions in the header.
        self.addFilter(u"First filter")
        self.addFilter(u"Second filter")

        self.assertContentEqual([u"First filter", u"Second filter"],
                                self.getSubscriptionEmailHeaders())

    def test_header_other_subscriber_by_person(self):
        # Filters for a different subscribers are included only
        # in email messages relevant to them, even if they might
        # all be for the same notification.
        other_person = self.factory.makePerson()
        other_subscription = self.bug.default_bugtask.target.addSubscription(
            other_person, other_person)
        self.addFilter(u"Someone's filter", other_subscription)
        self.addNotificationRecipient(self.notification, other_person)

        self.addFilter(u"Test filter")

        the_subscriber = self.subscription.subscriber
        self.assertEqual(
            {other_person.preferredemail.email: [u"Someone's filter"],
             the_subscriber.preferredemail.email: [u"Test filter"]},
            self.getSubscriptionEmailHeaders(by_person=True))

    def test_body_empty(self):
        # An initial filter with no description doesn't cause any
        # text to be added to the email body.
        self.assertContentEqual([],
                                self.getSubscriptionEmailBody())

    def test_body_single(self):
        # A single filter with a description makes all emails
        # include that particular filter description in the body.
        self.addFilter(u"Test filter")

        self.assertContentEqual([u"Matching subscriptions: Test filter"],
                                self.getSubscriptionEmailBody())

    def test_body_multiple(self):
        # Multiple filters with description make all emails
        # include them in the email body.
        self.addFilter(u"First filter")
        self.addFilter(u"Second filter")

        self.assertContentEqual(
            [u"Matching subscriptions: First filter, Second filter"],
            self.getSubscriptionEmailBody())

    def test_muted(self):
        self.addFilter(u"Test filter")
        BugSubscriptionFilterMute(
            person=self.subscription.subscriber,
            filter=self.notification.bug_filters.one())
        filtered, omitted, messages = construct_email_notifications(
            [self.notification])
        self.assertEqual(list(messages), [])

    def test_header_multiple_one_muted(self):
        # Multiple filters with a description make all emails
        # include all filter descriptions in the header.
        self.addFilter(u"First filter")
        self.addFilter(u"Second filter")
        BugSubscriptionFilterMute(
            person=self.subscription.subscriber,
            filter=self.notification.bug_filters[0])

        self.assertContentEqual([u"Second filter"],
                                self.getSubscriptionEmailHeaders())

def fetch_notifications(subscriber, bug):
    return IStore(BugNotification).find(
        BugNotification,
        BugNotification.id == BugNotificationRecipient.bug_notificationID,
        BugNotificationRecipient.personID == subscriber.id,
        BugNotification.bug == bug)

class TestEmailNotificationsWithFiltersWhenBugCreated(TestCaseWithFactory):
    # See bug 720147.

    layer = LaunchpadZopelessLayer

    def setUp(self):
        super(TestEmailNotificationsWithFiltersWhenBugCreated, self).setUp()
        self.subscriber = self.factory.makePerson()
        self.submitter = self.factory.makePerson()
        self.product = self.factory.makeProduct(
            bug_supervisor=self.submitter)
        self.subscription = self.product.addSubscription(
            self.subscriber, self.subscriber)
        self.filter = self.subscription.bug_filters[0]
        self.filter.description = u'Needs triage'
        self.filter.statuses = [BugTaskStatus.NEW, BugTaskStatus.INCOMPLETE]

    def test_filters_match_when_bug_is_created(self):
        message = u"this is an unfiltered comment"
        params = CreateBugParams(
            title=u"crashes all the time",
            comment=message, owner=self.submitter,
            status=BugTaskStatus.NEW)
        bug = self.product.createBug(params)
        notification = fetch_notifications(self.subscriber, bug).one()
        self.assertEqual(notification.message.text_contents, message)

    def test_filters_do_not_match_when_bug_is_created(self):
        message = u"this is a filtered comment"
        params = CreateBugParams(
            title=u"crashes all the time",
            comment=message, owner=self.submitter,
            status=BugTaskStatus.TRIAGED,
            importance=BugTaskImportance.HIGH)
        bug = self.product.createBug(params)
        notifications = fetch_notifications(self.subscriber, bug)
        self.assertTrue(notifications.is_empty())

class TestManageNotificationsMessage(TestCaseWithFactory):

    layer = LaunchpadZopelessLayer

    def test_manage_notifications_message_is_included(self):
        # Set up a subscription to a product.
        subscriber = self.factory.makePerson()
        submitter = self.factory.makePerson()
        product = self.factory.makeProduct(
            bug_supervisor=submitter)
        product.addSubscription(subscriber, subscriber)
        # Create a bug that will match the subscription.
        bug = product.createBug(CreateBugParams(
            title=self.factory.getUniqueString(),
            comment=self.factory.getUniqueString(),
            owner=submitter))
        notification = fetch_notifications(subscriber, bug).one()
        _, _, (message,) = construct_email_notifications([notification])
        payload = message.get_payload()
        self.assertThat(payload, Contains(
            'To manage notifications about this bug go to:\nhttp://'))

class TestNotificationSignatureSeparator(TestCase):

    def test_signature_separator(self):
        # Email signatures are often separated from the body of a message by a
        # special separator so user agents can identify the signature for
        # special treatment (hiding, stripping when replying, colorizing,
        # etc.).  The bug notification messages follow the convention.
        names = ['bug-notification-verbose.txt', 'bug-notification.txt']
        for name in names:
            template = get_email_template(name, 'bugs')
            self.assertTrue(re.search('^-- $', template, re.MULTILINE))

class TestExpandedNotificationFooters(EmailNotificationTestBase):

    layer = LaunchpadZopelessLayer

    def test_expanded_footer(self):
        # Recipients with expanded_notification_footers receive an expanded
        # footer on messages, which is separated by the correct number of
        # newlines.
        with lp_dbuser(), person_logged_in(self.bug_subscriber):
            self.bug_subscriber.expanded_notification_footers = True
            expected_to = str(self.bug_subscriber.preferredemail.email)
        self.bug.addChange(BugTitleChange(
            self.ten_minutes_ago, self.person, "title",
            "Old summary", "New summary"))
        [payload] = [
            payload for message, payload in self.get_messages()
            if message["to"] == expected_to]
        self.assertThat(payload, MatchesRegex(
            r'.*To manage notifications about this bug go to:\n'
            r'http://.*\+subscriptions\n'
            r'\n'
            r'Launchpad-Notification-Type: bug\n', re.S))

class TestDeferredNotifications(TestCaseWithFactory):

    layer = LaunchpadZopelessLayer

    def setUp(self):
        super(TestDeferredNotifications, self).setUp()
        self.notification_set = getUtility(IBugNotificationSet)
        # Ensure there are no outstanding notifications.
        for notification in self.notification_set.getNotificationsToSend():
            notification.destroySelf()
        self.ten_minutes_ago = datetime.now(pytz.UTC) - timedelta(minutes=10)

    def _make_deferred_notification(self):
        bug = self.factory.makeBug()
        empty_recipients = BugNotificationRecipients()
        message = getUtility(IMessageSet).fromText(
            'subject', 'a comment.', bug.owner,
            datecreated=self.ten_minutes_ago)
        self.notification_set.addNotification(
            bug, False, message, empty_recipients, None, deferred=True)

    def test_deferred_notifications(self):
        # Create some deferred notifications and show that processing them
        # puts then in the state where they are ready to send.
        num = 5
        for i in range(num):
            self._make_deferred_notification()
        deferred = self.notification_set.getDeferredNotifications()
        self.assertEqual(num, deferred.count())
        process_deferred_notifications(deferred)
        # Now that are all in the PENDING state.
        ready_to_send = self.notification_set.getNotificationsToSend()
        self.assertEqual(num, len(ready_to_send))
        # And there are no longer any deferred.
        deferred = self.notification_set.getDeferredNotifications()
        self.assertEqual(0, deferred.count())

class BrokenMailer(TestMailer):
    """A mailer that raises an exception for certain recipient addresses."""

    def __init__(self, broken):
        self.broken = broken

    def send(self, from_addr, to_addrs, message):
        if any(to_addr in self.broken for to_addr in to_addrs):
            raise SMTPException("test requested delivery failure")
        return super(BrokenMailer, self).send(from_addr, to_addrs, message)

class TestSendBugNotifications(TestCaseWithFactory):

    layer = LaunchpadZopelessLayer

    def setUp(self):
        super(TestSendBugNotifications, self).setUp()
        self.notification_set = getUtility(IBugNotificationSet)
        # Ensure there are no outstanding notifications.
        for notification in self.notification_set.getNotificationsToSend():
            notification.destroySelf()
        self.ten_minutes_ago = datetime.now(pytz.UTC) - timedelta(minutes=10)

    def test_oops_on_failed_delivery(self):
        # If one notification fails to send, it logs an OOPS and doesn't get
        # in the way of sending other notifications.
        set_immediate_mail_delivery(False)
        subscribers = []
        for i in range(3):
            bug = self.factory.makeBug()
            subscriber = self.factory.makePerson()
            subscribers.append(subscriber.preferredemail.email)
            bug.default_bugtask.target.addSubscription(subscriber, subscriber)
            message = getUtility(IMessageSet).fromText(
                "subject", "a comment.", bug.owner,
                datecreated=self.ten_minutes_ago)
            bug.addCommentNotification(message)

        notifications = list(get_email_notifications(
            self.notification_set.getNotificationsToSend()))
        self.assertEqual(3, len(notifications))

        mailer = BrokenMailer([subscribers[1]])
        with ZopeUtilityFixture(mailer, IMailDelivery, "Mail"):
            switch_dbuser(config.malone.bugnotification_dbuser)
            script = SendBugNotifications(
                "send-bug-notifications", config.malone.bugnotification_dbuser,
                ["-q"])
            script.txn = self.layer.txn
            script.main()

        self.assertEqual(1, len(self.oopses))
        self.assertIn(
            "SMTPException: test requested delivery failure",
            self.oopses[0]["tb_text"])
        for i, (bug_notifications, _, messages) in enumerate(notifications):
            for bug_notification in bug_notifications:
                if i == 1:
                    self.assertEqual(
                        BugNotificationStatus.PENDING, bug_notification.status)
                else:
                    self.assertEqual(
                        BugNotificationStatus.SENT, bug_notification.status)