# Example showing how to fit a 2d line with kornia / pytorch
import matplotlib.pyplot as plt
import torch
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import gradio as gr
from kornia.geometry.line import ParametrizedLine, fit_line


def inference(point1, point2, point3, point4):
    std = 1.2  # standard deviation for the points
    num_points = 50  # total number of points

    # create a baseline
    p0 = torch.tensor([point1, point2], dtype=torch.float32)
    p1 = torch.tensor([point3, point4], dtype=torch.float32)
    l1 = ParametrizedLine.through(p0, p1)
    
    # sample some points and weights
    pts, w = [], []
    for t in torch.linspace(-10, 10, num_points):
        p2 = l1.point_at(t)
        p2_noise = torch.rand_like(p2) * std
        p2 += p2_noise
        pts.append(p2)
        w.append(1 - p2_noise.mean())
    pts = torch.stack(pts)
    w = torch.stack(w)
    if len(pts.shape) == 2:
        pts = pts[None]
    if len(w.shape) == 1:
        w = w[None]
    
    l2 = fit_line(pts, w)
    
    # project some points along the estimated line
    p3 = l2.point_at(-10)
    p4 = l2.point_at(10)

    X = torch.stack((p3, p4)).detach().numpy()
    X_pts = pts.detach().numpy()

    fig = plt.figure()
    plt.plot(X_pts[:, 0], X_pts[:, 1], 'ro')
    plt.plot(X[:, 0], X[:, 1])
    return fig

inputs = [
    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 1"),
    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 2"),
    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 3"),
    gr.inputs.Slider(0.0, 10.0, default=0.0, label="Point 4"),
]
outputs = gr.Plot()

examples = [
    [0.0, 0.0, 1.0, 1.0],
]

title = 'Line Fitting'

demo = gr.Interface(
    fn=inference,
    inputs=inputs,
    outputs=outputs,
    title=title,
    cache_examples=True,
    theme='huggingface',
    live=True,
    examples=examples,

)
demo.launch()