import gradio as gr
import torch
import torch.nn as nn
import time
import math

class Lreg(nn.Module):
    def __init__(self):
        super(Lreg,self).__init__()
        self.l = nn.Linear(1,1)

    def forward(self,x):
        y = self.l(x)
        return y

model = Lreg()
crit = nn.MSELoss()
optim = torch.optim.SGD(model.parameters(),lr=0.01)

def weightsbias(weight,bias,predicted,progress=gr.Progress()):
    weight = float(weight)
    bias = float(bias)
    predicted = float(predicted)
    x = torch.randn(1000)
    y = (x*weight)+bias
    progress(0,desc="Starting")
    model.train()
    for i in progress.tqdm(range(1,100)):
        for idx,ele in enumerate(x):
            ele = ele.unsqueeze(0)
            out = model(ele)
            optim.zero_grad()
            loss = crit(out,y[idx].unsqueeze(0))
            loss.backward()
            optim.step()
    model.eval()
    t = torch.tensor([predicted],dtype=torch.float32)
    rout = model(t)
    c = model.state_dict()
    return [rout.item(),c['l.weight'].item(),c['l.bias'].item()]


with gr.Blocks() as iface:
    weights = gr.Textbox(label="Weights")
    bias = gr.Textbox(label="Bias")
    predicted = gr.Textbox(label="Prediction Number")
    output = gr.Number(label="Output Predicted")
    w = gr.Number(label="Calculated Weight ")
    b = gr.Number(label="Calculated Bias ")
    train = gr.Button("Train")
    train.click(fn=weightsbias, inputs=[weights,bias,predicted], outputs=[output,w,b])

iface.queue(concurrency_count=10).launch()