import torch
import copy
from torch.testing._internal.common_methods_invocations import op_db
from functorch_additional_op_db import additional_op_db
from enum import Enum
import functorch._src.top_operators_github_usage as top_ops
import pprint
import unittest
import enum
from torch.testing._internal.common_device_type import toleranceOverride

# Importing these files make modifications to the op_db that we need
import test_ops  # noqa: F401
import test_vmap  # noqa: F401

all_overridable = list(torch.overrides.get_testing_overrides().keys())

public_docs = [
    (torch.nn.functional, 'torch.nn.functional', 'docs/source/nn.functional.rst'),
    (torch.fft, 'torch.fft', 'docs/source/fft.rst'),
    (torch.special, 'torch.special', 'docs/source/special.rst'),
    (torch.linalg, 'torch.linalg', 'docs/source/linalg.rst'),
    (torch, 'torch', 'docs/source/torch.rst'),
    (torch.Tensor, 'torch.Tensor', 'docs/source/tensors.rst'),
]

# torch.abs, Tensor.abs, Tensor.abs_ are all considered to be different


def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/debug-cpu'):
    results = {}
    all_overridable_apis = set(torch.overrides.get_testing_overrides().keys())
    for module, module_name, src in public_docs:
        with open(f'{pytorch_root}/{src}') as f:
            lines = f.readlines()
        # APIs eitehr begin with 4 spaces or ".. autofunction::"
        api_lines1 = [line.strip() for line in lines if line.startswith(' ' * 4)]
        api_lines2 = [line.strip()[len('.. autofunction:: '):]
                      for line in lines if line.startswith('.. autofunction::')]
        lines = api_lines1 + api_lines2
        lines = [line[7:] if line.startswith('Tensor.') else line for line in lines]
        lines = [line for line in lines if hasattr(module, line)]
        for line in lines:
            api = getattr(module, line)
            if api in all_overridable_apis:
                results[f'{module_name}.{line}'] = api
    return results


denylist = {
    'torch.Tensor.data_ptr',
    'torch.Tensor.dim',
    'torch.Tensor.element_size',
    'torch.Tensor.backward',
    'torch.Tensor.as_strided',
    'torch.Tensor.register_hook',
    'torch.Tensor.record_stream',
    'torch.Tensor.qscheme',
    'torch.Tensor.ndimension',
    'torch.Tensor.smm',
    'torch.Tensor.sspaddmm',
    'torch.Tensor.retain_grad',
    'torch.Tensor.sparse_mask',
    'torch.Tensor.sparse_dim',
    'torch.Tensor.dense_dim',
    'torch.Tensor.values',
    'torch.Tensor.indices',
    'torch.Tensor.numel',
    'torch.Tensor.size',
    'torch.Tensor.nelement',
    'torch.Tensor.q_scale',
    'torch.Tensor.q_zero_point',
    'torch.Tensor.q_per_channel_scales',
    'torch.Tensor.q_per_channel_zero_points',
    'torch.Tensor.q_per_channel_axis',
    'torch.Tensor.int_repr',
    'torch.Tensor.to_sparse',
    'torch.Tensor.is_inference',
    'torch.Tensor.storage',
    'torch.Tensor.storage_type',
}


def get_method_only_ops_we_care_about():
    apis = get_public_overridable_apis()
    result = []
    for key, _ in apis.items():
        if not key.startswith('torch.Tensor'):
            continue
        if key in denylist:
            continue
        api = key.split('.')[2]
        # filter out in-place
        if api.endswith('_'):
            continue
        if f'torch.{api}' not in apis.keys():
            result.append(api)
    return result

# Deduplicates torch.abs and Tensor.abs


def get_public_overridable_ops():
    results = get_public_overridable_apis()
    cpy = copy.deepcopy(results)
    for key, _ in cpy.items():
        if not key.startswith('torch.Tensor'):
            continue
        api = key.split('.')[2]
        if f'torch.{api}' in results.keys():
            del results[key]
    return results


def get_public_overridable_outplace_ops():
    results = get_public_overridable_ops()
    cpy = copy.deepcopy(results)
    for key, _ in cpy.items():
        # NB: there are no dunder methods bcs we don't document those
        if key.endswith('_'):
            del results[key]
    return results


def get_public_overridable_outplace_we_care_about():
    results = get_public_overridable_outplace_ops()
    cpy = copy.deepcopy(results)
    for key, _ in cpy.items():
        # quantization
        if 'quant' in key or '.q_' in key:
            del results[key]

        # is_cpu, etc. It doesn't make sense to have OpInfos for these
        if '.is_' in key:
            del results[key]

        if key in denylist and key in results:
            del results[key]
    return results

# e.g. nn.functional.softmax


def get_op(dotted_name):
    names = dotted_name.split('.')
    mod = torch
    for name in names:
        if not hasattr(mod, name):
            return None
        mod = getattr(mod, name)
    return mod

# Maps function -> [OpInfo]


def get_ops_covered_by_opinfos():
    ops = {}

    def safe_append(dct, key, val):
        if key in dct:
            dct[key].append(val)
        else:
            dct[key] = [val]

    for opinfo in op_db:
        func_op = get_op(opinfo.name)
        if func_op:
            safe_append(ops, func_op, opinfo)
        if opinfo.method_variant:
            safe_append(ops, opinfo.method_variant, opinfo)
        if opinfo.inplace_variant:
            safe_append(ops, opinfo.inplace_variant, opinfo)
        for alias in opinfo.aliases:
            safe_append(ops, alias.op, opinfo)
    return ops


factory_fns = {
    'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'randperm',
    'linspace', 'logspace', 'hann_window', 'full', 'eye', 'blackman_window',
    'barlett_window', 'randint', 'range', 'arange',
}


def get_top_ops(torch_threshold, nn_fn_threshold, with_counts=False):
    denylist = set({
        # These are either not real "operators", factory functions
        # that trivially work, or not-documented ops.
        'load', 'no_grad', 'save', 'from_numpy',
        'manual_seed', 'set_grad_enabled',
        'set_default_tensor_type', 'set_num_threads',
        'set_printoptions', 'numel',
        'set_default_dtype', 'sparse_coo_tensor', 'set_rng_state',
        'get_rng_state', 'get_default_dtype', 'initial_seed',
        'get_num_threads', 'quantize_per_tensor',
        'hann_window', 'is_tensor', 'as_tensor',
        'equal', 'enable_grad', 'seed', 'is_storage',
        'is_floating_point', 'nn.functional.torch',
        'set_flush_denormal', 'set_num_interop_threads', 'dequantize',
        'get_num_interop_threads', 'nn.functional.math',
        'nn.functional.threshold_',
        'nn.functional.selu_',
        'nn.functional.elu_',
        'nn.functional.rrelu_',
        'nn.functional.leaky_relu_',
        'nn.functional.hardtanh_',
        'nn.functional.has_torch_function',
        'nn.functional.has_torch_function_unary',
        'nn.functional.has_torch_function_variadic',
        'nn.functional.handle_torch_function',
        'nn.functional.adaptive_max_pool1d_with_indices',
        'nn.functional.adaptive_max_pool2d_with_indices',
        'nn.functional.adaptive_max_pool3d_with_indices',
        'nn.functional.fractional_max_pool2d_with_indices',
        'nn.functional.fractional_max_pool3d_with_indices',
        'is_complex',
        'grad',
        'quantize_per_channel',
        'nn.functional.max_pool2d_with_indices',
        'nn.functional.max_pool3d_with_indices',
        'nn.functional.max_pool1d_with_indices',
        'nn.functional.celu_',
        'nn.functional.grad',
        'nn.functional.relu_',
        'nn.functional.boolean_dispatch',
        'nn.functional.assert_int_or_pair',
        'fft',  # is namespace
    })

    torch_ops = top_ops.top_torch
    nn_fn_ops = top_ops.get_nn_functional_top_list()
    torch_ops = [op for op in torch_ops if op[0] not in denylist]
    nn_fn_ops = [op for op in nn_fn_ops if op[0] not in denylist]

    ops = torch_ops[:torch_threshold] + nn_fn_ops[:nn_fn_threshold]

    # Now, sort by priority
    ops.sort(reverse=True, key=lambda op: op[1])
    if not with_counts:
        ops = [op[0] for op in ops]
    return ops


def get_ops_percentage(torch_threshold, nn_fn_threshold):
    data = top_ops.top_torch + top_ops.get_nn_functional_top_list()

    def get_num_usages(opname):
        # Ignore this, this is heavily inflated
        if opname == 't':
            return 0
        result = [op[1] for op in data if op[0] == opname]
        assert len(result) == 1
        return result[0]

    # get all operators that are not in the denylist
    all_ops = get_top_ops(999999, 999999)
    total_op_usages = sum([get_num_usages(op) for op in all_ops])

    # get subset of all operators
    subset_ops = get_top_ops(torch_threshold, nn_fn_threshold)
    subset_op_usages = sum([get_num_usages(op) for op in subset_ops])
    return subset_op_usages / total_op_usages


def get_top_ops_not_covered_by_opinfo(torch_threshold=0, nn_fn_threshold=0):
    ops = get_top_ops(torch_threshold, nn_fn_threshold)

    ops_with_opinfo = []
    for op in op_db:
        ops_with_opinfo.append(op.name)
        ops_with_opinfo.extend([op.name for op in op.aliases])
    ops_with_opinfo = set(ops_with_opinfo)

    result = [op for op in ops if op not in ops_with_opinfo]
    result = [op for op in result if op not in denylist]
    result = [op for op in result if op not in factory_fns]
    return result


def get_covered_ops(ops_list, invert=False):
    ops_covered_by_opinfo = get_ops_covered_by_opinfos()
    overridable_outplace_ops = ops_list
    results = {}
    for key, op in overridable_outplace_ops.items():
        cond = op in ops_covered_by_opinfo
        if invert:
            cond = not cond
        if cond:
            results[key] = op
    return results


class Status(Enum):
    Correct = 0
    Fast = 1


tests = {
    'test_vmap_exhaustive',
    'test_op_has_batch_rule',
    'test_vjp',
    'test_vmapvjp',
    'test_vmapvjp_has_batch_rule',
    'test_jvp',
    'test_vmapjvp',
}


def is_decorateinfo_skip_or_xfail(decorateinfo):
    assert len(decorateinfo.decorators) == 1
    actual_decorator = decorateinfo.decorators[0]
    if isinstance(actual_decorator, toleranceOverride):
        return False
    if actual_decorator == unittest.expectedFailure:
        return True
    # Assume the rest are skips
    return True


def get_all_tested_ops():
    overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
    op_to_opinfo = get_ops_covered_by_opinfos()
    result = set({})
    for name, op in get_covered_ops(overridable_outplace_we_care_about).items():
        opinfos = op_to_opinfo[op]
        for opinfo in opinfos:
            result.add(opinfo.name)
    return result


def get_skipped_or_xfailed_ops_for(test_name):
    overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
    op_to_opinfo = get_ops_covered_by_opinfos()
    result = set({})
    for name, op in get_covered_ops(overridable_outplace_we_care_about).items():
        opinfos = op_to_opinfo[op]
        for opinfo in opinfos:
            for decorator in opinfo.decorators:
                if not hasattr(decorator, 'test_name'):
                    continue
                if decorator.test_name != test_name:
                    continue
                if is_decorateinfo_skip_or_xfail(decorator):
                    result.add(opinfo.name)
    return result


def get_statuses(for_subset=None, invert=False):
    overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
    if for_subset is not None:
        overridable_outplace_we_care_about = {
            k: v
            for k, v in overridable_outplace_we_care_about.items()
            # Removes "torch."
            if k[6:] in for_subset
        }
    op_to_opinfo = get_ops_covered_by_opinfos()
    result = {}
    _ = get_covered_ops(overridable_outplace_we_care_about)

    def get_covered_tests(op):
        opinfos = op_to_opinfo[op]
        result = copy.deepcopy(tests)
        for opinfo in opinfos:
            for decorator in opinfo.decorators:
                if not hasattr(decorator, 'test_name'):
                    continue
                if decorator.test_name in tests and decorator.test_name in result:
                    result.remove(decorator.test_name)
        return result

    def get_all_aliases(op):
        opinfos = op_to_opinfo[op]
        result = []
        for opinfo in opinfos:
            result.append(opinfo.name)
            result.extend(opinfo.aliases)
        return set(result)

    for name, op in get_covered_ops(overridable_outplace_we_care_about).items():
        successful_tests = get_covered_tests(op)
        failed_tests = tests - successful_tests
        result[name] = failed_tests if invert else successful_tests
    return result


def transpose_statuses(for_subset=None, invert=False):
    statuses = get_statuses(for_subset, invert=invert)
    result = {}
    for test in tests:
        result[test] = set({})
    for op, supported in statuses.items():
        for test in supported:
            result[test].add(op)
    return result


overridable_apis = get_public_overridable_apis()

overridable_ops = get_public_overridable_ops()

overridable_outplace_ops = get_public_overridable_outplace_ops()

overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()

tested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about)
untested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about, invert=True)

# print("List of OpInfos we need:")
# for key in untested_overridable_outplace_ops.keys():
#     print(key)
# print("-" * 80)
# print("")

print(f'Overridable public APIs: {len(overridable_apis)}')
print(f'Overridable public ops: {len(overridable_ops)}')
print(f'Overridable public outplace ops: {len(overridable_outplace_ops)}')
print(f'Overridable public outplace ops we care about: {len(overridable_outplace_we_care_about)}')
print(f'OpInfo-tested overridable public outplace ops: {len(tested_overridable_outplace_ops)}')


def remove_torch(name):
    assert name[:6] == 'torch.'
    return name[6:]


def get_list_of_all_tests():
    all_tests = list(tested_overridable_outplace_ops.keys())
    return set([remove_torch(test) for test in all_tests])


mytest = {
    'test_vmap_exhaustive',
    'test_op_has_batch_rule',
    'test_vjp',
    'test_vmapvjp',
    'test_vmapvjp_has_batch_rule',
}

print('*' * 80)
all_tests = get_list_of_all_tests()
for test in mytest:
    result = get_skipped_or_xfailed_ops_for(test)
    diff = len(all_tests - result)
    print(f'{test}: {diff}')


def get_jvp_coverage(subset=None):
    # - number that support autograd
    # - number that support forward_ad (in pytorch core)
    # - number that support functorch.jvp
    op_to_opinfo = get_ops_covered_by_opinfos()
    ops_dct = tested_overridable_outplace_ops
    if subset is not None:
        ops_dct = {name: op for name, op in ops_dct.items()
                   if remove_torch(name) in subset}
    supports_autograd_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items()
                                 if op_to_opinfo[fn][0].supports_autograd}
    supports_forwardad_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items()
                                  if op_to_opinfo[fn][0].supports_forward_ad}

    ops = set([remove_torch(test) for test in list(ops_dct.keys())])
    supports_autograd = set([remove_torch(test)
                             for test in list(supports_autograd_ops_dct.keys())])
    supports_forward_ad = set([remove_torch(test)
                               for test in list(supports_forwardad_ops_dct.keys())])
    assert supports_forward_ad.issubset(supports_autograd)
    assert supports_autograd.issubset(ops)

    failed_ops = get_skipped_or_xfailed_ops_for('test_jvp')

    coverage = len(supports_forward_ad - failed_ops)
    no_forward_ad = len(supports_autograd) - len(supports_forward_ad)
    print(f'test_jvp, {coverage}, {no_forward_ad}, {len(ops)}')


get_jvp_coverage()
get_jvp_coverage(get_top_ops(100, 25))
for op in get_top_ops(100, 25):
    print(op)
print('*' * 80)

# result = get_skipped_or_xfailed_ops_for('test_vmap_exhaustive')
# result = get_skipped_or_xfailed_ops_for('test_op_has_batch_rule')
# result = get_skipped_or_xfailed_ops_for('test_vjp')
# result = get_skipped_or_xfailed_ops_for('test_vmapvjp')
# result = get_skipped_or_xfailed_ops_for('test_vmapvjp_has_batch_rule')
# import pdb; pdb.set_trace()

statuses = transpose_statuses()
for test in tests:
    print(f'{test} coverage {len(statuses[test])}')

method_only_ops = get_method_only_ops_we_care_about()
# for op in method_only_ops:
#     print(f'    {op},')

top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(100, 25)
print('=' * 80)
for op in top_ops_not_covered_by_opinfo:
    print(f'{op}, {top_ops.usage_count[op]}')

# print("top ops not covered by opinfo: ")
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(200, 50)
# for op in top_ops_not_covered_by_opinfo:
#     print(f'{op}, {top_ops.usage_count[op]}')

# print("top ops not covered by opinfo: ")
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(220, 92)
# for op in top_ops_not_covered_by_opinfo:
#    print(f'{op}, {top_ops.usage_count[op]}')

# print("top ops not covered by opinfo: ")
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(999, 999)
# for op in top_ops_not_covered_by_opinfo:
#     print(f'{op}, {top_ops.usage_count[op]}')


def remove_from_set(parent, to_remove):
    for to_remove_elt in to_remove:
        if to_remove_elt in parent:
            parent.remove(to_remove_elt)


def print_coverage_info(th=100, nn=25):
    print('=' * 80)
    print(f"top {th}, {nn} coverage")
    statuses = transpose_statuses(get_top_ops(th, nn), invert=True)
    top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(th, nn)

    # testing problems
    exemptions = {
        'torch.nn.functional.dropout',  # randomness
    }

    # Allowed exemptions
    vmap_exemptions = {
        'torch.randn_like',  # randomness
        'torch.rand_like',  # randomness
        'torch.allclose',  # number output
        'torch.unique',  # dynamic
        'torch.nonzero',  # dynamic
        'torch.masked_select',  # dynamic
        'torch.prod',  # dynamic (backward)
        'torch.norm',  # norm with nuc is not commonly used; we support the other cases.
        'torch.svd',  # There isn't a bug, it is just nondeterministic so we can't test it.
        'torch.nn.functional.embedding',  # We support everything except the sparse option.
    }
    remove_from_set(statuses['test_vmap_exhaustive'], vmap_exemptions)
    remove_from_set(statuses['test_vmapvjp'], vmap_exemptions)
    remove_from_set(statuses['test_vmapvjp_has_batch_rule'], vmap_exemptions)
    remove_from_set(statuses['test_op_has_batch_rule'], vmap_exemptions)
    remove_from_set(statuses['test_vmapjvp'], vmap_exemptions)
    for test in tests:
        remove_from_set(statuses[test], exemptions)

    print(f"total ops in set: {th + nn}")
    print(f"tested by OpInfo: {th + nn - len(top_ops_not_covered_by_opinfo)}")
    for test in tests:
        if test in {'test_jvp', 'test_vmapjvp'}:
            continue
        print(f'{test} failing coverage {len(statuses[test])}')

    # We don't care about these yet
    del statuses['test_jvp']
    del statuses['test_vmapjvp']

    pprint.pprint(statuses)


def get_name_to_opinfo_map():
    dct = {}
    for op in (op_db + additional_op_db):
        def add(name, op):
            if name not in dct:
                dct[name] = []
            dct[name].append(op)
        add(op.name, op)
        for alias in op.aliases:
            add(alias.name, op)
    return dct


NAME_TO_OPINFO = get_name_to_opinfo_map()


class Support(enum.Enum):
    NO = 0
    YES = 1
    UNKNOWN = 2


FACTORY_FNS = {
    'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'range',
    'full', 'randperm', 'eye', 'randint', 'linspace', 'logspace',
}

VJP_EXEMPTIONS = {
    'nn.functional.dropout',  # not actually problem, randomness testing artifact
    'nn.functional.dropout2d',  # not actually problem, randomness testing artifact
    'nn.functional.rrelu',  # not actually problem, randomness testing artifact
    'bernoulli',  # not actually problem, randomness testing artifact
    'normal',  # not actually problem, randomness testing artifact
}

VMAP_EXEMPTIONS = {
    'randn_like',  # randomness
    'rand_like',  # randomness
    'allclose',  # number output
    'unique',  # dynamic
    'nonzero',  # dynamic
    'masked_select',  # dynamic
    'prod',  # dynamic (backward)
    'norm',  # norm with nuc is not commonly used; we support the other cases.
    'svd',  # There isn't a bug, it is just nondeterministic so we can't test it.
    'nn.functional.embedding',  # We support everything except the sparse option.
    'nn.functional.dropout',  # randomness
    'nn.functional.dropout2d',  # randomness
    'bernoulli',  # randomness
    'multinomial',  # randomness
    'normal',  # randomness
}

JVP_EXEMPTIONS = {
    'nn.functional.dropout',  # not actually problem, randomness testing artifact
    'nn.functional.dropout2d',  # not actually problem, randomness testing artifact
    'nn.functional.rrelu',  # not actually problem, randomness testing artifact
    'normal',  # not actually problem, randomness testing artifact
    'bernoulli',  # not actually problem, randomness testing artifact
}


class Operator:
    def __init__(self, name):
        self.name = name
        self.opinfos = NAME_TO_OPINFO.get(name, None)
        assert self.opinfos is None or len(self.opinfos) > 0

    def has_opinfo(self):
        return self.opinfos is not None

    def __repr__(self):
        return f'Operator("{self.name}")'

    def __hash__(self):
        return hash(self.name)

    def no_opinfos_skip_test(self, test_name):
        """Returns NO if any opinfos have a skip or xfail for the test"""
        if not self.has_opinfo():
            return Support.UNKNOWN
        for opinfo in self.opinfos:
            for decorator in opinfo.decorators:
                if not hasattr(decorator, 'test_name'):
                    continue
                if decorator.test_name != test_name:
                    continue
                if is_decorateinfo_skip_or_xfail(decorator):
                    return Support.NO
        return Support.YES

    def any_opinfo_attr(self, attr):
        if not self.has_opinfo():
            raise RuntimeError()
        return any([getattr(opinfo, attr) for opinfo in self.opinfos])

    def all_opinfo_attr(self, attr):
        if not self.has_opinfo():
            raise RuntimeError()
        return all([getattr(opinfo, attr) for opinfo in self.opinfos])

    def supports_vjp(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in VJP_EXEMPTIONS:
            return Support.YES
        return self.no_opinfos_skip_test('test_vjp')

    def supports_vmap(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in VMAP_EXEMPTIONS:
            return Support.YES
        return self.no_opinfos_skip_test('test_vmap_exhaustive')

    def supports_fast_vmap(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in VMAP_EXEMPTIONS:
            return Support.YES
        return self.no_opinfos_skip_test('test_op_has_batch_rule')

    def supports_vmapvjp(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in VMAP_EXEMPTIONS:
            return Support.YES
        return self.no_opinfos_skip_test('test_vmapvjp')

    def supports_fast_vmapvjp(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in VMAP_EXEMPTIONS:
            return Support.YES
        return self.no_opinfos_skip_test('test_vmapvjp_has_batch_rule')

    def supports_jvp(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        if self.name in JVP_EXEMPTIONS:
            return Support.YES
        if not self.has_opinfo():
            return Support.UNKNOWN
        if self.any_opinfo_attr('supports_autograd') and \
                not self.all_opinfo_attr('supports_forward_ad'):
            return Support.NO
        return self.no_opinfos_skip_test('test_jvp')

    def supports_jvpvjp(self):
        if self.name in FACTORY_FNS:
            return Support.YES
        exemptions = {
            # we have support (see OpInfo), testing artifact
            'nn.functional.dropout2d',
            'nn.functional.dropout',
            # exception: we dont even support double backward for this
            'nn.functional.hardswish',
            'bernoulli',  # this isn't differentiable
            'normal',  # not differentiable
        }
        if self.name in exemptions:
            return Support.YES
        return self.no_opinfos_skip_test('test_jvpvjp')

    def _supports_vmapjvp_base(self, test):
        if self.name in FACTORY_FNS:
            return Support.YES
        VMAPJVP_EXEMPTIONS = {
            'prod',  # dynamic (backward)
            'nn.functional.batch_norm',  # testing problem
            'normal',  # not actually problem, randomness testing artifact
            'bernoulli',  # not actually problem, randomness testing artifact
            'nn.functional.dropout2d',  # not actually problem, randomness testing artifact
            'nn.functional.dropout',  # not actually problem, randomness testing artifact
            # Not a problem.
            # It's just that the max_norm testing mutates inputs...
            # (we have our own functorch variant of the OpInfo without max_norm)
            'nn.functional.embedding',
        }
        if self.name in VMAPJVP_EXEMPTIONS:
            return Support.YES
        if not self.has_opinfo():
            return Support.UNKNOWN
        if self.any_opinfo_attr('supports_autograd') and \
                not self.all_opinfo_attr('supports_forward_ad'):
            return Support.NO
        return self.no_opinfos_skip_test(test)

    def supports_vmapjvp(self):
        return self._supports_vmapjvp_base('test_vmapjvpall')

    def supports_fast_vmapjvp(self):
        return self._supports_vmapjvp_base('test_vmapjvpall_has_batch_rule')


class OperatorSet:
    def __init__(self, operators):
        self.data = set(operators)

    @classmethod
    def from_names(cls, names):
        return OperatorSet([Operator(name) for name in names])

    @classmethod
    def from_top_ops_threshold(cls, torch_threshold, nn_fn_threshold):
        names = get_top_ops(torch_threshold, nn_fn_threshold)
        return cls.from_names(names)

    @classmethod
    def from_top125(cls):
        return cls.from_top_ops_threshold(100, 25)

    @classmethod
    def from_top160(cls):
        return cls.from_top_ops_threshold(107, 53)

    @classmethod
    def all(cls):
        dct = get_public_overridable_outplace_we_care_about()
        names = dct.keys()
        names_sanitized = []
        for n in names:
            torch_tensor = 'torch.Tensor.'
            torch_dot = 'torch.'
            if n.startswith(torch_tensor):
                names_sanitized.append(n[len(torch_tensor):])
            elif n.startswith(torch_dot):
                names_sanitized.append(n[len(torch_dot):])
            else:
                raise AssertionError()
        return cls.from_names(names_sanitized)

    def query(self, operator_method, filter=(Support.NO, Support.YES, Support.UNKNOWN)):
        result = {}
        for key in filter:
            result[key] = set([])
        for op in self.data:
            support_status = operator_method(op)
            if support_status in filter:
                result[support_status].add(op)
        return result

    def summary(self):
        checks = [
            'supports_vjp',
            'supports_vmap',
            'supports_fast_vmap',
            'supports_vmapvjp',
            'supports_fast_vmapvjp',
            'supports_jvp',
            'supports_vmapjvp',
            'supports_fast_vmapjvp',
            'supports_jvpvjp',
        ]
        result = ['test, yes, no, unknown']
        for check in checks:
            accessor = getattr(Operator, check)
            all_results = self.query(accessor)
            yes_amt = len(all_results[Support.YES])
            no_amt = len(all_results[Support.NO])
            unknown_amt = len(all_results[Support.UNKNOWN])
            result.append(f'{check}, {yes_amt}, {no_amt}, {unknown_amt}')
        return '\n'.join(result)


opset = OperatorSet.all()
has_no_opinfo = opset.query(Operator.has_opinfo, (False,))

print("=" * 30 + " Summary " + "=" * 30)
print(f'% of usages on github: {get_ops_percentage(99999, 99999)}')
print(opset.summary())

# sanity checks
result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)

print("=" * 30 + " Top 60 Summary " + "=" * 30)
print(f'% of usages on github: {get_ops_percentage(35, 25)}')
opset = OperatorSet.from_top_ops_threshold(35, 25)
# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# kresult = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
# kpprint.pprint(result)
# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# pprint.pprint(result)
print(opset.summary())

print("=" * 30 + " Top 125 Summary " + "=" * 30)
print(f'% of usages on github: {get_ops_percentage(100, 25)}')
opset = OperatorSet.from_top125()
# result = opset.query(Operator.supports_vmap, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
print("supports_vjp")
result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN))
pprint.pprint(result)
print("supports_jvp")
result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
pprint.pprint(result)
print("supports_vmapjvp")
result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
pprint.pprint(result)
print("supports_jvpvjp")
result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
pprint.pprint(result)
# result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# pprint.pprint(result)
print(opset.summary())

# print("=" * 30 + " Top 160 Summary " + "=" * 30)
# opset = OperatorSet.from_top160()
# result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
# pprint.pprint(result)
# print(opset.summary())

# Print list of everything in order
# all_ops = get_top_ops(999999, 999999, with_counts=True)
# for op, count in all_ops:
#     print(f'{op}, {count}')