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 generate_table(n1,dim1,n2,dim2,out_shape,n_dim):
    makdown_table = ""
    makdown_table += "| <!-- --> "*n_dim + "|\n" # header
    makdown_table += "|---"*n_dim + "|\n" # line

    # tensor 1
    for index in range(n_dim):
        if index < (n_dim - n1):
            makdown_table += "| "
        else:
            makdown_table += f"| {dim1[index - (n_dim- n1)]}"
    makdown_table +="|\n"
    # tensor 2
    for index in range(n_dim):
        if index < (n_dim - n2):
            makdown_table += "| "
        else:
            makdown_table += f"| {dim2[index - (n_dim- n2)]}"
    makdown_table +="|\n"
    return makdown_table

def predict(dim1, dim2):
    dim1 = sanitize_dimention(dim1)
    dim2 = sanitize_dimention(dim2)
    n1,dim1,n2,dim2,out_shape =  generate_example(dim1, dim2)
    code =  EXAMPLE_MD.format(
        n1=str(n1), dim1=str(dim1), n2=str(n2), dim2=str(dim2), out_shape=str(out_shape)
    )
    n1 = len(dim1)
    n2 = len(dim2)
    n_dim = max(n1,n2)
    table1 = generate_table(n1,dim1,n2,dim2,out_shape,n_dim)

    # TODO 
    # add code exmplanation here
    
    
    out = code
    out+= "\n# Step1 (alignment)" + "\n" + table1
    return out


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

demo.launch(debug=True)