import numpy as np


config = {
    "data_dir": r"C:\Users\admin\Projects\autopilot\datasets\ya_go_maps\images",
    "image_size": (256, 256),
    "batch_size": 32,
    "train_split": 0.8,
    "num_workers": 0,
    "epochs": 100,
    "learning_rate": 2e-4,
    "dropout_rate": 0.3,
    "backbone": "resnet18",
    "output_dir": r"C:\Users\admin\Projects\autopilot\models\SiaN\runs",
}


def get_camera_matrix(w, h):
    return np.array([[w / 2, 0, w / 2], [0, h / 2, h / 2], [0, 0, 1]], dtype=np.float32)


def generate_random_homography_params(angle_range=10, translation_range=0.1, scale_range=(0.9, 1.1)):
    scale = np.random.uniform(*scale_range)
    tx = np.random.uniform(-translation_range, translation_range)
    ty = np.random.uniform(-translation_range, translation_range)
    rx = np.radians(np.random.uniform(-angle_range, angle_range))
    ry = np.radians(np.random.uniform(-angle_range, angle_range))
    rz = np.radians(np.random.uniform(-angle_range, angle_range))
    return np.array([rx, ry, rz, tx, ty, scale])


def homography_params_to_matrix(params, K):
    rx, ry, rz, tx, ty, scale = params
    cy, sy = np.cos(rz), np.sin(rz)
    cp, sp = np.cos(ry), np.sin(ry)
    cr, sr = np.cos(rx), np.sin(rx)
    Rz = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]], dtype=np.float32)
    Ry = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]], dtype=np.float32)
    Rx = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]], dtype=np.float32)
    T = np.array([[1, 0, tx], [0, 1, ty], [0, 0, scale]], dtype=np.float32)
    return K @ Rx @ Ry @ Rz @ T @ np.linalg.inv(K)


def matrix_to_homography_params(H, K):
    K_inv = np.linalg.inv(K)
    E = K_inv @ H @ K
    scale = np.sqrt(np.linalg.det(E[:2, :2]))
    R = E[:2, :2] / scale
    tx, ty = E[0, 2], E[1, 2]
    rz = np.arctan2(R[1, 0], R[0, 0])
    r20, r21 = E[2, 0], E[2, 1]
    ry = np.arctan2(r20, r21)
    rx = np.arctan2(-E[1, 2], E[1, 1])
    return np.array([rx, ry, rz, tx, ty, scale], dtype=np.float32)