import cv2
import numpy as np


def intersect_lines(line_a, line_b):
    """Intersect two lines given as (x1, y1, x2, y2), or None if parallel."""
    x1, y1, x2, y2 = line_a
    x3, y3, x4, y4 = line_b

    p1 = np.array([x1, y1, 1.0])
    p2 = np.array([x2, y2, 1.0])
    p3 = np.array([x3, y3, 1.0])
    p4 = np.array([x4, y4, 1.0])

    l1 = np.cross(p1, p2)
    l2 = np.cross(p3, p4)
    x = np.cross(l1, l2)

    if abs(x[2]) < 1e-8:
        return None

    return x[0] / x[2], x[1] / x[2]


def build_point_map(landmarks):
    """Index a list of landmark dicts (each with a 'name' key) by name."""
    return {lm["name"]: lm for lm in landmarks}


def set_midpoint(pt_map, name, a, b):
    """Store the midpoint of two named points in pt_map under `name`."""
    if (a in pt_map) and (b in pt_map):
        pt_map[name] = {
            "orig_x": (pt_map[a]["orig_x"] + pt_map[b]["orig_x"]) / 2,
            "orig_y": (pt_map[a]["orig_y"] + pt_map[b]["orig_y"]) / 2,
        }
    return pt_map


def get_angle(pt_map, p1, p2, p3):
    """Angle in degrees at vertex p2, between rays p2->p1 and p2->p3."""
    if (p1 not in pt_map) or (p2 not in pt_map) or (p3 not in pt_map):
        return None

    v21 = np.array([
        pt_map[p1]["orig_x"] - pt_map[p2]["orig_x"],
        pt_map[p1]["orig_y"] - pt_map[p2]["orig_y"],
    ])
    v23 = np.array([
        pt_map[p3]["orig_x"] - pt_map[p2]["orig_x"],
        pt_map[p3]["orig_y"] - pt_map[p2]["orig_y"],
    ])

    dot_product = np.dot(v21, v23)
    cross_product_z = abs(v21[0] * v23[1] - v21[1] * v23[0])

    angle_radians = np.arctan2(cross_product_z, dot_product)
    angle_degrees = np.degrees(angle_radians)

    if angle_degrees < 0:
        angle_degrees += 360

    return angle_degrees


def plot_line(image, pt_map, a, b, color, skip_display=False):
    """Draw a line between two named points in pt_map; returns (image, line)."""
    if (a not in pt_map) or (b not in pt_map):
        return image, None

    px1, py1 = pt_map[a]["orig_x"], pt_map[a]["orig_y"]
    px2, py2 = pt_map[b]["orig_x"], pt_map[b]["orig_y"]

    if not skip_display:
        cv2.line(
            image,
            (int(round(px1)), int(round(py1))),
            (int(round(px2)), int(round(py2))),
            color, 2,
        )

    return image, [px1, py1, px2, py2]


def plot_angle(image, pt_map, a, b, c, color):
    """Draw the angle arc at vertex B formed by A-B and B-C."""
    if (a not in pt_map) or (b not in pt_map) or (c not in pt_map):
        return image

    xA, yA = pt_map[a]["orig_x"], pt_map[a]["orig_y"]
    xB, yB = pt_map[b]["orig_x"], pt_map[b]["orig_y"]
    xC, yC = pt_map[c]["orig_x"], pt_map[c]["orig_y"]

    m_ba = (yB - yA) / (xB - xA)
    b_ba = yB - m_ba * xB
    ba = np.linspace(xB, xA, 10)
    nx_a = ba[1]
    ny_a = m_ba * nx_a + b_ba

    pts = np.array([[
        [round(nx_a), round(ny_a)],
        [round(xB), round(yB)],
        [round(xC), round(yC)],
    ]]).astype(int)

    return cv2.polylines(image, [pts], isClosed=True, color=color, thickness=2)


def plot_landmarks(image, landmarks, color=(255, 255, 0), angle_color=(255, 0, 0)):
    """Draw landmark points plus the derived axes/joint-lines/angles used
    for CPAK measurement on top of image (modified in place, also returned).
    """
    pt_map = {}
    for pt in landmarks:
        pt_map[pt["name"]] = pt

        if pt["name"] in ("FR", "FL"):
            continue

        cv2.circle(image, (int(pt["orig_x"]), int(pt["orig_y"])), 5, color, -1)

    # FHR-FR, LFR-MFR LTR-MTR [FTR]-AR
    # FHL-FL, LFL-MFL LTL-MTL [FTL]-AL

    #  FTR = (LTR+MTR)/2
    #  FTL = (LTL+MTL)/2
    
    # Right femur: mechanical axis (FHR-FR) x joint line (LFR-MFR) -> FFR
    image, femoral_axis_r = plot_line(image, pt_map, "FHR", "FR", color, skip_display=True)
    image, joint_line_distal_femur_r = plot_line(image, pt_map, "LFR", "MFR", color)
    if femoral_axis_r is not None and joint_line_distal_femur_r is not None:
        intersection = intersect_lines(femoral_axis_r, joint_line_distal_femur_r)
        if intersection is not None:
            pt_map["FFR"] = {"orig_x": intersection[0], "orig_y": intersection[1]}
            image, _ = plot_line(image, pt_map, "FHR", "FFR", color)

    # Right tibia: joint line (LTR-MTR), midpoint FTR -> ankle center (AR)
    image, _ = plot_line(image, pt_map, "LTR", "MTR", color)
    set_midpoint(pt_map, "FTR", "LTR", "MTR")
    image, _ = plot_line(image, pt_map, "FTR", "AR", color)

    image = plot_angle(image, pt_map, "FHR", "FFR", "LFR", angle_color)
    image = plot_angle(image, pt_map, "AR", "FTR", "MTR", angle_color)

    # Left femur
    image, femoral_axis_l = plot_line(image, pt_map, "FHL", "FL", color, skip_display=True)
    image, joint_line_distal_femur_l = plot_line(image, pt_map, "LFL", "MFL", color)
    if femoral_axis_l is not None and joint_line_distal_femur_l is not None:
        intersection = intersect_lines(femoral_axis_l, joint_line_distal_femur_l)
        if intersection is not None:
            pt_map["FFL"] = {"orig_x": intersection[0], "orig_y": intersection[1]}
            image, _ = plot_line(image, pt_map, "FHL", "FFL", color)

    # Left tibia
    image, _ = plot_line(image, pt_map, "LTL", "MTL", color)
    set_midpoint(pt_map, "FTL", "LTL", "MTL")
    image, _ = plot_line(image, pt_map, "FTL", "AL", color)

    image = plot_angle(image, pt_map, "FHL", "FFL", "LFL", angle_color)
    image = plot_angle(image, pt_map, "AL", "FTL", "MTL", angle_color)

    return image