import numpy as np
import pymc as pm
from pytensor.tensor import TensorVariable
import pytensor
import pytensor.tensor as T
from pymc.pytensorf import collect_default_updates


N_CUSTOMERS = 10


def dist(
    alpha: TensorVariable,
    size: TensorVariable,
) -> TensorVariable:

    # Find out when new tables are created
    first_roll = pm.Uniform.dist(size=size)
    customers = T.cumsum(T.ones(size)) - 1
    newtable_chances = 1.0 * alpha / (alpha + customers)
    newtables = first_roll < newtable_chances

    # Function to assign old tables whenever a new table is not created
    def assign_table(isnew, index, prior):
        counts = T.unique(prior[:index], return_counts=True)[1]
        probs = counts / counts.sum()
        choice = pm.Categorical.dist(p=probs) + 1
        zeros_subtensor = prior[index]
        switch = isnew.astype("int")
        value1 = choice * (1 - switch)  # Assign an old table
        value2 = (prior[:index].max() + 1) * switch  # Add a new table
        output = T.set_subtensor(zeros_subtensor, value1 + value2)
        return output, collect_default_updates(output)

    indexes = pytensor.tensor.arange(newtables.shape[0])
    result, _ = pytensor.scan(
        fn=assign_table,
        outputs_info=T.ones_like(newtables, dtype='int64'),
        sequences=[newtables[1:], indexes[1:]],
    )
    return result[-1].astype("int")


print(f"{N_CUSTOMERS} customers with low alpha:", dist(1, N_CUSTOMERS).eval())
print(f"{N_CUSTOMERS} customers with very high alpha:", dist(100, N_CUSTOMERS).eval())

alpha = T.scalar(dtype='int64')
size = T.vector(shape=(0,), dtype='int64')
dist(alpha, size)

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[4], line 1
----> 1 dist(alpha, size)

Cell In[1], line 36, in dist(alpha, size)
     33     return output, collect_default_updates(output)
     35 indexes = pytensor.tensor.arange(newtables.shape[0])
---> 36 result, _ = pytensor.scan(
     37     fn=assign_table,
     38     outputs_info=T.ones_like(newtables, dtype='int64'),
     39     sequences=[newtables[1:], indexes[1:]],
     40 )
     41 return result[-1].astype("int")

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/scan/basic.py:744, in scan(fn, sequences, outputs_info, non_sequences, n_steps, truncate_gradient, go_backwards, mode, name, profile, allow_gc, strict, return_list)
    738     arg.name = init_out["initial"].name + "[t-1]"
    740 # We need now to allocate space for storing the output and copy
    741 # the initial state over. We do this using the expand function
    742 # defined in scan utils
    743 sit_sot_scan_inputs.append(
--> 744     expand_empty(
    745         unbroadcast(shape_padleft(actual_arg), 0),
    746         actual_n_steps,
    747     )
    748 )
    750 sit_sot_inner_slices.append(actual_arg)
    751 if i in return_steps:

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/scan/utils.py:235, in expand_empty(tensor_var, size)
    233 shapes = [tensor_var.shape[x] for x in range(tensor_var.ndim)]
    234 new_shape = [size + shapes[0]] + shapes[1:]
--> 235 empty = AllocEmpty(tensor_var.dtype)(*new_shape)
    237 ret = set_subtensor(empty[: shapes[0]], tensor_var)
    238 ret.tag.nan_guard_mode_check = False

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/op.py:295, in Op.__call__(self, *inputs, **kwargs)
    253 r"""Construct an `Apply` node using :meth:`Op.make_node` and return its outputs.
    254 
    255 This method is just a wrapper around :meth:`Op.make_node`.
   (...)
    292 
    293 """
    294 return_list = kwargs.pop("return_list", False)
--> 295 node = self.make_node(*inputs, **kwargs)
    297 if config.compute_test_value != "off":
    298     compute_test_value(node)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/tensor/basic.py:3835, in AllocEmpty.make_node(self, *_shape)
   3834 def make_node(self, *_shape):
-> 3835     _shape, static_shape = infer_static_shape(_shape)
   3836     otype = TensorType(dtype=self.dtype, shape=static_shape)
   3837     output = otype()

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/tensor/basic.py:1386, in infer_static_shape(shape)
   1382     raise TypeError(f"Shapes must be scalar integers; got {s_as_str}")
   1384 sh = [check_type(as_tensor_variable(s, ndim=0)) for s in shape]
-> 1386 shape_fg = FunctionGraph(
   1387     outputs=sh,
   1388     features=[ShapeFeature()],
   1389     clone=True,
   1390 )
   1391 folded_shape = rewrite_graph(shape_fg, custom_rewrite=topo_constant_folding).outputs
   1392 static_shape = tuple(
   1393     s.data.item() if isinstance(s, Constant) else None for s in folded_shape
   1394 )

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/fg.py:157, in FunctionGraph.__init__(self, inputs, outputs, features, clone, update_mapping, **clone_kwds)
    154     self.add_input(in_var, check=False)
    156 for output in outputs:
--> 157     self.add_output(output, reason="init")
    159 self.profile = None
    160 self.update_mapping = update_mapping

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/fg.py:167, in FunctionGraph.add_output(self, var, reason, import_missing)
    165 """Add a new variable as an output to this `FunctionGraph`."""
    166 self.outputs.append(var)
--> 167 self.import_var(var, reason=reason, import_missing=import_missing)
    168 self.clients[var].append(("output", len(self.outputs) - 1))

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/fg.py:308, in FunctionGraph.import_var(self, var, reason, import_missing)
    306 # Imports the owners of the variables
    307 if var.owner and var.owner not in self.apply_nodes:
--> 308     self.import_node(var.owner, reason=reason, import_missing=import_missing)
    309 elif (
    310     var.owner is None
    311     and not isinstance(var, AtomicVariable)
    312     and var not in self.inputs
    313 ):
    314     from pytensor.graph.null_type import NullType

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/fg.py:389, in FunctionGraph.import_node(self, apply_node, check, reason, import_missing)
    387         self.variables.add(input)
    388     self.add_client(input, (node, i))
--> 389 self.execute_callbacks("on_import", node, reason)

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/graph/fg.py:727, in FunctionGraph.execute_callbacks(self, name, *args, **kwargs)
    725         continue
    726     tf0 = time.perf_counter()
--> 727     fn(self, *args, **kwargs)
    728     self.execute_callbacks_times[feature] += time.perf_counter() - tf0
    729 self.execute_callbacks_time += time.perf_counter() - t0

File ~/miniconda3/envs/pymc_env/lib/python3.11/site-packages/pytensor/tensor/rewriting/shape.py:569, in ShapeFeature.on_import(self, fgraph, node, reason)
    562 for i, d in enumerate(sh):
    563     # Note: we ignore any shape element that is not typed (i.e.,
    564     # does not have a 'dtype' attribute). This means there may
    565     # still remain int elements that are int32 on 32-bit platforms,
    566     # but this works with `local_useless_subtensor`, so for now we
    567     # keep it this way. See #266 for a better long-term fix.
    568     if getattr(d, "dtype", "int64") != "int64":
--> 569         assert d.dtype in discrete_dtypes, (node, d.dtype)
    570         assert str(d.dtype) != "uint64", node
    571         new_shape += sh[len(new_shape) : i + 1]

AssertionError: (CumOp{None, add}(Alloc.0), 'float64')