import os
import zipfile

import numpy as np
import pytest

from cellquant.annotations import (
    ImageAnnotation,
    Shape,
    discover,
    image_key,
    load_annotation_file,
    parse_cvat_xml,
    rasterize,
    to_rgb_uint8,
)

XML = """<?xml version="1.0" encoding="utf-8"?>
<annotations>
  <version>1.1</version>
  <meta><task><size>1</size></task></meta>
  <image id="0" name="K1 demo - RGB copy.tif" width="100" height="80">
    <polyline label="Cell, Antibody +" source="manual" occluded="0" points="10.0,10.0;30.0,10.0;30.0,30.0;10.0,30.0;10.0,10.0" z_order="0"></polyline>
    <polygon label="Cell, Nostril, Antibody -" source="manual" occluded="0" points="50,50;70,50;70,70;50,70" z_order="0"></polygon>
    <ellipse label="Nostril in Nucleus" source="manual" occluded="0" cx="60" cy="60" rx="3" ry="2" z_order="0"></ellipse>
    <polyline label="Nostril External of Nucleus" source="manual" occluded="0" points="90.2,5.1;90.4,5.3;90.6,5.1" z_order="0"></polyline>
    <box label="Cell, Antibody -" source="manual" occluded="0" xtl="80" ytl="60" xbr="95" ybr="75" z_order="0"></box>
  </image>
</annotations>
"""


def test_parse_xml_shapes_and_counts():
    anns = parse_cvat_xml(XML.encode())
    assert len(anns) == 1
    a = anns[0]
    assert (a.width, a.height) == (100, 80)
    assert a.key == "K1"
    c = a.counts()
    assert c["Cell, Antibody +"] == 1 and c["Cell, Nostril, Antibody -"] == 1 and c["Cell, Antibody -"] == 1
    assert c["Nostril in Nucleus"] == 1 and c["Nostril External of Nucleus"] == 1


def test_rasterize_areas_and_attributes():
    a = parse_cvat_xml(XML.encode())[0]
    R = rasterize(a)
    assert R.cell_inst.shape == (80, 100)
    assert R.n_cells == 3 and R.n_nostrils == 2
    areas = np.bincount(R.cell_inst.ravel())[1:]
    # 20x20 squares rasterise to ~21x21 pixels (inclusive edges) -> 400..441
    assert all(380 <= x <= 450 for x in areas[:2]) or all(380 <= x <= 450 for x in sorted(areas)[-2:])
    # tiny polyline (sub-pixel) keeps its centroid pixel
    nareas = np.bincount(R.nostril_inst.ravel())[1:]
    assert nareas.min() >= 1
    # ellipse rx=3, ry=2 -> pi*6 ~ 19 px
    assert 10 <= nareas.max() <= 30
    ab = R.cell_antibody()
    hn = R.cell_has_nostril()
    labels = R.cell_labels
    for i, lab in enumerate(labels):
        if i == 0:
            continue
        assert ab[i] == lab.endswith("+")
        assert hn[i] == ("Nostril" in lab)
    assert R.nostril_inside().sum() == 1


def test_image_key_handles_copies_and_double_spaces():
    assert image_key("F114  F 35D SC Iba W002 - Focused EDF Document_RGB.tif") == "F114"
    assert image_key("C110 F 35D SC Iba W001 - Focused EDF Document_RGB copy (1).tif") == "C110"


def test_discover_pairs_by_internal_name(tmp_path):
    from PIL import Image
    Image.fromarray(np.zeros((80, 100, 3), np.uint8)).save(tmp_path / "K1 demo - RGB copy (1).png")
    Image.fromarray(np.zeros((80, 100, 3), np.uint8)).save(tmp_path / "K2 other.png")  # unannotated
    with zipfile.ZipFile(tmp_path / "WRONGNAME_Annotated.zip", "w") as zf:
        zf.writestr("annotations.xml", XML)
    samples, unann = discover(str(tmp_path), verbose=False)
    assert [s.key for s in samples] == ["K1"]
    assert len(unann) == 1 and unann[0].endswith("K2 other.png")
    assert samples[0].annotation.count("Cell, Antibody +") == 1


def test_load_annotation_file_zip_and_xml(tmp_path):
    (tmp_path / "a.xml").write_text(XML)
    with zipfile.ZipFile(tmp_path / "a.zip", "w") as zf:
        zf.writestr("annotations.xml", XML)
    assert len(load_annotation_file(str(tmp_path / "a.xml"))[0].shapes) == 5
    assert len(load_annotation_file(str(tmp_path / "a.zip"))[0].shapes) == 5
    with pytest.raises(ValueError):
        load_annotation_file(str(tmp_path / "nothing.txt"))


def test_to_rgb_uint8_variants():
    g = np.full((4, 5), 7, np.uint8)
    assert to_rgb_uint8(g).shape == (4, 5, 3)
    chw = np.zeros((3, 4, 5), np.uint8)
    assert to_rgb_uint8(chw).shape == (4, 5, 3)
    rgba = np.zeros((4, 5, 4), np.uint8)
    assert to_rgb_uint8(rgba).shape == (4, 5, 3)
    u16 = np.full((4, 5, 3), 4000, np.uint16)
    out = to_rgb_uint8(u16)
    assert out.dtype == np.uint8 and out.max() == 255
    f01 = np.full((4, 5, 3), 0.5, np.float32)
    assert to_rgb_uint8(f01).max() in (127, 128)
    with pytest.raises(ValueError):
        to_rgb_uint8(np.zeros((40, 50, 2), np.uint8))
    with pytest.raises(ValueError):
        to_rgb_uint8(np.zeros((4, 50, 2), np.uint8))   # 4 rows is not "RGBA channel-first"


def test_shape_area_shoelace():
    s = Shape("Cell, Antibody +", "polygon", np.array([[0, 0], [10, 0], [10, 10], [0, 10]], np.float32))
    assert abs(s.area - 100) < 1e-6