"""Template placement must follow face pose, not unequal left/right bangs."""
import numpy as np
import pytest

from hairline import constants as C
from hairline.hairline_2d import (
    align_hairline_for_template,
    clamp_hairline_to_silhouette,
)


def _scene(angle=0.0, shape=(600, 900)):
    h, w = shape
    scale = np.array([w, h])
    lm = np.zeros((468, 3), np.float32)
    xy = np.full((468, 2), [450.0, 300.0])
    for i, (left, right) in enumerate(C.MP_TOP_ANCHOR_PAIRS):
        xy[left] = [300.0 + i * 12, 210.0 - i * 8]
        xy[right] = [600.0 - i * 12, 210.0 - i * 8]
    xy[10] = [450, 150]
    xy[152] = [450, 450]
    theta = np.deg2rad(angle)
    rotation = np.array([[np.cos(theta), -np.sin(theta)],
                         [np.sin(theta), np.cos(theta)]])
    xy = (xy - [450, 300]) @ rotation.T + [450, 300]
    lm[:, :2] = xy / scale
    up = rotation @ [0.0, -1.0]
    anchors = lm[C.MP_TOP_ANCHORS, :2] * scale
    distances = np.full(C.N_ANCHORS, 50.0)
    indices = {k: i for i, k in enumerate(C.MP_TOP_ANCHORS)}
    distances[indices[67]] = 80.0
    distances[indices[297]] = 20.0
    hairline = ((anchors + distances[:, None] * up) / scale).astype(np.float32)
    return lm, hairline, np.ones(C.N_ANCHORS, bool), anchors, up


@pytest.mark.parametrize("angle,shape", [(0, (600, 900)), (23, (600, 900)),
                                         (-17, (1200, 900))])
def test_uneven_bangs_do_not_shear_template(angle, shape):
    lm, raw, valid, anchors, up = _scene(angle, shape)
    raw_copy, lm_copy = raw.copy(), lm.copy()
    aligned = align_hairline_for_template(lm, raw, valid, shape)
    displacement = aligned * shape[::-1] - anchors
    # Every anchor extends 50 px along the rotated face axis, even when
    # one side's hair boundary is 60 px higher than its counterpart.
    np.testing.assert_allclose(displacement, np.tile(50 * up, (C.N_ANCHORS, 1)), atol=1e-4)
    np.testing.assert_array_equal(raw, raw_copy)
    np.testing.assert_array_equal(lm, lm_copy)


def test_anatomical_pairs_are_not_reversed_anchor_order():
    shape = (600, 900)
    lm, raw, valid, anchors, up = _scene(shape=shape)
    # Endpoints of MP_TOP_ANCHORS are not mirror-ordered.
    distances = {127: 10, 356: 30, 234: 60, 454: 100}
    for mp_idx, distance in distances.items():
        i = C.MP_TOP_ANCHORS.index(mp_idx)
        raw[i] = (anchors[i] + distance * up) / shape[::-1]
    aligned = align_hairline_for_template(lm, raw, valid, shape)
    actual = (aligned * shape[::-1] - anchors) @ up
    for mp_idx, expected in [(127, 20), (356, 20), (234, 80), (454, 80)]:
        assert actual[C.MP_TOP_ANCHORS.index(mp_idx)] == pytest.approx(expected, abs=1e-4)


def test_missing_ray_hit_does_not_pull_valid_partner():
    shape = (600, 900)
    lm, raw, valid, _anchors, _up = _scene(shape=shape)
    left, right = C.MP_TOP_ANCHORS.index(67), C.MP_TOP_ANCHORS.index(297)
    valid[right] = False
    raw[right, 1] = -0.5  # An extrapolated fallback outside the head.
    aligned = align_hairline_for_template(lm, raw, valid, shape)
    np.testing.assert_allclose(aligned[[left, right]], raw[[left, right]], atol=1e-7)


def test_template_alignment_still_allows_silhouette_clamping():
    shape = (600, 900)
    lm, raw, valid, _anchors, _up = _scene(shape=shape)
    aligned = align_hairline_for_template(lm, raw, valid, shape)
    parse_map = np.zeros(shape, np.int32)
    parse_map[140:450, 250:650] = C.PARSE_SKIN
    clamped = clamp_hairline_to_silhouette(aligned, parse_map)
    assert np.min(clamped[:, 1] * shape[0]) >= 142 - 1e-4