import matplotlib.pyplot as plt
import torch
import matplotlib
matplotlib.use('Agg')
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.unsqueeze(0)
    if len(w.shape) == 1:
        w = w.unsqueeze(0)

    l2 = fit_line(pts, w)

    # project some points along the estimated line
    p3 = l2.point_at(torch.tensor(-10.0))
    p4 = l2.point_at(torch.tensor(10.0))
    X = torch.stack((p3, p4)).squeeze().detach().numpy()
    X_pts = pts.squeeze().detach().numpy()

    fig, ax = plt.subplots()
    ax.plot(X_pts[:, 0], X_pts[:, 1], 'ro')
    ax.plot(X[:, 0], X[:, 1])
    ax.set_xlim(X_pts[:, 0].min() - 1, X_pts[:, 0].max() + 1)
    ax.set_ylim(X_pts[:, 1].min() - 1, X_pts[:, 1].max() + 1)
    return fig

inputs = [
    gr.Slider(0.0, 10.0, value=0.0, label="Point 1 X"),
    gr.Slider(0.0, 10.0, value=0.0, label="Point 1 Y"),
    gr.Slider(0.0, 10.0, value=10.0, label="Point 2 X"),
    gr.Slider(0.0, 10.0, value=10.0, label="Point 2 Y"),
]

outputs = gr.Plot()

examples = [
    [0.0, 0.0, 10.0, 10.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()