"""Spherical projection and gradient-domain distance fusion for an isolated trial."""

import numpy as np
from scipy.ndimage import map_coordinates
from scipy.sparse import coo_matrix, diags
from scipy.sparse.linalg import spsolve


def panorama_directions(width, height):
    u, v = np.meshgrid((np.arange(width) + 0.5) / width,
                       (np.arange(height) + 0.5) / height)
    return uv_directions(u, v)


def uv_directions(u, v):
    theta, phi = 2 * np.pi * u, np.pi * v
    return np.stack((np.sin(phi) * np.cos(theta),
                     np.sin(phi) * np.sin(theta), np.cos(phi)), axis=-1)


def camera_bases():
    ratio = (1 + np.sqrt(5)) / 2
    vertices = []
    for a in (-1, 1):
        for b in (-ratio, ratio):
            vertices.extend(((0, a, b), (a, b, 0), (b, 0, a)))
    bases = []
    for forward in np.asarray(vertices):
        forward = forward / np.linalg.norm(forward)
        right = np.cross(forward, (0, 0, 1))
        right /= np.linalg.norm(right)
        down = np.cross(forward, right)
        bases.append(np.stack((right, down, forward)))
    return bases


def view_rays(size):
    x, y = np.meshgrid(2 * (np.arange(size) + 0.5) / size - 1,
                       2 * (np.arange(size) + 0.5) / size - 1)
    return np.stack((x, y, np.ones_like(x)), axis=-1)


def sample_panorama(image, directions):
    directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True)
    u = np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1
    v = np.arccos(np.clip(directions[..., 2], -1, 1)) / np.pi
    height, width = image.shape[:2]
    padded = np.concatenate((image[:, -1:], image, image[:, :1]), axis=1)
    coords = (np.clip(v * height - 0.5, 0, height - 1), u * width + 0.5)
    if image.ndim == 2:
        return map_coordinates(padded, coords, order=1, mode="nearest")
    return np.stack([map_coordinates(padded[..., c], coords, order=1,
                                    mode="nearest") for c in range(image.shape[2])], -1)


def room_reference(directions):
    """Trace a furnished rectangular room from a camera 1.6 m above its floor."""
    directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True)
    bounds = np.where(directions >= 0, (3.0, 4.0, 1.4), (-3.0, -4.0, -1.6))
    hits = bounds / np.where(np.abs(directions) > 1e-12, directions, 1e-12)
    hits = np.where(hits > 0, hits, np.inf)
    surface = np.argmin(hits, axis=-1)
    distance = np.min(hits, axis=-1)
    point = directions * distance[..., None]
    colors = np.array(((0.68, 0.76, 0.80), (0.81, 0.73, 0.61), (0.9, 0.88, 0.82)))
    color = colors[surface].copy()
    floor = (surface == 2) & (directions[..., 2] < 0)
    tiles = (np.floor(point[..., 0] * 2) + np.floor(point[..., 1] * 2)) % 2
    color[floor] = np.stack((0.40 + 0.08 * tiles, 0.27 + 0.06 * tiles,
                            0.15 + 0.04 * tiles), -1)[floor]
    lines = (np.abs(np.sin(point[..., 0] * np.pi * 2)) < 0.04) | (
        np.abs(np.sin(point[..., 1] * np.pi * 2)) < 0.04)
    color[floor & lines] *= 0.6
    window = (surface == 1) & (directions[..., 1] > 0) & (
        np.abs(point[..., 0] + 0.4) < 1.05) & (np.abs(point[..., 2] - 0.25) < 0.7)
    color[window] = (0.48, 0.72, 0.88)
    frame = window & ((np.abs(point[..., 0] + 0.4) < 0.045) |
                      (np.abs(point[..., 2] - 0.25) < 0.04))
    color[frame] = 0.93
    for lower, upper, base_color in (
        ((0.8, 1.0, -1.6), (2.0, 2.2, -0.6), (0.28, 0.43, 0.34)),
        ((-2.8, -1.8, -1.6), (-2.1, 0.5, 0.7), (0.43, 0.26, 0.15)),
    ):
        safe = np.where(np.abs(directions) > 1e-12, directions, 1e-12)
        a, b = np.asarray(lower) / safe, np.asarray(upper) / safe
        near, far = np.max(np.minimum(a, b), axis=-1), np.min(np.maximum(a, b), axis=-1)
        visible = (near > 0) & (near < far) & (near < distance)
        distance[visible] = near[visible]
        color[visible] = base_color
    color *= (0.75 + 0.25 * np.clip(directions[..., 2] + 0.6, 0, 1))[..., None]
    return distance, color


def warp_distances(distances, masks, bases, width, height):
    directions = panorama_directions(width, height)
    logs, validities = [], []
    for distance, mask, basis in zip(distances, masks, bases):
        rays = directions @ basis.T
        projected = rays[..., :2] / np.maximum(rays[..., 2:], 1e-8)
        valid = (rays[..., 2] > 0) & (np.abs(projected) < 0.98).all(axis=-1)
        size = distance.shape[0]
        coords = ((projected[..., 1] + 1) * size / 2 - 0.5,
                  (projected[..., 0] + 1) * size / 2 - 0.5)
        valid &= map_coordinates(mask.astype(float), coords, order=0, mode="nearest") > 0.5
        logs.append(map_coordinates(np.log(np.maximum(distance, 1e-6)), coords,
                                    order=1, mode="nearest"))
        validities.append(valid)
    return np.asarray(logs), np.asarray(validities)


def merge_distances(distances, masks, bases, width, height):
    """Fuse log-distance derivatives with a cyclic longitude and a scale anchor.

    Algorithm reference: microsoft/MoGe, moge/utils/panorama.py.
    This independent implementation solves normal equations on a small trial grid.
    """
    logs, valid = warp_distances(distances, masks, bases, width, height)
    count = width * height
    index = np.arange(count).reshape(height, width)
    start = np.concatenate((index.ravel(), index[:-1].ravel()))
    end = np.concatenate((np.roll(index, -1, axis=1).ravel(), index[1:].ravel()))
    rows = np.arange(len(start))
    gradient = coo_matrix((np.concatenate((-np.ones(len(start)), np.ones(len(start)))),
                          (np.tile(rows, 2), np.concatenate((start, end)))),
                         shape=(len(start), count)).tocsr()
    laplacian = -(gradient.T @ gradient).tocsr()
    edge_valid = valid.reshape(12, -1)[:, start] & valid.reshape(12, -1)[:, end]
    grad_samples = logs.reshape(12, -1)[:, end] - logs.reshape(12, -1)[:, start]
    grad_target = np.sum(grad_samples * edge_valid, axis=0) / np.maximum(edge_valid.sum(0), 1)
    lap_valid = valid.copy()
    lap_valid &= np.roll(valid, 1, axis=2) & np.roll(valid, -1, axis=2)
    lap_valid[:, 1:] &= valid[:, :-1]
    lap_valid[:, :-1] &= valid[:, 1:]
    lap_valid = lap_valid.reshape(12, -1)
    lap_samples = np.asarray((laplacian @ logs.reshape(12, -1).T).T)
    lap_target = np.sum(lap_samples * lap_valid, axis=0) / np.maximum(lap_valid.sum(0), 1)
    g = gradient[edge_valid.any(0)]
    l = laplacian[lap_valid.any(0)]
    coverage = valid.any(0)
    if not coverage.any():
        raise ValueError("No valid panorama predictions")
    baseline = np.sum(logs * valid, axis=0) / np.maximum(valid.sum(0), 1)
    matrix = g.T @ g + l.T @ l + diags(np.full(count, 1e-7))
    target = g.T @ grad_target[edge_valid.any(0)] + l.T @ lap_target[lap_valid.any(0)]
    solution = spsolve(matrix.tocsc(), target + 1e-7 * baseline.ravel()).reshape(height, width)
    solution += np.median(baseline[coverage] - solution[coverage])
    return np.exp(solution).astype(np.float32), coverage