"""Mechanisms for inferring function types based on callsites.

Currently works by collecting all argument types at callsites,
synthesizing a list of possible function types from that, trying them
all, and picking the one with the fewest errors that we think is the
"best".

Can return JSON that pyannotate can use to apply the annotations to code.

There are a bunch of TODOs here:
 * Maybe want a way to surface the choices not selected??
 * We can generate an exponential number of type suggestions, and probably want
   a way to not always need to check them all.
 * Our heuristics for what types to try are primitive and not yet
   supported by real practice.
 * More!

Other things:
 * This is super brute force. Could we integrate with the typechecker
   more to understand more about what is going on?
 * Like something with tracking constraints/unification variables?
 * No understanding of type variables at *all*
"""

from __future__ import annotations

import itertools
import json
import os
from contextlib import contextmanager
from typing import Callable, Iterator, NamedTuple, TypeVar, cast
from typing_extensions import TypedDict

from mypy.argmap import map_actuals_to_formals
from mypy.build import Graph, State
from mypy.checkexpr import has_any_type
from mypy.find_sources import InvalidSourceList, SourceFinder
from mypy.join import join_type_list
from mypy.meet import meet_type_list
from mypy.modulefinder import PYTHON_EXTENSIONS
from mypy.nodes import (
    ARG_STAR,
    ARG_STAR2,
    ArgKind,
    CallExpr,
    Decorator,
    Expression,
    FuncDef,
    MypyFile,
    RefExpr,
    ReturnStmt,
    SymbolNode,
    SymbolTable,
    TypeInfo,
    reverse_builtin_aliases,
)
from mypy.options import Options
from mypy.plugin import FunctionContext, MethodContext, Plugin
from mypy.server.update import FineGrainedBuildManager
from mypy.state import state
from mypy.traverser import TraverserVisitor
from mypy.typeops import make_simplified_union
from mypy.types import (
    AnyType,
    CallableType,
    FunctionLike,
    Instance,
    NoneType,
    ProperType,
    TupleType,
    Type,
    TypeAliasType,
    TypedDictType,
    TypeOfAny,
    TypeStrVisitor,
    TypeTranslator,
    TypeVarType,
    UninhabitedType,
    UnionType,
    get_proper_type,
)
from mypy.types_utils import is_overlapping_none, remove_optional
from mypy.util import split_target

class PyAnnotateSignature(TypedDict):
    return_type: str
    arg_types: list[str]

class Callsite(NamedTuple):
    path: str
    line: int
    arg_kinds: list[list[ArgKind]]
    callee_arg_names: list[str | None]
    arg_names: list[list[str | None]]
    arg_types: list[list[Type]]

class SuggestionPlugin(Plugin):
    """Plugin that records all calls to a given target."""

    def __init__(self, target: str) -> None:
        if target.endswith((".__new__", ".__init__")):
            target = target.rsplit(".", 1)[0]

        self.target = target
        # List of call sites found by dmypy suggest:
        # (path, line, <arg kinds>, <arg names>, <arg types>)
        self.mystery_hits: list[Callsite] = []

    def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
        if fullname == self.target:
            return self.log
        else:
            return None

    def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None:
        if fullname == self.target:
            return self.log
        else:
            return None

    def log(self, ctx: FunctionContext | MethodContext) -> Type:
        self.mystery_hits.append(
            Callsite(
                ctx.api.path,
                ctx.context.line,
                ctx.arg_kinds,
                ctx.callee_arg_names,
                ctx.arg_names,
                ctx.arg_types,
            )
        )
        return ctx.default_return_type

# NOTE: We could make this a bunch faster by implementing a StatementVisitor that skips
# traversing into expressions
class ReturnFinder(TraverserVisitor):
    """Visitor for finding all types returned from a function."""

    def __init__(self, typemap: dict[Expression, Type]) -> None:
        self.typemap = typemap
        self.return_types: list[Type] = []

    def visit_return_stmt(self, o: ReturnStmt) -> None:
        if o.expr is not None and o.expr in self.typemap:
            self.return_types.append(self.typemap[o.expr])

    def visit_func_def(self, o: FuncDef) -> None:
        # Skip nested functions
        pass

def get_return_types(typemap: dict[Expression, Type], func: FuncDef) -> list[Type]:
    """Find all the types returned by return statements in func."""
    finder = ReturnFinder(typemap)
    func.body.accept(finder)
    return finder.return_types

class ArgUseFinder(TraverserVisitor):
    """Visitor for finding all the types of arguments that each arg is passed to.

    This is extremely simple minded but might be effective anyways.
    """

    def __init__(self, func: FuncDef, typemap: dict[Expression, Type]) -> None:
        self.typemap = typemap
        self.arg_types: dict[SymbolNode, list[Type]] = {arg.variable: [] for arg in func.arguments}

    def visit_call_expr(self, o: CallExpr) -> None:
        if not any(isinstance(e, RefExpr) and e.node in self.arg_types for e in o.args):
            return

        typ = get_proper_type(self.typemap.get(o.callee))
        if not isinstance(typ, CallableType):
            return

        formal_to_actual = map_actuals_to_formals(
            o.arg_kinds,
            o.arg_names,
            typ.arg_kinds,
            typ.arg_names,
            lambda n: AnyType(TypeOfAny.special_form),
        )

        for i, args in enumerate(formal_to_actual):
            for arg_idx in args:
                arg = o.args[arg_idx]
                if isinstance(arg, RefExpr) and arg.node in self.arg_types:
                    self.arg_types[arg.node].append(typ.arg_types[i])

def get_arg_uses(typemap: dict[Expression, Type], func: FuncDef) -> list[list[Type]]:
    """Find all the types of arguments that each arg is passed to.

    For example, given
      def foo(x: int) -> None: ...
      def bar(x: str) -> None: ...
      def test(x, y):
          foo(x)
          bar(y)

    this will return [[int], [str]].
    """
    finder = ArgUseFinder(func, typemap)
    func.body.accept(finder)
    return [finder.arg_types[arg.variable] for arg in func.arguments]

class SuggestionFailure(Exception):
    pass

def is_explicit_any(typ: AnyType) -> bool:
    # Originally I wanted to count as explicit anything derived from an explicit any, but that
    # seemed too strict in some testing.
    # return (typ.type_of_any == TypeOfAny.explicit
    #         or (typ.source_any is not None and typ.source_any.type_of_any == TypeOfAny.explicit))
    # Important question: what should we do with source_any stuff? Does that count?
    # And actually should explicit anys count at all?? Maybe not!
    return typ.type_of_any == TypeOfAny.explicit

def is_implicit_any(typ: Type) -> bool:
    typ = get_proper_type(typ)
    return isinstance(typ, AnyType) and not is_explicit_any(typ)

class SuggestionEngine:
    """Engine for finding call sites and suggesting signatures."""

    def __init__(
        self,
        fgmanager: FineGrainedBuildManager,
        *,
        json: bool,
        no_errors: bool = False,
        no_any: bool = False,
        flex_any: float | None = None,
        use_fixme: str | None = None,
        max_guesses: int | None = None,
    ) -> None:
        self.fgmanager = fgmanager
        self.manager = fgmanager.manager
        self.plugin = self.manager.plugin
        self.graph = fgmanager.graph
        self.finder = SourceFinder(self.manager.fscache, self.manager.options)

        self.give_json = json
        self.no_errors = no_errors
        self.flex_any = flex_any
        if no_any:
            self.flex_any = 1.0

        self.max_guesses = max_guesses or 64
        self.use_fixme = use_fixme

    def suggest(self, function: str) -> str:
        """Suggest an inferred type for function."""
        mod, func_name, node = self.find_node(function)

        with self.restore_after(mod):
            with self.with_export_types():
                suggestion = self.get_suggestion(mod, node)

        if self.give_json:
            return self.json_suggestion(mod, func_name, node, suggestion)
        else:
            return self.format_signature(suggestion)

    def suggest_callsites(self, function: str) -> str:
        """Find a list of call sites of function."""
        mod, _, node = self.find_node(function)
        with self.restore_after(mod):
            callsites, _ = self.get_callsites(node)

        return "\n".join(
            dedup(
                [
                    f"{path}:{line}: {self.format_args(arg_kinds, arg_names, arg_types)}"
                    for path, line, arg_kinds, _, arg_names, arg_types in callsites
                ]
            )
        )

    @contextmanager
    def restore_after(self, module: str) -> Iterator[None]:
        """Context manager that reloads a module after executing the body.

        This should undo any damage done to the module state while mucking around.
        """
        try:
            yield
        finally:
            self.reload(self.graph[module])

    @contextmanager
    def with_export_types(self) -> Iterator[None]:
        """Context manager that enables the export_types flag in the body.

        This causes type information to be exported into the manager's all_types variable.
        """
        old = self.manager.options.export_types
        self.manager.options.export_types = True
        try:
            yield
        finally:
            self.manager.options.export_types = old

    def get_trivial_type(self, fdef: FuncDef) -> CallableType:
        """Generate a trivial callable type from a func def, with all Anys"""
        # The Anys are marked as being from the suggestion engine
        # since they need some special treatment (specifically,
        # constraint generation ignores them.)
        return CallableType(
            [AnyType(TypeOfAny.suggestion_engine) for _ in fdef.arg_kinds],
            fdef.arg_kinds,
            fdef.arg_names,
            AnyType(TypeOfAny.suggestion_engine),
            self.named_type("builtins.function"),
        )

    def get_starting_type(self, fdef: FuncDef) -> CallableType:
        if isinstance(fdef.type, CallableType):
            return make_suggestion_anys(fdef.type)
        else:
            return self.get_trivial_type(fdef)

    def get_args(
        self,
        is_method: bool,
        base: CallableType,
        defaults: list[Type | None],
        callsites: list[Callsite],
        uses: list[list[Type]],
    ) -> list[list[Type]]:
        """Produce a list of type suggestions for each argument type."""
        types: list[list[Type]] = []
        for i in range(len(base.arg_kinds)):
            # Make self args Any but this will get overridden somewhere in the checker
            if i == 0 and is_method:
                types.append([AnyType(TypeOfAny.suggestion_engine)])
                continue

            all_arg_types = []
            for call in callsites:
                for typ in call.arg_types[i - is_method]:
                    # Collect all the types except for implicit anys
                    if not is_implicit_any(typ):
                        all_arg_types.append(typ)
            all_use_types = []
            for typ in uses[i]:
                # Collect all the types except for implicit anys
                if not is_implicit_any(typ):
                    all_use_types.append(typ)
            # Add in any default argument types
            default = defaults[i]
            if default:
                all_arg_types.append(default)
                if all_use_types:
                    all_use_types.append(default)

            arg_types = []

            if all_arg_types and all(
                isinstance(get_proper_type(tp), NoneType) for tp in all_arg_types
            ):
                arg_types.append(
                    UnionType.make_union([all_arg_types[0], AnyType(TypeOfAny.explicit)])
                )
            elif all_arg_types:
                arg_types.extend(generate_type_combinations(all_arg_types))
            else:
                arg_types.append(AnyType(TypeOfAny.explicit))

            if all_use_types:
                # This is a meet because the type needs to be compatible with all the uses
                arg_types.append(meet_type_list(all_use_types))

            types.append(arg_types)
        return types

    def get_default_arg_types(self, fdef: FuncDef) -> list[Type | None]:
        return [
            self.manager.all_types[arg.initializer] if arg.initializer else None
            for arg in fdef.arguments
        ]

    def get_guesses(
        self,
        is_method: bool,
        base: CallableType,
        defaults: list[Type | None],
        callsites: list[Callsite],
        uses: list[list[Type]],
    ) -> list[CallableType]:
        """Compute a list of guesses for a function's type.

        This focuses just on the argument types, and doesn't change the provided return type.
        """
        options = self.get_args(is_method, base, defaults, callsites, uses)

        # Take the first `max_guesses` guesses.
        product = itertools.islice(itertools.product(*options), 0, self.max_guesses)
        return [refine_callable(base, base.copy_modified(arg_types=list(x))) for x in product]

    def get_callsites(self, func: FuncDef) -> tuple[list[Callsite], list[str]]:
        """Find all call sites of a function."""
        new_type = self.get_starting_type(func)

        collector_plugin = SuggestionPlugin(func.fullname)

        self.plugin._plugins.insert(0, collector_plugin)
        try:
            errors = self.try_type(func, new_type)
        finally:
            self.plugin._plugins.pop(0)

        return collector_plugin.mystery_hits, errors

    def filter_options(
        self, guesses: list[CallableType], is_method: bool, ignore_return: bool
    ) -> list[CallableType]:
        """Apply any configured filters to the possible guesses.

        Currently the only option is filtering based on Any prevalance."""
        return [
            t
            for t in guesses
            if self.flex_any is None
            or any_score_callable(t, is_method, ignore_return) >= self.flex_any
        ]

    def find_best(self, func: FuncDef, guesses: list[CallableType]) -> tuple[CallableType, int]:
        """From a list of possible function types, find the best one.

        For best, we want the fewest errors, then the best "score" from score_callable.
        """
        if not guesses:
            raise SuggestionFailure("No guesses that match criteria!")
        errors = {guess: self.try_type(func, guess) for guess in guesses}
        best = min(guesses, key=lambda s: (count_errors(errors[s]), self.score_callable(s)))
        return best, count_errors(errors[best])

    def get_guesses_from_parent(self, node: FuncDef) -> list[CallableType]:
        """Try to get a guess of a method type from a parent class."""
        if not node.info:
            return []

        for parent in node.info.mro[1:]:
            pnode = parent.names.get(node.name)
            if pnode and isinstance(pnode.node, (FuncDef, Decorator)):
                typ = get_proper_type(pnode.node.type)
                # FIXME: Doesn't work right with generic tyeps
                if isinstance(typ, CallableType) and len(typ.arg_types) == len(node.arguments):
                    # Return the first thing we find, since it probably doesn't make sense
                    # to grab things further up in the chain if an earlier parent has it.
                    return [typ]

        return []

    def get_suggestion(self, mod: str, node: FuncDef) -> PyAnnotateSignature:
        """Compute a suggestion for a function.

        Return the type and whether the first argument should be ignored.
        """
        graph = self.graph
        callsites, orig_errors = self.get_callsites(node)
        uses = get_arg_uses(self.manager.all_types, node)

        if self.no_errors and orig_errors:
            raise SuggestionFailure("Function does not typecheck.")

        is_method = bool(node.info) and not node.is_static

        with state.strict_optional_set(graph[mod].options.strict_optional):
            guesses = self.get_guesses(
                is_method,
                self.get_starting_type(node),
                self.get_default_arg_types(node),
                callsites,
                uses,
            )
        guesses += self.get_guesses_from_parent(node)
        guesses = self.filter_options(guesses, is_method, ignore_return=True)
        best, _ = self.find_best(node, guesses)

        # Now try to find the return type!
        self.try_type(node, best)
        returns = get_return_types(self.manager.all_types, node)
        with state.strict_optional_set(graph[mod].options.strict_optional):
            if returns:
                ret_types = generate_type_combinations(returns)
            else:
                ret_types = [NoneType()]

        guesses = [best.copy_modified(ret_type=refine_type(best.ret_type, t)) for t in ret_types]
        guesses = self.filter_options(guesses, is_method, ignore_return=False)
        best, errors = self.find_best(node, guesses)

        if self.no_errors and errors:
            raise SuggestionFailure("No annotation without errors")

        return self.pyannotate_signature(mod, is_method, best)

    def format_args(
        self,
        arg_kinds: list[list[ArgKind]],
        arg_names: list[list[str | None]],
        arg_types: list[list[Type]],
    ) -> str:
        args: list[str] = []
        for i in range(len(arg_types)):
            for kind, name, typ in zip(arg_kinds[i], arg_names[i], arg_types[i]):
                arg = self.format_type(None, typ)
                if kind == ARG_STAR:
                    arg = "*" + arg
                elif kind == ARG_STAR2:
                    arg = "**" + arg
                elif kind.is_named():
                    if name:
                        arg = f"{name}={arg}"
            args.append(arg)
        return f"({', '.join(args)})"

    def find_node(self, key: str) -> tuple[str, str, FuncDef]:
        """From a target name, return module/target names and the func def.

        The 'key' argument can be in one of two formats:
        * As the function full name, e.g., package.module.Cls.method
        * As the function location as file and line separated by column,
          e.g., path/to/file.py:42
        """
        # TODO: Also return OverloadedFuncDef -- currently these are ignored.
        node: SymbolNode | None = None
        if ":" in key:
            if key.count(":") > 1:
                raise SuggestionFailure(
                    "Malformed location for function: {}. Must be either"
                    " package.module.Class.method or path/to/file.py:line".format(key)
                )
            file, line = key.split(":")
            if not line.isdigit():
                raise SuggestionFailure(f"Line number must be a number. Got {line}")
            line_number = int(line)
            modname, node = self.find_node_by_file_and_line(file, line_number)
            tail = node.fullname[len(modname) + 1 :]  # add one to account for '.'
        else:
            target = split_target(self.fgmanager.graph, key)
            if not target:
                raise SuggestionFailure(f"Cannot find module for {key}")
            modname, tail = target
            node = self.find_node_by_module_and_name(modname, tail)

        if isinstance(node, Decorator):
            node = self.extract_from_decorator(node)
            if not node:
                raise SuggestionFailure(f"Object {key} is a decorator we can't handle")

        if not isinstance(node, FuncDef):
            raise SuggestionFailure(f"Object {key} is not a function")

        return modname, tail, node

    def find_node_by_module_and_name(self, modname: str, tail: str) -> SymbolNode | None:
        """Find symbol node by module id and qualified name.

        Raise SuggestionFailure if can't find one.
        """
        tree = self.ensure_loaded(self.fgmanager.graph[modname])

        # N.B. This is reimplemented from update's lookup_target
        # basically just to produce better error messages.

        names: SymbolTable = tree.names

        # Look through any classes
        components = tail.split(".")
        for i, component in enumerate(components[:-1]):
            if component not in names:
                raise SuggestionFailure(
                    "Unknown class {}.{}".format(modname, ".".join(components[: i + 1]))
                )
            node: SymbolNode | None = names[component].node
            if not isinstance(node, TypeInfo):
                raise SuggestionFailure(
                    "Object {}.{} is not a class".format(modname, ".".join(components[: i + 1]))
                )
            names = node.names

        # Look for the actual function/method
        funcname = components[-1]
        if funcname not in names:
            key = modname + "." + tail
            raise SuggestionFailure(
                "Unknown {} {}".format("method" if len(components) > 1 else "function", key)
            )
        return names[funcname].node

    def find_node_by_file_and_line(self, file: str, line: int) -> tuple[str, SymbolNode]:
        """Find symbol node by path to file and line number.

        Find the first function declared *before or on* the line number.

        Return module id and the node found. Raise SuggestionFailure if can't find one.
        """
        if not any(file.endswith(ext) for ext in PYTHON_EXTENSIONS):
            raise SuggestionFailure("Source file is not a Python file")
        try:
            modname, _ = self.finder.crawl_up(os.path.normpath(file))
        except InvalidSourceList as e:
            raise SuggestionFailure("Invalid source file name: " + file) from e
        if modname not in self.graph:
            raise SuggestionFailure("Unknown module: " + modname)
        # We must be sure about any edits in this file as this might affect the line numbers.
        tree = self.ensure_loaded(self.fgmanager.graph[modname], force=True)
        node: SymbolNode | None = None
        closest_line: int | None = None
        # TODO: Handle nested functions.
        for _, sym, _ in tree.local_definitions():
            if isinstance(sym.node, (FuncDef, Decorator)):
                sym_line = sym.node.line
            # TODO: add support for OverloadedFuncDef.
            else:
                continue

            # We want the closest function above the specified line
            if sym_line <= line and (closest_line is None or sym_line > closest_line):
                closest_line = sym_line
                node = sym.node
        if not node:
            raise SuggestionFailure(f"Cannot find a function at line {line}")
        return modname, node

    def extract_from_decorator(self, node: Decorator) -> FuncDef | None:
        for dec in node.decorators:
            typ = None
            if isinstance(dec, RefExpr) and isinstance(dec.node, FuncDef):
                typ = dec.node.type
            elif (
                isinstance(dec, CallExpr)
                and isinstance(dec.callee, RefExpr)
                and isinstance(dec.callee.node, FuncDef)
                and isinstance(dec.callee.node.type, CallableType)
            ):
                typ = get_proper_type(dec.callee.node.type.ret_type)

            if not isinstance(typ, FunctionLike):
                return None
            for ct in typ.items:
                if not (
                    len(ct.arg_types) == 1
                    and isinstance(ct.arg_types[0], TypeVarType)
                    and ct.arg_types[0] == ct.ret_type
                ):
                    return None

        return node.func

    def try_type(self, func: FuncDef, typ: ProperType) -> list[str]:
        """Recheck a function while assuming it has type typ.

        Return all error messages.
        """
        old = func.unanalyzed_type
        # During reprocessing, unanalyzed_type gets copied to type (by aststrip).
        # We set type to None to ensure that the type always changes during
        # reprocessing.
        func.type = None
        func.unanalyzed_type = typ
        try:
            res = self.fgmanager.trigger(func.fullname)
            # if res:
            #     print('===', typ)
            #     print('\n'.join(res))
            return res
        finally:
            func.unanalyzed_type = old

    def reload(self, state: State) -> list[str]:
        """Recheck the module given by state."""
        assert state.path is not None
        self.fgmanager.flush_cache()
        return self.fgmanager.update([(state.id, state.path)], [])

    def ensure_loaded(self, state: State, force: bool = False) -> MypyFile:
        """Make sure that the module represented by state is fully loaded."""
        if not state.tree or state.tree.is_cache_skeleton or force:
            self.reload(state)
        assert state.tree is not None
        return state.tree

    def named_type(self, s: str) -> Instance:
        return self.manager.semantic_analyzer.named_type(s)

    def json_suggestion(
        self, mod: str, func_name: str, node: FuncDef, suggestion: PyAnnotateSignature
    ) -> str:
        """Produce a json blob for a suggestion suitable for application by pyannotate."""
        # pyannotate irritatingly drops class names for class and static methods
        if node.is_class or node.is_static:
            func_name = func_name.split(".", 1)[-1]

        # pyannotate works with either paths relative to where the
        # module is rooted or with absolute paths. We produce absolute
        # paths because it is simpler.
        path = os.path.abspath(self.graph[mod].xpath)

        obj = {
            "signature": suggestion,
            "line": node.line,
            "path": path,
            "func_name": func_name,
            "samples": 0,
        }
        return json.dumps([obj], sort_keys=True)

    def pyannotate_signature(
        self, cur_module: str | None, is_method: bool, typ: CallableType
    ) -> PyAnnotateSignature:
        """Format a callable type as a pyannotate dict"""
        start = int(is_method)
        return {
            "arg_types": [self.format_type(cur_module, t) for t in typ.arg_types[start:]],
            "return_type": self.format_type(cur_module, typ.ret_type),
        }

    def format_signature(self, sig: PyAnnotateSignature) -> str:
        """Format a callable type in a way suitable as an annotation... kind of"""
        return f"({', '.join(sig['arg_types'])}) -> {sig['return_type']}"

    def format_type(self, cur_module: str | None, typ: Type) -> str:
        if self.use_fixme and isinstance(get_proper_type(typ), AnyType):
            return self.use_fixme
        return typ.accept(TypeFormatter(cur_module, self.graph, self.manager.options))

    def score_type(self, t: Type, arg_pos: bool) -> int:
        """Generate a score for a type that we use to pick which type to use.

        Lower is better, prefer non-union/non-any types. Don't penalize optionals.
        """
        t = get_proper_type(t)
        if isinstance(t, AnyType):
            return 20
        if arg_pos and isinstance(t, NoneType):
            return 20
        if isinstance(t, UnionType):
            if any(isinstance(get_proper_type(x), AnyType) for x in t.items):
                return 20
            if any(has_any_type(x) for x in t.items):
                return 15
            if not is_overlapping_none(t):
                return 10
        if isinstance(t, CallableType) and (has_any_type(t) or is_tricky_callable(t)):
            return 10
        return 0

    def score_callable(self, t: CallableType) -> int:
        return sum(self.score_type(x, arg_pos=True) for x in t.arg_types) + self.score_type(
            t.ret_type, arg_pos=False
        )

def any_score_type(ut: Type, arg_pos: bool) -> float:
    """Generate a very made up number representing the Anyness of a type.

    Higher is better, 1.0 is max
    """
    t = get_proper_type(ut)
    if isinstance(t, AnyType) and t.type_of_any != TypeOfAny.suggestion_engine:
        return 0
    if isinstance(t, NoneType) and arg_pos:
        return 0.5
    if isinstance(t, UnionType):
        if any(isinstance(get_proper_type(x), AnyType) for x in t.items):
            return 0.5
        if any(has_any_type(x) for x in t.items):
            return 0.25
    if isinstance(t, CallableType) and is_tricky_callable(t):
        return 0.5
    if has_any_type(t):
        return 0.5

    return 1.0

def any_score_callable(t: CallableType, is_method: bool, ignore_return: bool) -> float:
    # Ignore the first argument of methods
    scores = [any_score_type(x, arg_pos=True) for x in t.arg_types[int(is_method) :]]
    # Return type counts twice (since it spreads type information), unless it is
    # None in which case it does not count at all. (Though it *does* still count
    # if there are no arguments.)
    if not isinstance(get_proper_type(t.ret_type), NoneType) or not scores:
        ret = 1.0 if ignore_return else any_score_type(t.ret_type, arg_pos=False)
        scores += [ret, ret]

    return sum(scores) / len(scores)

def is_tricky_callable(t: CallableType) -> bool:
    """Is t a callable that we need to put a ... in for syntax reasons?"""
    return t.is_ellipsis_args or any(k.is_star() or k.is_named() for k in t.arg_kinds)

class TypeFormatter(TypeStrVisitor):
    """Visitor used to format types"""

    # TODO: Probably a lot
    def __init__(self, module: str | None, graph: Graph, options: Options) -> None:
        super().__init__(options=options)
        self.module = module
        self.graph = graph

    def visit_any(self, t: AnyType) -> str:
        if t.missing_import_name:
            return t.missing_import_name
        else:
            return "Any"

    def visit_instance(self, t: Instance) -> str:
        s = t.type.fullname or t.type.name or None
        if s is None:
            return "<???>"
        if s in reverse_builtin_aliases:
            s = reverse_builtin_aliases[s]

        mod_obj = split_target(self.graph, s)
        assert mod_obj
        mod, obj = mod_obj

        # If a class is imported into the current module, rewrite the reference
        # to point to the current module. This helps the annotation tool avoid
        # inserting redundant imports when a type has been reexported.
        if self.module:
            parts = obj.split(".")  # need to split the object part if it is a nested class
            tree = self.graph[self.module].tree
            if tree and parts[0] in tree.names:
                mod = self.module

        if (mod, obj) == ("builtins", "tuple"):
            mod, obj = "typing", "Tuple[" + t.args[0].accept(self) + ", ...]"
        elif t.args:
            obj += f"[{self.list_str(t.args)}]"

        if mod_obj == ("builtins", "unicode"):
            return "Text"
        elif mod == "builtins":
            return obj
        else:
            delim = "." if "." not in obj else ":"
            return mod + delim + obj

    def visit_tuple_type(self, t: TupleType) -> str:
        if t.partial_fallback and t.partial_fallback.type:
            fallback_name = t.partial_fallback.type.fullname
            if fallback_name != "builtins.tuple":
                return t.partial_fallback.accept(self)
        s = self.list_str(t.items)
        return f"Tuple[{s}]"

    def visit_uninhabited_type(self, t: UninhabitedType) -> str:
        return "Any"

    def visit_typeddict_type(self, t: TypedDictType) -> str:
        return t.fallback.accept(self)

    def visit_union_type(self, t: UnionType) -> str:
        if len(t.items) == 2 and is_overlapping_none(t):
            return f"Optional[{remove_optional(t).accept(self)}]"
        else:
            return super().visit_union_type(t)

    def visit_callable_type(self, t: CallableType) -> str:
        # TODO: use extended callables?
        if is_tricky_callable(t):
            arg_str = "..."
        else:
            # Note: for default arguments, we just assume that they
            # are required.  This isn't right, but neither is the
            # other thing, and I suspect this will produce more better
            # results than falling back to `...`
            args = [typ.accept(self) for typ in t.arg_types]
            arg_str = f"[{', '.join(args)}]"

        return f"Callable[{arg_str}, {t.ret_type.accept(self)}]"

TType = TypeVar("TType", bound=Type)

def make_suggestion_anys(t: TType) -> TType:
    """Make all anys in the type as coming from the suggestion engine.

    This keeps those Anys from influencing constraint generation,
    which allows us to do better when refining types.
    """
    return cast(TType, t.accept(MakeSuggestionAny()))

class MakeSuggestionAny(TypeTranslator):
    def visit_any(self, t: AnyType) -> Type:
        if not t.missing_import_name:
            return t.copy_modified(type_of_any=TypeOfAny.suggestion_engine)
        else:
            return t

    def visit_type_alias_type(self, t: TypeAliasType) -> Type:
        return t.copy_modified(args=[a.accept(self) for a in t.args])

def generate_type_combinations(types: list[Type]) -> list[Type]:
    """Generate possible combinations of a list of types.

    mypy essentially supports two different ways to do this: joining the types
    and unioning the types. We try both.
    """
    joined_type = join_type_list(types)
    union_type = make_simplified_union(types)
    if joined_type == union_type:
        return [joined_type]
    else:
        return [joined_type, union_type]

def count_errors(msgs: list[str]) -> int:
    return len([x for x in msgs if " error: " in x])

def refine_type(ti: Type, si: Type) -> Type:
    """Refine `ti` by replacing Anys in it with information taken from `si`

    This basically works by, when the types have the same structure,
    traversing both of them in parallel and replacing Any on the left
    with whatever the type on the right is. If the types don't have the
    same structure (or aren't supported), the left type is chosen.

    For example:
      refine(Any, T) = T,  for all T
      refine(float, int) = float
      refine(List[Any], List[int]) = List[int]
      refine(Dict[int, Any], Dict[Any, int]) = Dict[int, int]
      refine(Tuple[int, Any], Tuple[Any, int]) = Tuple[int, int]

      refine(Callable[[Any], Any], Callable[[int], int]) = Callable[[int], int]
      refine(Callable[..., int], Callable[[int, float], Any]) = Callable[[int, float], int]

      refine(Optional[Any], int) = Optional[int]
      refine(Optional[Any], Optional[int]) = Optional[int]
      refine(Optional[Any], Union[int, str]) = Optional[Union[int, str]]
      refine(Optional[List[Any]], List[int]) = List[int]

    """
    t = get_proper_type(ti)
    s = get_proper_type(si)

    if isinstance(t, AnyType):
        # If s is also an Any, we return if it is a missing_import Any
        return t if isinstance(s, AnyType) and t.missing_import_name else s

    if isinstance(t, Instance) and isinstance(s, Instance) and t.type == s.type:
        return t.copy_modified(args=[refine_type(ta, sa) for ta, sa in zip(t.args, s.args)])

    if (
        isinstance(t, TupleType)
        and isinstance(s, TupleType)
        and t.partial_fallback == s.partial_fallback
        and len(t.items) == len(s.items)
    ):
        return t.copy_modified(items=[refine_type(ta, sa) for ta, sa in zip(t.items, s.items)])

    if isinstance(t, CallableType) and isinstance(s, CallableType):
        return refine_callable(t, s)

    if isinstance(t, UnionType):
        return refine_union(t, s)

    # TODO: Refining of builtins.tuple, Type?

    return t

def refine_union(t: UnionType, s: ProperType) -> Type:
    """Refine a union type based on another type.

    This is done by refining every component of the union against the
    right hand side type (or every component of its union if it is
    one). If an element of the union is successfully refined, we drop it
    from the union in favor of the refined versions.
    """
    # Don't try to do any union refining if the types are already the
    # same.  This prevents things like refining Optional[Any] against
    # itself and producing None.
    if t == s:
        return t

    rhs_items = s.items if isinstance(s, UnionType) else [s]

    new_items = []
    for lhs in t.items:
        refined = False
        for rhs in rhs_items:
            new = refine_type(lhs, rhs)
            if new != lhs:
                new_items.append(new)
                refined = True
        if not refined:
            new_items.append(lhs)

    # Turn strict optional on when simplifying the union since we
    # don't want to drop Nones.
    with state.strict_optional_set(True):
        return make_simplified_union(new_items)

def refine_callable(t: CallableType, s: CallableType) -> CallableType:
    """Refine a callable based on another.

    See comments for refine_type.
    """
    if t.fallback != s.fallback:
        return t

    if t.is_ellipsis_args and not is_tricky_callable(s):
        return s.copy_modified(ret_type=refine_type(t.ret_type, s.ret_type))

    if is_tricky_callable(t) or t.arg_kinds != s.arg_kinds:
        return t

    return t.copy_modified(
        arg_types=[refine_type(ta, sa) for ta, sa in zip(t.arg_types, s.arg_types)],
        ret_type=refine_type(t.ret_type, s.ret_type),
    )

T = TypeVar("T")

def dedup(old: list[T]) -> list[T]:
    new: list[T] = []
    for x in old:
        if x not in new:
            new.append(x)
    return new


# file: rfcomm-client.py
# auth: Albert Huang <albert@csail.mit.edu>
# desc: simple demonstration of a client application that uses RFCOMM sockets
#       intended for use with rfcomm-server
#
# $Id: rfcomm-client.py 424 2006-08-24 03:35:54Z albert $

from bluetooth import *
import subprocess
import obd
import time

host = '00:1D:A5:00:0F:B6'
port = 1

#retcode = subprocess.call("rfcomm connect /dev/rfcomm0 00:1D:A5:00:0F:B6 1 &",shell=True)
#print("Return code is:", retcode)

proc = subprocess.Popen(["rfcomm", "connect", "/dev/rfcomm0", "00:1D:A5:00:0F:B6", "1","&"])
time.sleep(10)
print("Slept for 10 secs: Command should have completed")
conn = obd.OBD("/dev/rfcomm0")
cmd = obd.commands.DISTANCE_SINCE_DTC_CLEAR
print("querying OBD")
resp = conn.query(cmd)
print(resp.value,resp.unit)
proc2 = subprocess.Popen(["rfcomm", "release", "/dev/rfcomm0"])
time.sleep(5)
print("Disconnected")


import sys
import time

start_secs = time.time()

l=[]
my_file = open("inp.txt", "r")
Lines = my_file.readlines()
for line in Lines:
    l.append(line.strip())

d=dict()
for e in l:
    arr=e.split('bags contain')
    c1=arr[0].strip()
    temp=arr[1].strip()
    temp=temp.replace('bags','').replace('bag','').replace('.','')
    arr2=temp.split(',')
    arr3=[]
    for e2 in arr2:
        temp2=e2.strip()
        if temp2=='no other':
            arr3.append(temp2)
        else:
            idx=temp2.find(' ')
            num=temp2[0:idx]
            col=temp2[idx+1:]
            arr3.append(col)
    d[c1]=arr3

def findGold(c,seen):
    if c == 'shiny gold':
        return 1
    if c == 'no other':
        return 0
    if c in seen:
        return 0
    seen.append(c)
    arr = d[c]
    num=0
    for c2 in arr:
        num =num+findGold(c2,seen)
    return num

tot=0
for x, y in d.items():
    sarr =[]
    if x == 'shiny gold':
        continue
    n2 = findGold(x,sarr)
    if n2 > 0:
        tot = tot + 1
print(tot)
end_secs = time.time()
print(str(end_secs-start_secs))


'''
Problem Description

Given a sorted array A containing N integers both positive and negative.

You need to create another array containing the squares of all the elements in A and return it in non-decreasing order.

Try to this in O(N) time.

Problem Constraints
1 <= N <= 105.

-103 <= A[i] <= 103

Input Format
First and only argument is an integer array A.

Output Format
Return a integer array as described in the problem above.

Example Input
Input 1:

 A = [-6, -3, -1, 2, 4, 5]
Input 2:

 A = [-5, -4, -2, 0, 1]

Example Output
Output 1:

 [1, 4, 9, 16, 25, 36]
Output 2:

 [0, 1, 4, 16, 25]
'''

def solve(A):
    '''
    :param A: -> List of numbers
    :return: -> List of sorted squared integers
    '''
    start = 0
    end = len(A) - 1
    res = []

    for _ in range(len(A)):
        if abs(A[start]) > abs(A[end]):
            res.append(A[start] * A[start])
            start += 1
        else:
            res.append(A[end] * A[end])
            end -= 1
    return list(reversed(res))

print(solve([-6, -3, -1, 2, 4, 5]))
print(solve([-5, -4, -2, 0, 1]))
print(solve([1, 2, 3, 4, 5]))
print(solve([-5, -4, -3, -2, -1]))


from copy import deepcopy
from datetime import datetime, timedelta

import pytest

from boto3.session import Session
from botocore.exceptions import ClientError, NoCredentialsError
from dateutil.tz import tzlocal
from mock.mock import patch

from ecs_deploy.ecs import EcsService, EcsTaskDefinition, \
    UnknownContainerError, EcsTaskDefinitionDiff, EcsClient, \
    EcsAction, EcsConnectionError, DeployAction, ScaleAction, RunAction, \
    EcsTaskDefinitionCommandError, UnknownTaskDefinitionError, LAUNCH_TYPE_EC2

CLUSTER_NAME = u'test-cluster'
CLUSTER_ARN = u'arn:aws:ecs:eu-central-1:123456789012:cluster/%s' % CLUSTER_NAME
SERVICE_NAME = u'test-service'
SERVICE_ARN = u'ecs-svc/12345678901234567890'
DESIRED_COUNT = 2
TASK_DEFINITION_FAMILY_1 = u'test-task'
TASK_DEFINITION_REVISION_1 = 1
TASK_DEFINITION_ROLE_ARN_1 = u'arn:test:role:1'
TASK_DEFINITION_ARN_1 = u'arn:aws:ecs:eu-central-1:123456789012:task-definition/%s:%s' % (TASK_DEFINITION_FAMILY_1,
                                                                                TASK_DEFINITION_REVISION_1)
TASK_DEFINITION_VOLUMES_1 = []
TASK_DEFINITION_CONTAINERS_1 = [
    {u'name': u'webserver', u'image': u'webserver:123', u'command': u'run',
     u'environment': ({"name": "foo", "value": "bar"}, {"name": "lorem", "value": "ipsum"}, {"name": "empty", "value": ""}),
     u'secrets': ({"name": "baz", "valueFrom": "qux"}, {"name": "dolor", "valueFrom": "sit"})},
    {u'name': u'application', u'image': u'application:123', u'command': u'run', u'environment': ()}
]

TASK_DEFINITION_FAMILY_2 = u'test-task'
TASK_DEFINITION_REVISION_2 = 2
TASK_DEFINITION_ARN_2 = u'arn:aws:ecs:eu-central-1:123456789012:task-definition/%s:%s' % (TASK_DEFINITION_FAMILY_2,
                                                                                TASK_DEFINITION_REVISION_2)
TASK_DEFINITION_VOLUMES_2 = []
TASK_DEFINITION_CONTAINERS_2 = [
    {u'name': u'webserver', u'image': u'webserver:123', u'command': u'run',
     u'environment': ({"name": "foo", "value": "bar"}, {"name": "lorem", "value": "ipsum"}, {"name": "empty", "value": ""}),
     u'secrets': ({"name": "baz", "valueFrom": "qux"}, {"name": "dolor", "valueFrom": "sit"})},
    {u'name': u'application', u'image': u'application:123', u'command': u'run', u'environment': ()}
]

TASK_DEFINITION_REVISION_3 = 3
TASK_DEFINITION_ARN_3 = u'arn:aws:ecs:eu-central-1:123456789012:task-definition/%s:%s' % (TASK_DEFINITION_FAMILY_1,
                                                                                TASK_DEFINITION_REVISION_3)
TASK_DEFINITION_VOLUMES_3 = []
TASK_DEFINITION_CONTAINERS_3 = [
    {u'name': u'webserver', u'image': u'webserver:456', u'command': u'execute',
     u'environment': ({"name": "foo", "value": "foobar"}, {"name": "newvar", "value": "new value"}),
     u'secrets': ({"name": "baz", "valueFrom": "foobaz"}, {"name": "dolor", "valueFrom": "loremdolor"})},
    {u'name': u'application', u'image': u'application:123', u'command': u'run', u'environment': ()}
]
TASK_DEFINITION_ROLE_ARN_3 = u'arn:test:another-role:1'

PAYLOAD_TASK_DEFINITION_1 = {
    u'taskDefinitionArn': TASK_DEFINITION_ARN_1,
    u'family': TASK_DEFINITION_FAMILY_1,
    u'revision': TASK_DEFINITION_REVISION_1,
    u'taskRoleArn': TASK_DEFINITION_ROLE_ARN_1,
    u'executionRoleArn': TASK_DEFINITION_ROLE_ARN_1,
    u'volumes': deepcopy(TASK_DEFINITION_VOLUMES_1),
    u'containerDefinitions': deepcopy(TASK_DEFINITION_CONTAINERS_1),
    u'status': u'active',
    u'requiresAttributes': {},
    u'networkMode': u'host',
    u'placementConstraints': {},
    u'unknownProperty': u'lorem-ipsum',
    u'compatibilities': [u'EC2'],
}

PAYLOAD_TASK_DEFINITION_2 = {
    u'taskDefinitionArn': TASK_DEFINITION_ARN_2,
    u'family': TASK_DEFINITION_FAMILY_2,
    u'revision': TASK_DEFINITION_REVISION_2,
    u'volumes': deepcopy(TASK_DEFINITION_VOLUMES_2),
    u'containerDefinitions': deepcopy(TASK_DEFINITION_CONTAINERS_2),
    u'status': u'active',
    u'unknownProperty': u'lorem-ipsum',
    u'compatibilities': [u'EC2'],
}

PAYLOAD_TASK_DEFINITION_3 = {
    u'taskDefinitionArn': TASK_DEFINITION_ARN_3,
    u'family': TASK_DEFINITION_FAMILY_1,
    u'revision': TASK_DEFINITION_REVISION_3,
    u'taskRoleArn': TASK_DEFINITION_ROLE_ARN_3,
    u'executionRoleArn': TASK_DEFINITION_ROLE_ARN_3,
    u'volumes': deepcopy(TASK_DEFINITION_VOLUMES_3),
    u'containerDefinitions': deepcopy(TASK_DEFINITION_CONTAINERS_3),
    u'status': u'active',
    u'requiresAttributes': {},
    u'networkMode': u'host',
    u'placementConstraints': {},
    u'unknownProperty': u'lorem-ipsum',
    u'compatibilities': [u'EC2'],
}

TASK_ARN_1 = u'arn:aws:ecs:eu-central-1:123456789012:task/12345678-1234-1234-1234-123456789011'
TASK_ARN_2 = u'arn:aws:ecs:eu-central-1:123456789012:task/12345678-1234-1234-1234-123456789012'

PAYLOAD_TASK_1 = {
    u'taskArn': TASK_ARN_1,
    u'clusterArn': CLUSTER_ARN,
    u'taskDefinitionArn': TASK_DEFINITION_ARN_1,
    u'containerInstanceArn': u'arn:aws:ecs:eu-central-1:123456789012:container-instance/12345678-123456-123456-123456',
    u'overrides': {u'containerOverrides': []},
    u'lastStatus': u'RUNNING',
    u'desiredStatus': u'RUNNING',
    u'containers': TASK_DEFINITION_CONTAINERS_1,
    u'startedBy': SERVICE_ARN
}

PAYLOAD_TASK_2 = {
    u'taskArn': TASK_ARN_2,
    u'clusterArn': CLUSTER_ARN,
    u'taskDefinitionArn': TASK_DEFINITION_ARN_1,
    u'containerInstanceArn': u'arn:aws:ecs:eu-central-1:123456789012:container-instance/12345678-123456-123456-123456',
    u'overrides': {u'containerOverrides': []},
    u'lastStatus': u'RUNNING',
    u'desiredStatus': u'RUNNING',
    u'containers': TASK_DEFINITION_CONTAINERS_1,
    u'startedBy': SERVICE_ARN
}

PAYLOAD_DEPLOYMENTS = [
    {
        u'status': u'PRIMARY',
        u'pendingCount': 0,
        u'desiredCount': DESIRED_COUNT,
        u'runningCount': DESIRED_COUNT,
        u'taskDefinition': TASK_DEFINITION_ARN_1,
        u'createdAt': datetime(2016, 3, 11, 12, 0, 0, 000000, tzinfo=tzlocal()),
        u'updatedAt': datetime(2016, 3, 11, 12, 5, 0, 000000, tzinfo=tzlocal()),
        u'id': u'ecs-svc/0000000000000000002',
    }
]

PAYLOAD_EVENTS = [
    {
        u'id': u'error',
        u'createdAt': datetime.now(tz=tzlocal()),
        u'message': u'Service was unable to Lorem Ipsum'
    },
    {
        u'id': u'older_error',
        u'createdAt': datetime(2016, 3, 11, 12, 0, 10, 000000, tzinfo=tzlocal()),
        u'message': u'Service was unable to Lorem Ipsum'
    }
]

PAYLOAD_SERVICE = {
    u'serviceName': SERVICE_NAME,
    u'desiredCount': DESIRED_COUNT,
    u'taskDefinition': TASK_DEFINITION_ARN_1,
    u'deployments': PAYLOAD_DEPLOYMENTS,
    u'events': []
}

PAYLOAD_SERVICE_WITH_ERRORS = {
    u'serviceName': SERVICE_NAME,
    u'desiredCount': DESIRED_COUNT,
    u'taskDefinition': TASK_DEFINITION_ARN_1,
    u'deployments': PAYLOAD_DEPLOYMENTS,
    u'events': PAYLOAD_EVENTS
}

PAYLOAD_SERVICE_WITHOUT_DEPLOYMENTS = {
    u'serviceName': SERVICE_NAME,
    u'desiredCount': DESIRED_COUNT,
    u'taskDefinition': TASK_DEFINITION_ARN_1,
    u'deployments': [],
    u'events': []
}

RESPONSE_SERVICE = {
    u"service": PAYLOAD_SERVICE
}

RESPONSE_SERVICE_WITH_ERRORS = {
    u"service": PAYLOAD_SERVICE_WITH_ERRORS
}

RESPONSE_DESCRIBE_SERVICES = {
    u"services": [PAYLOAD_SERVICE]
}

RESPONSE_TASK_DEFINITION = {
    u"taskDefinition": PAYLOAD_TASK_DEFINITION_1
}

RESPONSE_TASK_DEFINITION_2 = {
    u"taskDefinition": PAYLOAD_TASK_DEFINITION_2
}

RESPONSE_TASK_DEFINITION_3 = {
    u"taskDefinition": PAYLOAD_TASK_DEFINITION_3
}

RESPONSE_TASK_DEFINITIONS = {
    TASK_DEFINITION_ARN_1: RESPONSE_TASK_DEFINITION,
    TASK_DEFINITION_ARN_2: RESPONSE_TASK_DEFINITION_2,
    TASK_DEFINITION_ARN_3: RESPONSE_TASK_DEFINITION_3,
    u'test-task:1': RESPONSE_TASK_DEFINITION,
    u'test-task:2': RESPONSE_TASK_DEFINITION_2,
    u'test-task:3': RESPONSE_TASK_DEFINITION_3,
    u'test-task': RESPONSE_TASK_DEFINITION_2,
}

RESPONSE_LIST_TASKS_2 = {
    u"taskArns": [TASK_ARN_1, TASK_ARN_2]
}

RESPONSE_LIST_TASKS_1 = {
    u"taskArns": [TASK_ARN_1]
}

RESPONSE_LIST_TASKS_0 = {
    u"taskArns": []
}

RESPONSE_DESCRIBE_TASKS = {
    u"tasks": [PAYLOAD_TASK_1, PAYLOAD_TASK_2]
}

@pytest.fixture()
def task_definition():
    return EcsTaskDefinition(**deepcopy(PAYLOAD_TASK_DEFINITION_1))

@pytest.fixture
def task_definition_revision_2():
    return EcsTaskDefinition(**deepcopy(PAYLOAD_TASK_DEFINITION_2))

@pytest.fixture
def service():
    return EcsService(CLUSTER_NAME, deepcopy(PAYLOAD_SERVICE))

@pytest.fixture
def service_with_errors():
    return EcsService(CLUSTER_NAME, deepcopy(PAYLOAD_SERVICE_WITH_ERRORS))

@pytest.fixture
def service_without_deployments():
    return EcsService(CLUSTER_NAME, deepcopy(PAYLOAD_SERVICE_WITHOUT_DEPLOYMENTS))

def test_service_init(service):
    assert isinstance(service, dict)
    assert service.cluster == CLUSTER_NAME
    assert service[u'desiredCount'] == DESIRED_COUNT
    assert service[u'taskDefinition'] == TASK_DEFINITION_ARN_1

def test_service_set_task_definition(service, task_definition):
    assert service.task_definition == TASK_DEFINITION_ARN_1
    service.set_task_definition(task_definition)
    assert service.task_definition == task_definition.arn

def test_service_name(service):
    assert service.name == SERVICE_NAME

def test_service_deployment_created_at(service):
    assert service.deployment_created_at == datetime(2016, 3, 11, 12, 00, 00, 000000, tzinfo=tzlocal())

def test_service_deployment_updated_at(service):
    assert service.deployment_updated_at == datetime(2016, 3, 11, 12, 5, 00, 000000, tzinfo=tzlocal())

def test_service_deployment_created_at_without_deployments(service_without_deployments):
    now = datetime.now()
    assert service_without_deployments.deployment_created_at >= now
    assert service_without_deployments.deployment_created_at <= datetime.now()

def test_service_deployment_updated_at_without_deployments(service_without_deployments):
    now = datetime.now()
    assert service_without_deployments.deployment_updated_at >= now
    assert service_without_deployments.deployment_updated_at <= datetime.now()

def test_service_errors(service_with_errors):
    assert len(service_with_errors.errors) == 1

def test_service_older_errors(service_with_errors):
    assert len(service_with_errors.older_errors) == 1

def test_task_family(task_definition):
    assert task_definition.family == TASK_DEFINITION_FAMILY_1

def test_task_containers(task_definition):
    assert task_definition.containers == TASK_DEFINITION_CONTAINERS_2

def test_task_container_names(task_definition):
    assert u'webserver' in task_definition.container_names
    assert u'application' in task_definition.container_names
    assert u'foobar' not in task_definition.container_names

def test_task_volumes(task_definition):
    assert task_definition.volumes == TASK_DEFINITION_VOLUMES_2

def test_task_revision(task_definition):
    assert task_definition.revision == TASK_DEFINITION_REVISION_1

def test_task_no_diff(task_definition):
    assert task_definition.diff == []

def test_task_image_diff(task_definition):
    task_definition.set_images(u'foobar')
    assert len(task_definition.diff) == 2

    for diff in task_definition.diff:
        assert isinstance(diff, EcsTaskDefinitionDiff)

def test_task_set_tag(task_definition):
    task_definition.set_images(u'foobar')
    for container in task_definition.containers:
        assert container[u'image'].endswith(u':foobar')

def test_task_set_image(task_definition):
    task_definition.set_images(webserver=u'new-image:123', application=u'app-image:latest')
    for container in task_definition.containers:
        if container[u'name'] == u'webserver':
            assert container[u'image'] == u'new-image:123'
        if container[u'name'] == u'application':
            assert container[u'image'] == u'app-image:latest'

def test_task_set_environment(task_definition):
    assert len(task_definition.containers[0]['environment']) == 3

    task_definition.set_environment(((u'webserver', u'foo', u'baz'), (u'webserver', u'some-name', u'some-value')))

    assert len(task_definition.containers[0]['environment']) == 4

    assert {'name': 'lorem', 'value': 'ipsum'} in task_definition.containers[0]['environment']
    assert {'name': 'foo', 'value': 'baz'} in task_definition.containers[0]['environment']
    assert {'name': 'some-name', 'value': 'some-value'} in task_definition.containers[0]['environment']

def test_task_set_environment_exclusively(task_definition):
    assert len(task_definition.containers[0]['environment']) == 3
    assert len(task_definition.containers[1]['environment']) == 0

    task_definition.set_environment(((u'application', u'foo', u'baz'), (u'application', u'new-var', u'new-value')), exclusive=True)

    assert len(task_definition.containers[0]['environment']) == 0
    assert len(task_definition.containers[1]['environment']) == 2

    assert task_definition.containers[0]['environment'] == []
    assert {'name': 'foo', 'value': 'baz'} in task_definition.containers[1]['environment']
    assert {'name': 'new-var', 'value': 'new-value'} in task_definition.containers[1]['environment']

def test_task_set_secrets_exclusively(task_definition):
    assert len(task_definition.containers[0]['secrets']) == 2

    task_definition.set_secrets(((u'webserver', u'new-secret', u'another-place'), ), exclusive=True)

    assert len(task_definition.containers[0]['secrets']) == 1
    assert {'name': 'new-secret', 'valueFrom': 'another-place'} in task_definition.containers[0]['secrets']

def test_task_set_secrets(task_definition):
    task_definition.set_secrets(((u'webserver', u'foo', u'baz'), (u'webserver', u'some-name', u'some-value')))

    assert {'name': 'dolor', 'valueFrom': 'sit'} in task_definition.containers[0]['secrets']
    assert {'name': 'foo', 'valueFrom': 'baz'} in task_definition.containers[0]['secrets']
    assert {'name': 'some-name', 'valueFrom': 'some-value'} in task_definition.containers[0]['secrets']

def test_task_set_image_for_unknown_container(task_definition):
    with pytest.raises(UnknownContainerError):
        task_definition.set_images(foobar=u'new-image:123')

def test_task_set_command(task_definition):
    task_definition.set_commands(webserver=u'run-webserver', application=u'run-application')
    for container in task_definition.containers:
        if container[u'name'] == u'webserver':
            assert container[u'command'] == [u'run-webserver']
        if container[u'name'] == u'application':
            assert container[u'command'] == [u'run-application']

def test_task_set_command_with_multiple_arguments(task_definition):
    task_definition.set_commands(webserver=u'run-webserver arg1 arg2', application=u'run-application arg1 arg2')
    for container in task_definition.containers:
        if container[u'name'] == u'webserver':
            assert container[u'command'] == [u'run-webserver', u'arg1', u'arg2']
        if container[u'name'] == u'application':
            assert container[u'command'] == [u'run-application', u'arg1', u'arg2']

def test_task_set_command_with_empty_argument(task_definition):
    empty_argument = " "
    task_definition.set_commands(webserver=empty_argument + u'run-webserver arg1 arg2')
    for container in task_definition.containers:
        if container[u'name'] == u'webserver':
            assert container[u'command'] == [u'run-webserver', u'arg1', u'arg2']

def test_task_set_command_as_json_list(task_definition):
    task_definition.set_commands(webserver=u'["run-webserver", "arg1", "arg2"]', application=u'["run-application", "arg1", "arg2"]')
    for container in task_definition.containers:
        if container[u'name'] == u'webserver':
            assert container[u'command'] == [u'run-webserver', u'arg1', u'arg2']
        if container[u'name'] == u'application':
            assert container[u'command'] == [u'run-application', u'arg1', u'arg2']

def test_task_set_command_as_invalid_json_list(task_definition):
    with pytest.raises(EcsTaskDefinitionCommandError):
        task_definition.set_commands(webserver=u'["run-webserver, "arg1" arg2"]', application=u'["run-application" "arg1 "arg2"]')

def test_task_set_command_for_unknown_container(task_definition):
    with pytest.raises(UnknownContainerError):
        task_definition.set_images(foobar=u'run-foobar')

def test_task_get_overrides(task_definition):
    assert task_definition.get_overrides() == []

def test_task_get_overrides_with_command(task_definition):
    task_definition.set_commands(webserver='/usr/bin/python script.py')
    overrides = task_definition.get_overrides()
    assert len(overrides) == 1
    assert overrides[0]['command'] == ['/usr/bin/python','script.py']

def test_task_get_overrides_with_environment(task_definition):
    task_definition.set_environment((('webserver', 'foo', 'baz'),))
    overrides = task_definition.get_overrides()
    assert len(overrides) == 1
    assert overrides[0]['name'] == 'webserver'
    assert dict(name='foo', value='baz') in overrides[0]['environment']

def test_task_get_overrides_with_secrets(task_definition):
    task_definition.set_secrets((('webserver', 'foo', 'baz'),))
    overrides = task_definition.get_overrides()
    assert len(overrides) == 1
    assert overrides[0]['name'] == 'webserver'
    assert dict(name='foo', valueFrom='baz') in overrides[0]['secrets']

def test_task_get_overrides_with_command_environment_and_secrets(task_definition):
    task_definition.set_commands(webserver='/usr/bin/python script.py')
    task_definition.set_environment((('webserver', 'foo', 'baz'),))
    task_definition.set_secrets((('webserver', 'bar', 'qux'),))
    overrides = task_definition.get_overrides()
    assert len(overrides) == 1
    assert overrides[0]['name'] == 'webserver'
    assert overrides[0]['command'] == ['/usr/bin/python','script.py']
    assert dict(name='foo', value='baz') in overrides[0]['environment']
    assert dict(name='bar', valueFrom='qux') in overrides[0]['secrets']

def test_task_get_overrides_with_command_secrets_and_environment_for_multiple_containers(task_definition):
    task_definition.set_commands(application='/usr/bin/python script.py')
    task_definition.set_environment((('webserver', 'foo', 'baz'),))
    task_definition.set_secrets((('webserver', 'bar', 'qux'),))
    overrides = task_definition.get_overrides()
    assert len(overrides) == 2
    assert overrides[0]['name'] == 'application'
    assert overrides[0]['command'] == ['/usr/bin/python','script.py']
    assert overrides[1]['name'] == 'webserver'
    assert dict(name='foo', value='baz') in overrides[1]['environment']
    assert dict(name='bar', valueFrom='qux') in overrides[1]['secrets']

def test_task_get_overrides_command(task_definition):
    command = task_definition.get_overrides_command('/usr/bin/python script.py')
    assert isinstance(command, list)
    assert command == ['/usr/bin/python','script.py']

def test_task_get_overrides_environment(task_definition):
    environment = task_definition.get_overrides_env(dict(foo='bar'))
    assert isinstance(environment, list)
    assert len(environment) == 1
    assert environment[0] == dict(name='foo', value='bar')

def test_task_get_overrides_secrets(task_definition):
    secrets = task_definition.get_overrides_secrets(dict(foo='bar'))
    assert isinstance(secrets, list)
    assert len(secrets) == 1
    assert secrets[0] == dict(name='foo', valueFrom='bar')

def test_task_definition_diff():
    diff = EcsTaskDefinitionDiff(u'webserver', u'image', u'new', u'old')
    assert str(diff) == u'Changed image of container "webserver" to: "new" (was: "old")'

@patch.object(Session, 'client')
@patch.object(Session, '__init__')
def test_client_init(mocked_init, mocked_client):
    mocked_init.return_value = None

    EcsClient(u'access_key_id', u'secret_access_key', u'region', u'profile', u'session_token')

    mocked_init.assert_called_once_with(aws_access_key_id=u'access_key_id',
                                        aws_secret_access_key=u'secret_access_key',
                                        profile_name=u'profile',
                                        region_name=u'region',
                                        aws_session_token=u'session_token')
    mocked_client.assert_any_call(u'ecs')
    mocked_client.assert_any_call(u'events')

@pytest.fixture
@patch.object(Session, 'client')
@patch.object(Session, '__init__')
def client(mocked_init, mocked_client):
    mocked_init.return_value = None
    return EcsClient(u'access_key_id', u'secret_access_key', u'region', u'profile', u'session_token')

def test_client_describe_services(client):
    client.describe_services(u'test-cluster', u'test-service')
    client.boto.describe_services.assert_called_once_with(cluster=u'test-cluster', services=[u'test-service'])

def test_client_describe_task_definition(client):
    client.describe_task_definition(u'task_definition_arn')
    client.boto.describe_task_definition.assert_called_once_with(taskDefinition=u'task_definition_arn')

def test_client_describe_unknown_task_definition(client):
    error_response = {u'Error': {u'Code': u'ClientException', u'Message': u'Unable to describe task definition.'}}
    client.boto.describe_task_definition.side_effect = ClientError(error_response, u'DescribeServices')
    with pytest.raises(UnknownTaskDefinitionError):
        client.describe_task_definition(u'task_definition_arn')

def test_client_list_tasks(client):
    client.list_tasks(u'test-cluster', u'test-service')
    client.boto.list_tasks.assert_called_once_with(cluster=u'test-cluster', serviceName=u'test-service')

def test_client_describe_tasks(client):
    client.describe_tasks(u'test-cluster', u'task-arns')
    client.boto.describe_tasks.assert_called_once_with(cluster=u'test-cluster', tasks=u'task-arns')

def test_client_register_task_definition(client):
    containers = [{u'name': u'foo'}]
    volumes = [{u'foo': u'bar'}]
    role_arn = 'arn:test:role'
    execution_role_arn = 'arn:test:role'
    task_definition = EcsTaskDefinition(
        containerDefinitions=containers,
        volumes=volumes,
        family=u'family',
        revision=1,
        taskRoleArn=role_arn,
        executionRoleArn=execution_role_arn,
        status='active',
        taskDefinitionArn='arn:task',
        requiresAttributes={},
        unkownProperty='foobar'
    )

    client.register_task_definition(
        family=task_definition.family,
        containers=task_definition.containers,
        volumes=task_definition.volumes,
        role_arn=task_definition.role_arn,
        execution_role_arn=execution_role_arn,
        additional_properties=task_definition.additional_properties
    )

    client.boto.register_task_definition.assert_called_once_with(
        family=u'family',
        containerDefinitions=containers,
        volumes=volumes,
        taskRoleArn=role_arn,
        executionRoleArn=execution_role_arn,
        unkownProperty='foobar'
    )

def test_client_deregister_task_definition(client):
    client.deregister_task_definition(u'task_definition_arn')
    client.boto.deregister_task_definition.assert_called_once_with(taskDefinition=u'task_definition_arn')

def test_client_update_service(client):
    client.update_service(u'test-cluster', u'test-service', 5, u'task-definition')
    client.boto.update_service.assert_called_once_with(
        cluster=u'test-cluster',
        service=u'test-service',
        desiredCount=5,
        taskDefinition=u'task-definition'
    )

def test_client_update_service_without_desired_count(client):
    client.update_service(u'test-cluster', u'test-service', None, u'task-definition')
    client.boto.update_service.assert_called_once_with(
        cluster=u'test-cluster',
        service=u'test-service',
        taskDefinition=u'task-definition'
    )

def test_client_run_task(client):
    client.run_task(
        cluster=u'test-cluster',
        task_definition=u'test-task',
        count=2,
        started_by='test',
        overrides=dict(foo='bar')
    )

    client.boto.run_task.assert_called_once_with(
        cluster=u'test-cluster',
        taskDefinition=u'test-task',
        count=2,
        startedBy='test',
        overrides=dict(foo='bar')
    )

def test_ecs_action_init(client):
    action = EcsAction(client, u'test-cluster', u'test-service')
    assert action.client == client
    assert action.cluster_name == u'test-cluster'
    assert action.service_name == u'test-service'
    assert isinstance(action.service, EcsService)

def test_ecs_action_init_with_invalid_cluster():
    with pytest.raises(EcsConnectionError) as excinfo:
        client = EcsTestClient(u'access_key',  u'secret_key')
        EcsAction(client, u'invliad-cluster', u'test-service')
    assert str(excinfo.value) == u'An error occurred (ClusterNotFoundException) when calling the DescribeServices ' \
                                 u'operation: Cluster not found.'

def test_ecs_action_init_with_invalid_service():
    with pytest.raises(EcsConnectionError) as excinfo:
        client = EcsTestClient(u'access_key',  u'secret_key')
        EcsAction(client, u'test-cluster', u'invalid-service')
    assert str(excinfo.value) == u'An error occurred when calling the DescribeServices operation: Service not found.'

def test_ecs_action_init_without_credentials():
    with pytest.raises(EcsConnectionError) as excinfo:
        client = EcsTestClient()
        EcsAction(client, u'test-cluster', u'invalid-service')
    assert str(excinfo.value) == u'Unable to locate credentials. Configure credentials by running "aws configure".'

def test_ecs_action_get_service():
    client = EcsTestClient(u'access_key', u'secret_key')
    action = EcsAction(client, u'test-cluster', u'test-service')
    service = action.get_service()
    assert service.name == u'test-service'
    assert service.cluster == u'test-cluster'

@patch.object(EcsClient, '__init__')
def test_ecs_action_get_current_task_definition(client, service):
    client.describe_task_definition.return_value = RESPONSE_TASK_DEFINITION

    action = EcsAction(client, u'test-cluster', u'test-service')
    task_definition = action.get_current_task_definition(service)

    client.describe_task_definition.assert_called_once_with(
        task_definition_arn=service.task_definition
    )

    assert isinstance(task_definition, EcsTaskDefinition)
    assert task_definition.family == u'test-task'
    assert task_definition.revision == 1
    assert task_definition.arn == u'arn:aws:ecs:eu-central-1:123456789012:task-definition/test-task:1'

@patch.object(EcsClient, '__init__')
def test_update_task_definition(client, task_definition):
    client.register_task_definition.return_value = RESPONSE_TASK_DEFINITION

    action = EcsAction(client, u'test-cluster', u'test-service')
    new_task_definition = action.update_task_definition(task_definition)

    assert isinstance(new_task_definition, EcsTaskDefinition)
    client.register_task_definition.assert_called_once_with(
        family=task_definition.family,
        containers=task_definition.containers,
        volumes=task_definition.volumes,
        role_arn=task_definition.role_arn,
        execution_role_arn=task_definition.execution_role_arn,
        additional_properties={
            u'networkMode': u'host',
            u'placementConstraints': {},
            u'unknownProperty': u'lorem-ipsum'
        }
    )

@patch.object(EcsClient, '__init__')
def test_deregister_task_definition(client, task_definition):
    action = EcsAction(client, u'test-cluster', u'test-service')
    action.deregister_task_definition(task_definition)

    client.deregister_task_definition.assert_called_once_with(
        task_definition.arn
    )

@patch.object(EcsClient, '__init__')
def test_update_service(client, service):
    client.update_service.return_value = RESPONSE_SERVICE

    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    new_service = action.update_service(service)

    assert isinstance(new_service, EcsService)
    client.update_service.assert_called_once_with(
        cluster=service.cluster,
        service=service.name,
        desired_count=None,
        task_definition=service.task_definition
    )

@patch.object(EcsClient, '__init__')
def test_is_deployed(client, service):
    client.list_tasks.return_value = RESPONSE_LIST_TASKS_1
    client.describe_tasks.return_value = RESPONSE_DESCRIBE_TASKS

    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    is_deployed = action.is_deployed(service)

    assert is_deployed is True
    client.list_tasks.assert_called_once_with(
        cluster_name=service.cluster,
        service_name=service.name
    )

@patch.object(EcsClient, '__init__')
def test_is_not_deployed_with_more_than_one_deployment(client, service):
    service['deployments'].append(service['deployments'][0])

    client.list_tasks.return_value = RESPONSE_LIST_TASKS_1
    client.describe_tasks.return_value = RESPONSE_DESCRIBE_TASKS

    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    is_deployed = action.is_deployed(service)

    assert is_deployed is False

@patch.object(EcsClient, '__init__')
def test_is_deployed_if_no_tasks_should_be_running(client, service):
    client.list_tasks.return_value = RESPONSE_LIST_TASKS_0
    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    service[u'desiredCount'] = 0
    is_deployed = action.is_deployed(service)
    assert is_deployed is True

@patch.object(EcsClient, '__init__')
def test_is_not_deployed_if_no_tasks_running(client, service):
    client.list_tasks.return_value = RESPONSE_LIST_TASKS_0
    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    is_deployed = action.is_deployed(service)
    assert is_deployed is False

@patch.object(EcsClient, '__init__')
def test_get_running_tasks_count(client, service):
    client.describe_tasks.return_value = RESPONSE_DESCRIBE_TASKS
    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    running_count = action.get_running_tasks_count(service, [TASK_ARN_1, TASK_ARN_2])
    assert running_count == 2

@patch.object(EcsClient, '__init__')
def test_get_running_tasks_count_new_revision(client, service, task_definition_revision_2):
    client.describe_tasks.return_value = RESPONSE_DESCRIBE_TASKS
    action = EcsAction(client, CLUSTER_NAME, SERVICE_NAME)
    service.set_task_definition(task_definition_revision_2)
    running_count = action.get_running_tasks_count(service, [TASK_ARN_1, TASK_ARN_2])
    assert running_count == 0

@patch.object(EcsClient, '__init__')
def test_deploy_action(client, task_definition_revision_2):
    action = DeployAction(client, CLUSTER_NAME, SERVICE_NAME)
    updated_service = action.deploy(task_definition_revision_2)

    assert action.service.task_definition == task_definition_revision_2.arn
    assert isinstance(updated_service, EcsService)

    client.describe_services.assert_called_once_with(
        cluster_name=CLUSTER_NAME,
        service_name=SERVICE_NAME
    )
    client.update_service.assert_called_once_with(
        cluster=action.service.cluster,
        service=action.service.name,
        desired_count=action.service.desired_count,
        task_definition=task_definition_revision_2.arn
    )

@patch.object(EcsClient, '__init__')
def test_scale_action(client):
    action = ScaleAction(client, CLUSTER_NAME, SERVICE_NAME)
    updated_service = action.scale(5)

    assert isinstance(updated_service, EcsService)

    client.describe_services.assert_called_once_with(
        cluster_name=CLUSTER_NAME,
        service_name=SERVICE_NAME
    )
    client.update_service.assert_called_once_with(
        cluster=action.service.cluster,
        service=action.service.name,
        desired_count=5,
        task_definition=action.service.task_definition
    )

@patch.object(EcsClient, '__init__')
def test_run_action(client):
    action = RunAction(client, CLUSTER_NAME)
    assert len(action.started_tasks) == 0

@patch.object(EcsClient, '__init__')
def test_run_action_run(client, task_definition):
    action = RunAction(client, CLUSTER_NAME)
    client.run_task.return_value = dict(tasks=[dict(taskArn='A'), dict(taskArn='B')])
    action.run(task_definition, 2, 'test', LAUNCH_TYPE_EC2, (), (), False)

    client.run_task.assert_called_once_with(
        cluster=CLUSTER_NAME,
        task_definition=task_definition.family_revision,
        count=2,
        started_by='test',
        overrides=dict(containerOverrides=task_definition.get_overrides()),
        launchtype=LAUNCH_TYPE_EC2,
        subnets=(),
        security_groups=(),
        public_ip=False
    )

    assert len(action.started_tasks) == 2

def test_ecs_server_get_warnings():
    since = datetime.now() - timedelta(hours=1)
    until = datetime.now() + timedelta(hours=1)

    event_unable = {
        u'createdAt': datetime.now(),
        u'message': u'unable to foo',
    }

    event_unknown = {
        u'createdAt': datetime.now(),
        u'message': u'unkown foo',
    }

    service = EcsService('foo', {
        u'deployments': [],
        u'events': [event_unable, event_unknown],
    })

    assert len(service.get_warnings(since, until)) == 1

class EcsTestClient(object):
    def __init__(self, access_key_id=None, secret_access_key=None, region=None,
                 profile=None, deployment_errors=False, client_errors=False,
                 wait=0):
        super(EcsTestClient, self).__init__()
        self.access_key_id = access_key_id
        self.secret_access_key = secret_access_key
        self.region = region
        self.profile = profile
        self.deployment_errors = deployment_errors
        self.client_errors = client_errors
        self.wait_until = datetime.now() + timedelta(seconds=wait)

    def describe_services(self, cluster_name, service_name):
        if not self.access_key_id or not self.secret_access_key:
            raise NoCredentialsError()
        if cluster_name != u'test-cluster':
            error_response = {u'Error': {u'Code': u'ClusterNotFoundException', u'Message': u'Cluster not found.'}}
            raise ClientError(error_response, u'DescribeServices')
        if service_name != u'test-service':
            return {u'services': []}
        if self.deployment_errors:
            return {
                u"services": [PAYLOAD_SERVICE_WITH_ERRORS],
                u"failures": []
            }
        return {
            u"services": [PAYLOAD_SERVICE],
            u"failures": []
        }

    def describe_task_definition(self, task_definition_arn):
        if not self.access_key_id or not self.secret_access_key:
            raise EcsConnectionError(u'Unable to locate credentials. Configure credentials by running "aws configure".')
        if task_definition_arn in RESPONSE_TASK_DEFINITIONS:
            return deepcopy(RESPONSE_TASK_DEFINITIONS[task_definition_arn])
        raise UnknownTaskDefinitionError('Unknown task definition arn: %s' % task_definition_arn)

    def list_tasks(self, cluster_name, service_name):
        if self.wait_until <= datetime.now():
            return deepcopy(RESPONSE_LIST_TASKS_2)
        return deepcopy(RESPONSE_LIST_TASKS_0)

    def describe_tasks(self, cluster_name, task_arns):
        return deepcopy(RESPONSE_DESCRIBE_TASKS)

    def register_task_definition(self, family, containers, volumes, role_arn,
                                 execution_role_arn, additional_properties):
        if not self.access_key_id or not self.secret_access_key:
            raise EcsConnectionError(u'Unable to locate credentials. Configure credentials by running "aws configure".')
        return deepcopy(RESPONSE_TASK_DEFINITION_2)

    def deregister_task_definition(self, task_definition_arn):
        return deepcopy(RESPONSE_TASK_DEFINITION)

    def update_service(self, cluster, service, desired_count, task_definition):
        if self.client_errors:
            error = dict(Error=dict(Code=123, Message="Something went wrong"))
            raise ClientError(error, 'fake_error')
        if self.deployment_errors:
            return deepcopy(RESPONSE_SERVICE_WITH_ERRORS)
        return deepcopy(RESPONSE_SERVICE)

    def run_task(self, cluster, task_definition, count, started_by, overrides,
                 launchtype='EC2', subnets=(), security_groups=(),
                 public_ip=False):
        if not self.access_key_id or not self.secret_access_key:
            raise EcsConnectionError(u'Unable to locate credentials. Configure credentials by running "aws configure".')
        if cluster == 'unknown-cluster':
            raise EcsConnectionError(u'An error occurred (ClusterNotFoundException) when calling the RunTask operation: Cluster not found.')
        if self.deployment_errors:
            error = dict(Error=dict(Code=123, Message="Something went wrong"))
            raise ClientError(error, 'fake_error')
        return dict(tasks=[dict(taskArn='arn:foo:bar'), dict(taskArn='arn:lorem:ipsum')])

    def update_rule(self, cluster, rule, task_definition):
        if not self.access_key_id or not self.secret_access_key:
            raise EcsConnectionError(u'Unable to locate credentials. Configure credentials by running "aws configure".')
        if cluster == 'unknown-cluster':
            raise EcsConnectionError(u'An error occurred (ClusterNotFoundException) when calling the RunTask operation: Cluster not found.')


import PySimpleGUI as sg
import socket
import threading
import warnings

warnings.filterwarnings("ignore")
sg.theme("DarkTeal5")
sg.set_options(font=("Helvetica", 13))
layout = [
    [sg.Text("Name:"), sg.InputText(key="-NAME-", size=(30, 1)),
     sg.Button("Connect", key="-CONNECT-", size=(15, 1)),
     sg.Button("Stop", key="-STOP1-", size=(15, 1))],
    [sg.Multiline(key="-DISPLAY-", size=(70, 23), disabled=True,
                  autoscroll=True)],
    [sg.InputText(key="-MESSAGE-", size=(55, 1)),
     sg.Button("Send", key="-SEND-", bind_return_key=True, size=(15, 1))],
]

window = sg.Window("Client", layout, size=(600, 425))

username = ""
client = None
HOST_ADDR = "127.0.0.1"
HOST_PORT = 8080
msgs = []

def connect():
    global username, client
    if len(values["-NAME-"]) < 1:
        sg.popup_error("Enter a valid name.")
    else:
        username = values["-NAME-"]
        connect_to_server(username)

def connect_to_server(name):
    global client, HOST_PORT, HOST_ADDR
    try:
        client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client.connect((HOST_ADDR, HOST_PORT))
        client.send(name.encode())

        window["-NAME-"].update(disabled=True)
        window["-CONNECT-"].update(disabled=True)
        window["-MESSAGE-"].update(disabled=False)

        thread = threading.Thread(target=receive_message_from_server,
                                  args=(client,))
        thread.start()
    except Exception as e:
        sg.popup_error(
            "Cannot connect to host: " + HOST_ADDR + " on port: " + str(
                HOST_PORT) + ". Server may be unavailable. Try again later.")

def receive_message_from_server(sck):
    try:
        while True:
            from_server = sck.recv(4096).decode()

            if not from_server:
                break

            if '->' in from_server:
                from_server = from_server.replace('->', ' -> ')

            msgs.append(from_server)
            text = "\n".join(msgs)
            window["-DISPLAY-"].update(f'{text}\n')
    except OSError:
        exit()

    sck.close()
    window.close()

while True:
    event, values = window.read()

    if event == sg.WIN_CLOSED:
        break
    elif event == "-CONNECT-":
        connect()
    elif event == "-STOP1-":
        if client:
            client.close()
        window.close()
        exit()
    elif event == "-SEND-":
        msg = values["-MESSAGE-"].replace('\n', '')
        if msg:
            client.send(msg.encode())
            msgs.append('You -> '+msg)
            text = "\n".join(msgs)
            window["-DISPLAY-"].update(f'{text}\n')
            window["-MESSAGE-"].update(value="")
            window["-MESSAGE-"].set_focus()

window.close()


import logging

from fuel_agent_ci.objects import Object

LOG = logging.getLogger(__name__)

class Artifact(Object):
    __typename__ = 'artifact'

    def __init__(self, env, name, url, path, unpack=None, clean=None):
        self.env = env
        self.name = name
        self.url = url
        self.path = path
        self.unpack = unpack
        self.clean = clean

    def get(self):
        if not self.status():
            LOG.debug('Getting artifact %s' % self.name)
            self.env.driver.artifact_get(self)

    def clean(self):
        if self.status():
            LOG.debug('Cleaning artifact %s' % self.name)
            self.env.driver.artifact_clean(self)

    def status(self):
        status = self.env.driver.artifact_status(self)
        LOG.debug('Artifact %s status %s' % (self.name, status))
        return status


#!/usr/bin/env python3

from newrelic.agenttest.app.testdecorators import *
from newrelic.agenttest.app.java.testdecorators import *
import newrelic.agenttest.unittest

class RulesTest(newrelic.agenttest.unittest.TestCase):

    @collector(
        custom={
            'connect_options' : {
                'transaction_name_rules' : [{
                    'eval_order'       : 2,
                    'each_segment'     : False,
                    'match_expression' : r'^WebTransaction/Custom/java_test_webapp/NewRelicApiServlet/SOMETHING/(.*)$',
                    'replacement'      : r'WebTransaction/NewRelicApiServlet/SOMETHING_ELSE/\1',
                    'terminate_chain'  : True,
                    'replace_all'      : False,
                    'ignore'           : False
                }],
            }
        }
    )
    @java_agent
    @java
    @java_test_webapp
    @tomcat(select={'version' : 7})
    def test_transaction_rename(self):
        self.assertTrue(self.test_env.tomcat.get_path_and_shutdown(
                '/java_test_webapp/NewRelicApiServlet',
                api='setTransactionName',
                api_args='/java_test_webapp/NewRelicApiServlet/SOMETHING/ETC'))

        self.assertIterAttributes(
            self.test_env.collector.get_unscoped_metrics('WebTransaction/NewRelicApiServlet/SOMETHING_ELSE/ETC'),
            { 'count' : 1 })

        self.assertIterAttributes(
            self.test_env.collector.get_scoped_metrics('WebTransaction/NewRelicApiServlet/SOMETHING_ELSE/ETC'),
            {
                'name' : 'Java/com.nr.test.servlet.NewRelicApiServlet/init',
                'count' : 2
            },
            {
                'name' : 'Java/javax.servlet.ServletRequestListener/requestInitialized',
                'count' : 1
            },
            {
                'name' : 'Servlet/com.nr.test.servlet.NewRelicApiServlet/service',
                'count' : 1
            })

    @collector(
        custom={
            'connect_options' : {
                'metric_name_rules' : [{
                    'eval_order'       : 2,
                    'each_segment'     : False,
                    'match_expression' : r'^Custom/RequestTime/(.*)$',
                    'replacement'      : r'Custom/\1/RequestTime',
                    'terminate_chain'  : True,
                    'replace_all'      : False,
                    'ignore'           : False
                }]
            }
        }
    )
    @java_agent
    @java
    @java_test_webapp
    @tomcat(select={'version' : 7})
    def test_metric_rename(self):
        self.assertTrue(self.test_env.tomcat.get_path_and_shutdown(
                '/java_test_webapp/NewRelicApiServlet',
                api='recordMetric',
                api_args='Custom/RequestTime/Time,1'))

        self.assertIterAttributes(
            self.test_env.collector.get_unscoped_metrics('WebTransaction/Servlet/NewRelicApiServlet'),
            { 'count' : 1 })

        self.assertIterAttributes(
            self.test_env.collector.get_unscoped_metrics('Custom/Time/RequestTime'),
            {
                'name' : 'Custom/Time/RequestTime',
                'count' : 1
            })

if __name__ == '__main__':
    newrelic.agenttest.unittest.initenv(__file__)
    newrelic.agenttest.unittest.main()


import os
today = '22-04-19_'

def parse_args():
    import argparse
    parser = argparse.ArgumentParser(description='Correlation of reserved stars')

    parser.add_argument('--piff_cat',
                        default='/home2/dfa/sobreira/alsina/catalogs/y3a1-v29',
                        help='Full Path to the Only stars Piff catalog')
    parser.add_argument('--exps_file',
                        default='/home/dfa/sobreira/alsina/DESWL/psf/ally3.grizY',
                        #default='/home/dfa/sobreira/alsina/DESWL/psf/testexp',
                        help='list of exposures (in lieu of separate exps)')
    parser.add_argument('--bands', default='riz', type=str,
                         help='Limit to the given bands')
    parser.add_argument('--use_reserved', default=True,
                        action='store_const', const=True,
                        help='just use the objects with the RESERVED flag')
    parser.add_argument('--frac', default=1., type=float,
                        help='Choose a random fraction of the input stars')
    parser.add_argument('--mod', default=True,
                        action='store_const', const=True,
                        help='If true it substracts the mean to each field before calculate correlations')
    parser.add_argument('--obs', default=False,
                        action='store_const', const=True,
                        help='Use e_obs instead of e_piff to calculate modified rho stats')
    parser.add_argument('--outpath', default='/home/dfa/sobreira/alsina/alpha-beta-gamma/code/correlations/',
                        help='location of the output of the files')

    args = parser.parse_args()

    return args

def measure_rho(data, max_sep=300, sep_units='arcmin',  tag=None, prefix='piff', mod=True,  obs=False ):
    """Compute the rho statistics
    """
    import treecorr
    import numpy as np

    e1 = data['obs_e1']
    e2 = data['obs_e2']
    p_e1 = data[prefix+'_e1']
    p_e2 = data[prefix+'_e2']
    T = data['obs_T']
    p_T = data[prefix+'_T']

    de1 = e1-p_e1
    de2 = e2-p_e2
    dt = (T-p_T)/T
    w1 = p_e1*dt
    w2 = p_e2*dt
    w1obs = e1*dt
    w2obs = e2*dt

    #Modified ellipticities
    if(mod):
        e1 = e1 - np.array(np.mean(e1))
        e2 = e2 - np.array(np.mean(e2))
        p_e1 = p_e1 - np.array(np.mean(p_e1))
        p_e2 = p_e2 - np.array(np.mean(p_e2))
        de1 = de1 - np.array(np.mean(de1))
        de2 = de2 - np.array(np.mean(de2))
        w1 = w1 - np.array(np.mean(w1))
        w2 = w2 - np.array(np.mean(w2))
        w1obs = w1obs - np.array(np.mean(w1obs))
        w2obs = w2obs - np.array(np.mean(w2obs))

    ra = data['ra']
    dec = data['dec']
    print('ra = ',ra)
    print('dec = ',dec)
    if(obs):
        ecat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=e1, g2=e2)
        decat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=de1, g2=de2)
        wcat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=w1obs, g2=w2obs)
    else:
        ecat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=p_e1, g2=p_e2)
        decat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=de1, g2=de2)
        #wcat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=w1, g2=w2)
        wcat = treecorr.Catalog(ra=ra, dec=dec, ra_units='deg', dec_units='deg', g1=w1obs, g2=w2obs)
    ecat.name = 'ecat'
    decat.name = 'decat'
    wcat.name = 'wcat'
    if tag is not None:
        for cat in [ ecat, decat, wcat ]:
            cat.name = tag + ":"  + cat.name

    bin_config = dict(
        sep_units = sep_units,
        nbins = 20,
        min_sep = 2.5,
        max_sep = 250,)

    #sep_units = 'degrees',
    '''
    bin_slop = 0.1,
    min_sep = 0.5,
    max_sep = max_sep,
    bin_size = 0.2,
    '''

    results = []
    for (cat1, cat2) in [(ecat, ecat),
                         (decat, decat),
                          (decat, ecat),
                          (wcat, wcat),
                          (decat, wcat),
                          (ecat, wcat) ]:
        print('Doing correlation of %s vs %s'%(cat1.name, cat2.name))

        rho = treecorr.GGCorrelation(bin_config, verbose=2)

        if cat1 is cat2:
            rho.process(cat1)
        else:
            rho.process(cat1, cat2)
        print('mean xi+ = ',rho.xip.mean())
        print('mean xi- = ',rho.xim.mean())
        results.append(rho)

    return results

def main():
    import sys
    sys.path.insert(0, '/home/dfa/sobreira/alsina/alpha-beta-gamma/code/src')
    #sys.path.insert(0, '/global/cscratch1/sd/alsina/alpha-beta-gamma/code/src')

    import numpy as np
    from read_psf_cats import read_data,  toList
    from astropy.io import fits

    args = parse_args()

    #Make directory where the ouput data will be
    outpath = os.path.expanduser(args.outpath)
    try:
        if not os.path.exists(outpath):
            os.makedirs(outpath)
    except OSError:
        if not os.path.exists(outpath): raise

    #STATISTIC USING ONLY RESERVED STARS
    keys = ['ra', 'dec','obs_e1', 'obs_e2', 'obs_T',
            'piff_e1', 'piff_e2', 'piff_T',  'mag']

    exps = toList(args.exps_file)
    data, bands, tilings = read_data(exps, args.piff_cat , keys,
                                     limit_bands=args.bands,
                                     use_reserved=args.use_reserved)
    print("Objects",  len(data))
    data = data[data['mag']<20]
    print("Objects with magnitude <20",  len(data))

    names=['BIN1', 'BIN2','ANGBIN', 'VALUE', 'ANG']
    forms = ['i4', 'i4', 'i4',  'f8',  'f8']
    dtype = dict(names = names, formats=forms)
    nrows = 20
    outdata = np.recarray((nrows, ), dtype=dtype)

    namesout=['RHO0P','RHO1P', 'RHO2P', 'RHO3P', 'RHO4P', 'RHO5P',
              'RHO0M','RHO1M', 'RHO2M', 'RHO3M', 'RHO4M', 'RHO5M']

    rho0, rho1, rho2, rho3, rho4, rho5 = measure_rho(data,  mod=args.mod, obs=args.obs)
    angarr = np.arange(nrows)
    thetaarr = np.exp(rho0.meanlogr)

    rho0parr = rho0.xip; rho1parr = rho1.xip; rho2parr = rho2.xip
    rho3parr = rho3.xip; rho4parr = rho4.xip; rho5parr = rho5.xip
    rho0marr = rho0.xim; rho1marr = rho1.xim; rho2marr = rho2.xim
    rho3marr = rho3.xim; rho4marr = rho4.xim; rho5marr = rho5.xim
    varrho0arr = 2*rho0.varxi; varrho1arr = 2*rho1.varxi; varrho2arr = 2*rho2.varxi;
    varrho3arr = 2*rho3.varxi; varrho4arr = 2*rho4.varxi; varrho5arr = 2*rho5.varxi;

    rhos = [rho0parr, rho1parr, rho2parr, rho3parr, rho4parr,
            rho5parr, rho0marr, rho1marr, rho2marr, rho3marr, rho4marr,
            rho5marr]
    vares = [varrho0arr, varrho1arr, varrho2arr, varrho3arr,
             varrho4arr, varrho5arr, varrho0arr, varrho1arr, varrho2arr,
             varrho3arr, varrho4arr, varrho5arr]
    for i, nam in enumerate(namesout):
        covmat = np.diag(vares[i])
        hdu = fits.PrimaryHDU()
        hdul = fits.HDUList([hdu])
        covmathdu = fits.ImageHDU(covmat, name='COVMAT')
        hdul.insert(1, covmathdu)

        angarray = np.exp(rho0.meanlogr)
        valuearray =  np.array(rhos[i])
        bin1array = np.array([ -999]*nrows)
        bin2array = np.array([ -999]*nrows)
        angbinarray = np.arange(nrows)
        array_list = [bin1array, bin2array, angbinarray, valuearray,  angarray ]
        for array, name in zip(array_list, names): outdata[name] = array

        corrhdu = fits.BinTableHDU(outdata, name=nam)
        hdul.insert(2, corrhdu)

        hdul.writeto(outpath + nam + '.fits', overwrite=True)

if __name__ == "__main__":
    main()


import os
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from networks.network import Actor, V_critic, Q_critic, MixNet

class OMIGA(object):
    def __init__(self, observation_spec, action_spec, num_agent, eval_env, config):
        self._alpha = 10
        self._gamma = config['gamma']
        self._tau = config['tau']
        self._hidden_sizes = config['hidden_sizes']
        self._mix_hidden_sizes = config['mix_hidden_sizes']
        self._batch_size = config['batch_size']
        self._lr = config['lr']
        self._grad_norm_clip = config['grad_norm_clip']
        self._num_agent = num_agent
        self._device = config['device']
        self._eval_env = eval_env
        self._iteration = 0
        self._optimizers = dict()

        # v-network
        self._v_network = V_critic(observation_spec, num_agent, self._hidden_sizes, self._device).to(self._device)
        self._v_target_network = copy.deepcopy(self._v_network)
        self._optimizers['v'] = torch.optim.Adam(self._v_network.parameters(), self._lr)

        # q-network and mix-network
        self._q_network = Q_critic(observation_spec, action_spec, num_agent, self._hidden_sizes, self._device).to(self._device)
        self._q_target_network = copy.deepcopy(self._q_network)
        self._mix_network = MixNet(observation_spec, action_spec, num_agent, self._mix_hidden_sizes, self._device).to(self._device)
        self._mix_target_network = copy.deepcopy(self._mix_network)
        self._q_param = list(self._q_network.parameters()) + list(self._mix_network.parameters())
        self._optimizers['q'] = torch.optim.Adam(self._q_param, self._lr)

        # policy-network
        self._policy_network = Actor(observation_spec, action_spec, num_agent, self._hidden_sizes, self._device).to(self._device)
        self._optimizers['policy'] = torch.optim.Adam(self._policy_network.parameters(), self._lr)

    def q_loss(self, o_with_a_id, s, o_next_with_id, s_next, r, mask, result={}):
        q_values = self._q_network(o_with_a_id)
        w, b = self._mix_network(s)
        q_total = (w * q_values).sum(dim=-2) + b.squeeze(dim=-1)

        v_next_target_values = self._v_target_network(o_next_with_id)
        w_next_target, b_next_target = self._mix_target_network(s_next)
        v_next_target_total = (w_next_target * v_next_target_values).sum(dim=-2) + b_next_target.squeeze(dim=-1)

        expected_q_total = r+ self._gamma * mask * v_next_target_total.detach()
        q_loss = ((q_total - expected_q_total.detach())**2).mean()

        result.update({
            'q_loss': q_loss,
            'expected_q_total': expected_q_total.mean(),
            'q_total': q_total.mean(),
            'w1': w[:,0,:].mean(),
            'w2': w[:,1,:].mean(),
            'b': b.mean(),
            'q_values1': q_values[:,0,:].mean(),
            'q_values2': q_values[:,1,:].mean(),
        })

        return result

    def v_loss(self, z, w_target, v_values, result={}):
        max_z = torch.max(z)
        max_z = torch.where(max_z < -1.0, torch.tensor(-1.0).to(self._device), max_z)
        max_z = max_z.detach()

        v_loss = torch.exp(z - max_z) + torch.exp(-max_z) * w_target * v_values / self._alpha
        v_loss = v_loss.mean()

        result.update({
            'v_loss': v_loss,
        })
        return result

    def policy_loss(self, exp_a, a, o_with_id, result={}):
        log_probs = self._policy_network.get_log_density(o_with_id, a)
        policy_loss = -(exp_a * log_probs).mean()

        result.update({
            'policy_loss': policy_loss,
        })
        return result

    def train_step(self, o, s, a, r, mask, s_next, o_next, a_next):
        # Shared network values
        one_hot_agent_id = torch.eye(self._num_agent).expand(o.shape[0], -1, -1).to(self._device)
        o_with_id = torch.cat((o, one_hot_agent_id), dim=-1)
        o_with_a_id = torch.cat((o, a, one_hot_agent_id), dim=-1)
        o_next_with_id = torch.cat((o_next, one_hot_agent_id), dim=-1)

        # q_loss
        loss_result = self.q_loss(o_with_a_id, s, o_next_with_id, s_next, r, mask, result={})

        # v and policy shared values
        q_target_values = self._q_target_network(o_with_a_id)
        w_target, b_target = self._mix_target_network(s)

        v_values = self._v_network(o_with_id)

        z = 1/self._alpha * (w_target.detach() * q_target_values.detach() - w_target.detach() * v_values)
        z = torch.clamp(z, min=-10.0, max=10.0)

        exp_a = torch.exp(z).detach().squeeze(-1)

        # v_loss
        loss_result = self.v_loss(z, w_target.detach(), v_values, result=loss_result)
        # policy_loss
        loss_result = self.policy_loss(exp_a, a, o_with_id, result=loss_result)

        self._optimizers['policy'].zero_grad()
        loss_result['policy_loss'].backward()
        nn.utils.clip_grad_norm_(self._policy_network.parameters(), self._grad_norm_clip)
        self._optimizers['policy'].step()

        self._optimizers['q'].zero_grad()
        loss_result['q_loss'].backward()
        nn.utils.clip_grad_norm_(self._q_param, self._grad_norm_clip)
        self._optimizers['q'].step()

        self._optimizers['v'].zero_grad()
        loss_result['v_loss'].backward()
        nn.utils.clip_grad_norm_(self._v_network.parameters(), self._grad_norm_clip)
        self._optimizers['v'].step()

        # soft update
        for param, target_param in zip(self._q_network.parameters(), self._q_target_network.parameters()):
            target_param.data.copy_(self._tau * param.data + (1 - self._tau) * target_param.data)
        for param, target_param in zip(self._mix_network.parameters(), self._mix_target_network.parameters()):
            target_param.data.copy_(self._tau * param.data + (1 - self._tau) * target_param.data)
        for param, target_param in zip(self._v_network.parameters(), self._v_target_network.parameters()):
            target_param.data.copy_(self._tau * param.data + (1 - self._tau) * target_param.data)

        self._iteration += 1

        loss_result.update({
            'v_values1': v_values[:,0,:].mean(),
            'v_values2': v_values[:,0,:].mean(),
            'q_target_values1': q_target_values[:,0,:].mean(),
            'q_target_values2': q_target_values[:,1,:].mean(),
        })
        return loss_result

    def step(self, o):
        o = torch.from_numpy(o).to(self._device)
        one_hot_agent_id = torch.eye(self._num_agent).expand(o.shape[0], -1, -1).to(self._device)
        o_with_id = torch.cat((o, one_hot_agent_id), dim=-1)
        action = self._policy_network.get_deterministic_action(o_with_id)

        return action.detach().cpu()