import gradio as gr
import torch


EXAMPLE_MD = """
```python
import torch 

t1 = torch.arange({n1}).view({dim1})

t2 = torch.arange({n2}).view({dim2})

(t1 @ t2).shape = {out_shape} 

```
"""


def generate_example(dim1: list, dim2: list):
    n1 = 1
    n2 = 1
    for i in dim1:
        n1 *= i
    for i in dim2:
        n2 *= i

    t1 = torch.arange(n1).view(dim1)
    t2 = torch.arange(n2).view(dim2)
    try:
        out_shape = list((t1 @ t2).shape)
    except RuntimeError:
        out_shape = "error"

    return n1,dim1,n2,dim2,out_shape


def sanitize_dimention(dim):
    if dim is None:
        gr.Error("one of the dimentions is empty, please fill it")
    if "[" in dim:
        dim = dim.replace("[", "")
    if "]" in dim:
        dim = dim.replace("]", "")
    if "," in dim:
        dim = dim.replace(",", " ").strip()
        out = [int(i.strip()) for i in dim.split()]
    else:
        out = [int(dim.strip())]
    if 0 in out:
        gr.Error(
            "Found the number 0 in one of the dimensions which is not allowed, consider using 1 instead"
        )
    return out


def predict(dim1, dim2):
    dim1 = sanitize_dimention(dim1)
    dim2 = sanitize_dimention(dim2)
    n1,dim1,n2,dim2,out_shape =  generate_example(dim1, dim2)

    # TODO 
    # add code exmplanation here
    
    return EXAMPLE_MD.format(
        n1=str(n1), dim1=str(dim1), s2=str(n2), dim2=str(dim2), out_shape=str(out_shape)
    )


demo = gr.Interface(
    predict,
    inputs=["text", "text"],
    outputs=["markdown"],
    examples=[["1,2,3", "5,3,7"], ["1,2,3", "5,2,7"]],
)

demo.launch(debug=True)