# AUTOGENERATED! DO NOT EDIT! File to edit: diffdrive.ipynb (unless otherwise specified).

__all__ = ['axes', 'plane', 'ray', 'show_3d', 'read_image', 'conv2']

# Cell
import gtsam
import math
import PIL
import numpy as np
import plotly.express as px
import plotly.graph_objects as go

# Cell
def axes(length=1.5):
    """Add axes to a plotly figure"""
    return [
        go.Scatter3d(x=[0,length], y=[0,0], z=[0,0], mode="lines", line=dict(color="red")),
        go.Scatter3d(x=[0,0], y=[0,length], z=[0,0], mode="lines", line=dict(color="green")),
        go.Scatter3d(x=[0,0], y=[0,0], z=[0,length], mode="lines", line=dict(color="blue")),
        go.Scatter3d(x=[0], y=[0], z=[0], mode="markers", marker=dict(size=3, color="cyan")),
    ]

def plane(Z:float):
    return [
            go.Scatter3d(x=[-1.2, 1.2], y=[-1,-1], z=[Z, Z], mode="lines", line=dict(color="blue")),
            go.Scatter3d(x=[-1.2, 1.2], y=[ 1, 1], z=[Z, Z], mode="lines", line=dict(color="blue")),
            go.Scatter3d(x=[-1.2,-1.2], y=[-1, 1], z=[Z, Z], mode="lines", line=dict(color="blue")),
            go.Scatter3d(x=[ 1.2, 1.2], y=[-1, 1], z=[Z, Z], mode="lines", line=dict(color="blue")),
    ]

def ray(point3, F, color="orange"):
    X,Y,Z = point3
    x,y = F*X/Z, F*Y/Z
    return go.Scatter3d(x=[X,0,x], y=[Y,0,y], z=[Z,0,F], marker=dict(size=3, color=color))

def show_3d(fig):
    camera = dict(up=dict(x=0, y=-1, z=0), center=dict(x=0, y=0, z=0), eye=dict(x=1.7, y=-0.3, z=-1.1))
    fig.update_layout(scene_camera=camera, margin=dict(t=0, r=0, l=0, b=0), title="Pinhole Camera Model", showlegend=False)
    fig.show()


# Cell
import PIL
import requests

def read_image(image_name):
    """Read image from a the book repo"""
    url = f"https://raw.githubusercontent.com/gtbook/robotics/main/Figures5/{image_name}"
    return PIL.Image.open(requests.get(url, stream=True).raw)

# Cell
import torch

def conv2(input, filter):
    """Convolve input image of shape (iW,iH) with filter of shape (kW,kH)"""
    iW, iH = input.shape
    kW, kH = filter.shape
    _input = torch.reshape(input, (1, 1, iW, iH))
    _filter = torch.reshape(filter, (1, 1, kW, kH))
    _output = torch.conv2d(_input, _filter, padding='same')
    return torch.reshape(_output, (iW, iH))