from pathlib import Path
import re
import xml.etree.ElementTree as ET

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image


OUT = Path("exports")
OUT.mkdir(exist_ok=True)

width_in = 3.50
height_in = 2.40
dpi = 300

plt.rcParams.update(
    {
        "font.size": 8,
        "axes.labelsize": 8,
        "axes.titlesize": 9,
        "legend.fontsize": 7,
        "xtick.labelsize": 7,
        "ytick.labelsize": 7,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
        "svg.fonttype": "none",
    }
)

days = np.arange(1, 7)
control = np.array([2.1, 2.4, 2.8, 3.0, 3.4, 3.7])
treatment = np.array([2.0, 2.7, 3.4, 4.1, 4.6, 5.0])

fig, ax = plt.subplots(figsize=(width_in, height_in), dpi=dpi, layout="constrained")
ax.plot(days, control, marker="o", linewidth=1.4, label="Control")
ax.plot(days, treatment, marker="s", linewidth=1.4, label="Treatment")
ax.set_title("Response over time")
ax.set_xlabel("Day")
ax.set_ylabel("Mean response")
ax.grid(True, linewidth=0.4, alpha=0.35)
ax.legend(frameon=False)

metadata = {
    "Title": "Publication figure export",
    "Author": "Data Team",
    "Creator": "Matplotlib publication export script",
}

pdf = OUT / "publication-figure.pdf"
svg = OUT / "publication-figure.svg"
png = OUT / "publication-figure.png"

fig.savefig(pdf, metadata=metadata)
fig.savefig(svg, metadata={"Title": metadata["Title"], "Creator": metadata["Creator"]})
fig.savefig(png, dpi=dpi, metadata={"Title": metadata["Title"]})
plt.close(fig)


def pdf_size_inches(path):
    match = re.search(
        rb"/MediaBox\s*\[\s*0\s+0\s+([0-9.]+)\s+([0-9.]+)\s*\]",
        path.read_bytes(),
    )
    if not match:
        raise RuntimeError(f"MediaBox not found in {path}")
    return float(match.group(1)) / 72, float(match.group(2)) / 72


def svg_unit_to_inches(value):
    if value.endswith("pt"):
        return float(value[:-2]) / 72
    if value.endswith("in"):
        return float(value[:-2])
    raise RuntimeError(f"Unsupported SVG unit: {value}")


def svg_size_inches(path):
    root = ET.parse(path).getroot()
    return svg_unit_to_inches(root.attrib["width"]), svg_unit_to_inches(
        root.attrib["height"]
    )


def print_inches(label, path, size):
    print(f"{path.name:<24} {label:<3} {size[0]:.2f} x {size[1]:.2f} in")


with Image.open(png) as image:
    png_size = image.size
    png_title = image.text.get("Title", "")

print(f"matplotlib {matplotlib.__version__} ({matplotlib.get_backend()} backend)")
print(f"target                   {width_in:.2f} x {height_in:.2f} in at {dpi} dpi")
print_inches("PDF", pdf, pdf_size_inches(pdf))
print_inches("SVG", svg, svg_size_inches(svg))
print(f"{png.name:<24} PNG {png_size[0]} x {png_size[1]} px")
print(f"png title metadata      {png_title}")