import cv2
import numpy as np

class StereoCalc:

    def __init__(self, alpha=0.5, cal_file='dspl_qhd.xml'):
        # just like before we specify an enum flag, but this time it is
        # FILE_STORAGE_READ
        cv_file = cv2.FileStorage(cal_file, cv2.FILE_STORAGE_READ)
        # for some reason __getattr__ doesn't work for FileStorage object in python
        # however in the C++ documentation, getNode, which is also available,
        # does the same thing
        # note we also have to specify the type to retrieve other wise we only get a
        # FileNode object back instead of a matrix
        self.alpha = alpha
        self.B = cv_file.getNode("B").mat()
        # convert B to vector
        self.B = np.squeeze(self.B)
        self.k1 = cv_file.getNode("k1").mat()
        self.k2 = cv_file.getNode("k2").mat()
        self.Fmatrix = cv_file.getNode("Fmatrix").mat()
        self.R12 = cv_file.getNode("R12").mat()

        cv_file.release()

    def calc_position(self, left_x, left_y, right_x, right_y):

        imagepoints1d = np.transpose(np.array([left_x, left_y, 1.0]))
        imagepoints2d = np.transpose(np.array([right_x, right_y, 1.0]))

        R = np.transpose(self.R12)
        thetarec = np.arctan2(-self.B[2], self.B[0])
        ctr_ = np.cos(thetarec)
        str_ = np.sin(thetarec)

        Rrectheta = np.array([
            [ctr_, np.array(0.0), -str_],
            [0.0, 1.0, 0.0],
            [str_, 0.0, ctr_]
        ])

        Bprime = Rrectheta @ self.B

        psirec = np.arctan2(Bprime[1], Bprime[0])
        spsir = np.sin(psirec)
        cpsir = np.cos(psirec)

        Rrecpsi = np.array([
            [cpsir, spsir, 0.0],
            [-spsir, cpsir, 0.0],
            [0.0, 0.0, 1.0]
        ])

        Rrec = Rrecpsi @ Rrectheta
        Bcheck = Rrec * self.B

        R1 = Rrec

        # Left image
        pud1 = self.k1 @ imagepoints1d
        vv1 = np.transpose(np.array([pud1[0], pud1[1], 1.0]))
        rad2 = vv1[0]**2 + vv1[1]**2

        vv3 = (1 + self.alpha * rad2) * vv1
        vv3[2] = 1.0

        vv2 = np.linalg.inv(self.k1) @ vv3

        imagepoints1 = vv2

        # right image
        pud2 = self.k2 @ imagepoints2d
        vv4 = np.transpose(np.array([pud2[0], pud2[1], 1.0]))
        rad2 = vv4[0]**2 + vv4[1]**2

        vv6 = (1 + self.alpha * rad2) * vv4
        vv6[2] = 1.0

        vv5 = np.linalg.inv(self.k2) @ vv6

        imagepoints2 = vv5

        p1 = imagepoints1
        p2 = imagepoints2

        epiline = self.Fmatrix @ p1
        epiline = epiline / (epiline[0]**2 + epiline[1]**2)

        doterror_inpixels = np.transpose(p2) @ epiline

        vv1 = np.linalg.inv(self.k2) @ R @ self.k1 @ p1
        vv2 = np.linalg.inv(self.k2) @ R @ self.B

        print(vv2)

        M = np.array([
            [vv1[0], -vv2[0]],
            [vv1[1], -vv2[1]],
            [vv1[2], -vv2[2]]
        ])

        print(M)

        MtM = np.transpose(M) @ M
        vv0 = np.linalg.inv(MtM) @ np.transpose(M) @ p2

        rho1 = vv0[0] / vv0[1]

        Position = rho1 * self.k1 @ p1

        Position_perp = R1 @ Position

        return Position_perp, doterror_inpixels

    def calibration_info(self):
        print("B:")
        print(str(self.B))
        print("k1:")
        print(str(self.k1))
        print("k2:")
        print(str(self.k2))
        print("FMatrix:")
        print(str(self.Fmatrix))
        print("R12:")
        print(str(self.R12))


if __name__=="__main__":

    s_calc = StereoCalc()
    s_calc.calibration_info()

    # test some points
    print(s_calc.calc_position(100,100,120,145))