""" main api point """
import logging
from typing import List

from . import dal, mappers
from .model import PriceModel, SecuritySymbol

class PriceDbApplication:
    """ Contains the main public interfaces """

    def __init__(self, session=None):
        self.logger = logging.getLogger(__name__)
        self.price_repo = None
        self.security_repo = None
        self.__session = session

    def add_price(self, price: PriceModel):
        """ Creates a new price record """
        # assert isinstance(price, PriceModel)

        if not price:
            raise ValueError("Cannot add price. The received model is null!")

        mapper = mappers.PriceMapper()
        entity = mapper.map_model(price)

        self.add_price_entity(entity)

    def add_price_entity(self, price: dal.Price):
        """ Adds the price """
        from decimal import Decimal

        # check if the price already exists in db.
        repo = self.get_price_repository()
        existing = (
            repo.query
            .filter(dal.Price.namespace == price.namespace)
            .filter(dal.Price.symbol == price.symbol)
            .filter(dal.Price.date == price.date)
            .filter(dal.Price.time == price.time)
            .first()
        )
        if existing:
            # Update existing price.
            new_value = Decimal(price.value) / Decimal(price.denom)
            self.logger.info(f"Exists: {price}")
            if price.currency != existing.currency:
                raise ValueError(
                    f"The currency is different for price {price}!")
            if existing.value != price.value:
                existing.value = price.value
                self.logger.info(f"Updating to {new_value}.")
            if existing.denom != price.denom:
                existing.denom = price.denom
        else:
            # Insert new price
            self.session.add(price)
            self.logger.info(f"Added {price}")

    def download_price(self, symbol: str, currency: str, agent: str) -> PriceModel:
        """ Download and save price online """
        price = self.__download_price(symbol, currency, agent)
        self.save()
        return price

    def download_prices(self, **kwargs):
        """ Downloads all the prices that are listed in the Security table.
        Accepts filter arguments: currency, agent, symbol, namespace.
        """
        currency: str = kwargs.get('currency', None)
        if currency:
            currency = currency.upper()
        agent: str = kwargs.get('agent', None)
        if agent:
            agent = agent.upper()
        symbol: str = kwargs.get('symbol', None)
        if symbol:
            symbol = symbol.upper()
        namespace: str = kwargs.get('namespace', None)
        if namespace:
            namespace = namespace.upper()
        securities = self.__get_securities(currency, agent, symbol, namespace)
        #self.logger.debug(securities)

        for sec in securities:
            symbol = f"{sec.namespace}:{sec.symbol}"
            currency = sec.currency
            agent = sec.updater
            #self.logger.debug(f"Initiating download for {symbol} {currency} with {agent}...")
            try:
                self.__download_price(symbol.strip(), currency, agent)
            except Exception as e:
                self.logger.error(str(e))
        self.save()

    def import_prices(self, file_path: str, currency_symbol: str):
        """ Incomplete """
        from .csv import CsvParser

        assert isinstance(file_path, str)
        assert isinstance(currency_symbol, str)

        self.logger.debug(f"Importing {file_path}")
        parser = CsvParser()
        prices = parser.parse_file(file_path, currency_symbol)

        counter = 0
        session = self.session
        # Create insert statements
        mapper = mappers.PriceMapper()
        for price in prices:
            new_price = mapper.map_model(price)
            self.add_price_entity(new_price)
            counter += 1
        # Save all to database
        session.commit()
        print(f"{counter} records inserted.")

    def get_latest_price(self, symbol: SecuritySymbol) -> PriceModel:
        """ Returns the latest price for the given symbol """
        from .repositories import PriceRepository

        # TODO should include the currency? Need a public model for exposing the result.
        assert isinstance(symbol, SecuritySymbol)

        session = self.session
        repo = PriceRepository(session)
        query = (
            repo.query
            .filter(dal.Price.symbol == symbol.mnemonic)
            .order_by(dal.Price.date.desc(), dal.Price.time.desc())
        )
        if symbol.namespace:
            query = query.filter(dal.Price.namespace == symbol.namespace)

        latest = query.first()

        # map
        mapper = mappers.PriceMapper()
        result = mapper.map_entity(latest)

        return result

    @property
    def session(self):
        """ Returns the current db session """
        if not self.__session:
            self.__session = dal.get_default_session()
        return self.__session

    def get_prices(self, date: str, currency: str) -> List[PriceModel]:
        """ Fetches all the prices for the given arguments """
        from .repositories import PriceRepository

        session = self.session
        repo = PriceRepository(session)
        query = repo.query
        if date:
            query = query.filter(dal.Price.date == date)
        if currency:
            query = query.filter(dal.Price.currency == currency)
        # Sort by symbol.
        query = query.order_by(dal.Price.namespace, dal.Price.symbol)
        price_entities = query.all()

        mapper = mappers.PriceMapper()
        result = []
        for entity in price_entities:
            model = mapper.map_entity(entity)
            result.append(model)
        return result

    def get_prices_on(self, on_date: str, namespace: str, symbol: str):
        """ Returns the latest price on the date """
        repo = self.get_price_repository()
        query = (
            repo.query.filter(dal.Price.namespace == namespace)
            .filter(dal.Price.symbol == symbol)
            .filter(dal.Price.date == on_date)
            .order_by(dal.Price.time.desc())
        )
        result = query.first()
        # logging.debug(result)
        return result

    def get_latest_prices(self):
        """ Fetches the latest prices for all symbols """
        # get all symbols first, for which we have prices available
        from .repositories import PriceRepository

        repo = PriceRepository(self.__session)

        query = (
            repo.session.query(dal.Price.namespace, dal.Price.symbol)
            .order_by(dal.Price.namespace, dal.Price.symbol)
            .distinct()
        )
        all_symbols = query.all()
        # self.logger.debug(all_symbols)

        # Get the latest prices for these symbols
        latest_prices = []
        for symbol_price in all_symbols:
            symbol = SecuritySymbol(symbol_price.namespace, symbol_price.symbol)
            latest = self.get_latest_price(symbol)
            latest_prices.append(latest)
        return latest_prices

    def get_price_repository(self):
        """ Price repository """
        from .repositories import PriceRepository

        if not self.price_repo:
            self.price_repo = PriceRepository(self.session)
        return self.price_repo

    def get_security_repository(self):
        """ Security repository """
        from .repositories import SecurityRepository

        if not self.security_repo:
            self.security_repo = SecurityRepository(self.session)
        return self.security_repo

    def prune_all(self) -> int:
        """
        Prune historical prices for all symbols, leaving only the latest.
        Returns the number of items removed.
        """
        from .repositories import PriceRepository

        # get all symbols that have prices
        repo = PriceRepository()
        items = repo.query.distinct(dal.Price.namespace, dal.Price.symbol).all()
        # self.logger.debug(items)
        count = 0

        for item in items:
            symbol = SecuritySymbol(item.namespace, item.symbol)
            deleted = self.prune(symbol)
            if deleted:
                count += 1

        return count

    def prune(self, symbol: SecuritySymbol):
        """
        Delete all but the latest available price for the given symbol.
        Returns the number of items removed.
        """
        from .repositories import PriceRepository

        assert isinstance(symbol, SecuritySymbol)

        self.logger.debug(f"pruning prices for {symbol}")

        repo = PriceRepository()
        query = (
            repo.query.filter(dal.Price.namespace == symbol.namespace)
            .filter(dal.Price.symbol == symbol.mnemonic)
            .order_by(dal.Price.date.desc())
            .order_by(dal.Price.time.desc())
        )
        all_prices = query.all()
        # self.logger.debug(f"fetched {all_prices}")

        deleted = False
        first = True
        for single in all_prices:
            if not first:
                repo.query.filter(dal.Price.id == single.id).delete()
                deleted = True
                self.logger.debug(f"deleting {single.id}")
            else:
                first = False

        repo.save()

        return deleted

    def save(self):
        """ Save changes """
        if self.__session:
            self.session.commit()
        else:
            self.logger.warning("Save called but no session open.")

    def __download_price(self, symbol: str, currency: str, agent: str):
        """ Downloads and parses the price """
        from finance_quote_python import Quote

        assert isinstance(symbol, str)
        assert isinstance(currency, str)
        assert isinstance(agent, str)

        if not symbol:
            return None

        #self.logger.info(f"Downloading {symbol}... ")

        dl = Quote()
        dl.logger = self.logger

        dl.set_source(agent)
        dl.set_currency(currency)

        result = dl.fetch(agent, [symbol])

        if not result:
            raise ValueError(f"Did not receive a response for {symbol}.")

        price = result[0]

        if not price:
            raise ValueError(f"Price not downloaded/parsed for {symbol}.")
        else:
            # Create price data entity, to be inserted.
            self.add_price(price)

        return price

    def __get_securities(self, currency: str, agent: str, symbol: str,
                         namespace: str) -> List[dal.Security]:
        """ Fetches the securities that match the given filters """
        repo = self.get_security_repository()
        query = repo.query

        if currency is not None:
            query = query.filter(dal.Security.currency == currency)

        if agent is not None:
            query = query.filter(dal.Security.updater == agent)

        if symbol is not None:
            query = query.filter(dal.Security.symbol == symbol)

        if namespace is not None:
            query = query.filter(dal.Security.namespace == namespace)

        # Sorting
        query = query.order_by(dal.Security.namespace, dal.Security.symbol)

        securities = query.all()
        return securities


import pprint as pp
import itertools
import inspect
from functools import partial
from collections import OrderedDict
from pythonwhat import test_funcs
from pythonwhat.State import State
from pythonwhat.checks.check_wrappers import state_partial

TEST_NAMES = [
    "test_mc",
    "test_or",
    "test_with",
    "test_import",
    "test_object",
    "test_correct",
    "test_if_else",
    "test_for_loop",
    "test_function",
    "test_list_comp",
    "test_function_v2",
    "test_data_frame",
    "test_while_loop",
    "test_student_typed",
    "test_object_accessed",
    "test_output_contains",
    "test_expression_result",
    "test_expression_output",
    "test_function_definition",
    "test_object_after_expression",
]

SUB_TESTS = {
    "test_if_else": ["test", "body", "orelse"],
    "test_list_comp": ["comp_iter", "body", "ifs"],
    "check_correct": ["check", "diagnose"],
    "test_for_loop": ["for_iter", "body", "orelse"],
    "test_while_loop": ["test", "body", "orelse"],
    "test_with": ["context_tests", "body"],
    "test_function_definition": ["body"],
    "check_or": ["tests"],
}

class Tree(object):
    def __init__(self):
        """
        Represent a tree of nodes, by holding the currently active node.

        This class is necessary to put sub-tests onto a graph, because they may
        exist inside of function calls. By getting the currently active node
        from the tree, Probe instances may add subtests as children on that node,
        and update it when recursing over sub-tests.

        """
        self.root = Node(name="root")
        self.crnt_node = self.root

    @classmethod
    def str_branch(cls, node, str_func=lambda s: str(s)):
        # dict(getattr(s.data.get("bound_args", {}), 'arguments', {}))
        f = node.data.get("func")
        this_node = (
            "  " * node.depth + "("
            + getattr(f, "__name__", node.name)
            + str_func(node)
            + ")\n"
        )
        return this_node + "".join(
            map(lambda x: cls.str_branch(x, str_func), node.child_list)
        )

    def __str__(self):
        return self.str_branch(self.crnt_node)

    def descend(self, node=None):
        node = self.crnt_node if node is None else node
        children = map(self.descend, node.child_list)
        base = [node] if node.name != "root" else []
        return sum(children, base)

    def __iter__(self):
        for ii in self.descend(self.crnt_node):
            yield ii

class Node(object):
    def __init__(self, child_list=None, data=None, name="unnamed", arg_name=""):
        """
        Hold a function call with its bound arguments, along with child nodes.

        """
        self.parent = None
        self.name = name
        self.arg_name = arg_name
        self.child_list = [] if child_list is None else child_list
        self.data = {} if data is None else data
        self.updated = False
        # hacky way to add their argument name when a function of tests was given

    def __call__(self, state=None):
        """Call original function with its arguments, and optional state"""
        ba = self.data["bound_args"]
        if state:
            func = self.data["func"]
            ba = inspect.signature(func).bind(state, *ba.args[1:], **ba.kwargs)
            func(*ba.args, **ba.kwargs)
            return state
        else:
            self.data["func"](*ba.args, **ba.kwargs)
            ba.apply_defaults()
            return ba.arguments["state"]

    def __str__(self):
        return pp.pformat(dict(getattr(self.data.get("bound_args", {}), 'arguments', {})))

    def __iter__(self):
        for c in self.child_list:
            yield c

    def partial(self):
        """Return partial of original function call"""
        ba = self.data["bound_args"]
        return state_partial(self.data["func"], *ba.args[1:], **ba.kwargs)

    def update_child_calls(self):
        """Replace child nodes on original function call with their partials"""

        for node in filter(lambda n: len(n.arg_name), self.child_list):
            self.data["bound_args"].arguments[node.arg_name] = node.partial()
        self.updated = True

    def remove_child(self, node):
        index = self.child_list.index(node)
        del self.child_list[index]
        return index

    def add_child(self, child):
        # since it is a tree, there is only one parent
        # note this means we do not allow edges between same layer units
        if child.parent:
            child.parent.remove_child(child)
        child.parent = self
        self.child_list.append(child)

    def descend(self, include_me=True):
        """Descend depth first into all child nodes"""
        if include_me:
            yield self

        for child in self.child_list:
            yield child
            yield from child.descend()

    @property
    def depth(self):
        if self.parent:
            return self.parent.depth + 1
        else:
            return 0

class NodeList(Node):
    def partial(self):
        return [node.partial() for node in self.child_list]

    def update_child_calls(self):
        pass

class Probe(object):
    def __init__(self, tree, f, eval_on_call=False):
        self.tree = tree
        self.f = f
        self.test_name = f.__name__
        self.eval_on_call = eval_on_call
        # TODO: auto sub_test detection
        self.sub_tests = SUB_TESTS.get(self.test_name) or []

    def __call__(self, *args, **kwargs):
        """Bind arguments to original function signature, and store in a Node

        This is used to discover what tests and sub-tests the SCT would like to
        call, and defer them for later execution via their node instance. Node
        instances are assembled into a tree.

        """
        if (len(args) > 0 and not isinstance(args[0], State)) or len(args) == 0:
            # no state placeholder if a state is passed
            args = ['state_placeholder'] + list(args)

        bound_args = inspect.signature(self.f).bind(*args, **kwargs)

        data = dict(bound_args=bound_args, func=self.f)
        this_node = Node(data=data, name=self.test_name)
        if self.tree is not None:
            self.tree.crnt_node.add_child(this_node)

        # First pass to set up branches off node
        arguments = bound_args.arguments
        for subtest in self.sub_tests:  # TODO: auto sub test detection
            if subtest in arguments and arguments[subtest]:
                self.build_sub_test_nodes(arguments[subtest], self.tree, this_node, subtest)

        # Second pass to build node and all its children into a subtest
        for node in this_node.descend(include_me=True):
            if node.updated:  # already built, e.g. node used multiple times
                continue
            else:
                node.update_child_calls()

        if self.eval_on_call:
            return this_node()
        else:
            return this_node

    @staticmethod
    def build_sub_test_nodes(test, tree, node, arg_name):
        # note that I've made the strong assumption that
        # if not a function, then test is a dict, list or tuple of them
        if isinstance(test, (list, tuple)):
            nl = NodeList(name="List", arg_name=arg_name)
            node.add_child(nl)
            for ii, f in enumerate(test):
                Probe.build_sub_test_nodes(f, tree, nl, str(ii))
        elif isinstance(test, Node):
            # test was a lambdaless subtest call, which produced a node
            # so need to tell it what its arg_name was on parent test
            test.arg_name = arg_name
            node.add_child(test)
        elif callable(test):
            # test was inside a lambda, function containing subtests or v2 F() chain object with subtests
            # since either may contain multiple subtests, we put them in a node list
            nl = NodeList(name="ListDeferred", arg_name=arg_name)
            node.add_child(nl)
            if tree is not None:
                prev_node, tree.crnt_node = tree.crnt_node, nl
                test()
                tree.crnt_node = prev_node
            else:
                test()
        elif test is not None:
            raise Exception("Expected a function or list/tuple/dict of functions")

def build_probe_context():
    tree = Tree()
    probe_context = {s: Probe(tree, getattr(test_funcs, s)) for s in TEST_NAMES}
    return tree, probe_context


from protowhat.sct_syntax import Chain as ProtoChain, F as ProtoF, state_dec_gen
from pythonwhat.checks.check_wrappers import scts
from pythonwhat.State import State
from pythonwhat.probe import Node, Probe, TEST_NAMES
from pythonwhat.utils import include_v1
from pythonwhat import test_funcs
from functools import wraps

# TODO: could define scts for check_wrappers at the module level
ATTR_SCTS = scts.copy()

def multi_dec(f):
    """Decorator for multi to remove nodes for original test functions from root node"""

    @wraps(f)
    def wrapper(*args, **kwargs):
        args = (
            args[0] if len(args) == 1 and isinstance(args[0], (list, tuple)) else args
        )
        for arg in args:
            if isinstance(arg, Node) and arg.parent.name is "root":
                arg.parent.remove_child(arg)
                arg.update_child_calls()
        return f(*args, **kwargs)

    return wrapper

state_dec = state_dec_gen(State, ATTR_SCTS)

class Chain(ProtoChain):
    def __init__(self, state, attr_scts=ATTR_SCTS):
        super().__init__(state, attr_scts)

class F(ProtoF):
    def __init__(self, stack=None, attr_scts=ATTR_SCTS):
        super().__init__(stack, attr_scts)

def Ex(state=None):
    return Chain(state or State.root_state)

if include_v1():
    # Prepare SCTs that may be chained attributes ----------------------
    # decorate functions that may try to run test_* function nodes as subtests
    # so they remove those nodes from the tree
    for k in ["multi", "with_context"]:
        ATTR_SCTS[k] = multi_dec(ATTR_SCTS[k])

    # allow test_* functions as chained attributes
    for k in TEST_NAMES:
        ATTR_SCTS[k] = Probe(tree=None, f=getattr(test_funcs, k), eval_on_call=True)

    # original logical test_* functions behave like multi
    # this is necessary to allow them to take check_* funcs as args
    # since probe behavior will try to call all SCTs passed (assuming they're also probes)
    for k in ["test_or", "test_correct"]:
        ATTR_SCTS[k] = multi_dec(getattr(test_funcs, k))

# Prepare check_funcs to be used alone (e.g. test = check_with().check_body())
v2_check_functions = {k: state_dec(v) for k, v in scts.items()}


from pythonwhat.tasks import (
    getResultInProcess,
    getOutputInProcess,
    getErrorInProcess,
    ReprFail,
    isDefinedInProcess,
    getOptionFromProcess,
    ReprFail,
    UndefinedValue,
)
from protowhat.Test import Test
from pythonwhat.Test import EqualTest, DefinedCollTest
from protowhat.Feedback import InstructorError
from pythonwhat.Feedback import Feedback
from pythonwhat import utils
from functools import partial
import re
import copy
import ast

evalCalls = {
    "value": getResultInProcess,
    "output": getOutputInProcess,
    "error": getErrorInProcess,
}

def has_part(state, name, msg, fmt_kwargs=None, index=None):
    d = {
        "sol_part": state.solution_parts,
        "stu_part": state.student_parts,
        **fmt_kwargs,
    }

    def verify(part, index):
        if index is not None:
            if isinstance(index, list):
                for ind in index:
                    part = part[ind]
            else:
                part = part[index]
        if part is None:
            raise KeyError

    # Chceck if it's there in the solution
    _msg = state.build_message(msg, d)
    _err_msg = "SCT fails on solution: " + _msg
    try:
        verify(state.solution_parts[name], index)
    except (KeyError, IndexError):
        raise InstructorError(_err_msg)

    try:
        verify(state.student_parts[name], index)
    except (KeyError, IndexError):
        state.report(Feedback(_msg, state))

    return state

def has_equal_part(state, name, msg):
    d = {
        "stu_part": state.student_parts,
        "sol_part": state.solution_parts,
        "name": name,
    }

    _msg = state.build_message(msg, d)
    state.do_test(
        EqualTest(d["stu_part"][name], d["sol_part"][name], Feedback(_msg, state))
    )

    return state

# TODO: shouldn't have to hardcode message
def has_equal_part_len(state, name, unequal_msg):
    """Verify that a part that is zoomed in on has equal length.

    Typically used in the context of ``check_function_def()``

    Arguments:
        name (str): name of the part for which to check the length to the corresponding part in the solution.
        unequal_msg (str): Message in case the lengths do not match.
        state (State): state as passed by the SCT chain. Don't specify this explicitly.

    :Examples:

        Student and solution code::

            def shout(word):
                return word + '!!!'

        SCT that checks number of arguments::

            Ex().check_function_def('shout').has_equal_part_len('args', 'not enough args!')
    """
    d = dict(
        stu_len=len(state.student_parts[name]), sol_len=len(state.solution_parts[name])
    )

    if d["stu_len"] != d["sol_len"]:
        _msg = state.build_message(unequal_msg, d)
        state.report(Feedback(_msg, state))

    return state

# Expression tests -----------------------------------------------------------

def has_equal_ast(state, incorrect_msg=None, code=None, exact=True, append=None):
    """Test whether abstract syntax trees match between the student and solution code.

    ``has_equal_ast()`` can be used in two ways:

    * As a robust version of ``has_code()``. By setting ``code``, you can look for the AST representation of ``code`` in the student's submission.
      But be aware that ``a`` and ``a = 1`` won't match, as reading and assigning are not the same in an AST.
      Use ``ast.dump(ast.parse(code))`` to see an AST representation of ``code``.
    * As an expression-based check when using more advanced SCT chain, e.g. to compare the equality of expressions to set function arguments.

    Args:
        incorrect_msg: message displayed when ASTs mismatch. When you specify ``code`` yourself, you have to specify this.
        code: optional code to use instead of the solution AST.
        exact: whether the representations must match exactly. If false, the solution AST
               only needs to be contained within the student AST (similar to using test student typed).
               Defaults to ``True``, unless the ``code`` argument has been specified.

    :Example:

        Student and Solution Code::

            dict(a = 'value').keys()

        SCT::

            # all pass
            Ex().has_equal_ast()
            Ex().has_equal_ast(code = "dict(a = 'value').keys()")
            Ex().has_equal_ast(code = "dict(a = 'value')", exact = False)

        Student and Solution Code::

            import numpy as np
            arr = np.array([1, 2, 3, 4, 5])
            np.mean(arr)

        SCT::

            # Check underlying value of arugment a of np.mean:
            Ex().check_function('numpy.mean').check_args('a').has_equal_ast()

            # Only check AST equality of expression used to specify argument a:
            Ex().check_function('numpy.mean').check_args('a').has_equal_ast()

    """
    if utils.v2_only():
        state.assert_is_not(["object_assignments"], "has_equal_ast", ["check_object"])
        state.assert_is_not(["function_calls"], "has_equal_ast", ["check_function"])

    if code and incorrect_msg is None:
        raise InstructorError(
            "If you manually specify the code to match inside has_equal_ast(), "
            "you have to explicitly set the `incorrect_msg` argument."
        )

    if (
        append is None
    ):  # if not specified, set to False if incorrect_msg was manually specified
        append = incorrect_msg is None
    if incorrect_msg is None:
        incorrect_msg = "Expected `{{sol_str}}`, but got `{{stu_str}}`."

    def parse_tree(tree):
        # get contents of module.body if only 1 element
        crnt = (
            tree.body[0]
            if isinstance(tree, ast.Module) and len(tree.body) == 1
            else tree
        )

        # remove Expr if it exists
        return ast.dump(crnt.value if isinstance(crnt, ast.Expr) else crnt)

    stu_rep = parse_tree(state.student_ast)
    sol_rep = parse_tree(state.solution_ast if not code else ast.parse(code))

    fmt_kwargs = {
        "sol_str": state.solution_code if not code else code,
        "stu_str": state.student_code,
    }

    _msg = state.build_message(incorrect_msg, fmt_kwargs, append=append)

    if exact and not code:
        state.do_test(EqualTest(stu_rep, sol_rep, Feedback(_msg, state)))
    elif not sol_rep in stu_rep:
        state.report(Feedback(_msg, state))

    return state

DEFAULT_INCORRECT_MSG = "Expected {{test_desc}}`{{sol_eval}}`, but got `{{stu_eval}}`."
DEFAULT_ERROR_MSG = "Running {{'it' if parent['part'] else 'the higlighted expression'}} generated an error: `{{stu_str}}`."
DEFAULT_ERROR_MSG_INV = "Running {{'it' if parent['part'] else 'the higlighted expression'}} didn't generate an error, but it should!"
DEFAULT_UNDEFINED_NAME_MSG = "Running {{'it' if parent['part'] else 'the higlighted expression'}} should define a variable `{{name}}` without errors, but it doesn't."
DEFAULT_INCORRECT_NAME_MSG = (
    "Are you sure you assigned the correct value to `{{name}}`?"
)
DEFAULT_INCORRECT_EXPR_CODE_MSG = (
    "Running the expression `{{expr_code}}` didn't generate the expected result."
)

args_string = """

    Args:
        incorrect_msg (str): feedback message if the {0} of the expression in the solution
          doesn't match the one of the student. This feedback message will be expanded if it is used
          in the context of another check function, like ``check_if_else``.
        error_msg (str): feedback message if there was an error when running the targeted student code.
          Note that when testing for an error, this message is displayed when none is raised.
        undefined_msg (str): feedback message if the ``name`` argument is defined, but a variable
          with that name doesn't exist after running the targeted student code.
        extra_env (dict): set variables to the extra environment. They will update the student and solution environment in
          the active state before the student/solution code in the active state is ran. This argument should contain a
          dictionary with the keys the names of the variables you want to set, and the values are the values of these variables.
          You can also use ``set_env()`` for this.
        context_vals (list): set variables which are bound in a ``for`` loop to certain values.
          This argument is only useful when checking a for loop (or list comprehensions).
          It contains a list with the values of the bound variables.
          You can also use ``set_context()`` for this.
        pre_code (str): the code in string form that should be executed before the expression is executed.
          This is the ideal place to set a random seed, for example.
        expr_code (str): if this argument is set, the expression in the student/solution code will not
          be ran. Instead, the given piece of code will be ran in the student as well as the solution environment
          and the result will be compared.
        name (str): If this is specified, the {0} of running this expression after running the focused expression
          is returned, instead of the {0} of the focussed expression in itself. This is typically used to inspect the
          {0} of an object after executing the body of e.g. a ``for`` loop.
        copy (bool): whether to try to deep copy objects in the environment, such as lists, that could
          accidentally be mutated. Disable to speed up SCTs. Disabling may lead to cryptic mutation issues.
        func: custom binary function of form f(stu_result, sol_result), for equality testing.
        override: If specified, this avoids the execution of the targeted code in the solution process. Instead, it
          will compare the {0} of the expression in the student process with the value specified in ``override``.
          Typically used in a ``SingleProcessExercise`` or if you want to allow for different solutions other than
          the one coded up in the solution.
    """

def has_expr(
    state,
    incorrect_msg=None,
    error_msg=None,
    undefined_msg=None,
    append=None,
    extra_env=None,
    context_vals=None,
    pre_code=None,
    expr_code=None,
    name=None,
    copy=True,
    func=None,
    override=None,
    test=None,  # todo: default or arg before state
):

    if (
        append is None
    ):  # if not specified, set to False if incorrect_msg was manually specified
        append = incorrect_msg is None
    if incorrect_msg is None:
        if name:
            incorrect_msg = DEFAULT_INCORRECT_NAME_MSG
        elif expr_code:
            incorrect_msg = DEFAULT_INCORRECT_EXPR_CODE_MSG
        else:
            incorrect_msg = DEFAULT_INCORRECT_MSG
    if undefined_msg is None:
        undefined_msg = DEFAULT_UNDEFINED_NAME_MSG
    if error_msg is None:
        if test == "error":
            error_msg = DEFAULT_ERROR_MSG_INV
        else:
            error_msg = DEFAULT_ERROR_MSG

    get_func = partial(
        evalCalls[test],
        extra_env=extra_env,
        context_vals=context_vals,
        pre_code=pre_code,
        expr_code=expr_code,
        name=name,
        copy=copy,
    )

    if override is not None:
        # don't bother with running expression and fetching output/value
        # eval_sol, str_sol = eval
        eval_sol, str_sol = override, str(override)
    else:
        eval_sol, str_sol = get_func(
            tree=state.solution_ast,
            process=state.solution_process,
            context=state.solution_context,
            env=state.solution_env,
        )

        if (test == "error") ^ isinstance(eval_sol, Exception):
            raise InstructorError(
                "Evaluating expression raised error in solution process (or didn't raise if testing for one). "
                "Error: {} - {}".format(type(eval_sol), str_sol)
            )
        if isinstance(eval_sol, ReprFail):
            raise InstructorError(
                "Couldn't extract the value for the highlighted expression from the solution process: "
                + eval_sol.info
            )

    eval_stu, str_stu = get_func(
        tree=state.student_ast,
        process=state.student_process,
        context=state.student_context,
        env=state.student_env,
    )

    # kwargs ---
    fmt_kwargs = {
        "stu_part": state.student_parts,
        "sol_part": state.solution_parts,
        "name": name,
        "test": test,
        "test_desc": "" if test == "value" else "the %s " % test,
        "expr_code": expr_code,
    }

    fmt_kwargs["stu_eval"] = utils.shorten_str(str(eval_stu))
    fmt_kwargs["sol_eval"] = utils.shorten_str(str(eval_sol))
    if incorrect_msg == DEFAULT_INCORRECT_MSG and (
        fmt_kwargs["stu_eval"] is None
        or fmt_kwargs["sol_eval"] is None
        or fmt_kwargs["stu_eval"] == fmt_kwargs["sol_eval"]
    ):
        incorrect_msg = "Expected something different."

    # tests ---
    # error in process
    if (test == "error") ^ isinstance(eval_stu, Exception):
        fmt_kwargs["stu_str"] = str_stu
        _msg = state.build_message(error_msg, fmt_kwargs, append=append)
        feedback = Feedback(_msg, state)
        state.report(feedback)

    # name is undefined after running expression
    if isinstance(eval_stu, UndefinedValue):
        _msg = state.build_message(undefined_msg, fmt_kwargs, append=append)
        state.report(Feedback(_msg, state))

    # test equality of results
    _msg = state.build_message(incorrect_msg, fmt_kwargs, append=append)
    state.do_test(EqualTest(eval_stu, eval_sol, Feedback(_msg, state), func))

    return state

has_equal_value = partial(has_expr, test="value")
has_equal_value.__doc__ = (
    """Run targeted student and solution code, and compare returned value.

    When called on an SCT chain, ``has_equal_value()`` will execute the student and solution
    code that is 'zoomed in on' and compare the returned values.
    """
    + args_string.format("returned value", "value")
    + """
    :Example:

        Student code and solution code::

            import numpy as np
            arr = np.array([1, 2, 3, 4, 5])
            np.mean(arr)

        SCT::

            # Verify equality of arr:
            Ex().check_object('arr').has_equal_value()

            # Verify whether arr was correctly set in np.mean
            Ex().check_function('numpy.mean').check_args('a').has_equal_value()

            # Verify whether np.mean(arr) produced the same result
            Ex().check_function('numpy.mean').has_equal_value()

    """
)

has_equal_output = partial(has_expr, test="output")
has_equal_output.__doc__ = """Run targeted student and solution code, and compare output.

    When called on an SCT chain, ``has_equal_output()`` will execute the student and solution
    code that is 'zoomed in on' and compare the output.
    """ + args_string.format(
    "output"
)

has_equal_error = partial(has_expr, test="error")
has_equal_error.__doc__ = """Run targeted student and solution code, and compare generated errors.

    When called on an SCT chain, ``has_equal_error()`` will execute the student and solution
    code that is 'zoomed in on' and compare the errors that they generate.
    """ + args_string.format(
    "error"
)

## Various has tests ----------------------------------------------------------

from pythonwhat.Test import StringContainsTest

def has_code(state, text, pattern=True, not_typed_msg=None):
    """Test the student code.

    Tests if the student typed a (pattern of) text. It is advised to use ``has_equal_ast()`` instead of ``has_code()``,
    as it is more robust to small syntactical differences that don't change the code's behavior.

    Args:
        text (str): the text that is searched for
        pattern (bool): if True (the default), the text is treated as a pattern. If False, it is treated as plain text.
        not_typed_msg (str): feedback message to be displayed if the student did not type the text.

    :Example:

        Student code and solution code::

            y = 1 + 2 + 3

        SCT::

            # Verify that student code contains pattern (not robust!!):
            Ex().has_code(r"1\\s*\\+2\\s*\\+3")

    """
    if not not_typed_msg:
        if pattern:
            not_typed_msg = "Could not find the correct pattern in your code."
        else:
            not_typed_msg = "Could not find the following text in your code: %r" % text

    student_code = state.student_code

    _msg = state.build_message(not_typed_msg)
    state.do_test(
        StringContainsTest(student_code, text, pattern, Feedback(_msg, state))
    )

    return state

def has_import(
    state,
    name,
    same_as=False,
    not_imported_msg="Did you import `{{pkg}}`?",
    incorrect_as_msg="Did you import `{{pkg}}` as `{{alias}}`?",
):
    """Checks whether student imported a package or function correctly.

    Python features many ways to import packages.
    All of these different methods revolve around the ``import``, ``from`` and ``as`` keywords.
    ``has_import()`` provides a robust way to check whether a student correctly imported a certain package.

    By default, ``has_import()`` allows for different ways of aliasing the imported package or function.
    If you want to make sure the correct alias was used to refer to the package or function that was imported,
    set ``same_as=True``.

    Args:
        name (str): the name of the package that has to be checked.
        same_as (bool): if True, the alias of the package or function has to be the same. Defaults to False.
        not_imported_msg (str): feedback message when the package is not imported.
        incorrect_as_msg (str): feedback message if the alias is wrong.

    :Example:

        Example 1, where aliases don't matter (defaut): ::

            # solution
            import matplotlib.pyplot as plt

            # sct
            Ex().has_import("matplotlib.pyplot")

            # passing submissions
            import matplotlib.pyplot as plt
            from matplotlib import pyplot as plt
            import matplotlib.pyplot as pltttt

            # failing submissions
            import matplotlib as mpl

        Example 2, where the SCT is coded so aliases do matter: ::

            # solution
            import matplotlib.pyplot as plt

            # sct
            Ex().has_import("matplotlib.pyplot", same_as=True)

            # passing submissions
            import matplotlib.pyplot as plt
            from matplotlib import pyplot as plt

            # failing submissions
            import matplotlib.pyplot as pltttt

    """
    student_imports = state.ast_dispatcher("imports", state.student_ast)
    solution_imports = state.ast_dispatcher("imports", state.solution_ast)

    if name not in solution_imports:
        raise InstructorError(
            "`has_import()` couldn't find an import of the package %s in your solution code."
            % name
        )

    fmt_kwargs = {"pkg": name, "alias": solution_imports[name]}

    _msg = state.build_message(not_imported_msg, fmt_kwargs)
    state.do_test(DefinedCollTest(name, student_imports, _msg))

    if same_as:
        _msg = state.build_message(incorrect_as_msg, fmt_kwargs)
        state.do_test(EqualTest(solution_imports[name], student_imports[name], _msg))

    return state

def has_output(state, text, pattern=True, no_output_msg=None):
    """Search student output for a pattern.

    Among the student and solution process, the student submission and solution code as a string,
    the ``Ex()`` state also contains the output that a student generated with his or her submission.

    With ``has_output()``, you can access this output and match it against a regular or fixed expression.

    Args:
        text (str): the text that is searched for
        pattern (bool): if True (default), the text is treated as a pattern. If False, it is treated as plain text.
        no_output_msg (str): feedback message to be displayed if the output is not found.

    :Example:

        As an example, suppose we want a student to print out a sentence: ::

            # Print the "This is some ... stuff"
            print("This is some weird stuff")

        The following SCT tests whether the student prints out ``This is some weird stuff``: ::

            # Using exact string matching
            Ex().has_output("This is some weird stuff", pattern = False)

            # Using a regular expression (more robust)
            # pattern = True is the default
            msg = "Print out ``This is some ... stuff`` to the output, " + \\
                  "fill in ``...`` with a word you like."
            Ex().has_output(r"This is some \w* stuff", no_output_msg = msg)

    """
    if not no_output_msg:
        no_output_msg = "You did not output the correct things."

    _msg = state.build_message(no_output_msg)
    state.do_test(StringContainsTest(state.raw_student_output, text, pattern, _msg))

    return state

def has_printout(
    state, index, not_printed_msg=None, pre_code=None, name=None, copy=False
):
    """Check if the right printouts happened.

    ``has_printout()`` will look for the printout in the solution code that you specified with ``index`` (0 in this case), rerun the ``print()`` call in
    the solution process, capture its output, and verify whether the output is present in the output of the student.

    This is more robust as ``Ex().check_function('print')`` initiated chains as students can use as many
    printouts as they want, as long as they do the correct one somewhere.

    Args:
        index (int): index of the ``print()`` call in the solution whose output you want to search for in the student output.
        not_printed_msg (str): if specified, this overrides the default message that is generated when the output
          is not found in the student output.
        pre_code (str): Python code as a string that is executed before running the targeted student call.
          This is the ideal place to set a random seed, for example.
        copy (bool): whether to try to deep copy objects in the environment, such as lists, that could
          accidentally be mutated. Disabled by default, which speeds up SCTs.
        state (State): state as passed by the SCT chain. Don't specify this explicitly.

    :Example:

        Suppose you want somebody to print out 4: ::

            print(1, 2, 3, 4)

        The following SCT would check that: ::

            Ex().has_printout(0)

        All of the following SCTs would pass: ::

            print(1, 2, 3, 4)
            print('1 2 3 4')
            print(1, 2, '3 4')
            print("random"); print(1, 2, 3, 4)

    :Example:

        Watch out: ``has_printout()`` will effectively **rerun** the ``print()`` call in the solution process after the entire solution script was executed.
        If your solution script updates the value of `x` after executing it, ``has_printout()`` will not work.

        Suppose you have the following solution: ::

            x = 4
            print(x)
            x = 6

        The following SCT will not work: ::

            Ex().has_printout(0)

        Why? When the ``print(x)`` call is executed, the value of ``x`` will be 6, and pythonwhat will look for the output `'6`' in the output the student generated.
        In cases like these, ``has_printout()`` cannot be used.

    :Example:

        Inside a for loop ``has_printout()``

        Suppose you have the following solution: ::

            for i in range(5):
                print(i)

        The following SCT will not work: ::

            Ex().check_for_loop().check_body().has_printout(0)

        The reason is that ``has_printout()`` can only be called from the root state. ``Ex()``.
        If you want to check printouts done in e.g. a for loop, you have to use a `check_function('print')` chain instead: ::

            Ex().check_for_loop().check_body().\\
                set_context(0).check_function('print').\\
                check_args(0).has_equal_value()

    """

    extra_msg = "If you want to check printouts done in e.g. a for loop, you have to use a `check_function('print')` chain instead."
    state.assert_root("has_printout", extra_msg=extra_msg)

    if not_printed_msg is None:
        not_printed_msg = (
            "Have you used `{{sol_call}}` to do the appropriate printouts?"
        )

    try:
        sol_call_ast = state.ast_dispatcher("function_calls", state.solution_ast)[
            "print"
        ][index]["node"]
    except (KeyError, IndexError):
        raise InstructorError(
            "`has_printout({})` couldn't find the {} print call in your solution.".format(
                index, utils.get_ord(index + 1)
            )
        )

    out_sol, str_sol = getOutputInProcess(
        tree=sol_call_ast,
        process=state.solution_process,
        context=state.solution_context,
        env=state.solution_env,
        pre_code=pre_code,
        copy=copy,
    )

    sol_call_str = state.solution_ast_tokens.get_text(sol_call_ast)

    if isinstance(str_sol, Exception):
        raise InstructorError(
            "Evaluating the solution expression {} raised error in solution process."
            "Error: {} - {}".format(sol_call_str, type(out_sol), str_sol)
        )

    _msg = state.build_message(not_printed_msg, {"sol_call": sol_call_str})

    has_output(state, out_sol.strip(), pattern=False, no_output_msg=_msg)

    return state

def has_no_error(
    state,
    incorrect_msg="Have a look at the console: your code contains an error. Fix it and try again!",
):
    """Check whether the submission did not generate a runtime error.

    If all SCTs for an exercise pass, before marking the submission as correct pythonwhat will automatically check whether
    the student submission generated an error. This means it is not needed to use ``has_no_error()`` explicitly.

    However, in some cases, using ``has_no_error()`` explicitly somewhere throughout your SCT execution can be helpful:

    - If you want to make sure people didn't write typos when writing a long function name.
    - If you want to first verify whether a function actually runs, before checking whether the arguments were specified correctly.
    - More generally, if, because of the content, it's instrumental that the script runs without
      errors before doing any other verifications.

    Args:
        incorrect_msg: if specified, this overrides the default message if the student code generated an error.

    :Example:

        Suppose you're verifying an exercise about model training and validation: ::

            # pre exercise code
            import numpy as np
            from sklearn.model_selection import train_test_split
            from sklearn import datasets
            from sklearn import svm

            iris = datasets.load_iris()
            iris.data.shape, iris.target.shape

            # solution
            X_train, X_test, y_train, y_test = train_test_split(
                iris.data, iris.target, test_size=0.4, random_state=0)

        If you want to make sure that ``train_test_split()`` ran without errors,
        which would check if the student typed the function without typos and used
        sensical arguments, you could use the following SCT: ::

            Ex().has_no_error()
            Ex().check_function('sklearn.model_selection.train_test_split').multi(
                check_args(['arrays', 0]).has_equal_value(),
                check_args(['arrays', 0]).has_equal_value(),
                check_args(['options', 'test_size']).has_equal_value(),
                check_args(['options', 'random_state']).has_equal_value()
            )

        If, on the other hand, you want to fall back onto pythonwhat's built in behavior,
        that checks for an error before marking the exercise as correct, you can simply
        leave of the ``has_no_error()`` step.

    """
    state.assert_root("has_no_error")

    if state.reporter.errors:
        _msg = state.build_message(
            incorrect_msg, {"error": str(state.reporter.errors[0])}
        )
        state.report(Feedback(_msg, state))

    return state

MC_VAR_NAME = "selected_option"

def has_chosen(state, correct, msgs):
    """Test multiple choice exercise.

    Test for a MultipleChoiceExercise. The correct answer (as an integer) and feedback messages
    are passed to this function.

    Args:
        correct (int): the index of the correct answer (should be an instruction). Starts at 1.
        msgs (list(str)): a list containing all feedback messages belonging to each choice of the
                          student. The list should have the same length as the number of options.
    """
    if not issubclass(type(correct), int):
        raise InstructorError(
            "Inside `has_chosen()`, the argument `correct` should be an integer."
        )

    student_process = state.student_process
    if not isDefinedInProcess(MC_VAR_NAME, student_process):
        raise InstructorError("Option not available in the student process")
    else:
        selected_option = getOptionFromProcess(student_process, MC_VAR_NAME)
        if not issubclass(type(selected_option), int):
            raise InstructorError("selected_option should be an integer")

        if selected_option < 1 or correct < 1:
            raise InstructorError(
                "selected_option and correct should be greater than zero"
            )

        if selected_option > len(msgs) or correct > len(msgs):
            raise InstructorError("there are not enough feedback messages defined")

        feedback_msg = msgs[selected_option - 1]

        state.reporter.success_msg = msgs[correct - 1]

        state.do_test(EqualTest(selected_option, correct, feedback_msg))


from pythonwhat.checks.check_funcs import part_to_child, StubState
from pythonwhat.tasks import getSignatureInProcess
from pythonwhat.utils import get_ord, get_times
from protowhat.Test import Test
from protowhat.Feedback import InstructorError
from pythonwhat.Feedback import Feedback
from pythonwhat.parsing import IndexedDict
from functools import partial

def bind_args(signature, args_part):
    pos_args = []
    kw_args = {}
    for k, arg in args_part.items():
        if isinstance(k, int):
            pos_args.append(arg)
        else:
            kw_args[k] = arg

    bound_args = signature.bind(*pos_args, **kw_args)
    return IndexedDict(bound_args.arguments)

def get_mapped_name(name, mappings):
    # get name by splitting on periods
    if "." in name:
        for orig, full_name in mappings.items():
            if name.startswith(full_name):
                return name.replace(full_name, orig)
    return name

MISSING_MSG = "Did you call `{{mapped_name}}()`{{' ' + times if index>0}}?"
SIG_ISSUE_MSG = (
    "Have you specified the arguments for `{{mapped_name}}()` using the right syntax?"
)
PREPEND_MSG = "Check your {{ord + ' ' if index>0}}call of `{{mapped_name}}()`. "

def check_function(
    state,
    name,
    index=0,
    missing_msg=None,
    params_not_matched_msg=None,
    expand_msg=None,
    signature=True,
):
    """Check whether a particular function is called.

    ``check_function()`` is typically followed by:

    - ``check_args()`` to check whether the arguments were specified.
      In turn, ``check_args()`` can be followed by ``has_equal_value()`` or ``has_equal_ast()``
      to assert that the arguments were correctly specified.
    - ``has_equal_value()`` to check whether rerunning the function call coded by the student
      gives the same result as calling the function call as in the solution.

    Checking function calls is a tricky topic. Please visit the
    `dedicated article <articles/checking_function_calls.html>`_ for more explanation,
    edge cases and best practices.

    Args:
        name (str): the name of the function to be tested. When checking functions in packages, always
            use the 'full path' of the function.
        index (int): index of the function call to be checked. Defaults to 0.
        missing_msg (str): If specified, this overrides an automatically generated feedback message in case
            the student did not call the function correctly.
        params_not_matched_msg (str): If specified, this overrides an automatically generated feedback message
            in case the function parameters were not successfully matched.
        expand_msg (str): If specified, this overrides any messages that are prepended by previous SCT chains.
        signature (Signature): Normally, check_function() can figure out what the function signature is,
            but it might be necessary to use ``sig_from_params()`` to manually build a signature and pass this along.
        state (State): State object that is passed from the SCT Chain (don't specify this).

    :Examples:

        Student code and solution code::

            import numpy as np
            arr = np.array([1, 2, 3, 4, 5])
            np.mean(arr)

        SCT::

            # Verify whether arr was correctly set in np.mean
            Ex().check_function('numpy.mean').check_args('a').has_equal_value()

            # Verify whether np.mean(arr) produced the same result
            Ex().check_function('numpy.mean').has_equal_value()
    """

    append_missing = missing_msg is None
    append_params_not_matched = params_not_matched_msg is None
    if missing_msg is None:
        missing_msg = MISSING_MSG
    if expand_msg is None:
        expand_msg = PREPEND_MSG
    if params_not_matched_msg is None:
        params_not_matched_msg = SIG_ISSUE_MSG

    stu_out = state.ast_dispatcher("function_calls", state.student_ast)
    sol_out = state.ast_dispatcher("function_calls", state.solution_ast)

    student_mappings = state.ast_dispatcher("mappings", state.student_ast)

    fmt_kwargs = {
        "times": get_times(index + 1),
        "ord": get_ord(index + 1),
        "index": index,
        "mapped_name": get_mapped_name(name, student_mappings),
    }

    # Get Parts ----
    # Copy, otherwise signature binding overwrites sol_out[name][index]['args']
    try:
        sol_parts = {**sol_out[name][index]}
    except KeyError:
        raise InstructorError(
            "`check_function()` couldn't find a call of `%s()` in the solution code. Make sure you get the mapping right!"
            % name
        )
    except IndexError:
        raise InstructorError(
            "`check_function()` couldn't find %s calls of `%s()` in your solution code."
            % (index + 1, name)
        )

    try:
        # Copy, otherwise signature binding overwrites stu_out[name][index]['args']
        stu_parts = {**stu_out[name][index]}
    except (KeyError, IndexError):
        _msg = state.build_message(missing_msg, fmt_kwargs, append=append_missing)
        state.report(Feedback(_msg, state))

    # Signatures -----
    if signature:
        signature = None if isinstance(signature, bool) else signature
        get_sig = partial(
            getSignatureInProcess,
            name=name,
            signature=signature,
            manual_sigs=state.get_manual_sigs(),
        )

        try:
            sol_sig = get_sig(
                mapped_name=sol_parts["name"], process=state.solution_process
            )
            sol_parts["args"] = bind_args(sol_sig, sol_parts["args"])
        except Exception as e:
            raise InstructorError(
                "`check_function()` couldn't match the %s call of `%s` to its signature:\n%s "
                % (get_ord(index + 1), name, e)
            )

        try:
            stu_sig = get_sig(
                mapped_name=stu_parts["name"], process=state.student_process
            )
            stu_parts["args"] = bind_args(stu_sig, stu_parts["args"])
        except Exception:
            _msg = state.build_message(
                params_not_matched_msg, fmt_kwargs, append=append_params_not_matched
            )
            state.report(
                Feedback(
                    _msg, StubState(stu_parts["node"], state.highlighting_disabled)
                )
            )

    # three types of parts: pos_args, keywords, args (e.g. these are bound to sig)
    append_message = {"msg": expand_msg, "kwargs": fmt_kwargs}
    child = part_to_child(
        stu_parts, sol_parts, append_message, state, node_name="function_calls"
    )
    return child


from pythonwhat import utils
import dill
import pickle
import pythonwhat
import ast
import inspect
from copy import deepcopy
from pickle import PicklingError
from pythonwhat.utils_env import set_context_vals, assign_from_ast
from contextlib import contextmanager
from functools import partial, wraps
from protowhat.Feedback import InstructorError

# Shell is passed as a parameter to partially applied functions in executeTask
# Process is passed as a parameter in SCT function

def process_task(f):
    """Decorator to (optionally) run function in a process."""
    sig = inspect.signature(f)

    @wraps(f)
    def wrapper(*args, **kwargs):
        # get bound arguments for call
        ba = sig.bind_partial(*args, **kwargs)
        # when process is specified, remove from args and use to execute
        process = ba.arguments.get("process")
        if process:
            ba.arguments["process"] = None
            # partial function since shell argument may have been left
            # unspecified, as it will be passed when the process executes
            pf = partial(wrapper, *ba.args, **ba.kwargs)
            return process.executeTask(pf)
        # otherwise, run original function
        return f(*ba.args, **ba.kwargs)

    return wrapper

def get_env(ns):
    if "__env__" in ns:
        return ns["__env__"]
    else:
        return ns

@contextmanager
def capture_output():
    import sys
    from io import StringIO

    oldout, olderr = sys.stdout, sys.stderr
    out = [StringIO(), StringIO()]
    sys.stdout, sys.stderr = out
    yield out
    sys.stdout, sys.stderr = oldout, olderr
    out[0] = out[0].getvalue()
    out[1] = out[1].getvalue()

# MC
@process_task
def getOptionFromProcess(process, name, shell):
    return shell.user_ns[name]

# Is a variable is defined in the process?
@process_task
def isDefinedInProcess(name, process, shell):
    return name in get_env(shell.user_ns)

# Is a variable is of a certain class in the process?
@process_task
def isInstanceInProcess(name, klass, process, shell):
    return isinstance(get_env(shell.user_ns)[name], klass)

# Get the columns of a Pandas data frame in the process
@process_task
def getColumnsInProcess(name, process, shell):
    return list(get_env(shell.user_ns)[name].columns)

# Is a key defined in a collection in the process?
@process_task
def isDefinedCollInProcess(name, key, process, shell):
    return key in get_env(shell.user_ns)[name]

# Get the signature of a function inside the process

def get_signature(name, mapped_name, signature, manual_sigs, env):
    if isinstance(signature, str):
        if signature in manual_sigs:
            signature = inspect.Signature(manual_sigs[signature])
        else:
            raise InstructorError("signature error - specified signature not found")

    if signature is None:
        # establish function
        try:
            fun = eval(mapped_name, env)
        except:
            raise InstructorError("%s() was not found." % mapped_name)

        # first go through manual sigs
        # try to get signature
        try:
            if name in manual_sigs:
                signature = inspect.Signature(manual_sigs[name])
            else:
                # it might be a method, and we have to find the general method name
                if "." in mapped_name:
                    els = name.split(".")
                    try:
                        els[0] = type(eval(els[0], env)).__name__
                        generic_name = ".".join(els[:])
                    except:
                        raise InstructorError("signature error - cannot convert call")
                    if generic_name in manual_sigs:
                        signature = inspect.Signature(manual_sigs[generic_name])
                    else:
                        raise InstructorError(
                            "signature error - %s not in builtins" % generic_name
                        )
                else:
                    raise InstructorError("manual signature not found")
        except Exception as e:
            try:
                signature = inspect.signature(fun)
            except:
                raise InstructorError(e.args[0] + " and cannot determine signature")

    return signature

# Get the signature of a function based on an object inside the process
@process_task
def getSignatureInProcess(name, mapped_name, signature, manual_sigs, process, shell):
    return get_signature(
        name=name,
        mapped_name=mapped_name,
        signature=signature,
        manual_sigs=manual_sigs,
        env=get_env(shell.user_ns),
    )

@process_task
def getSignatureFromObjInProcess(obj_char, process, shell):
    try:
        return inspect.signature(eval(obj_char, get_env(shell.user_ns)))
    except:
        return None

# Stuff for test_with

from contextlib import ExitStack

def context_env_update(context_list, env):
    es = ExitStack()
    for item in context_list:
        # create context manager and enter
        tmp_name = "__pw_cm"
        cm_code = compile(ast.Expression(item.context_expr), "<context_eval>", "eval")
        env[tmp_name] = es.enter_context(eval(cm_code, env))

        # assign to its optional_vars in separte dict
        if item.optional_vars:
            code = assign_from_ast(item.optional_vars, tmp_name)
            exec(code, env)

    return es

@process_task
def setUpNewEnvInProcess(context, process, shell):
    shell.user_ns["__env__"] = utils.copy_env(shell.user_ns)
    try:
        es = context_env_update(context, shell.user_ns["__env__"])
        shell.user_ns["__exit_stack__"] = es
        return True
    except Exception as e:
        return e

# break down environment
def context_objs_exit(es):
    try:
        es.close()
        return True
    except Exception as e:
        return e

@process_task
def breakDownNewEnvInProcess(process, shell):
    try:
        res = context_objs_exit(shell.user_ns["__exit_stack__"])
        del shell.user_ns["__exit_stack__"]
        del shell.user_ns["__env__"]
        return res
    except:
        return False

# Tasks that may need to serialize across processes ===========================

# Get a bytes or string representation of an object in the process
@process_task
def getClass(name, process, shell):
    try:
        obj = get_env(shell.user_ns)[name]
        obj_type = type(obj)
        return obj_type.__module__ + "." + obj_type.__name__
    except:
        return None

@process_task
def convert(name, converter, process, shell):
    return dill.loads(converter)(get_env(shell.user_ns)[name])

@process_task
def getStreamPickle(name, process, shell):
    try:
        return pickle.dumps(get_env(shell.user_ns)[name])
    except:
        return None

@process_task
def getStreamDill(name, process, shell):
    try:
        return dill.dumps(get_env(shell.user_ns)[name])
    except:
        return None

class ReprFail:
    def __init__(self, info):
        self.info = info

def getRepresentation(name, process):
    obj_class = getClass(name, process)
    converters = pythonwhat.State.State.root_state.converters
    if obj_class in converters:
        repres = convert(name, dill.dumps(converters[obj_class]), process)
        if errored(repres):
            return ReprFail("manual conversion failed: {}".format(repres))
        else:
            return repres
    else:
        # first try to pickle
        try:
            stream = getStreamPickle(name, process)
            if not errored(stream):
                return pickle.loads(stream)
        except:
            pass

        # if it failed, try to dill
        try:
            stream = getStreamDill(name, process)
            if not errored(stream):
                return dill.loads(stream)
            return ReprFail(
                "dilling inside process failed for %s - write manual converter"
                % obj_class
            )
        except PicklingError:
            return ReprFail(
                "undilling of bytestream failed with PicklingError - write manual converter"
            )
        except Exception as e:
            return ReprFail(
                "undilling of bytestream failed for class %s - write manual converter."
                "Error: %s - %s" % (obj_class, type(e), e)
            )

def errored(el):
    return el is None or (isinstance(el, list) and "backend-error" in str(el))

# Make wrapper for getting an object representation from process --------------
class UndefinedValue:
    pass

def getResultFromProcess(res, tempname, process):
    """Get a value from process, return tuple of value, res if succesful"""
    if not isinstance(res, (UndefinedValue, Exception)):
        value = getRepresentation(tempname, process)
        return value, res
    else:
        return res, str(res)

# decorator to automatically get value after running process task function
def get_rep(f):
    sig = inspect.signature(f)

    @wraps(f)
    def wrapper(*args, **kwargs):
        # get bound arguments for call
        ba = sig.bind_partial(*args, **kwargs)
        ba.apply_defaults()
        # get tempname, process arg values
        tempname = ba.arguments["tempname"]
        process = ba.arguments["process"]
        # run process task
        res = f(*args, **kwargs)
        # get result from task
        return getResultFromProcess(res, tempname, process)

    return wrapper

## Get the output of a tree (with setting envs, pre_code and/er expr_code)
@process_task
def get_output(f, process, shell, *args, **kwargs):
    with capture_output() as out:
        res = f(*args, process=process, shell=shell, **kwargs)

    out_str = out[0].strip()
    if not isinstance(res, Exception):
        toret = out_str or "no printouts"
        return toret, toret
    else:
        return res, str(res)

@process_task
def get_error(f, *args, **kwargs):
    res = f(*args, **kwargs)
    return (res, str(res)) if isinstance(res, Exception) else (None, res)

# General tasks to eval or exec code, with decorated counterparts -------------

# Eval an expression tree or node (with setting envs, pre_code and/or expr_code)
@process_task
def taskRunEval(
    tree,
    process,
    shell,
    env=None,
    extra_env=None,
    context=None,
    context_vals=None,
    pre_code="",
    expr_code="",
    name="",
    copy=True,
    tempname="_evaluation_object_",
    call=None,
):
    try:
        # Prepare code and mode -----------------------------------------------
        if (expr_code and name) or (not expr_code and isinstance(tree, ast.Module)):
            mode = "exec"
        else:
            mode = "eval"
            tree = ast.Expression(tree)

        # Expression code takes precedence over tree code
        if expr_code:
            code = expr_code
            tree = ast.parse(code, mode=mode)
        else:
            code = compile(tree, "<script>", mode)

        # Set up environment --------------------------------------------------
        # avoid deepy copy if specified, or just looking up variable by name
        if isinstance(tree, ast.Expression):
            tree = tree.body
        if not copy or (
            isinstance(tree, (ast.Name, ast.Subscript))
            and isinstance(tree.ctx, ast.Load)
        ):
            new_env = dict(get_env(shell.user_ns))
        else:
            # might raise an error if object refuses pickle interface
            # used by deepcopy to restore class
            new_env = utils.copy_env(get_env(shell.user_ns))

        if env is not None:
            new_env.update(deepcopy(env))
        if extra_env is not None:
            new_env.update(deepcopy(extra_env))
        if context is not None:
            set_context_vals(new_env, context, context_vals)

        # Execute code --------------------------------------------------------
        # Run pre_code if specified
        if pre_code:
            exec(pre_code, new_env)

        if mode == "eval":
            obj = eval(code, new_env)
        else:
            exec(code, new_env)
            obj = "exec only"

        # If name given, get from new_env
        if name:
            try:
                obj = eval(name, new_env)
            except NameError:
                return UndefinedValue()

        # call object if dict with args and kwargs was passed
        if call is not None:
            obj = obj(*call["args"], **call["kwargs"])

        # Set object as temp variable in original environment, so we can
        # later get its class, etc.., in order to extract it from process
        get_env(shell.user_ns)[tempname] = obj
        return str(obj)

    except Exception as e:
        return e

getResultInProcess = get_rep(taskRunEval)
getOutputInProcess = partial(get_output, taskRunEval)
getErrorInProcess = partial(get_error, taskRunEval)


import ast

def assign_from_ast(node, expr):
    """
    Creates code to assign name (or tuple of names) node from expr

    This is useful for recreating destructuring assignment behavior, like
    a, *b = [1,2,3].
    """
    if isinstance(expr, str):
        expr = ast.Name(id=expr, ctx=ast.Load())
    mod = ast.Module([ast.Assign(targets=[node], value=expr)])
    ast.fix_missing_locations(mod)
    return compile(mod, "<assignment_script>", "exec")

def set_context_vals(env, context, context_vals=None):
    env.update(**context)
    # support old test_* functions with context_val arguments
    if context_vals:
        crnt_ctx = context.context
        if len(crnt_ctx) == 1 != len(context_vals):
            context_vals = [context_vals]

        env.update({key: value for (key, value) in zip(crnt_ctx, context_vals)})


from protowhat.checks.check_logic import multi, check_not, check_or, check_correct, disable_highlighting, fail
from protowhat.Feedback import InstructorError
import ast

multi.__doc__ = str(multi.__doc__) + """
    :Example:

        Suppose we want to verify the following function call: ::

            round(1.2345, ndigits=2)

        The following SCT would verify this, using ``multi`` to
        'branch out' the state to two sub-SCTs: ::

            Ex().check_function('round').multi(
                check_args(0).has_equal_value(),
                check_args('ndigits').has_equal_value()
            )
    """

check_not.__doc__ = str(check_not.__doc__) + """
    :Example:
        The SCT fails with feedback for a specific incorrect value, defined using an override: ::

            Ex().check_object('result').multi(
                check_not(
                    has_equal_value(override=100),
                    msg='100 is incorrect for reason xyz.'
                ),
                has_equal_value()
            )

        Notice that ``check_not`` comes before the ``has_equal_value`` test
        that checks if the student value is equal to the solution value.

    :Example:
        The SCT below runs two ``has_code`` cases: ::

            Ex().check_not(
                has_code('mean'),
                has_code('median'),
                msg='Check your code'
            )

        If students use ``mean`` or ``median`` anywhere in their code, this SCT will fail.

    Note:
        - This function is not yet tested with all checks, please report unexpected behaviour.
        - This function can be thought as a NOT(x OR y OR ...) statement, since all tests it runs must fail
        - This function can be considered a direct counterpart of multi.

    """

check_or.__doc__ = str(check_or.__doc__) + """
    :Example:

        The SCT below tests that the student typed either 'mean' or 'median': ::

            Ex().check_or(
                has_code('mean'),
                has_code('median')
            )

        If the student didn't type either, the feedback message generated by ``has_code(mean)``,
        the first SCT, will be presented to the student.

    """

check_correct.__doc__ = str(check_correct.__doc__) + """
    :Example:

        The SCT below tests whether an object is correct. Only if the object is not correct, will
        the function calling checks be executed ::

            Ex().check_correct(
                check_object('x').has_equal_value(),
                check_function('round').check_args(0).has_equal_value()
            )

    """

# utility functions -----------------------------------------------------------

fail.__doc__ = str(fail.__doc__) + """
    :Example:

        As a trivial SCT example, ::

            Ex().check_for_loop().check_body().fail()

        This can also be helpful for debugging SCTs, as it can be used to stop testing as a given point.
    """

def override(state, solution):
    """Override the solution code with something arbitrary.

    There might be cases in which you want to temporarily override the solution code
    so you can allow for alternative ways of solving an exercise.
    When you use ``override()`` in an SCT chain, the remainder of that SCT chain will
    run as if the solution code you specified is the only code that was in the solution.

    Check the glossary for an example (pandas plotting)

    Args:
        solution: solution code as a string that overrides the original solution code.
        state: State instance describing student and solution code. Can be omitted if used with Ex().
    """

    # the old ast may be a number of node types, but generally either a
    # (1) ast.Module, or for single expressions...
    # (2) whatever was grabbed using module.body[0]
    # (3) module.body[0].value, when module.body[0] is an Expr node
    old_ast = state.solution_ast
    new_ast = ast.parse(solution)
    if not isinstance(old_ast, ast.Module) and len(new_ast.body) == 1:
        expr = new_ast.body[0]
        candidates = [expr, expr.value] if isinstance(expr, ast.Expr) else [expr]
        for node in candidates:
            if isinstance(node, old_ast.__class__):
                new_ast = node
                break

    kwargs = state.messages[-1] if state.messages else {}
    child = state.to_child(
        solution_ast=new_ast,
        student_ast=state.student_ast,
        highlight=state.highlight,
        append_message={"msg": "", "kwargs": kwargs},
    )

    return child

def set_context(state, *args, **kwargs):
    """Update context values for student and solution environments.

    When ``has_equal_x()`` is used after this, the context values (in ``for`` loops and function definitions, for example)
    will have the values specified through his function. It is the function equivalent of the ``context_vals`` argument of
    the ``has_equal_x()`` functions.

    - Note 1: excess args and unmatched kwargs will be unused in the student environment.
    - Note 2: When you try to set context values that don't match any target variables in the solution code,
      ``set_context()`` raises an exception that lists the ones available.
    - Note 3: positional arguments are more robust to the student using different names for context values.
    - Note 4: You have to specify arguments either by position, either by name. A combination is not possible.

    :Example:

        Solution code::

            total = 0
            for i in range(10):
                print(i ** 2)

        Student submission that will pass (different iterator, different calculation)::

            total = 0
            for j in range(10):
                print(j * j)

        SCT::

            # set_context is robust against different names of context values.
            Ex().check_for_loop().check_body().multi(
                set_context(1).has_equal_output(),
                set_context(2).has_equal_output(),
                set_context(3).has_equal_output()
            )

            # equivalent SCT, by setting context_vals in has_equal_output()
            Ex().check_for_loop().check_body().\\
                multi([s.has_equal_output(context_vals=[i]) for i in range(1, 4)])

    """

    stu_crnt = state.student_context.context
    sol_crnt = state.solution_context.context

    # for now, you can't specify both
    if len(args) > 0 and len(kwargs) > 0:
        raise InstructorError(
            "In `set_context()`, specify arguments either by position, either by name."
        )

    # set args specified by pos -----------------------------------------------
    if args:
        # stop if too many pos args for solution
        if len(args) > len(sol_crnt):
            raise InstructorError(
                "Too many positional args. There are {} context vals, but tried to set {}".format(
                    len(sol_crnt), len(args)
                )
            )
        # set pos args
        upd_sol = sol_crnt.update(dict(zip(sol_crnt.keys(), args)))
        upd_stu = stu_crnt.update(dict(zip(stu_crnt.keys(), args)))
    else:
        upd_sol = sol_crnt
        upd_stu = stu_crnt

    # set args specified by keyword -------------------------------------------
    if kwargs:
        # stop if keywords don't match with solution
        if set(kwargs) - set(upd_sol):
            raise InstructorError(
                "`set_context()` failed: context val names are {}, but you tried to set {}.".format(
                    upd_sol or "missing", sorted(list(kwargs.keys()))
                )
            )
        out_sol = upd_sol.update(kwargs)
        # need to match keys in kwargs with corresponding keys in stu context
        # in case they used, e.g., different loop variable names
        match_keys = dict(zip(sol_crnt.keys(), stu_crnt.keys()))
        out_stu = upd_stu.update(
            {match_keys[k]: v for k, v in kwargs.items() if k in match_keys}
        )
    else:
        out_sol = upd_sol
        out_stu = upd_stu

    return state.to_child(
        student_context=out_stu, solution_context=out_sol, highlight=state.highlight
    )

def set_env(state, **kwargs):
    """Update/set environemnt variables for student and solution environments.

    When ``has_equal_x()`` is used after this, the variables specified through this function will
    be available in the student and solution process. Note that you will not see these variables
    in the student process of the state produced by this function: the values are saved on the state
    and are only added to the student and solution processes when ``has_equal_ast()`` is called.

    :Example:

        Student and Solution Code::

            a = 1
            if a > 4:
                print('pretty large')

        SCT::

            # check if condition works with different values of a
            Ex().check_if_else().check_test().multi(
                set_env(a = 3).has_equal_value(),
                set_env(a = 4).has_equal_value(),
                set_env(a = 5).has_equal_value()
            )

            # equivalent SCT, by setting extra_env in has_equal_value()
            Ex().check_if_else().check_test().\\
                multi([has_equal_value(extra_env={'a': i}) for i in range(3, 6)])
    """

    stu_crnt = state.student_env.context
    sol_crnt = state.solution_env.context

    stu_new = stu_crnt.update(kwargs)
    sol_new = sol_crnt.update(kwargs)

    return state.to_child(
        student_env=stu_new, solution_env=sol_new, highlight=state.highlight
    )

disable_highlighting.__doc__ = str(disable_highlighting.__doc__) + """
    :Examples:

        SCT that will mark the 'number' portion if it is incorrect::

            Ex().check_function('round').check_args(0).has_equal_ast()

        SCT chains that will not mark certain mistakes. The earlier you put the function, the more types of mistakes will no longer be highlighted::

            Ex().disable_highlighting().check_function('round').check_args(0).has_equal_ast()
            Ex().check_function('round').disable_highlighting().check_args(0).has_equal_ast()
            Ex().check_function('round').check_args(0).disable_highlighting().has_equal_ast()
    """


from pythonwhat.parsing import ObjectAssignmentParser
from pythonwhat.Test import (
    DefinedProcessTest,
    InstanceProcessTest,
    DefinedCollProcessTest,
)
from protowhat.Feedback import InstructorError
from pythonwhat.Feedback import Feedback
from pythonwhat.tasks import (
    isDefinedInProcess,
    isInstanceInProcess,
    isDefinedCollInProcess,
)
from pythonwhat.checks.check_funcs import part_to_child
from pythonwhat.utils import v2_only
import pandas as pd
import ast

def check_object(
    state, index, missing_msg=None, expand_msg=None, typestr="variable"
):
    """Check object existence (and equality)

    Check whether an object is defined in the student's process, and zoom in on its value in both
    student and solution process to inspect quality (with has_equal_value().

    In ``pythonbackend``, both the student's submission as well as the solution code are executed, in separate processes.
    ``check_object()`` looks at these processes and checks if the referenced object is available in the student process.
    Next, you can use ``has_equal_value()`` to check whether the objects in the student and solution process correspond.

    Args:
        index (str): the name of the object which value has to be checked.
        missing_msg (str): feedback message when the object is not defined in the student process.
        expand_msg (str): If specified, this overrides any messages that are prepended by previous SCT chains.

    :Example:

        Suppose you want the student to create a variable ``x``, equal to 15: ::

            x = 15

        The following SCT will verify this: ::

            Ex().check_object("x").has_equal_value()

        - ``check_object()`` will check if the variable ``x`` is defined in the student process.
        - ``has_equal_value()`` will check whether the value of ``x`` in the solution process is the same as in the student process.

        Note that ``has_equal_value()`` only looks at **end result** of a variable in the student process.
        In the example, how the object ``x`` came about in the student's submission, does not matter.
        This means that all of the following submission will also pass the above SCT: ::

            x = 15
            x = 12 + 3
            x = 3; x += 12

    :Example:

        As the previous example mentioned, ``has_equal_value()`` only looks at the **end result**. If your exercise is
        first initializing and object and further down the script is updating the object, you can only look at the final value!

        Suppose you want the student to initialize and populate a list `my_list` as follows: ::

            my_list = []
            for i in range(20):
                if i % 3 == 0:
                    my_list.append(i)

        There is no robust way to verify whether `my_list = [0]` was coded correctly in a separate way.
        The best SCT would look something like this: ::

            msg = "Have you correctly initialized `my_list`?"
            Ex().check_correct(
                check_object('my_list').has_equal_value(),
                multi(
                    # check initialization: [] or list()
                    check_or(
                        has_equal_ast(code = "[]", incorrect_msg = msg),
                        check_function('list')
                    ),
                    check_for_loop().multi(
                        check_iter().has_equal_value(),
                        check_body().check_if_else().multi(
                            check_test().multi(
                                set_context(2).has_equal_value(),
                                set_context(3).has_equal_value()
                            ),
                            check_body().set_context(3).\\
                                set_env(my_list = [0]).\\
                                has_equal_value(name = 'my_list')
                        )
                    )
                )
            )

        - ``check_correct()`` is used to robustly check whether ``my_list`` was built correctly.
        - If ``my_list`` is not correct, **both** the initialization and the population code are checked.

    :Example:

        Because checking object correctness incorrectly is such a common misconception, we're adding another example: ::

            import pandas as pd
            df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
            df['c'] = [7, 8, 9]

        The following SCT would be **wrong**, as it does not factor in the possibility that the 'add column ``c``' step could've been wrong: ::

            Ex().check_correct(
                check_object('df').has_equal_value(),
                check_function('pandas.DataFrame').check_args(0).has_equal_value()
            )

        The following SCT would be better, as it is specific to the steps: ::

            # verify the df = pd.DataFrame(...) step
            Ex().check_correct(
                check_df('df').multi(
                    check_keys('a').has_equal_value(),
                    check_keys('b').has_equal_value()
                ),
                check_function('pandas.DataFrame').check_args(0).has_equal_value()
            )

            # verify the df['c'] = [...] step
            Ex().check_df('df').check_keys('c').has_equal_value()

    :Example:

        pythonwhat compares the objects in the student and solution process with the ``==`` operator.
        For basic objects, this ``==`` is operator is properly implemented, so that the objects can be effectively compared.
        For more complex objects that are produced by third-party packages, however, it's possible that this equality operator is not implemented in a way you'd expect.
        Often, for these object types the ``==`` will compare the actual object instances: ::

            # pre exercise code
            class Number():
                def __init__(self, n):
                    self.n = n

            # solution
            x = Number(1)

            # sct that won't work
            Ex().check_object().has_equal_value()

            # sct
            Ex().check_object().has_equal_value(expr_code = 'x.n')

            # submissions that will pass this sct
            x = Number(1)
            x = Number(2 - 1)

        The basic SCT like in the previous example will notwork here.
        Notice how we used the ``expr_code`` argument to _override_ which value `has_equal_value()` is checking.
        Instead of checking whether `x` corresponds between student and solution process, it's now executing the expression ``x.n``
        and seeing if the result of running this expression in both student and solution process match.

    """

    # Only do the assertion if PYTHONWHAT_V2_ONLY is set to '1'
    if v2_only():
        extra_msg = "If you want to check the value of an object in e.g. a for loop, use `has_equal_value(name = 'my_obj')` instead."
        state.assert_root("check_object", extra_msg=extra_msg)

    if missing_msg is None:
        missing_msg = "Did you define the {{typestr}} `{{index}}` without errors?"

    if expand_msg is None:
        expand_msg = "Did you correctly define the {{typestr}} `{{index}}`? "

    if (
        not isDefinedInProcess(index, state.solution_process)
        and state.has_different_processes()
    ):
        raise InstructorError(
            "`check_object()` couldn't find object `%s` in the solution process."
            % index
        )

    append_message = {"msg": expand_msg, "kwargs": {"index": index, "typestr": typestr}}

    # create child state, using either parser output, or create part from name
    fallback = lambda: ObjectAssignmentParser.get_part(index)
    stu_part = state.ast_dispatcher("object_assignments", state.student_ast).get(index, fallback())
    sol_part = state.ast_dispatcher("object_assignments", state.solution_ast).get(index, fallback())

    # test object exists
    _msg = state.build_message(missing_msg, append_message["kwargs"])
    state.do_test(DefinedProcessTest(index, state.student_process, Feedback(_msg)))

    child = part_to_child(
        stu_part, sol_part, append_message, state, node_name="object_assignments"
    )

    return child

def is_instance(state, inst, not_instance_msg=None):
    """Check whether an object is an instance of a certain class.

    ``is_instance()`` can currently only be used when chained from ``check_object()``, the function that is
    used to 'zoom in' on the object of interest.

    Args:
        inst (class): The class that the object should have.
        not_instance_msg (str): When specified, this overrides the automatically generated message in case
            the object does not have the expected class.
        state (State): The state that is passed in through the SCT chain (don't specify this).

    :Example:

        Student code and solution code::

            import numpy as np
            arr = np.array([1, 2, 3, 4, 5])

        SCT::

            # Verify the class of arr
            import numpy
            Ex().check_object('arr').is_instance(numpy.ndarray)
    """

    state.assert_is(["object_assignments"], "is_instance", ["check_object"])

    sol_name = state.solution_parts.get("name")
    stu_name = state.student_parts.get("name")

    if not_instance_msg is None:
        not_instance_msg = "Is it a {{inst.__name__}}?"

    if not isInstanceInProcess(sol_name, inst, state.solution_process):
        raise InstructorError(
            "`is_instance()` noticed that `%s` is not a `%s` in the solution process."
            % (sol_name, inst.__name__)
        )

    _msg = state.build_message(not_instance_msg, {"inst": inst})
    feedback = Feedback(_msg, state)
    state.do_test(InstanceProcessTest(stu_name, inst, state.student_process, feedback))

    return state

def check_df(
    state, index, missing_msg=None, not_instance_msg=None, expand_msg=None
):
    """Check whether a DataFrame was defined and it is the right type

    ``check_df()`` is a combo of ``check_object()`` and ``is_instance()`` that checks whether the specified object exists
    and whether the specified object is pandas DataFrame.

    You can continue checking the data frame with ``check_keys()`` function to 'zoom in' on a particular column in the pandas DataFrame:

    Args:
        index (str): Name of the data frame to zoom in on.
        missing_msg (str): See ``check_object()``.
        not_instance_msg (str): See ``is_instance()``.
        expand_msg (str): If specified, this overrides any messages that are prepended by previous SCT chains.

    :Example:

        Suppose you want the student to create a DataFrame ``my_df`` with two columns.
        The column ``a`` should contain the numbers 1 to 3,
        while the contents of column ``b`` can be anything: ::

            import pandas as pd
            my_df = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "n", "y"]})

        The following SCT would robustly check that: ::

            Ex().check_df("my_df").multi(
                check_keys("a").has_equal_value(),
                check_keys("b")
            )

        - ``check_df()`` checks if ``my_df`` exists (``check_object()`` behind the scenes) and is a DataFrame (``is_instance()``)
        - ``check_keys("a")`` zooms in on the column ``a`` of the data frame, and ``has_equal_value()`` checks if the columns correspond between student and solution process.
        - ``check_keys("b")`` zooms in on hte column ``b`` of the data frame, but there's no 'equality checking' happening

        The following submissions would pass the SCT above: ::

            my_df = pd.DataFrame({"a": [1, 1 + 1, 3], "b": ["a", "l", "l"]})
            my_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})

    """
    child = check_object(
        state,
        index,
        missing_msg=missing_msg,
        expand_msg=expand_msg,
        typestr="pandas DataFrame",
    )
    is_instance(child, pd.DataFrame, not_instance_msg=not_instance_msg)
    return child

def check_keys(state, key, missing_msg=None, expand_msg=None):
    """Check whether an object (dict, DataFrame, etc) has a key.

    ``check_keys()`` can currently only be used when chained from ``check_object()``, the function that is
    used to 'zoom in' on the object of interest.

    Args:
        key (str): Name of the key that the object should have.
        missing_msg (str): When specified, this overrides the automatically generated
            message in case the key does not exist.
        expand_msg (str): If specified, this overrides any messages that are prepended by previous SCT chains.
        state (State): The state that is passed in through the SCT chain (don't specify this).

    :Example:

        Student code and solution code::

            x = {'a': 2}

        SCT::

            # Verify that x contains a key a
            Ex().check_object('x').check_keys('a')

            # Verify that x contains a key a and a is correct.
            Ex().check_object('x').check_keys('a').has_equal_value()

    """

    state.assert_is(["object_assignments"], "is_instance", ["check_object", "check_df"])

    if missing_msg is None:
        missing_msg = "There is no {{ 'column' if 'DataFrame' in parent.typestr else 'key' }} `'{{key}}'`."
    if expand_msg is None:
        expand_msg = "Did you correctly set the {{ 'column' if 'DataFrame' in parent.typestr else 'key' }} `'{{key}}'`? "

    sol_name = state.solution_parts.get("name")
    stu_name = state.student_parts.get("name")

    if not isDefinedCollInProcess(sol_name, key, state.solution_process):
        raise InstructorError(
            "`check_keys()` couldn't find key `%s` in object `%s` in the solution process."
            % (key, sol_name)
        )

    # check if key available
    _msg = state.build_message(missing_msg, {"key": key})
    state.do_test(
        DefinedCollProcessTest(
            stu_name, key, state.student_process, Feedback(_msg, state)
        )
    )

    def get_part(name, key, highlight):
        if isinstance(key, str):
            slice_val = ast.Str(s=key)
        else:
            slice_val = ast.parse(str(key)).body[0].value
        expr = ast.Subscript(
            value=ast.Name(id=name, ctx=ast.Load()),
            slice=ast.Index(value=slice_val),
            ctx=ast.Load(),
        )
        ast.fix_missing_locations(expr)
        return {"node": expr, "highlight": highlight}

    stu_part = get_part(stu_name, key, state.student_parts.get("highlight"))
    sol_part = get_part(sol_name, key, state.solution_parts.get("highlight"))
    append_message = {"msg": expand_msg, "kwargs": {"key": key}}
    child = part_to_child(stu_part, sol_part, append_message, state)
    return child