import torch
import pypose as pp
from pypose.optim import LM
from pypose.optim.kernel import Huber
from pypose.optim.solver import Cholesky
from pypose.optim.strategy import TrustRegion
from pypose.optim.corrector import FastTriggs
from pypose.optim.scheduler import StopOnPlateau

EDN2NED = pp.from_matrix(torch.tensor(
       [[0., 0., 1., 0.],
        [1., 0., 0., 0.],
        [0., 1., 0., 0.],
        [0., 0., 0., 1.]], dtype=torch.float32),
    pp.SE3_type)

NED2EDN = EDN2NED.Inv()

class PoseGraph(torch.nn.Module):
    def __init__(self, **kwargs) -> None:
        super().__init__()
        self.register_buffer("K", torch.tensor([
                [320., 0.  , 320.],
                [0.  , 320., 240.],
                [0.  , 0.  , 1.  ]
            ]))
        self.register_buffer("pts1", torch.empty(1000,2))        # N x 2, uv coordinate
        self.register_buffer("pts2", torch.empty(1000,2))        # N x 2, uv coordinate
        self.register_buffer("NED2EDN", NED2EDN)
        self.register_buffer("EDN2NED", EDN2NED)
        self.fx, self.fy = self.K[0, 0], self.K[1, 1]
        self.cx, self.cy = self.K[0, 2], self.K[1, 2]

        self.T = pp.Parameter(pp.identity_SE3())
        self.depth = torch.nn.Parameter(torch.empty(1000,))

    @property
    def point3d(self) -> torch.Tensor:
        return pp.geometry.pixel2point(self.pts1, self.depth, self.K)

    def forward(self) -> torch.Tensor:
        return pp.reprojerr(
            self.point3d, self.pts2, self.K, (self.NED2EDN @ self.T @ self.EDN2NED).Inv() , reduction='none'
        )

graph = PoseGraph()
graph.load_state_dict(torch.load("./Results/step49.pkl"))
graph = graph.cuda()

kernel = Huber(delta=0.1)
corrector = FastTriggs(kernel)
optimizer = LM(graph, solver=Cholesky(),
                        strategy=TrustRegion(radius=1e3),
                        kernel=kernel,
                        corrector=corrector,
                        min=1e-8,
                        vectorize=True)
scheduler = StopOnPlateau(optimizer, steps=10,
                                        patience=4,
                                        decreasing=1e-6,
                                        verbose=True)
while scheduler.continual():
    loss = optimizer.step(input=())
    scheduler.step(loss)