from abc import ABC, ABCMeta, abstractmethod
from dataclasses import asdict, dataclass
from enum import Enum, EnumMeta, IntEnum, auto
import json
from typing import Callable, Dict, Mapping, Optional, TypeVar, Union, List

import cvc5 as pycvc5

from .sorts import *

class Kind(Enum):
    BVAdd           = auto()
    BVSub           = auto()
    BVMul           = auto()
    BVOr            = auto()
    BVAnd           = auto()
    BVNot           = auto()
    BVXor           = auto()
    BVExtract       = auto()
    BVConcat        = auto()
    BVOrr           = auto() # OR reduction
    BVXorr          = auto() # XOR reduction
    BVUle           = auto() # Unsigned less than/equal
    BVUlt           = auto() # Unsigned less than
    BVUge           = auto() # Unsigned greater than/equal
    BVUgt           = auto() # unsigned greater than
    BVZeroPad       = auto()
    BVSignExtend    = auto()
    BVSll           = auto() # Shift left (logical)
    BVSrl           = auto() # Shift right (logical)
    BVSra           = auto() # Shift right (arithmetic)
    Or              = auto()
    And             = auto()
    Not             = auto()
    Xor             = auto()
    Implies         = auto()
    Ite             = auto()
    Match           = auto()
    Lambda          = auto()
    Equal           = auto()
    Distinct        = auto()
    Exists          = auto()
    ForAll          = auto()
    Select          = auto()
    Store           = auto()

    def to_cvc5(self, _cvc5_ctx):
        try:
            return _OP_KIND_MAPPING[self]
        except KeyError:
            raise NotImplementedError(f"cannot convert kind {self}")

    @staticmethod
    def from_cvc5(cvc5_kind):
        try:
            return _OP_KIND_REV_MAPPING[cvc5_kind]
        except KeyError:
            raise NotImplementedError(f"cannot convert CVC5 kind {cvc5_kind}")

    def to_sygus2(self):
        raise NotImplementedError()

# https://cvc5.github.io/docs/api/cpp/kind.html
# Maps our Kind class to pycvc5.Kind...
_OP_KIND_MAPPING = {
    Kind.BVAdd:         pycvc5.Kind.BITVECTOR_ADD,
    Kind.BVSub:         pycvc5.Kind.BITVECTOR_SUB,
    Kind.BVMul:         pycvc5.Kind.BITVECTOR_MULT,

    Kind.BVOr:          pycvc5.Kind.BITVECTOR_OR,
    Kind.BVAnd:         pycvc5.Kind.BITVECTOR_AND,
    Kind.BVNot:         pycvc5.Kind.BITVECTOR_NOT,
    Kind.BVXor:         pycvc5.Kind.BITVECTOR_XOR,
    Kind.BVSll:         pycvc5.Kind.BITVECTOR_SHL,
    Kind.BVSrl:         pycvc5.Kind.BITVECTOR_LSHR,
    Kind.BVSra:         pycvc5.Kind.BITVECTOR_ASHR,

    Kind.BVExtract:     pycvc5.Kind.BITVECTOR_EXTRACT,
    Kind.BVConcat:      pycvc5.Kind.BITVECTOR_CONCAT,

    Kind.BVOrr:         pycvc5.Kind.BITVECTOR_REDOR,

    Kind.BVUle:         pycvc5.Kind.BITVECTOR_ULE,
    Kind.BVUlt:         pycvc5.Kind.BITVECTOR_ULT,
    Kind.BVUge:         pycvc5.Kind.BITVECTOR_UGE,
    Kind.BVUgt:         pycvc5.Kind.BITVECTOR_UGT,

    Kind.BVZeroPad:     pycvc5.Kind.BITVECTOR_ZERO_EXTEND,
    Kind.BVSignExtend:  pycvc5.Kind.BITVECTOR_SIGN_EXTEND,

    Kind.Equal:         pycvc5.Kind.EQUAL,
    Kind.Distinct:      pycvc5.Kind.DISTINCT,
    Kind.Or:            pycvc5.Kind.OR,
    Kind.And:           pycvc5.Kind.AND,
    Kind.Not:           pycvc5.Kind.NOT,
    Kind.Xor:           pycvc5.Kind.XOR,
    Kind.Implies:       pycvc5.Kind.IMPLIES,
    Kind.Ite:           pycvc5.Kind.ITE,
    Kind.Match:         pycvc5.Kind.MATCH,

    Kind.Lambda:        pycvc5.Kind.LAMBDA,
    Kind.Exists:        pycvc5.Kind.EXISTS,
    Kind.ForAll:        pycvc5.Kind.FORALL,

    Kind.Select:        pycvc5.Kind.SELECT,
    Kind.Store:         pycvc5.Kind.STORE
}
# ...and vice versa
_OP_KIND_REV_MAPPING = {v: k for k, v in _OP_KIND_MAPPING.items()}

_OP_SYGUS_SYMBOLS = {
    Kind.BVAdd: "bvadd",
    Kind.BVSub: "bvsub",
    Kind.BVMul: "bvmult",
    Kind.BVOr: "bvor",
    Kind.BVAnd: "bvand",
    Kind.BVNot: "bvnot",
    Kind.BVXor: "bvxor",
    Kind.BVSll: "bvshl",
    Kind.BVSrl: "bvashr",
    Kind.BVSra: "bvlshr",
    Kind.BVUlt: "bvult",
    Kind.BVUle: "bvule",
    Kind.BVUgt: "bvugt",
    Kind.BVUge: "bvuge",
    # extract is a special case
    Kind.BVConcat: "concat",
    Kind.Or: "or",
    Kind.And: "and",
    Kind.Not: "not",
    Kind.Xor: "xor",
    Kind.Equal: "=",
    Kind.Ite: "ite",
    Kind.Implies: "=>",
    Kind.Exists: "exists",
    Kind.ForAll: "forall",
    Kind.Select: "select",
    Kind.Store: "store",
}

# === BEGIN SMT TERMS ===

_T = TypeVar("_T")

class Term(Translatable, ABC):
    """
    Representation for an SMT term. Operator overrides perform type checking and in some cases,
    implicit type conversions between arguments; construct an `OpTerm` manually if this behavior
    is undesired.
    """

    def _binop_type_check(self, other, sext=False, zpad=True, cast_int=True) -> Tuple["Term", "Term"]:
        """
        Checks that two operands have the same sort.

        If sext or zpad are set, if both are bitvectors and one is shorter than the other,
        then the one with the smaller width is zero-padded/sign-extended to the appropriate width.
        If sext/zpad are set and one arg is a boolean but the other is a bitvector, the boolean will
        be upcasted with an ITE statement.

        When `cast_int` is set, if `other` is an int, it will automatically be wrapped in
        a constant of the appropriate sort.
        """
        assert not (zpad and sext), "zpad and sext cannot both be set"
        if cast_int and isinstance(other, int):
            if isinstance(other, BoolConst):
                other = other.value
            if self.is_bool_expr():
                assert other in (0, 1), f"cannot coerce int {other} to {self.sort}"
                return self, (BoolConst.T if other else BoolConst.F) # type: ignore
            elif self.is_bv_expr():
                bitwidth = self.c_bitwidth()
                assert other < (2 ** bitwidth), f"int constant {other} exceeds max value of bv{bitwidth}"
                return self, BVConst(other, bitwidth)
        assert isinstance(other, Term), f"cannot combine {self} with {other}"
        assert hasattr(self, "sort"), repr(self)
        assert hasattr(other, "sort"), repr(other)
        s_bw = self.c_bitwidth()
        o_bw = other.c_bitwidth()
        if s_bw == o_bw:
            return self, other
        if sext:
            if s_bw > o_bw:
                if other.is_bool_expr():
                    return self, other.ite(BVConst((1 << o_bw) - 1, o_bw), BVConst(0, o_bw))
                return self, other.sign_extend(BVConst(s_bw - o_bw, o_bw))
            else:
                if self.is_bool_expr():
                    return self.ite(BVConst((1 << o_bw) - 1, o_bw), BVConst(0, o_bw)), other
                return self.sign_extend(BVConst(o_bw - s_bw, s_bw)), other
        if zpad:
            if s_bw > o_bw:
                if other.is_bool_expr():
                    return self, other.ite(BVConst(1, o_bw), BVConst(0, o_bw))
                return self, other.zero_pad(BVConst(s_bw - o_bw, o_bw))
            else:
                if self.is_bool_expr():
                    return self.ite(BVConst(1, o_bw), BVConst(0, o_bw)), other
                return self.zero_pad(BVConst(o_bw - s_bw, s_bw)), other
        assert self.sort == other.sort, f"cannot combine value {self} of sort {self.sort} to {other} of sort {other.sort}"
        return self, other

    # === OPTIMIZATIONS AND TREE TRAVERSALS ===

    def replace_vars(self, m: Mapping["Variable", "Term"]) -> "Term":
        """
        Returns a new term tree with all references to `var` replaced by `new_term`.
        """
        if isinstance(self, Variable) and self in m:
            return m[self]
        t = self
        for i, child in enumerate(self._children):
            t = t._replace_child(child.replace_vars(m), i)
        return t

    def zpad_all_children(self, extra_bits: int) -> "Term":
        """
        Zero pads all immediate children by `extra_bits`.
        This is necessary to implement verilog carry bit idioms.
        """
        t = self
        for i, child in enumerate(self._children):
            t = t._replace_child(child.zpad(extra_bits), i)
        return t

    def preorder_visit_tree(self, visit_fn: Callable[["Term"], _T], shortcircuit=True) -> _T:
        """
        Calls `visit_fn` on this node, then recursively on all children.
        Returns whatever `visit_fn(self)` returns.

        If `shortcircuit` is True, then this function will return without visiting
        its children if  `visit_fn(self)` returns a falsey value.
        """
        rv = visit_fn(self)
        if not shortcircuit or bool(rv):
            for s in self._children:
                assert isinstance(s, Term), s
                s.preorder_visit_tree(visit_fn, shortcircuit)
        return rv

    def get_vars(self) -> List["Variable"]:
        """
        Returns the list of variables referenced in this expression tree.
        """
        visited = set()
        child_vars = []
        def visitor(term):
            if term in visited:
                return
            if isinstance(term, Variable):
                visited.add(term)
                child_vars.append(term)
        self.preorder_visit_tree(visitor, False)
        return child_vars

    def optimize(self) -> "Term":
        if isinstance(self, OpTerm):
            # This method must also traverse children!
            return self.optimize()
        t = self
        for i, child in enumerate(self._children):
            t = t._replace_child(child.optimize(), i)
        return t

    def const_fold(self) -> "Term":
        if isinstance(self, OpTerm):
            # This method must also traverse children!
            return self.const_fold()
        t = self
        for i, child in enumerate(self._children):
            t = t._replace_child(child.const_fold(), i)
        return t

    def typecheck(self) -> bool:
        """
        Performs rudimentary type checking (no scope/semantic checking).
        Returns True on success, False on failure. Short circuits.

        TODO for now, this just ensures everything in the tree is a term
        TODO error reporting
        """
        for s in self._children:
            if not isinstance(s, Term):
                print(f"type error: child term {s} is not a Term object")
                return False
            if not s.typecheck():
                return False
        return True

    def eval(self, values: Dict[str, int]) -> Union["BVConst", "BoolConst"]:
        """
        Evaluates this expression with variables replaced by the values specified in `values`.
        """
        raise NotImplementedError()

    # === PROPERTIES AND HELPER FUNCTIONS ===

    def is_const(self) -> bool:
        return isinstance(self, BVConst) or isinstance(self, BoolConst)

    def is_bool_expr(self) -> bool:
        return isinstance(self.sort, BoolSort)

    def is_bv_expr(self) -> bool:
        return isinstance(self.sort, BVSort)

    def is_array_expr(self) -> bool:
        return isinstance(self.sort, ArraySort)

    def is_bv_or_bool_expr(self) -> bool:
        return isinstance(self.sort, BVSort) or isinstance(self.sort, BoolSort)

    def c_bitwidth(self) -> int:
        """
        Gets the bitwidth of this term's sort, erroring if the sort does not have a bitwidth.
        """
        if not self.is_bv_or_bool_expr():
            raise Exception(f"cannot get bitwidth of non-bv/bool term {self}")
        return self.sort.bitwidth

    def _maybe_bool2bv1(self):
        if isinstance(self, BoolConst):
            return BVConst(int(self), 1)
        elif self.is_bool_expr():
            return self.ite(BVConst(1, 1), BVConst(0, 1))
        else:
            return self

    # === OPERATOR OVERRIDES AND SYNTACTIC SUGAR ===

    def ite(self, t_term, f_term):
        """
        Constructs an ITE expression with this variable as condition.

        If this term is of sort BV1, then an expression is automatically added to
        check if this is equal to the constant 1bv1.
        """
        if self.is_bv_expr():
            assert self.c_bitwidth() == 1
            cond = OpTerm(Kind.Equal, (self, BVConst(1, self.c_bitwidth())))
        else:
            assert self.is_bool_expr()
            cond = self
        t_term, f_term = t_term._binop_type_check(f_term)
        return OpTerm(Kind.Ite, (cond, t_term, f_term))

    def match_const(
        self,
        cases: Dict[Union[int, "BVConst"], "Term"],
        # default: Optional["Term"]=None TODO implement default
    ):
        """
        A more primitive form of an SMT match expression. SMTLIB match expressions
        perform destructuring akin to functional programming pattern matching, but
        we only support matching against constants.

        When this term takes on a certain value, the expression should evaluate to
        the corresponding value in `cases`. All keys in `cases` must be of the same
        sort as this term, or are casted if they are int.

        This match must be exhaustive, so if cases do not provide all possible values
        of this sort, then `default` must be specified. An Exception is raised if the match
        is not exhaustive.
        """
        assert self.is_bv_expr()
        sort_max = (1 << self.c_bitwidth()) - 1
        covered = 0 # Bit set representing which cases have been covered
        args = [self] # OpTerm sees this as pairs of (case, term) occurring consecutively
        for c, term in cases.items():
            if isinstance(c, int):
                c = BVConst(c, self.c_bitwidth())
            assert c.val <= sort_max
            if (covered & (1 << c.val)) != 0:
                # Check if the case was already set
                raise Exception(f"Match case {c.val} occurred multiple times")
            covered |= 1 << c.val
            args.append(c)
            args.append(term)
        all_covered = (1 << sort_max) - 1 # 1 bit for each possible value
        default = None
        if not default:
            assert covered != all_covered, "Match default was provided, but cases were already exhaustive"
        else:
            assert covered == all_covered, "Match was non-exhaustive"
        return OpTerm(Kind.Match, tuple(args))

    def implies(self, other):
        return OpTerm(Kind.Implies, self._binop_type_check(other))

    def __lt__(self, other):
        return OpTerm(Kind.BVUlt, self._binop_type_check(other))

    def __le__(self, other):
        return OpTerm(Kind.BVUle, self._binop_type_check(other))

    def op_eq(self, other):
        """
        We can't override `__eq__` without breaking a decent amount of stuff, so
        `op_eq` is syntactic sugar for an equality expression instead.
        """
        return OpTerm(Kind.Equal, self._binop_type_check(other))

    def op_ne(self, other):
        return OpTerm(Kind.Distinct, self._binop_type_check(other))

    def __gt__(self, other):
        return OpTerm(Kind.BVUgt, self._binop_type_check(other))

    def __ge__(self, other):
        return OpTerm(Kind.BVUge, self._binop_type_check(other))

    def __add__(self, other):
        return OpTerm(Kind.BVAdd, self._binop_type_check(other))

    def __and__(self, other):
        if isinstance(other, Term) and other.is_bool_expr():
            op = Kind.And
        else:
            op = Kind.BVAnd
        return OpTerm(op, self._binop_type_check(other))

    def __invert__(self):
        if self.is_bool_expr():
            op = Kind.Not
        else:
            op = Kind.BVNot
        return OpTerm(op, (self,))

    def __mul__(self, other):
        return OpTerm(Kind.BVMul, self._binop_type_check(other))

    def __neg__(self):
        raise NotImplementedError()

    def __or__(self, other):
        if isinstance(other, Term) and other.is_bool_expr():
            op = Kind.Or
        else:
            op = Kind.BVOr
        return OpTerm(op, self._binop_type_check(other))

    def __lshift__(self, other):
        return self.sll(other)

    def sll(self, other):
        # TODO typecheck all shifts
        return OpTerm(Kind.BVSll, self._binop_type_check(other))

    def __rshift__(self, other):
        # Python right shift is technically arithmetic since integers
        # don't have a fixed size
        return self.sra(other)

    def srl(self, other):
        return OpTerm(Kind.BVSrl, self._binop_type_check(other))

    def sra(self, other):
        return OpTerm(Kind.BVSra, self._binop_type_check(other))

    def __sub__(self, other):
        return OpTerm(Kind.BVSub, self._binop_type_check(other))

    def __xor__(self, other):
        if isinstance(other, Term) and other.is_bool_expr():
            op = Kind.Xor
        else:
            op = Kind.BVXor
        return OpTerm(op, self._binop_type_check(other))

    def __getitem__(self, key):
        if self.is_array_expr():
            # Array indexing
            if isinstance(key, Term):
                if not key.sort == self.sort.idx_sort:
                    raise TypeError(f"cannot index {self} (sort {self.sort}) with {key} (sort {key.sort})")
                return OpTerm(Kind.Select, (self, key))
            elif isinstance(key, int):
                # Convert key to appropriate index sort
                i_key = type(self.sort.idx_sort)(key)
                return OpTerm(Kind.Select, (self, i_key))
            else:
                raise TypeError(key)
        # Bitvector indexing
        if self.is_bool_expr():
            width = 1
        else:
            assert self.is_bv_expr(), "only BV terms support indexing, instead term was " + str(self)
            width = self.c_bitwidth()
        wrap = lambda i: BVConst(i, width)
        def shift_and_mask(term):
            """
            Encodes extracts of a variable index as a shift + mask.
            """
            if term.c_bitwidth() == width:
                shamt = term
            else:
                shamt = term.zpad(width - term.c_bitwidth())
            return OpTerm(Kind.BVAnd, (OpTerm(Kind.BVSrl, (self, shamt)), BVConst(1, width)))
        if isinstance(key, int):
            # Extracts a single bit
            # ite will auto-cast to a bool
            return OpTerm(Kind.BVExtract, (self, wrap(key), wrap(key))).ite(BoolConst.T, BoolConst.F)
        elif isinstance(key, Term):
            return shift_and_mask(key.start)
        elif isinstance(key, slice):
            if isinstance(key.start, Term) and not key.start.is_const() and key.start == key.stop:
                return shift_and_mask(key.start)
            if isinstance(key.start, int):
                hi = wrap(max(key.start, key.stop))
            else:
                hi = key.start
            if isinstance(key.stop, int):
                lo = wrap(min(key.start, key.stop))
            else:
                lo = key.stop
            if isinstance(hi, BVConst):
                assert hi.val < width, f"extract upper index {hi.val} of '{self}' exceeds bounds of bv{width}"
            if isinstance(lo, BVConst):
                assert lo.val >= 0, f"extract lower index {lo.val} of '{self}' was negative"
                if hi == lo:
                    return OpTerm(Kind.BVExtract, (self, hi, lo)).ite(BoolConst.T, BoolConst.F)
            return OpTerm(Kind.BVExtract, (self, hi, lo))
        raise TypeError(f"cannot index {self} with {key}")

    def __setitem__(self, key, value):
        if self.is_array_expr():
            if isinstance(key, type(self.sort.idx_sort)):
                pass
            elif isinstance(key, int):
                # Convert key to appropriate index sort
                key = type(self.sort.idx_sort)(key)
            else:
                raise TypeError(key)
            if isinstance(value, type(self.sort.elem_sort)):
                pass
            elif isinstance(key, int):
                value = type(self.sort.elem_sort)(key)
            else:
                raise TypeError(key)
            return OpTerm(Kind.Store, (self, key, value))
        raise TypeError(key)

    def concat(self, *others):
        if len(others) == 0:
            return self
        else:
            return OpTerm(Kind.BVConcat, (self, *others))

    def zpad(self, extra_bits: Union["BVConst", int]):
        return self.zero_pad(extra_bits)

    def zero_pad(self, extra_bits: Union["BVConst", int]):
        """Zero pads this term with an addition `extra_bits` bits."""
        if isinstance(extra_bits, int):
            extra_bits = BVConst(extra_bits, self.c_bitwidth())
        if self.is_bool_expr():
            new_width = self.c_bitwidth() + extra_bits.val
            return self.ite(BVConst(1, new_width), BVConst(0, new_width))
        return OpTerm(Kind.BVZeroPad, (self, extra_bits))

    def sext(self, extra_bits: Union["BVConst", int]):
        return self.sign_extend(extra_bits)

    def sign_extend(self, extra_bits: Union["BVConst", int]):
        """Sign extends this term with an addition `extra_bits` bits."""
        if isinstance(extra_bits, int):
            extra_bits = BVConst(extra_bits, self.c_bitwidth())
        if self.is_bool_expr():
            new_width = self.c_bitwidth() + extra_bits.val
            return self.ite(BVSort(new_width).max_bv_const(), BVConst(0, new_width))
        return OpTerm(Kind.BVSignExtend, (self, extra_bits))

    def orr(self):
        return OpTerm(Kind.BVOrr, (self,))

    def xorr(self):
        return OpTerm(Kind.BVXorr, (self,))

    # === ABSTRACT AND SHARED STATIC METHODS ===

    @staticmethod
    def from_cvc5(cvc5_term):
        from cvc5 import Kind as k
        cvc5_kind = cvc5_term.getKind()
        if cvc5_kind == k.VARIABLE:
            return Variable.from_cvc5(cvc5_term)
        elif cvc5_kind == k.LAMBDA:
            return LambdaTerm.from_cvc5(cvc5_term)
        elif cvc5_kind in (k.EXISTS, k.FORALL):
            return QuantTerm.from_cvc5(cvc5_term)
        elif cvc5_kind == k.CONST_BITVECTOR:
            return BVConst.from_cvc5(cvc5_term)
        elif cvc5_kind == k.CONST_BOOLEAN:
            return BoolConst(cvc5_term.getBooleanValue())
        elif cvc5_kind in _OP_KIND_REV_MAPPING:
            return OpTerm.from_cvc5(cvc5_term)

        # TODO what corresponds to UFTerm?
        raise NotImplementedError("Cannot convert from CVC5 term of kind " + str(cvc5_kind))

    @property
    @abstractmethod
    def sort(self):
        raise NotImplementedError()

    @property
    def _should_paren(self):
        return len(self._children) > 0 and not isinstance(self, ApplyUF) \
            and not (isinstance(self, OpTerm) and self.kind == Kind.BVExtract)

    @property
    @abstractmethod
    def _children(self):
        raise NotImplementedError()

    @abstractmethod
    def _replace_child(self, new_term, index) -> "Term":
        """
        Returns a new term with the `index`th child replaced by new_term
        """
        raise NotImplementedError()

# Needed to be able to make BoolConst sealed
class _TermMeta(ABCMeta, EnumMeta):
    pass

@dataclass(frozen=True)
class Variable(Term):
    """
    This class represents variable references. To obtain a `Translatable` for a variable
    declaration, use the `get_decl()` method.
    """

    name: str
    _sort: Sort

    def __post_init__(self):
        assert isinstance(self.name, str) and len(self.name) > 0
        assert isinstance(self._sort, Sort), f"{self._sort} is not a Sort instance"

    def __str__(self):
        return f"{self.name}"

    def to_uf(self) -> "UFTerm":
        """Converts this variable to a 0-arity `UFTerm`."""
        return UFTerm(self.name, self.sort, ())

    def get_decl(self, init_value=None) -> "VarDecl":
        """Gets a `Translatable` declaration of this variable."""
        return VarDecl(self.name, self.sort, init_value=init_value)

    def add_prefix(self, prefix) -> "Variable":
        return Variable(prefix + self.name, self.sort)

    def get_base(self) -> "Variable":
        return Variable(self.name[self.name.rfind(".") + 1:], self.sort)

    @staticmethod
    def from_cvc5(cvc5_term) -> "Variable":
        if cvc5_term.getKind() != pycvc5.Kind.VARIABLE:
            raise TypeError("Variable must be translated from pycvc5.Kind.VARIABLE, instead got " + str(cvc5_term.getKind()))
        # it seems like cvc5 will wrap the variable name in |var| if there are square brackets inside
        # so we need to strip those off
        name = cvc5_term.getSymbol()
        # technically fails if NAME is length 1, but at that point we don't care anymore
        if name[0] == "|" and name[-1] == "|":
            name = name[1:-1]
        return Variable(name, Sort.from_cvc5(cvc5_term.getSort()))

    @property
    def sort(self):
        return self._sort

    @property
    def _children(self):
        return []

    def _replace_child(self, new_term, index):
        return self

    def eval(self, values):
        assert self.name in values, f"could not evaluate variable '{self.name}' from mapping {values}"
        if self.is_bv_expr():
            return BVConst(values[self.name], self.c_bitwidth())
        elif self.is_bool_expr():
            return BoolConst(values[self.name])
        else:
            raise NotImplementedError()

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            if self in cvc5_ctx.terms:
                return cvc5_ctx.terms[self]
            else:
                v = cvc5_ctx.solver.mkVar(self.sort.to_cvc5(cvc5_ctx), self.name)
                cvc5_ctx.terms[self] = v
                return v
        elif tgt in (TargetFormat.SYGUS2, TargetFormat.VERILOG):
            return self.name
        elif tgt == TargetFormat.UCLID:
            if kwargs.get("prime_vars", False):
                return self.name + "'"
            else:
                return self.name
        elif tgt == TargetFormat.JSON:
            return json.dumps(asdict(self))
        raise NotImplementedError("cannot convert Variable to " + str(tgt))

bool_variable = lambda s, **kwargs: Variable(s, BoolSort(), **kwargs)
"""Constructs a variable of boolean sort with the given name."""

bv_variable = lambda s, w, **kwargs: Variable(s, BVSort(w), **kwargs)
"""Constructs a variable of a bitvector sort with the given name and width."""

@dataclass(frozen=True)
class VarDecl(Translatable):
    name: str
    sort: Sort
    init_value: Optional["BVConst"] = None

    def __str__(self):
        if self.init_value is None:
            return f"{self.name} : {self.sort}"
        else:
            return f"{self.name} : {self.sort} = {self.init_value}"

    def c_bitwidth(self):
        return self.get_ref().c_bitwidth()

    def get_ref(self):
        return Variable(self.name, self.sort)

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        """
        Possible kwargs:
        - `cvc5_ctx`: a `rtl2model.lynth.Cvc5Ctx` object used to get a reference to a cvc5 solver
        - `is_reg`: a boolean used by verilog conversion to determine if this should be declared
                    as a reg (if false or not specified, the variable will be a wire)
        - `anyconst`: a boolean used by verilog conversion; when specified, it will add the
                    `(* anyconst *)` yosys attribute
        """
        if tgt == TargetFormat.CVC5:
            # The CVC5 interface doesn't really distnguish between declarations
            # and references of bound variables
            return self.get_ref().to_cvc5(cvc5_ctx=kwargs["cvc5_ctx"])
        elif tgt == TargetFormat.SYGUS2:
            return f"(declare-var {self.name} {self.sort.to_sygus2()})"
        elif tgt == TargetFormat.VERILOG:
            if isinstance(self.sort, ArraySort):
                raise NotImplementedError("VarDecl verilog array translation not supported yet")
            is_reg = kwargs.get("is_reg", False)
            decl = "reg" if is_reg else "wire"
            if self.c_bitwidth() != 1:
                decl += " " + self.sort.to_verilog_str()
            decl += f" {self.name}"
            if self.init_value is not None:
                decl += " = " + self.init_value.to_verilog_str()
            decl += ";"
            anyconst = kwargs.get("anyconst", False)
            if anyconst:
                return  "(* anyconst *) " + decl
            else:
                return decl
        elif tgt == TargetFormat.UCLID:
            # TODO need to identify inputs/outputs
            return f"var {self.name} : {self.sort.to_uclid()};"
        raise NotImplementedError("cannot convert VarDecl to " + str(tgt))

@dataclass(frozen=True)
class OpTerm(Term):
    kind: Kind
    args: Tuple[Term, ...]

    @property
    def sort(self):
        # TODO better type checking
        bvops = { Kind.BVAdd, Kind.BVSub, Kind.BVOr, Kind.BVAnd, Kind.BVXor, Kind.BVNot,
                  Kind.BVSll, Kind.BVSrl, Kind.BVSra }
        if self.kind in bvops:
            return self.args[0].sort
        elif self.kind == Kind.BVExtract:
            # TODO fix this hack: in the CVC5 API, BVExtract must first be turned into
            # an operator via Solver::mkOp(BVExtract, high, low)
            # As a workaround, we store the high/low parameters as BVConst arguments
            assert isinstance(self.args[1], BVConst), f"extract first index must be BVConst, instead was {self.args[1]}"
            assert isinstance(self.args[2], BVConst), f"extract second index must be BVConst, instead was {self.args[2]}"
            return BVSort(self.args[1].val - self.args[2].val + 1)
        elif self.kind in (Kind.BVZeroPad, Kind.BVSignExtend):
            assert self.args[0].is_bv_expr(), repr(self.args[0])
            assert isinstance(self.args[1], BVConst), repr(self.args[1])
            return BVSort(self.args[0].c_bitwidth() + self.args[1].val)
        elif self.kind == Kind.BVConcat:
            return BVSort(sum(a.c_bitwidth() for a in self.args))
        elif self.kind == Kind.Ite:
            return self.args[1].sort
        elif self.kind == Kind.Match:
            return self.args[2].sort # Match args are (cond, c1, v1, c2, v2, ...)
        bitops = { Kind.Or, Kind.And, Kind.Xor, Kind.Not, Kind.Equal, Kind.Distinct, Kind.Implies,
                   Kind.BVUle, Kind.BVUlt, Kind.BVUge, Kind.BVUgt}
        if self.kind in bitops:
            return BoolSort()
        if self.kind == Kind.Select:
            assert self.args[0].is_array_expr(), ArraySort
            return self.args[0].sort.elem_sort
        raise NotImplementedError(f"Cannot get Sort for kind {self.kind}")

    @property
    def _children(self):
        return list(self.args)

    def _replace_child(self, new_term, index):
        new_args = list(self.args)
        new_args[index] = new_term
        return OpTerm(self.kind, tuple(new_args))

    def __str__(self):
        return self.to_verilog_str()

    @staticmethod
    def from_cvc5(cvc5_term):
        """
        Translates the provided cvc5 Term object into an `OpTerm`.
        """
        cvc5_kind = cvc5_term.getKind()
        if cvc5_kind not in _OP_KIND_REV_MAPPING:
            raise TypeError("OpTerm cannot be translated from " + str(cvc5_kind))
        kind = Kind.from_cvc5(cvc5_kind)
        if kind == Kind.BVExtract:
            op = cvc5_term.getOp()
            base = Term.from_cvc5(cvc5_term[0])
            assert op[0].getKind() == pycvc5.Kind.CONST_RATIONAL
            upper = BVConst(op[0].getIntegerValue(), base.c_bitwidth())
            assert op[1].getKind() == pycvc5.Kind.CONST_RATIONAL
            lower = BVConst(op[1].getIntegerValue(), base.c_bitwidth())
            return OpTerm(kind, (base, upper, lower))
        if kind in (Kind.BVZeroPad, Kind.BVSignExtend):
            op = cvc5_term.getOp()
            base = Term.from_cvc5(cvc5_term[0])
            assert op[0].getKind() == pycvc5.Kind.CONST_RATIONAL
            padamt = BVConst(op[0].getIntegerValue(), base.c_bitwidth())
            return OpTerm(kind, (base, padamt))
        return OpTerm(kind, tuple([Term.from_cvc5(t) for t in cvc5_term]))

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            cvc5_kind = self.kind.to_cvc5(self)
            if self.kind == Kind.BVExtract:
                # TODO special case BVExtract, Select, and Store for from_cvc5?
                assert isinstance(self.args[1], BVConst)
                assert isinstance(self.args[2], BVConst)
                op = cvc5_ctx.solver.mkOp(cvc5_kind, self.args[1].val, self.args[2].val)
                return cvc5_ctx.solver.mkTerm(op, self.args[0].to_cvc5(cvc5_ctx))
            elif self.kind == Kind.BVConcat:
                # Special case for bool operands for a concat
                args = [o._maybe_bool2bv1().to_cvc5(cvc5_ctx) for o in self.args]
                return cvc5_ctx.solver.mkTerm(cvc5_kind, *args)
            elif self.kind in (Kind.BVZeroPad, Kind.BVSignExtend):
                assert isinstance(self.args[1], BVConst)
                op = cvc5_ctx.solver.mkOp(cvc5_kind, self.args[1].val)
                return cvc5_ctx.solver.mkTerm(op, self.args[0].to_cvc5(cvc5_ctx))
            elif self.kind == Kind.Match:
                # Jank workaround until I figure out how to pattern match vs constants
                a0 = self.args[0]
                # Last expression is the "else" case, so cond doesn't matter
                t = self.args[-1]
                for i in range(len(self.args) - 4, 0, -2):
                    t = a0.op_eq(self.args[i]).ite(self.args[i + 1], t)
                return t.to_cvc5(cvc5_ctx)
                # case_terms = []
                # for i in range(1, len(self.args), 2):
                #     case_terms.append(
                #         cvc5_ctx.solver.mkTerm(
                #             pycvc5.Kind.MatchCase,
                #             self.args[i].to_cvc5(cvc5_ctx),
                #             self.args[i + 1].to_cvc5(cvc5_ctx),
                #         )
                #     )
                # return cvc5_ctx.solver.mkTerm(cvc5_kind, self.args[0].to_cvc5(cvc5_ctx), *case_terms)
            t = cvc5_ctx.solver.mkTerm(cvc5_kind, *[v.to_cvc5(cvc5_ctx) for v in self.args])
            return t
        elif tgt == TargetFormat.SYGUS2:
            if self.kind == Kind.Match:
                # Convert to a giant ITE chain
                a0 = self.args[0]
                # Last expression is the "else" case, so cond doesn't matter
                t = self.args[-1]
                for i in range(len(self.args) - 4, 0, -2):
                    t = a0.op_eq(self.args[i]).ite(self.args[i + 1], t)
                return t.to_sygus2()
            # extract, zero_extend, and sign_extend all use decimals instead of hex
            if self.kind == Kind.BVExtract:
                assert isinstance(self.args[1], BVConst), self.args
                assert isinstance(self.args[2], BVConst), self.args
                return f"((_ extract {self.args[1].val} {self.args[2].val}) {self.args[0].to_sygus2()})"
            if self.kind == Kind.BVZeroPad:
                assert isinstance(self.args[1], BVConst), self.args
                return f"((_ zero_extend {self.args[1].val}) {self.args[0].to_sygus2()})"
            if self.kind == Kind.BVSignExtend:
                assert isinstance(self.args[1], BVConst), self.args
                return f"((_ sign_extend {self.args[1].val}) {self.args[0].to_sygus2()})"
            if self.kind == Kind.Distinct:
                return f"(not (= " +  " ".join(a.to_sygus2() for a in self.args) + "))"
            return "(" + _OP_SYGUS_SYMBOLS[self.kind] + " " + " ".join([a.to_sygus2() for a in self.args]) + ")"
        elif tgt == TargetFormat.VERILOG:
            v = self.kind
            def wrap(term):
                if term._should_paren:
                    return "(" + term.to_verilog_str() + ")"
                else:
                    return term.to_verilog_str()
            unops = {
                Kind.Not: "!",
                Kind.BVNot: "~",
                Kind.BVOrr: "|",
                Kind.BVXorr: "^",
            }
            if v in unops:
                a0_str = wrap(self.args[0])
                return f"{unops[v]}{a0_str}"
            binops = {
                Kind.BVAdd: "+",
                Kind.BVSub: "-",
                Kind.BVMul: "*",
                Kind.BVOr: "|",
                Kind.BVAnd: "&",
                Kind.BVXor: "^",
                Kind.BVSll: "<<",
                Kind.BVSra: ">>",
                Kind.BVSrl: ">>>",
                Kind.BVUlt: "<",
                Kind.BVUle: "<=",
                Kind.BVUgt: ">",
                Kind.BVUge: ">=",
                Kind.Or: "||",
                Kind.And: "&&",
                Kind.Xor: "^", # TODO check if this differs from bv xor
                Kind.Equal: "==",
                Kind.Distinct: "!=",
            }
            if v in binops:
                a0_str = wrap(self.args[0])
                a1_str = wrap(self.args[1])
                return f"{a0_str} {binops[v]} {a1_str}"
            if v == Kind.Implies:
                a0_str = wrap(self.args[0])
                a1_str = wrap(self.args[1])
                return f"!{a0_str} || {a1_str}"
            if v == Kind.Ite:
                a0_str = wrap(self.args[0])
                a1_str = wrap(self.args[1])
                a2_str = wrap(self.args[2])
                return f"{a0_str} ? {a1_str} : {a2_str}"
            if v == Kind.Match:
                # Verilog case statements aren't expressions, so just ITE it up
                a0 = self.args[0]
                # Last expression is the "else" case, so cond doesn't matter
                t = self.args[-1]
                for i in range(len(self.args) - 4, 0, -2):
                    t = a0.op_eq(self.args[i]).ite(self.args[i + 1], t)
                return t.to_verilog_str()
            if v == Kind.BVConcat:
                return "{" + ",".join(wrap(a) for a in self.args) + "}"
            if v == Kind.BVExtract:
                a0_str = wrap(self.args[0])
                if self.args[1].is_const():
                    a1_str = str(int(self.args[1].val))
                else:
                    a1_str = wrap(self.args[1])
                if self.args[2].is_const():
                    a2_str = str(int(self.args[2].val))
                else:
                    a2_str = wrap(self.args[2])
                return f"{a0_str}[{a1_str}:{a2_str}]"
            if v == Kind.Select:
                a0_str = wrap(self.args[0])
                a1_str = wrap(self.args[1])
                return f"{a0_str}[{a1_str}]"
            if v == Kind.BVSignExtend:
                # duplicate MSB, then concat
                # e.g. {{8{ext[7]}}, ext} where ext is bv8
                a0_str = self.args[0].to_verilog_str()
                msb_i = self.args[0].c_bitwidth() - 1
                # brace balancing:
                # c = outermost concat curly brace
                # r = outside repeat
                # m = outside msb
                #        c0r0                  m0                 m1r1          c1
                return f"{{{{{self.args[1].val}{{{a0_str}[{msb_i}]}}}}, {a0_str}}}"
            if v == Kind.BVZeroPad:
                a0_str = self.args[0].to_verilog_str()
                #        c0r0                     r1          c1
                return f"{{{{{self.args[1].val}'h0}}, {a0_str}}}"
            raise NotImplementedError(v)
        elif tgt == TargetFormat.UCLID:
            def wrap(term):
                if term._should_paren:
                    return "(" + term.to_uclid(**kwargs) + ")"
                else:
                    return term.to_uclid(**kwargs)
            # TODO use uclid library
            unops = {
                Kind.Not: "!",
                Kind.BVNot: "~"
            }
            if self.kind in unops:
                return unops[self.kind] + wrap(self.args[0])
            binops = {
                Kind.BVAdd: "+",
                Kind.BVSub: "-",
                Kind.BVOr: "|",
                Kind.BVAnd: "&",
                Kind.BVXor: "^",
                Kind.BVUlt: "<_u",
                Kind.BVUle: "<=_u",
                Kind.BVUgt: ">_u",
                Kind.BVUge: ">=_u",
                Kind.Equal: "==",
                Kind.Distinct: "!=",
                Kind.Or: "||",
                Kind.And: "&&",
                Kind.Xor: "^", # TODO check if this differs from bv xor
                Kind.Implies: "==>",
            }
            if self.kind in binops:
                return wrap(self.args[0]) + " " + binops[self.kind] + " " + wrap(self.args[1])
            shifts = {
                Kind.BVSll: "bv_left_shift",
                Kind.BVSra: "bv_a_right_shift",
                Kind.BVSrl: "bv_l_right_shift",
            }
            if self.kind in shifts:
                return shifts[self.kind] + "(" + self.args[0].to_uclid(**kwargs) + ", " + self.args[1].to_uclid(**kwargs) + ")"
            if self.kind == Kind.Ite:
                return "if (" + self.args[0].to_uclid(**kwargs) + ") then " + wrap(self.args[1]) + " else " + wrap(self.args[2])
            if self.kind == Kind.BVExtract:
                if not isinstance(self.args[1], BVConst):
                    raise Exception(f"non-constant BVExtract not supported ({self.args[1]})")
                if not isinstance(self.args[2], BVConst):
                    raise Exception(f"non-constant BVExtract not supported ({self.args[2]})")
                a1_str = str(self.args[1].val)
                a2_str = str(self.args[2].val)
                return wrap(self.args[0]) + "[" + a1_str + ":" + a2_str + "]"
            if self.kind == Kind.BVConcat:
                return " ++ ".join(wrap(a) for a in self.args)
            if self.kind == Kind.Select:
                return wrap(self.args[0]) + "[" + self.args[1].to_uclid(**kwargs) + "]"
            if self.kind == Kind.BVZeroPad:
                assert isinstance(self.args[1], BVConst)
                return f"bv_zero_extend({self.args[1].val}, {self.args[0].to_uclid(**kwargs)})"
            if self.kind == Kind.Match:
                # uclid case statements aren't expressions, so just ITE it up
                a0 = self.args[0]
                # Last expression is the "else" case, so cond doesn't matter
                t = self.args[-1]
                for i in range(len(self.args) - 4, 0, -2):
                    t = a0.op_eq(self.args[i]).ite(self.args[i + 1], t)
                return t.to_uclid(**kwargs)
            raise NotImplementedError(self.kind)
        raise NotImplementedError("cannot convert OpTerm to " + str(tgt))

    def optimize(self) -> Term:
        return self.const_fold()

    def const_fold(self) -> Term:
        # Optimize all children first
        t = self
        for i, child in enumerate(self._children):
            t = t._replace_child(child.const_fold(), i)
        assert isinstance(t, OpTerm), t
        assert self.kind == t.kind
        # Special case for constant folding: ITEs/match where the cond is constant
        if t.kind == Kind.Match:
            cond = t.args[0]
            if cond.is_const():
                # Return value for correct case
                # Recall that every other arg is cond, then value
                for i in range(1, len(t.args) - 1):
                    case = t.args[i]
                    if case == cond:
                        return t.args[i + 1]
            else:
                return t
        if t.kind == Kind.Ite:
            cond = t.args[0]
            if cond.is_const():
                return t.args[1] if bool(cond) else t.args[2]
            else:
                return t
        # Constant folding: all arguments must be constants
        # This technically isn't true for stuff like concats,
        # but this heuristic works for binops and unops
        args = t._children
        for child in args:
            if not child.is_const():
                return t
        return t._do_const_eval(args)

    def eval(self, values):
        # Optimize all children first
        children = [child.eval(values) for child in self._children]
        return self._do_const_eval(children)

    def _do_const_eval(self, args):
        # unary ops
        a0 = args[0]
        a0_bw = a0.c_bitwidth()
        a0_mask = (1 << a0.c_bitwidth()) - 1
        if self.kind == Kind.Not:
            return BoolConst.T if a0 == BoolConst.F else BoolConst.T
        if self.kind == Kind.BVNot:
            return BVConst((~a0.val) & a0_mask, a0_bw)
        # binary ops
        a1 = args[1]
        if self.kind == Kind.BVAdd:
            return BVConst((a0.val + a1.val) & a0_mask, a0_bw)
        if self.kind == Kind.BVSub:
            # TODO check wrapping behavior
            return BVConst((a0.val - a1.val) & a0_mask, a0_bw)
        if self.kind == Kind.BVOr:
            return BVConst(a0.val | a1.val, a0_bw)
        if self.kind == Kind.BVAnd:
            return BVConst(a0.val & a1.val, a0_bw)
        if self.kind == Kind.BVXor:
            return BVConst(a0.val ^ a1.val, a0_bw)
        if self.kind == Kind.BVUlt:
            return BoolConst(a0.val < a1.val)
        if self.kind == Kind.BVUle:
            return BoolConst(a0.val <= a1.val)
        if self.kind == Kind.BVUgt:
            return BoolConst(a0.val > a1.val)
        if self.kind == Kind.BVUge:
            return BoolConst(a0.val >= a1.val)
        if self.kind == Kind.Equal:
            return BoolConst(a0.val == a1.val)
        if self.kind == Kind.Distinct:
            return BoolConst(a0.val != a1.val)
        if self.kind == Kind.Or:
            return a0 or a1 # IntEnum provides this automatically
        if self.kind == Kind.And:
            return a0 and a1
        if self.kind == Kind.Xor:
            return BoolConst.T if a0 != a1 else BoolConst.F
        if self.kind == Kind.Implies:
            # p ==> q == !p || q
            return (a0 == BoolConst.F) or a1
        # Kind.BVSll: "bv_left_shift",
        # Kind.BVSra: "bv_a_right_shift",
        # Kind.BVSrl: "bv_l_right_shift",
        # if self.kind == Kind.Select:
        if self.kind == Kind.BVZeroPad:
            return BVConst(a0.val, a1.val + a0_bw)
        # if self.Kind == Kind.SignExtend:
        # ternary operators and higher
        a2 = args[2]
        if self.kind == Kind.BVExtract:
            # TODO check semantics for off-by-one in verilog/cvc?
            # also, what direction does the slice go?
            high = a1.val
            low = a2.val
            width = (high - low) + 1
            mask = ((1 << width) - 1) << low
            return BVConst((a0.val & mask) >> low, width)
        # if self.kind == Kind.BVConcat:
        # Ternary operator
        if self.kind == Kind.Ite:
            return a1 if bool(a0) else a2
        return t

# TODO distinguish between references and declarations
# variables and UFTerms should be referenced in the same way, but obviously declared
# in different manners
@dataclass(frozen=True)
class UFTerm(Term):
    """
    A term representing an uninterpreted function of arbitrary arity.
    """
    name: str
    _sort: Sort
    params: Tuple[Variable, ...]

    def apply(self, *args):
        return ApplyUF(self, tuple(args))

    def __str__(self):
        return self.name

    @property
    def arity(self):
        return len(self.params)

    @staticmethod
    def from_cvc5(cvc5_term):
        raise NotImplementedError("Cannot convert from CVC5 UF term")

    @property
    def sort(self):
        return self._sort

    @property
    def _children(self):
        return []

    def _replace_child(self, new_term, index):
        return self

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            return cvc5_ctx.solver.declareFun(
                self.name,
                [p.sort.to_cvc5(cvc5_ctx) for p in self.params],
                self.sort.to_cvc5(cvc5_ctx)
            )
        elif tgt == TargetFormat.UCLID:
            params_s = ",".join(f"{v.name} : {v.sort.to_uclid()}" for v in self.params)
            return f"synthesis function {self.name}({params_s}) : {self.sort.to_uclid()};"
        raise NotImplementedError("cannot convert UFTerm to " + str(tgt))

@dataclass(frozen=True)
class LambdaTerm(Term):
    params: Tuple[Variable, ...]
    body: Term

    def apply(self, *args):
        assert len(args) == len(self.params)
        # jank hack for inlining lambdas
        return self.body.replace_vars({self.params[i]: args[i] for i in range(len(args))})

    @property
    def sort(self):
        return FunctionSort(tuple([p.sort for p in self.params]), self.body.sort)

    @property
    def _children(self):
        return self.body

    def _replace_child(self, new_term, index):
        if index != 0:
            raise IndexError()
        return LambdaTerm(self.params, new_term)

    def __str__(self):
        return f"({','.join(str(s) for s in self.params)}) -> {self.body}"

    @staticmethod
    def from_cvc5(cvc5_term):
        if cvc5_term.getKind() == pycvc5.Kind.LAMBDA:
            c_params = cvc5_term[0]
            c_body = cvc5_term[1]
            return LambdaTerm(
                tuple([Variable.from_cvc5(c) for c in c_params]),
                Term.from_cvc5(c_body),
            )
        else:
            raise TypeError("LambdaTerm must be translated from pycvc5.Kind.LAMBDA, instead got " + str(cvc5_term.getKind()))

    def _to_sygus2_fn(self, name: str):
        """hacky workaround since these aren't named"""
        return f"(define-fun {name} (" + \
            " ".join([f"({p.name} {p.sort.to_sygus2()})" for p in self.params]) + ") " + \
            self.body.sort.to_sygus2() + " " + \
            self.body.to_sygus2() \
            + ")"

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            if self in cvc5_ctx.terms:
                return cvc5_ctx.terms[self]
            else:
                cvc5_kind = pycvc5.Kind.LAMBDA
                # TODO this needs to be tested
                vlist = cvc5_ctx.solver.mkTerm(pycvc5.Kind.VARIABLE_LIST, *[p.to_cvc5(cvc5_ctx) for p in self.params])
                t = cvc5_ctx.solver.mkTerm(cvc5_kind, vlist, self.body.to_cvc5(cvc5_ctx))
                cvc5_ctx.terms[self] = t
                return t
        elif tgt in (TargetFormat.SYGUS2, TargetFormat.UCLID):
            return "(define-fun (" + \
                " ".join([f"({p.name} {p.sort.to_sygus2()})" for p in self.params]) + ") " + \
                self.body.sort.to_sygus2() + " " + \
                self.body.to_sygus2() \
                + ")"
        # return f"(" + \
        #     ", ".join([f"{p.name} : {p.sort.to_uclid()}" for p in self.params]) + ") : " + \
        #     self.body.sort.to_uclid() + " = " + \
        #     self.body.to_uclid()
        raise NotImplementedError("cannot convert LambdaTerm to " + str(tgt))

@dataclass(frozen=True)
class QuantTerm(Term):
    kind: Kind
    bound_vars: Tuple[Variable, ...]
    body: Term

    @property
    def sort(self):
        return BoolSort()

    @property
    def _children(self):
        return self.body

    def _replace_child(self, new_term, index):
        if index != 0:
            raise IndexError()
        return QuantTerm(self.kind, self.bound_vars, new_term)

    @staticmethod
    def from_cvc5(cvc5_term):
        k = cvc5_term.getKind()
        if k not in [pycvc5.Kind.Exists, pycvc5.Kind.Forall]:
            c_params = cvc5_term[0]
            c_body = cvc5_term[1]
            return QuantTerm(
                Kind.from_cvc5(k),
                tuple([Variable.from_cvc5(c) for c in c_params]),
                Term.from_cvc5(c_body),
            )
        else:
            raise TypeError("QuantTerm must be translated from pycvc5.Kind.Exists or Forall, instead got " + str(k))

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            if self in cvc5_ctx.terms:
                return cvc5_ctx.terms[self]
            else:
                cvc5_kind = self.kind.to_cvc5(cvc5_ctx)
                # TODO this needs to be tested
                vlist = cvc5_ctx.solver.mkTerm(pycvc5.Kind.VARIABLE_LIST, [p.to_cvc5(cvc5_ctx) for p in self.bound_vars])
                t = cvc5_ctx.solver.mkTerm(cvc5_kind, vlist, [self.body.to_cvc5(cvc5_ctx)])
                cvc5_ctx.terms[self] = t
                return t
        elif tgt in (TargetFormat.SYGUS2, TargetFormat.UCLID):
            return "(" + self.kind.to_sygus2() + " (" + \
                " ".join([f"({p.name} {p.sort.to_sygus2()})" for p in self.bound_vars]) + ") " + \
                self.body.sort.to_sygus2() + " " + \
                self.body.to_sygus2() \
                + ")"
        raise NotImplementedError("cannot convert QuantTerm to " + str(tgt))

@dataclass(frozen=True)
class ApplyUF(Term):
    """
    Term representing application of an uninterpreted function on the specified inputs.
    """
    fun: UFTerm
    # without making cvc5 unhappy
    input_values: Tuple[Term, ...]

    # def __post_init__(self):
    #     assert self.fun.arity == len(self.input_values)

    def __str__(self):
        return self.to_uclid()

    @property
    def sort(self):
        return self.fun.sort

    @property
    def _children(self):
        return list(self.input_values)

    def _replace_child(self, new_term, index):
        return ApplyUF(self.fun, tuple(t._replace_child(new_term, index) for t in self.input_values))

    @staticmethod
    def from_cvc5(cvc5_term):
        if cvc5_term.getKind() == pycvc5.Kind.APPLY_UF:
            raise NotImplementedError()
        else:
            raise TypeError("ApplyUF must be translated from pycvc5.Kind.APPLY_UF, instead got " + str(cvc5_term.getKind()))

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            cvc5_kind = pycvc5.Kind.APPLY_UF
            # Distinguish between ordinary UFs and synthfuns
            if self.fun.name in cvc5_ctx.synthfuns:
                cvc5_fun = cvc5_ctx.synthfuns[self.fun.name]
            else:
                cvc5_fun = self.fun.to_cvc5(cvc5_ctx)
            t = cvc5_ctx.solver.mkTerm(cvc5_kind, cvc5_fun, *[v.to_cvc5(cvc5_ctx) for v in self.input_values])
            return t
        elif tgt == TargetFormat.SYGUS2:
            return f"({self.fun.name} " + " ".join(a.to_sygus2() for a in self.input_values) + ")"
        elif tgt == TargetFormat.UCLID:
            args = ", ".join(i.to_uclid(**kwargs) for i in self.input_values)
            return self.fun.name + "(" + args + ")"
        raise NotImplementedError("cannot convert ApplyUF to " + str(tgt))

class BoolConst(Term, IntEnum, metaclass=_TermMeta):
    """SMT boolean constants for true and false. Can be coerced to/from an `int`."""
    F = 0
    T = 1

    @property
    def val(self) -> int:
        return int(self.value)

    @property
    def sort(self):
        return BoolSort()

    @property
    def _children(self):
        return []

    def _replace_child(self, new_term, index):
        return self

    def __str__(self):
        return "true" if self == self.T else "false"

    def eval(self, values):
        return self

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            if self == self.F:
                return cvc5_ctx.solver.mkFalse()
            else:
                return cvc5_ctx.solver.mkTrue()
        elif tgt == TargetFormat.VERILOG:
            return "1" if self == self.T else "0"
        elif tgt in (TargetFormat.SYGUS2, TargetFormat.UCLID):
            return "true" if self == self.T else "false"
        raise NotImplementedError("cannot convert BoolConst to " + str(tgt))

@dataclass(frozen=True)
class BVConst(Term):
    val: int
    width: int

    def __post_init__(self):
        assert isinstance(self.val, int), type(self.val)
        assert isinstance(self.width, int), type(self.width)

    @property
    def sort(self):
        return BVSort(self.width)

    @property
    def _children(self):
        return []

    def _replace_child(self, new_term, index):
        return self

    def __str__(self):
        return f"{self.val}bv{self.width}"

    @staticmethod
    def from_cvc5(cvc5_term):
        kind = cvc5_term.getKind()
        if kind != pycvc5.Kind.CONST_BITVECTOR:
            raise TypeError("BVConst must be translated from pycvc5.Kind.CONST_BITVECTOR, instead got " + str(kind))
        return BVConst(int(cvc5_term.getBitVectorValue(base=10)), cvc5_term.getSort().getBitVectorSize())

    def eval(self, values):
        return self

    def to_target_format(self, tgt: TargetFormat, **kwargs):
        if tgt == TargetFormat.CVC5:
            cvc5_ctx = kwargs["cvc5_ctx"]
            if BVSort(self.width) in cvc5_ctx.sorts:
                return cvc5_ctx.sorts[self]
            new_sort = cvc5_ctx.solver.mkBitVector(self.width, self.val)
            return new_sort
        elif tgt == TargetFormat.SYGUS2:
            # sygus literal lengths determine their type
            if self.width % 4 == 0:
                return f"#x{self.val:0{self.width // 4}x}"
            else:
                return f"#b{self.val:0{self.width}b}"
        elif tgt == TargetFormat.VERILOG:
            return "{}'h{:x}".format(self.width, self.val)
        elif tgt == TargetFormat.UCLID:
            return "0x{:x}bv{}".format(self.val, self.width)
        raise NotImplementedError("cannot convert BVConst to " + str(tgt))

# === END SMT TERMS ===


def add(a, b):
  c=a+b
  print(c)

add(4,65)
add(5,51)

def add(a,b):
  c=a+b
  return c

d= add(3,5)
print(d)

def add(a,b):
  c = a+b
  d = a-b
  return c, d #리턴을 여러개 가능, 튜플로 반환

print(add(5,4))

def isPrime(x):
  for i in range(2, x):
    if x%i==0:
      return False
  return True

a=[12,13,6,7,16]

for x in a:
  if isPrime(x):
    print(x, end=' ')


import numpy as np   # to handle matrix and data operation
import pandas as pd  # to read csv and handle dataframe
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator, FormatStrFormatter, AutoMinorLocator
from tqdm import tqdm
import copy
from mpl_toolkits.mplot3d import Axes3D
import random

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
from torch.utils.data import DataLoader
from torch.autograd import Variable

import pytorch_lightning as pl
from pytorch_lightning.core.decorators import auto_move_data

import torchvision.transforms as transforms
import torchvision.datasets

# from sklearn.model_selection import train_test_split
# from sklearn.cluster import KMeans
# from sklearn.decomposition import PCA
# import scipy.stats as sps
# import scipy.signal as spy_sig
# import scipy.ndimage as ndimage
# import scipy.optimize as spy_opt
# import random
# import os
# import pickle

random_seed = 1
torch.backends.cudnn.enabled = False

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

# torch_xla.core.xla_model.set_rng_state(1, device=device)
np.random.seed(0)
torch.manual_seed(0)

class Model(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = torchvision.models.resnet18(num_classes=10)

        self.model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        # self.model.conv1.conv = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)  ## googlenet
        # self.model.features[0] = nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(2, 2))

        self.loss = nn.CrossEntropyLoss()

    @auto_move_data
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_no):
        x, y = batch
        logits = self(x)
        loss = self.loss(logits, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.model.parameters(), lr=0.001)

def load_dataset(download=False, dataset='cifar'):
    global train_dataset_
    dataset = dataset.lower()
    batch_size = 100

    trans_og = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    # trans_og = transforms.Compose([transforms.Resize(96),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    AlexTransform = transforms.Compose([
        transforms.Resize((227, 227)),
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))])

    imagenet_trans = transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # trans =  transforms.Compose(
    #   [transforms.ToTensor(),
    #    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # trans = AlexTransform

    if dataset == 'cifar':
        trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        print("CIFAR10")
        train_dataset_ = torchvision.datasets.CIFAR10(root=r"Dataset/CIFAR10", train=True,
                                                      transform=trans, download=download)
        test_dataset_ = torchvision.datasets.CIFAR10(root=r"Dataset/CIFAR10", train=False, transform=trans)

    elif dataset == 'mnist':
        trans_og = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        print("MNIST")
        train_dataset_ = torchvision.datasets.MNIST(root=r"Dataset/", train=True,
                                                    transform=trans_og, download=download)
        test_dataset_ = torchvision.datasets.MNIST(root=r"Dataset/", train=False, transform=trans_og)

    elif dataset == 'imagenet':
        print("ImageNet")
        trans = imagenet_trans
        train_dataset = torchvision.datasets.ImageNet(root=r"Dataset/", split='train')
        test_dataset_ = torchvision.datasets.ImageNet(root=r"Dataset/", split='test', target_transform=trans)

    else:
        trans = transforms.Compose([transforms.Resize(96),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        train_dataset_ = torchvision.datasets.CIFAR10(root=r"Dataset/CIFAR10",
                                                      transform=trans, download=download)
        test_dataset_ = torchvision.datasets.CIFAR10(root=r"Dataset/CIFAR10", train=False, transform=trans)

    train_loader_ = DataLoader(dataset=train_dataset_, num_workers=2, batch_size=batch_size, shuffle=True)
    test_loader_ = DataLoader(dataset=test_dataset_, batch_size=100, shuffle=True)

    return train_loader_, test_loader_

def model_tester(model_, test_set):
    # Testing Accuracy
    model_.model.eval()
    model_.model.to(device)
    criterion = nn.CrossEntropyLoss()

    correct, total, running_loss = 0, 0, 0
    with torch.no_grad():
        for data in test_set:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model_.model(images)

            test_loss = criterion(outputs, labels).to(device)

            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            running_loss += test_loss.item() * images.size(0)

    accuracy = correct / total
    epoch_loss = running_loss / len(test_set)

    return accuracy, epoch_loss

def get_test_names(ex_model_og):
    ex_model = copy.deepcopy(ex_model_og)
    # params, names = count_parameters(ex_model.model)

    model_obj = copy.deepcopy(ex_model.model.state_dict())
    new_names = []

    for p_name, module in list(ex_model.model.named_modules())[1:]:
        nn_ = p_name + ".weight"
        if isinstance(module, nn.Conv2d):
            new_names.append(nn_)
        elif isinstance(module, nn.Linear):
            new_names.append(nn_)

    return new_names, model_obj

def test_harness(model_, test_set):
    mod_names, mod_obj = get_test_names(model_)
    res_dict = {}
    acc, loss = model_tester(model_, test_set)
    res_dict['OG'] = (acc, loss)
    print("OG", (acc, loss))
    print("\n")

    for name in tqdm(mod_names):
        model_clone = copy.deepcopy(model_)
        new_obj = copy.deepcopy(mod_obj)

        weight = copy.deepcopy(mod_obj[name])
        filtered_weight = torch.ones_like(weight)
        new_obj[name] = filtered_weight
        model_clone.model.load_state_dict(new_obj)
        acc, loss = model_tester(model_clone, test_set)
        res_dict[name] = (acc, loss)

        print(name, (acc, loss))
        print("\n")

        del filtered_weight

    return res_dict

train_loader, test_loader = load_dataset(download=False, dataset='mnist')
# WRAPPED_MODEL = xmp.MpModelWrapper(Model())
dataset_name = "MNIST"
e_no = 10
path = 'Saved_models/'
file_name = 'Init_3/ResNet18_' + dataset_name + "_" + str(e_no) + 'epochs.ckpt'
checkpoint = torch.load(path+file_name, map_location=torch.device("cpu"))

model = Model()
model.load_state_dict(checkpoint)
# red_data = test_harness(model, test_loader)
weight = model.model.layer4[1].conv1.weight.detach()

print(weight[0][0])
weight = weight/weight
print(weight[0][0])


import sys

infile = sys.argv[1] if len(sys.argv) > 1 else "day8.txt"
data = open(infile).read().strip()
lines = [x for x in data.split("\n")]

GRID = []
for line in lines:
    row = line
    GRID.append(row)

DIRECTIONS = [(-1, 0), (0, 1), (1, 0), (0, -1)]
ROWS = len(GRID)
COLUMNS = len(GRID[0])

part1 = 0
for row in range(ROWS):
    for column in range(COLUMNS):
        visible = False
        for (dr, dc) in DIRECTIONS:
            rr = row
            cc = column
            ok = True
            while True:
                rr += dr
                cc += dc
                if not (0<=rr<ROWS and 0<=cc<COLUMNS):
                    break
                if GRID[rr][cc] >= GRID[row][column]:
                    ok = False
            if ok:
                visible = True
        if visible:
            part1 += 1
print(part1)

part2 = 0
for row in range(ROWS):
    for column in range(COLUMNS):
        score = 1
        for (dr, dc) in DIRECTIONS:
            dist = 1
            rr = row + dr
            cc = column + dc
            while True:
                if not (0<=rr<ROWS and 0<=cc<COLUMNS):
                    dist -= 1
                    break
                if GRID[rr][cc] >= GRID[row][column]:
                    break
                dist += 1
                rr += dr
                cc += dc
            score *= dist
        part2 = max(part2, score)
print(part2)


#섬의 개수 세기
#2차원 그리드에서 1로 표시된 영역이 섬, 주어진 그리드에서 섬의 개수 찾기

class Solution:
    def num_islands(self, grid: list[list[str]]) -> int:
        count = 0 #섬 개수 초기화
        # 반복문 돌리기 for문이 두개!!! (grid[0][0] -> grid[3][4])
        for i in range(len(grid)):
            for j in range(len(grid[0])):
                # grid[i][j]이 1인 경우 ->
                if grid[i][j] == '1':
                    count += 1
                    stack = [(i, j)] #해당 셀의 좌표를 stack에 추가하기
                    while stack:
                        (y, x) = stack.pop()
                        direction = [(y - 1, x), (y + 1, x), (y, x - 1), (y, x + 1)]
                        for new_y, new_x in direction:
                            if 0 <= new_y < len(grid) and 0 <= new_x < len(grid[0]) and grid[new_y][new_x] == "1":
                                grid[new_y][new_x] = "0" #해당 지점 이미 방문
                                stack.append((new_y, new_x))
        return count

sol = Solution()
grid = [
    ["1","1","0","0","0"],
    ["1","1","0","0","0"],
    ["0","0","1","0","0"],
    ["0","0","0","1","1"]
]
result = sol.num_islands(grid)
print(result)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argrecord
import argrecord.argreplay
import shutil
import os
import sys

def copy(argstring):
    parser = argrecord.ArgumentRecorder(os.path.realpath(__file__))
    parser.add_argument('-m', '--make', action='store_true', private=True)
    parser.add_argument('-l', '--log',  type=str, dest='log_file')
    parser.add_argument('-a', '--append', action='store_true')
    parser.add_argument('-b', '--backup', type=str)
    parser.add_argument('--output_file',  type=str, output=True)
    parser.add_argument('input_file', nargs='?',  type=str, input=True)
    args = parser.parse_args(argstring)

    if not args.make or parser.replay_required(args):
        if args.log_file:
            parser.write_comments(args, args.log_file, append=args.append, backup=args.backup)
        infile = open(args.input_file, 'r') if args.input_file else sys.stdin
        outfile = open(args.output_file, 'w') if args.output_file else sys.stdout
        copy_buffer = infile.read(1024)
        while copy_buffer:
            outfile.write(copy_buffer)
            copy_buffer = infile.read(1024)
        if args.input_file:
            infile.close()
        if args.output_file:
            outfile.close()

        return True
    else:
        return False

def test_copy():
    # Create new test input file
    inf = open("in.txt", "w")
    inf.write("test\n")
    inf.close()
    if os.path.isfile('out.txt'):
        os.remove('out.txt')

    print("Test file copies first time")
    assert(copy(['--make', '--log', 'copy.log', '--output_file', 'out.txt', 'in.txt']))
    print("Test file does not copy second time")
    assert(not copy(['--make', '--log', 'copy.log', '--output_file', 'out.txt', 'in.txt']))

    print("Test replay works")
    os.environ['PATH'] = 'tests' + os.pathsep + os.environ['PATH']
    os.remove('out.txt')
    # Redirect stdin so pytest doesn't freak out
    sys.stdin = open(os.devnull, 'r')
    argrecord.argreplay.main(['copy.log'])
    assert(os.path.isfile('out.txt'))

    print("Test piping works")
    os.remove('out.txt')
    argrecord.argreplay.main(['test_pipe.log'])
    assert(os.path.isfile('out.txt'))

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


'''
bfs는 가장 가까운 물고기를 찾기 위한 수단
bfs를 만들어 놓고, 계속 반복해서 돌려야 한다.
'''

import sys
from collections import deque
import heapq
input = sys.stdin.readline

n = int(input())

board = [[]*n for _ in range(n)]
dsize = 2
eat = 0
second = 0

dr = [-1,0,1,0]
dc = [0,-1,0,1]

now = (-1,-1)
for i in range(n):
    board[i] = list(map(int, input().split()))
    for j in range(n):
        if not board[i][j]: continue
        if board[i][j]==9:
            now = (i,j)
            board[i][j] = 0

def bfs(row,col,size):
    # 가장 가까운 물꼬기를 찾기 위한 bfs
    visited = [[False]*n for _ in range(n)]
    togo = deque([(row,col,0)])
    fishes = []
    while(togo):
        nowr, nowc, dist = togo.popleft()
        for i in range(4):
            nr = nowr + dr[i]
            nc = nowc + dc[i]
            if 0<=nr<n and 0<=nc<n:
                if visited[nr][nc]: continue
                if board[nr][nc]<=size:
                    # 빈칸 방문
                    visited[nr][nc] = True
                    togo.append((nr,nc,dist+1))
                    if 0<board[nr][nc]<size:
                        # 먹을 수 있음
                        heapq.heappush(fishes, (dist+1, nr, nc))

    return fishes

while(True):
    nowr, nowc = now
    fishes = bfs(nowr, nowc, dsize)
    if not fishes:
        break
    # 물고기 먹음
    (dist, fr, fc) = heapq.heappop(fishes)
    second += dist
    board[fr][fc], board[nowr][nowc] = 0,0
    now = (fr, fc)
    eat+=1
    if eat==dsize:
        eat = 0
        dsize+=1

print(second)


import coalpy.gpu as g
import numpy as np
import math
import functools
from . import prefix_sum as gpu_prefix_sum

def prefix_sum(input_data, is_exclusive = False):
    accum = 0
    output = []
    for i in range(0, len(input_data), 1):
        if is_exclusive:
            output.append(accum)
            accum += input_data[i]
        else:
            accum += input_data[i]
            output.append(accum)
    return output

def test_cluster_gen(is_exclusive = False):
    buffersz = 8529
    input_data = np.array([x  for x in range(0, buffersz, 1)], dtype='i')
    test_input_buffer = g.Buffer(format = g.Format.R32_UINT, element_count = buffersz)

    reduction_buffers = gpu_prefix_sum.allocate_args(buffersz)

    cmd_list = g.CommandList()
    cmd_list.upload_resource(source = input_data, destination = test_input_buffer)
    output = gpu_prefix_sum.run(cmd_list, test_input_buffer, reduction_buffers, is_exclusive)

    g.schedule(cmd_list)

    dr = g.ResourceDownloadRequest(resource = output)
    dr.resolve()

    result = np.frombuffer(dr.data_as_bytearray(), dtype='i')
    result = np.resize(result, buffersz)
    expected = prefix_sum(input_data, is_exclusive)
    correct_count = functools.reduce(lambda x, y: x + y, [1 if x == y else 0 for (x, y) in zip(result, expected)])
    return True if correct_count == len(input_data) else False

def run_test(nm, fn):
    result = fn()
    print(nm + " : " + ("PASS" if result else "FAIL"))

def test_cluster_gen_inclusive():
    return test_cluster_gen(is_exclusive = False)

def test_cluster_gen_exclusive():
    return test_cluster_gen(is_exclusive = True)

if __name__ == "__main__":
    run_test("test prefix sum inclusive", test_cluster_gen_inclusive)
    run_test("test prefix sum exclusive", test_cluster_gen_exclusive)


import json
import boto3
from urllib.parse import unquote_plus

def lambda_handler(event, context):
    # parse the S3 triggered event
    record = event['Records'][0]

    # extract bucket from event
    bucket = record['s3']['bucket']['name']

    # extract key from event
    key = unquote_plus(record['s3']['object']['key'])

    ## for debugging
    # bucket = "fundmapper"
    # key = "04-StagingTables/series_data/202006/part-00000-39604d54-0f41-44be-8850-fafdc5733262-c000.snappy.parquet"

    # create s3 resource
    s3_resource = boto3.resource('s3')

    dest_prefix = "05-FinalTables"
    prefix, dataset, date, filename = key.split("/")
    print(bucket)
    # copy file from S3
    (s3_resource.Object(bucket, f"{dest_prefix}/{dataset}/{date}/{dataset}_{date}.parquet")
     .copy_from(CopySource=f"{bucket}/{key}"))

    # Delete the former file
    (s3_resource.Object(bucket, key)
     .delete())


def solution(park, routes):
    answer = []
    row = len(park)
    column = len(park[0])
    for idx,p in enumerate(park):
        if 'S' in p:
            start = [idx, p.index('S')]
    for route in routes:
        way, count = route.split()
        count = int(count)
        if way == 'E':
            if start[1] + count >= column:
                continue
            if 'X' in park[start[0]][start[1]: (start[1] + count+1)]:
                continue
            start[1] += count
        if way == 'W':
            if start[1] - count < 0:
                continue
            if 'X' in park[start[0]][start[1] - count:start[1] + 1]:
                continue
            start[1] -= count
        if way == 'N':
            if start[0] - count < 0:
                continue
            if 'X' in [park[x][start[1]] for x in range(start[0]-count, start[0]+1) ]:
                continue
            start[0] -= count
        if way == 'S':
            if start[0] + count >= row:
                continue
            if 'X' in [park[x][start[1]] for x in range(start[0], start[0]+count+1) ]:
                continue
            start[0] += count
    answer = start
    return answer


from collections import defaultdict
from enum import Enum
import logging
import sys
from typing import Dict, Optional, Sequence

from gym import Env, spaces

import numpy as np

from pyhtzee import Pyhtzee
from pyhtzee.classes import Category, PyhtzeeException, Rule

log = logging.getLogger(__name__)

class GameType(Enum):
    SUDDEN_DEATH = 0,
    RETRY_ON_WRONG_ACTION = 1

def get_score(score: Optional[int]) -> int:
    return score if score is not None else -1

def get_dice_face_counts(dice: Sequence[int]) -> Dict[int, int]:
    faces: Dict[int, int] = defaultdict(int)
    for die in dice:
        faces[die] += 1
    return faces

class YahtzeeSingleEnv(Env):
    metadata = {'render.modes': ['human']}

    def __init__(self,
                 rule: Rule = Rule.YAHTZEE_FREE_CHOICE_JOKER,
                 game_type: GameType = GameType.RETRY_ON_WRONG_ACTION,
                 seed=None):
        self.pyhtzee = Pyhtzee(seed=seed)
        self.rule = rule
        self.game_type = game_type
        self.action_space = spaces.Discrete(44)
        self.observation_space = spaces.Tuple((
            spaces.Discrete(13),  # round
            spaces.Discrete(4),  # sub-round
            spaces.Box(low=1, high=6, shape=(1,), dtype=np.uint8),  # die 1
            spaces.Box(low=1, high=6, shape=(1,), dtype=np.uint8),  # die 2
            spaces.Box(low=1, high=6, shape=(1,), dtype=np.uint8),  # die 3
            spaces.Box(low=1, high=6, shape=(1,), dtype=np.uint8),  # die 4
            spaces.Box(low=1, high=6, shape=(1,), dtype=np.uint8),  # die 5
            spaces.Box(low=-1, high=5, shape=(1,), dtype=np.int16),  # aces
            spaces.Box(low=-1, high=10, shape=(1,), dtype=np.int16),  # twos
            spaces.Box(low=-1, high=15, shape=(1,), dtype=np.int16),  # threes
            spaces.Box(low=-1, high=20, shape=(1,), dtype=np.int16),  # fours
            spaces.Box(low=-1, high=25, shape=(1,), dtype=np.int16),  # fives
            spaces.Box(low=-1, high=30, shape=(1,), dtype=np.int16),  # sixes
            spaces.Box(low=-1, high=30, shape=(1,), dtype=np.int16),  # three of a kind
            spaces.Box(low=-1, high=30, shape=(1,), dtype=np.int16),  # four of a kind
            spaces.Box(low=-1, high=25, shape=(1,), dtype=np.int16),  # full house
            spaces.Box(low=-1, high=30, shape=(1,), dtype=np.int16),  # small straight
            spaces.Box(low=-1, high=40, shape=(1,), dtype=np.int16),  # large straight
            spaces.Box(low=-1, high=30, shape=(1,), dtype=np.int16),  # chance
            spaces.Box(low=-1, high=50, shape=(1,), dtype=np.int16),  # yahtzee
            spaces.Box(low=-1, high=35, shape=(1,), dtype=np.int16),  # upper bonus
            spaces.Box(low=-1, high=1200, shape=(1,), dtype=np.int16),  # yahtzee bonus
        ))

    def get_observation_space(self):
        pyhtzee = self.pyhtzee
        return (
            pyhtzee.round,
            pyhtzee.sub_round,
            pyhtzee.dice[0],
            pyhtzee.dice[1],
            pyhtzee.dice[2],
            pyhtzee.dice[3],
            pyhtzee.dice[4],
            get_score(pyhtzee.scores.get(Category.ACES)),
            get_score(pyhtzee.scores.get(Category.TWOS)),
            get_score(pyhtzee.scores.get(Category.THREES)),
            get_score(pyhtzee.scores.get(Category.FOURS)),
            get_score(pyhtzee.scores.get(Category.FIVES)),
            get_score(pyhtzee.scores.get(Category.SIXES)),
            get_score(pyhtzee.scores.get(Category.THREE_OF_A_KIND)),
            get_score(pyhtzee.scores.get(Category.FOUR_OF_A_KIND)),
            get_score(pyhtzee.scores.get(Category.FULL_HOUSE)),
            get_score(pyhtzee.scores.get(Category.SMALL_STRAIGHT)),
            get_score(pyhtzee.scores.get(Category.LARGE_STRAIGHT)),
            get_score(pyhtzee.scores.get(Category.CHANCE)),
            get_score(pyhtzee.scores.get(Category.YAHTZEE)),
            get_score(pyhtzee.scores.get(Category.UPPER_SECTION_BONUS)),
            get_score(pyhtzee.scores.get(Category.YAHTZEE_BONUS)),
        )

    def sample_action(self):
        action = self.pyhtzee.sample_action()
        log.info(f'Sampled action: {action}')
        return action

    def step(self, action: int):
        pyhtzee = self.pyhtzee
        try:
            reward = pyhtzee.take_action(action)
            finished = pyhtzee.is_finished()
            valid_move = True
        except PyhtzeeException:
            valid_move = False
            reward = 0
            if self.game_type == GameType.SUDDEN_DEATH:
                log.info('Invalid action, terminating round.')
                finished = True
            else:  # retry on wrong action
                log.info('Invalid action, step ignored.')
                finished = False

        log.info(f'Finished step. Reward: {reward}, Finished: {finished}')
        debug_info = {
            'valid_move': valid_move,
        }
        return self.get_observation_space(), reward, finished, debug_info

    def reset(self):
        self.pyhtzee = Pyhtzee()

    def render(self, mode='human', close=False):
        dice = self.pyhtzee.dice
        outfile = sys.stdout
        outfile.write(f'Dice: {dice[0]} {dice[1]} {dice[2]} {dice[3]} {dice[4]} '
                      f'Round: {self.pyhtzee.round}.{self.pyhtzee.sub_round} '
                      f'Score: {self.pyhtzee.get_total_score()}\n')


from django.core.urlresolvers import reverse
from django.utils.html import format_html

from wagtail.wagtailadmin.search import SearchArea
from wagtail.wagtailcore import hooks
from wagtail.wagtailsnippets.permissions import user_can_edit_snippet_type

from wagtail.contrib.modeladmin.options import (
    ModelAdmin,
    ThumbnailMixin,
    modeladmin_register
)

from people.models import Person

class PeopleSearchArea(SearchArea):
    def is_shown(self, request):
        return user_can_edit_snippet_type(request.user, Person)

@hooks.register('register_admin_search_area')
def register_people_search_area():
    return PeopleSearchArea(
        'Team Members',
        reverse('people_person_modeladmin_index'),
        classnames='icon icon-group',
        order=150
    )

class PersonThumbnailMixin(ThumbnailMixin):
    thumb_image_field_name = 'photo'

class PeopleAdmin(ModelAdmin):
    model = Person
    menu_icon = 'group'
    menu_label = 'Team'
    menu_order = 300
    list_display = ('profile_photo', 'full_name', 'role_and_production')
    list_filter = ('role', 'production')
    list_per_page = 50
    search_fields = [
        'first_name',
        'last_name',
        'role',
        'biography',
        'production__title',
    ]

    def profile_photo(self, obj): # pylint: disable=no-self-use
        return PersonThumbnailMixin().admin_thumb(obj)
    profile_photo.short_description = "photo"

    def role_and_production(self, obj): # pylint: disable=no-self-use
        if obj.production and obj.role:
            return format_html(
                '{} at {}',
                obj.role,
                reverse('wagtailadmin_pages:edit', args=(obj.production.id,)),
                obj.production
            )
        elif obj.production:
            return format_html(
                '{}',
                reverse('wagtailadmin_pages:edit', args=(obj.production.id,)),
                obj.production
            )
        elif obj.role:
            return obj.role
        else:
            return ""
    role_and_production.admin_order_field = '-production'

    def full_name(self, obj): # pylint: disable=no-self-use
        return "%s %s" % (
            obj.first_name,
            obj.last_name
        )
    full_name.short_description = 'name'
    full_name.admin_order_field = 'last_name'

modeladmin_register(PeopleAdmin)


import os

from core.service import CoreService, addservice
from core.misc.ipaddr import IPv4Prefix, IPv6Prefix

class MesherService(CoreService):
    # a unique name is required, without spaces
    _name = "MesherService"
    # you can create your own group here
    _group = "Mesh"
    # list of other services this service depends on
    _depends = ()
    # per-node directories
    _dirs = ()
    # generated files (without a full path this file goes in the node's dir,
    #  e.g. /tmp/pycore.12345/n1.conf/)
    _configs = ('mesher-start.sh', 'mesher-stop.sh', )
    # this controls the starting order vs other enabled services
    _startindex = 1
    # list of startup commands, also may be generated during startup
    _startup = ('bash mesher-start.sh', )
    # list of shutdown commands
    _shutdown = ('bash mesher-stop.sh', )

    @classmethod
    def generateconfig(cls, node, filename, services):
        if filename == "mesher-start.sh":
            cfg = '''#!/bin/sh
sleep 1
sched=/tmp/scheduler.js
logfile={}/mesher-{}.log
if [ -e $sched ]; then
    nohup mesher-experimental $sched &> $logfile &
else
    nohup mesher-experimental &> $logfile &
fi
echo $! > mesher.pid
'''.format(node.session.sessiondir, node.name)

        if filename == "mesher-stop.sh":
            cfg = '''#!/bin/sh
kill `cat mesher.pid`
rm mesher.pid
'''

        return cfg

    @staticmethod
    def subnetentry(x):
        ''' Generate a subnet declaration block given an IPv4 prefix string
            for inclusion in the config file.
        '''
        if x.find(":") >= 0:
            # this is an IPv6 address
            return ""
        else:
            net = IPv4Prefix(x)
            return 'echo "  network %s"' % (net)

# this line is required to add the above class to the list of available services
addservice(MesherService)


''' Test the module :py:mod:`bp.blocks`.
'''
import datetime
import unittest
import asn1
from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa, ec
from bp.encoding.fields import (EidField)
from bp.encoding.blocks import PrimaryBlock, CanonicalBlock
from bp.util import BundleContainer
from bp.encoding.bpsec import (TargetResultList, TypeValuePair,
                               BlockIntegrityBlock)
from bp.config import Config
from bp.agent import Agent
from bp.app.bpsec import Bpsec, BPSEC_COSE_CONTEXT_ID

class TestBpsecCoseSign(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        super().setUpClass()

        from dbus.mainloop.glib import DBusGMainLoop
        # Must run before connection or real main loop is constructed
        DBusGMainLoop(set_as_default=True)

    def setUp(self):
        super().setUp()

        config = Config()
        config.node_id = 'dtn://node/'
        self._bp = Agent(config)
        self._app = self._bp._app['bpsec']

    def _dummy_ca_cert(self, ca_key):
        ca_name = x509.Name([
            x509.NameAttribute(x509.oid.NameOID.COMMON_NAME, 'CA'),
        ])
        nowtime = datetime.datetime.now(datetime.timezone.utc)
        cert = x509.CertificateBuilder().subject_name(
            ca_name
        ).issuer_name(
            ca_name
        ).public_key(
            ca_key.public_key()
        ).serial_number(
            x509.random_serial_number()
        ).not_valid_before(
            nowtime
        ).not_valid_after(
            nowtime + datetime.timedelta(days=1)
        ).add_extension(
            x509.BasicConstraints(ca=True, path_length=1),
            critical=True,
        ).add_extension(
            x509.KeyUsage(
                digital_signature=False,
                content_commitment=False,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=True,
                crl_sign=True,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=False,
        ).add_extension(
            x509.SubjectKeyIdentifier.from_public_key(ca_key.public_key()),
            critical=False,
        ).add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()),
            critical=False,
        ).sign(ca_key, hashes.SHA256(), backend=default_backend())
        return cert

    def _dummy_end_cert(self, ca_key, ca_cert, end_key):
        nowtime = datetime.datetime.now(datetime.timezone.utc)
        eid_enc = asn1.Encoder()
        eid_enc.start()
        eid_enc.write(self._app._config.node_id.encode('ascii'), asn1.Numbers.IA5String)
        sans = [
            x509.OtherName(
                x509.oid.ObjectIdentifier('1.3.6.1.5.5.7.8.11'),  # id-on-bundleEID
                eid_enc.output()
            )
        ]
        cert = x509.CertificateBuilder().subject_name(
            x509.Name([
                x509.NameAttribute(x509.oid.NameOID.COMMON_NAME, 'end-entity'),
            ])
        ).issuer_name(
            ca_cert.issuer
        ).public_key(
            end_key.public_key()
        ).serial_number(
            x509.random_serial_number()
        ).not_valid_before(
            nowtime
        ).not_valid_after(
            nowtime + datetime.timedelta(days=1)
        ).add_extension(
            x509.BasicConstraints(ca=False, path_length=None),
            critical=True,
        ).add_extension(
            x509.SubjectAlternativeName(sans),
            critical=False,
        ).add_extension(
            x509.KeyUsage(
                digital_signature=True,
                content_commitment=False,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=False,
                crl_sign=False,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=False,
        ).add_extension(
            x509.ExtendedKeyUsage([
                x509.oid.ObjectIdentifier('1.3.6.1.5.5.7.3.35')  # id-kp-bundleSecurity
            ]),
            critical=False,
        ).add_extension(
            x509.SubjectKeyIdentifier.from_public_key(end_key.public_key()),
            critical=False,
        ).add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()),
            critical=False,
        ).sign(ca_key, hashes.SHA256(), backend=default_backend())
        return cert

    def _dummy_ctr(self):
        ''' Generate a dummy bundle container.
        '''
        ctr = BundleContainer()
        ctr.bundle.primary = PrimaryBlock(
            destination="dtn://other/svc",
        )
        ctr.add_block(CanonicalBlock(
            type_code=1,
            btsd=b'hi',
        ))
        self._bp._apply_primary(ctr)
        return ctr

    def test_apply_bib_ec2(self):
        ca_key = ec.generate_private_key(ec.SECP256R1, backend=default_backend())
        ca_cert = self._dummy_ca_cert(ca_key)
        end_key = ec.generate_private_key(ec.SECP256R1, backend=default_backend())  # Curve for COSE ES256
        end_cert = self._dummy_end_cert(ca_key, ca_cert, end_key)

        ctx = self._app._contexts[BPSEC_COSE_CONTEXT_ID]
        ctx._ca_certs = [ca_cert]
        ctx._cert_chain = [end_cert]
        ctx._priv_key = end_key

        ctr = self._dummy_ctr()
        self._app._apply_bib(ctr)

        bibs = ctr.block_type(BlockIntegrityBlock)
        self.assertNotEqual([], bibs)
        bib = bibs[0]
        self.assertEqual([1], bib.payload.targets)

        ctr.record_action('deliver')
        self.assertFalse(self._app._verify_bib(ctr))
        self.assertEqual({'deliver'}, ctr.actions.keys())

    def test_apply_bib_rsa(self):
        ca_key = rsa.generate_private_key(0x10001, 1024, backend=default_backend())
        ca_cert = self._dummy_ca_cert(ca_key)
        end_key = rsa.generate_private_key(0x10001, 1024, backend=default_backend())
        end_cert = self._dummy_end_cert(ca_key, ca_cert, end_key)

        ctx = self._app._contexts[BPSEC_COSE_CONTEXT_ID]
        ctx._ca_certs = [ca_cert]
        ctx._cert_chain = [end_cert]
        ctx._priv_key = end_key

        ctr = self._dummy_ctr()
        self._app._apply_bib(ctr)

        bibs = ctr.block_type(BlockIntegrityBlock)
        self.assertNotEqual([], bibs)
        bib = bibs[0]
        self.assertEqual([1], bib.payload.targets)

        ctr.record_action('deliver')
        self.assertFalse(self._app._verify_bib(ctr))
        self.assertEqual({'deliver'}, ctr.actions.keys())


from typing import List

"""
https://leetcode.cn/problems/max-dot-product-of-two-subsequences/

请你返回 nums1 和 nums2 中两个长度相同的 非空 子序列的最大点积。

数组的非空子序列是通过删除原数组中某些元素（可能一个也不删除）后剩余数字组成的序列，但不能改变数字间相对顺序。
比方说，[2,3,5] 是 [1,2,3,4,5] 的一个子序列而 [1,5,3] 不是。

"""
def maxDotProduct(nums1: List[int], nums2: List[int]) -> int:
    """
    f[i][j] 表示只考虑数组nums1的前i个元素以及数组nums2的前j个元素时,
    可以得到的两个长度相同的非空子序列的最大点积。

    可以考虑每个数组中的最后一个元素:
    1. 选择成为点积。
    2. 不选择成为点积。
    """
    m, n = len(nums1), len(nums2)
    f = [[0] * n for _ in range(m)]

    for i in range(m):
        for j in range(n):
            xij = nums1[i] * nums2[j]
            f[i][j] = xij
            if i > 0:
                f[i][j] = max(f[i][j], f[i - 1][j])
            if j > 0:
                f[i][j] = max(f[i][j], f[i][j - 1])
            if i > 0 and j > 0:
                f[i][j] = max(f[i][j], f[i - 1][j - 1] + xij)

    return f[m - 1][n - 1]

动态规划/两个子序列的最大点积.py

两个子序列的最大点积.py


'''
    Postprocess teensy ouput
    ~~~~~~~~~~~~~~~~~~~~~~~~
'''

import argparse
import csv
import time
import os
import sys

def create_filename (Kp, Ki, Kd, setpoint):
    parts = [
        'teensy-output',
        'Kp={}'.format(Kp),
        'Ki={}'.format(Ki),
        'Kd={}'.format(Kd),
        time.strftime('%Y-%m-%d_%H,%M')
    ]
    return '{}.log'.format('_'.join(parts))

def log_lines (log):
    for line in log:
        line = line.strip()
        if line == '' or line[0] == ':':
            continue
        yield [float(_) for _ in line.split('\t')]

def main ():
    parser = argparse.ArgumentParser()
    parser.add_argument('logfile', help='Log file from the teensy.')
    parser.add_argument('kp', type=float)
    parser.add_argument('ki', default=0, nargs='?', type=float)
    parser.add_argument('kd', default=0, nargs='?', type=float)
    parser.add_argument('setpoint', default=100, nargs='?', type=float)
    parser.add_argument('cycle', default=500, nargs='?', type=int)

    args = parser.parse_args()

    if not os.path.isfile(args.logfile):
        print('{}: no such file'.format(args.logfile))
        return -1

    Kp = args.kp
    Ki = args.ki
    Kd = args.kd
    setpoint = args.setpoint

    filename = create_filename(Kp, Ki, Kd, setpoint)

    if os.path.isfile(filename):
        print('Output file {} already exists.'.format(filename))
        return -2

    with open(args.logfile) as log, open(filename, 'w') as out:
        vals = log_lines(log)
        writer = csv.writer(out, delimiter='\t')

        writer.writerow([
            time.strftime('%Y-%m-%d %H:%M:%S'),
            'Kp={}'.format(Kp),
            'Ki={}'.format(Ki),
            'Kd={}'.format(Kd),
            'setpoint={}'.format(setpoint),
            'cycle={}'.format(args.cycle),
        ])
        writer.writerow(['On','T','Output', 'dt', 'cycle', 't'])
        writer.writerows(vals)

if __name__ == '__main__':
    sys.exit(main())