import torch
import numpy as np
import random
import pdb


def recover_metric_depth(pred, gt):
    if type(pred).__module__ == torch.__name__:
        pred = pred.cpu().numpy()
    if type(gt).__module__ == torch.__name__:
        gt = gt.cpu().numpy()
    gt = gt.squeeze()
    pred = pred.squeeze()
    mask = (gt > 1e-8) & (pred > 1e-8)

    gt_mask = gt[mask]
    pred_mask = pred[mask]
    a, b = np.polyfit(pred_mask, gt_mask, deg=1)
    #print("scale {}".format(a))
    #print("shift {}".format(b))
    pred_metric = a * pred + b
    return pred_metric


def recover_scale_shift(pred, gt):
    if type(pred).__module__ == torch.__name__:
        pred = pred.cpu().numpy()
    if type(gt).__module__ == torch.__name__:
        gt = gt.cpu().numpy()
    gt = gt.squeeze()
    pred = pred.squeeze()
    mask = (gt > 1e-8) & (pred > 1e-8)

    gt_mask = gt[mask]
    pred_mask = pred[mask]
    a, b = np.polyfit(pred_mask, gt_mask, deg=1)
    #print("scale {}".format(a))
    #print("shift {}".format(b))
    pred_metric = a * pred + b
    return a, b


def normal_from_depth(depth, focal_length):
    """
    Compute surface normal from depth map
    """
    def normalization(data):
        mo_chang = np.sqrt(
            np.multiply(data[:, :, 0], data[:, :, 0])
            + np.multiply(data[:, :, 1], data[:, :, 1])
            + np.multiply(data[:, :, 2], data[:, :, 2])
        )
        mo_chang = np.dstack((mo_chang, mo_chang, mo_chang))
        return data/mo_chang

    width = 1024
    height = 768
    K = np.array([
        [focal_length, 0, width / 2],
        [0, focal_length, height / 2],
        [0, 0, 1],
    ])
    x, y = np.meshgrid(np.arange(0, width), np.arange(0, height))
    x = x.reshape([-1])
    y = y.reshape([-1])
    xyz = np.vstack((x, y, np.ones_like(x)))
    pts_3d = np.dot(np.linalg.inv(K), xyz*depth.reshape([-1]))
    pts_3d_world = pts_3d.reshape((3, height, width))
    f = pts_3d_world[:, 1:height-1, 2:width] - pts_3d_world[:, 1:height-1, 1:width-1]
    t = pts_3d_world[:, 2:height, 1:width-1] - pts_3d_world[:, 1:height-1, 1:width-1]
    normal_map = np.cross(f, t, axisa=0, axisb=0)
    normal_map = normalization(normal_map)

    # pad it to 400x600
    normal_map = np.pad(normal_map, ((1, 1), (1, 1), (0, 0)), 'edge')

    return normal_map


def estimate_equation(points):
    # use linear system
    if points.shape[0] == points.shape[1]:
        normal = np.linalg.solve(points, np.ones(points.shape[0]))
    else:
        normal = np.linalg.lstsq(points, np.ones(points.shape[0]), rcond=None)[0]

    offset = -1 / np.linalg.norm(normal)
    normal /= np.linalg.norm(normal)
    if normal[2] > 0:  # make sure n_z is negative
        normal = -normal
    return normal, offset


def fit_plane(points, thres=0.01, debug=False):
    final_inliers = []
    final_equation = np.array([0, 0, 0])
    final_offset = 0.0

    for i in range(200):
        # Sample 3 points
        sample = np.array([random.choice(points) for _ in range(3)])

        # n_x * x + n_y * y + n_z * z = offset
        try:
            equation, offset = estimate_equation(sample)
        except np.linalg.LinAlgError:  # typically a singular matrix
            continue
        error = points @ equation - offset
        inliers = points[error < thres]

        if debug:
            print('n and inliers: {} {}'.format(equation, len(inliers)))
        if len(inliers) > len(final_inliers):
            final_inliers = inliers
            final_offset = offset
            final_equation = equation

    equation, offset = estimate_equation(final_inliers)

    if debug:
        print('Best Results:')
        print(final_equation)
        print(len(final_inliers))

        print('Final Fit:')
        print(equation)
        print(offset)

    return equation, offset