main.py

#!/usr/bin/env python3
import asyncio
import aiohttp
import argparse
import json
import logging
import re
import traceback
import random
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Dict, Optional, Tuple
from datetime import datetime, timezone
from urllib.parse import urljoin

from PIL import Image, ImageOps
from rich.console import Console
from rich.progress import (
    Progress, SpinnerColumn, BarColumn, TextColumn, TimeElapsedColumn,
    TimeRemainingColumn, MofNCompleteColumn, TaskProgressColumn
)
from rich.logging import RichHandler

BASE_URL = "https://www.proteinatlas.org/"
console = Console()

def now_utc_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")

def setup_logging(debug: bool) -> None:
    """
    Configure logging with RichHandler. Enable DEBUG level if debug=True.

    Args:
        debug (bool): If True, set logging level to DEBUG, else INFO.
    """
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(message)s",
        datefmt="[%X]",
        handlers=[RichHandler(console=console, rich_tracebacks=True, show_level=True, show_path=False)]
    )
    logging.getLogger("PIL").setLevel(logging.WARNING)
    logging.getLogger("aiohttp.access").setLevel(logging.WARNING)

def safepath(name: str) -> str:
    """
    Sanitize a string to be safe for use as a filesystem path.

    Args:
        name (str): The input string.

    Returns:
        str: The sanitized string.
    """
    return re.sub(r"[^\w\-.]+", "_", name.strip())

def ensure_dir(p: Path) -> None:
    """
    Ensure that the given directory exists, creating it if necessary.

    Args:
        p (Path): The directory path to ensure.
    """
    p.mkdir(parents=True, exist_ok=True)

def pad_to_square(im: Image.Image, fill: int = 0) -> Image.Image:
    """
    Pad an image to make it square, filling with the specified color.

    Args:
        im (Image.Image): The input image.
        fill (int): The fill color (default 0).

    Returns:
        Image.Image: The padded square image.
    """
    w, h = im.size
    if w == h:
        return im
    size = max(w, h)
    delta_w = size - w
    delta_h = size - h
    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
    return ImageOps.expand(im, padding, fill=fill)

def process_image_to_spec(src_path: Path, dst_path: Path) -> Tuple[int, int]:
    """
    Process an image to grayscale, resize and pad it to 384x384 pixels, then save as PNG.

    Args:
        src_path (Path): Source image path.
        dst_path (Path): Destination PNG path.

    Returns:
        Tuple[int, int]: The original (width, height) of the image.
    """
    with Image.open(src_path) as im:
        im = im.convert("L")
        w0, h0 = im.size
        scale = 384 / min(w0, h0)
        new_w = max(1, int(round(w0 * scale)))
        new_h = max(1, int(round(h0 * scale)))
        im = im.resize((new_w, new_h), Image.Resampling.LANCZOS)
        im = pad_to_square(im, fill=0)
        if im.size != (384, 384):
            im = im.resize((384, 384), Image.Resampling.LANCZOS)
        ensure_dir(dst_path.parent)
        im.save(dst_path, format="PNG")
        return (w0, h0)

class HPAClient:
    def __init__(self, session: aiohttp.ClientSession, delay: float, semaphore: asyncio.Semaphore, retries: int, timeout: int):
        self.session = session
        self.delay = delay
        self.semaphore = semaphore
        self.retries = retries
        self.timeout = timeout
        self.xml_cache: Dict[str, str] = {}
        self.image_map_cache: Dict[str, Dict[str, Dict[str, Dict[str, dict]]]] = {}

    async def _get(self, url: str) -> str:
        """
        Perform an HTTP GET request with retries and delay.

        Args:
            url (str): The URL to fetch.

        Returns:
            str: The response text.

        Raises:
            Exception: If all retries fail.
        """
        for attempt in range(1, self.retries + 1):
            async with self.semaphore:
                try:
                    async with self.session.get(url, timeout=aiohttp.ClientTimeout(total=self.timeout)) as resp:
                        resp.raise_for_status()
                        text = await resp.text()
                        await asyncio.sleep(self.delay)
                        return text
                except Exception as e:
                    logging.debug(f"[GET] attempt {attempt}/{self.retries} failed: {e}")
                    if attempt == self.retries:
                        raise
                    await asyncio.sleep(min(2 * attempt, 5))
        raise RuntimeError("Unreachable retry loop")

    async def fetch_protein_xml(self, protein_id: str) -> str:
        """
        Fetch the XML data for a given protein ID and cache the parsed image map.

        Args:
            protein_id (str): The protein identifier.

        Returns:
            str: The XML string content.
        """
        url = urljoin(BASE_URL, f"{safepath(protein_id)}.xml")
        logging.debug(f"Fetching protein XML: {url}")
        xml_str = await self._get(url)
        self.xml_cache[protein_id] = xml_str
        self.image_map_cache[protein_id] = self.parse_xml_for_images(xml_str)
        return xml_str

    def parse_xml_for_images(self, xml_str: str) -> Dict[str, Dict[str, Dict[str, Dict[str, dict]]]]:
        """
        Parse the protein XML to extract image URLs and metadata, organized by antibody, cell line, image id, and channel.

        Args:
            xml_str (str): XML content as a string.

        Returns:
            Dict[str, Dict[str, Dict[str, Dict[str, dict]]]]: Nested mapping of antibody -> cell line -> image id -> channel -> info dict.
        """
        result = {}
        try:
            root = ET.fromstring(xml_str)
        except Exception as e:
            logging.error(f"Failed to parse XML: {e}")
            return result
        for entry in root.findall("entry"):
            protein_id = entry.findtext("identifier[@db='Ensembl']")
            gene_name = ""
            gene_elem = entry.find("gene")
            if gene_elem is not None and gene_elem.text:
                gene_name = gene_elem.text.strip()
            if not gene_name:
                name_elem = entry.find("name")
                if name_elem is not None and name_elem.text:
                    gene_name = name_elem.text.strip()
            for antibody in entry.findall("antibody"):
                ab_id = antibody.get("id") or antibody.findtext("id")
                if not ab_id:
                    continue
                result.setdefault(ab_id, {})
                for cell_expression in antibody.findall("cellExpression"):
                    for subassay in cell_expression.findall("subAssay"):
                        for data in subassay.findall("data"):
                            cell_line_elem = data.find("cellLine")
                            cell_line = cell_line_elem.text.strip() if cell_line_elem is not None and cell_line_elem.text else "Unknown"
                            for assay_image in data.findall("assayImage"):
                                for image_idx, image in enumerate(assay_image.findall("image")):
                                    url_elem = image.find("imageUrl")
                                    multi_url = url_elem.text if url_elem is not None else None
                                    image_id = None
                                    if multi_url:
                                        m = re.search(r"/[^/]*?_([A-Za-z]+[0-9]+_[0-9]+)(?:_|\.jpg)", multi_url)
                                        if m:
                                            image_id = m.group(1)
                                        else:
                                            m2 = re.search(r"/[^/]*?_([A-Za-z]+[0-9]+)(?:_|\.jpg)", multi_url)
                                            if m2:
                                                image_id = m2.group(1)
                                    if not image_id:
                                        tmpid = image.get("id")
                                        if tmpid and tmpid.strip():
                                            image_id = tmpid.strip()
                                        else:
                                            image_id = f"sample_{image_idx+1}"
                                    for channel in image.findall("channel"):
                                        ch_color = channel.get("color") or channel.findtext("color")
                                        if not ch_color or not multi_url:
                                            continue
                                        ch_color_lc = ch_color.lower()
                                        single_url = re.sub(r"_blue_red_green\.jpg$", f"_{ch_color_lc}.jpg", multi_url, flags=re.IGNORECASE)
                                        result.setdefault(ab_id, {})
                                        result[ab_id].setdefault(cell_line, {})
                                        result[ab_id][cell_line].setdefault(image_id, {})
                                        result[ab_id][cell_line][image_id][ch_color_lc] = {
                                            "url": single_url,
                                            "protein_id": protein_id,
                                            "gene_name": gene_name,
                                            "antibody_id": ab_id,
                                            "cell_line": cell_line,
                                            "channel": ch_color_lc,
                                            "image_id": image_id
                                        }
        return result

    def build_channel_image_url(self, protein_id: str, antibody_id: str, cell_line: str, channel: str) -> Optional[str]:
        """
        Build the image URL for a specific protein, antibody, cell line, and channel.

        Args:
            protein_id (str): Protein identifier.
            antibody_id (str): Antibody identifier.
            cell_line (str): Cell line name.
            channel (str): Channel name.

        Returns:
            Optional[str]: The image URL if found, else None.
        """
        info = self.image_map_cache.get(protein_id, {}).get(antibody_id, {}).get(cell_line, {}).get(channel)
        if isinstance(info, dict):
            return info.get("url")
        return None

class Downloader:
    def __init__(self, outdir: Path, client: HPAClient, progress: Progress, stats: Dict[str, int], error_log: Path):
        self.outdir = outdir
        self.client = client
        self.progress = progress
        self.stats = stats
        self.error_log = error_log
        self.task_overall = self.progress.add_task("[bold]Overall Progress", total=0)
        self.task_images = self.progress.add_task("Image Processing", total=0)

    async def process_protein(self, protein_id: str) -> None:
        """
        Download, process, and organize all images and metadata for a given protein.

        Args:
            protein_id (str): Protein identifier.
        """
        protein_dir = self.outdir / safepath(protein_id)
        ensure_dir(protein_dir)
        metadata_path = protein_dir / "metadata.json"

        try:
            await self.client.fetch_protein_xml(protein_id)
        except Exception as e:
            self._log_error(f"{protein_id}\tFetch protein XML failed\t{e}")
            return

        xml_map = self.client.image_map_cache.get(protein_id, {})
        image_tasks = []
        total_images = 0
        for ab, cl_dict in xml_map.items():
            for cl, imgid_map in cl_dict.items():
                for imgid, ch_map in imgid_map.items():
                    for ch, info in ch_map.items():
                        image_tasks.append((ab, cl, imgid, ch, info))
                        total_images += 1
        # Do not update self.task_images total here; keep it as set at initialization or in main_async.

        gene_name = ""
        for ab, cl_dict in xml_map.items():
            for cl, imgid_map in cl_dict.items():
                for imgid, ch_map in imgid_map.items():
                    for ch, info in ch_map.items():
                        gene_name = info.get("gene_name") or info.get("protein_name") or ""
                        break
                    if gene_name:
                        break
                if gene_name:
                    break
            if gene_name:
                break
        if not gene_name:
            gene_name = ""
        gene_name = gene_name.replace("\n", "").replace("\r", "").strip()

        metadata = {
            "protein_id": protein_id,
            "gene_name": gene_name,
            "antibodies": []
        }
        ab_map = {}
        cl_map = {}
        imgid_map_meta = {}
        for ab, cl, imgid, ch, info in image_tasks:
            if ab not in ab_map:
                antibody_obj = {"antibody_id": ab, "cell_lines": []}
                ab_map[ab] = antibody_obj
                metadata["antibodies"].append(antibody_obj)
            else:
                antibody_obj = ab_map[ab]
            if (ab, cl) not in cl_map:
                cell_line_obj = {"name": cl, "images": []}
                cl_map[(ab, cl)] = cell_line_obj
                antibody_obj["cell_lines"].append(cell_line_obj)
            else:
                cell_line_obj = cl_map[(ab, cl)]
            if (ab, cl, imgid) not in imgid_map_meta:
                image_obj = {"image_id": imgid, "channels": []}
                imgid_map_meta[(ab, cl, imgid)] = image_obj
                cell_line_obj["images"].append(image_obj)

        queue = asyncio.Queue()
        for ab, cl, imgid, ch, info in image_tasks:
            local_name = f"{protein_id}_{ab}_{cl}_{imgid}_{ch}.png"
            local_name = safepath(local_name)
            dst_png = protein_dir / local_name
            if dst_png.exists():
                rel_path = dst_png.relative_to(self.outdir)
                logging.info(f"Skip existing: {rel_path}")
                await queue.put((ab, cl, imgid, ch, info, True, local_name, dst_png))
            else:
                await queue.put((ab, cl, imgid, ch, info, False, local_name, dst_png))

        def write_metadata():
            with open(metadata_path, "w", encoding="utf-8") as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)

        async def consumer():
            while True:
                try:
                    ab, cl, imgid, ch, info, exists, local_name, dst_png = await queue.get()
                except Exception:
                    break
                url = info.get("url")
                tmp_jpg = protein_dir / f"__tmp__{random.randint(100000,999999)}.img"
                if exists:
                    rel_path = dst_png.relative_to(self.outdir)
                    logging.info(f"[Resume] File already exists, skipping download: {rel_path}")
                    try:
                        try:
                            with Image.open(dst_png) as im:
                                w0, h0 = im.size
                        except Exception:
                            w0 = h0 = 0
                        channel_obj = {
                            "name": ch,
                            "url": url,
                            "resolution": f"{w0}x{h0}",
                            "local_path": local_name,
                            "download_time": now_utc_iso()
                        }
                        image_obj = imgid_map_meta[(ab, cl, imgid)]
                        image_obj["channels"].append(channel_obj)
                        self.stats["images_ok"] += 1
                        self.progress.update(self.task_images, advance=1)
                        write_metadata()
                    except Exception as e:
                        self._log_error(f"{protein_id}\t{ab}\t{cl}\t{imgid}\t{ch}\tMetadata resume failed\t{e}")
                        self.stats["images_failed"] += 1
                        self.progress.update(self.task_images, advance=1)
                        write_metadata()
                    finally:
                        queue.task_done()
                    continue
                rel_dst = dst_png.relative_to(self.outdir)
                logging.info(
                    f"Downloading image: protein_id='{protein_id}', antibody_id='{ab}', cell_line='{cl}', image_id='{imgid}', channel='{ch}', url='{url}', save='{rel_dst}'"
                )
                if url is None:
                    self._log_error(f"{protein_id}\t{ab}\t{cl}\t{imgid}\t{ch}\tNo URL found for image")
                    self.stats["images_failed"] += 1
                    self.progress.update(self.task_images, advance=1)
                    queue.task_done()
                    continue
                try:
                    w0, h0 = await self._download_and_process(url, tmp_jpg, dst_png)
                    rel_dst2 = dst_png.relative_to(self.outdir)
                    logging.info(f"Processed image: {rel_dst2} (original size: {w0}x{h0})")
                    channel_obj = {
                        "name": ch,
                        "url": url,
                        "resolution": f"{w0}x{h0}",
                        "local_path": local_name,
                        "download_time": now_utc_iso()
                    }
                    image_obj = imgid_map_meta[(ab, cl, imgid)]
                    image_obj["channels"].append(channel_obj)
                    self.stats["images_ok"] += 1
                    self.progress.update(self.task_images, advance=1)
                    write_metadata()
                except Exception as e:
                    self._log_error(f"{protein_id}\t{ab}\t{cl}\t{imgid}\t{ch}\tDownload/Process failed\t{e}")
                    self.stats["images_failed"] += 1
                    self.progress.update(self.task_images, advance=1)
                    write_metadata()
                finally:
                    queue.task_done()

        num_workers = self.client.semaphore._value if hasattr(self.client.semaphore, "_value") else 6
        num_workers = max(1, num_workers)
        consumers = [asyncio.create_task(consumer()) for _ in range(num_workers)]
        await queue.join()
        for c in consumers:
            c.cancel()
        write_metadata()
        rel_meta = metadata_path.relative_to(self.outdir)
        logging.info(f"Wrote metadata: {rel_meta}")

    async def _download_and_process(self, url: str, tmp_path: Path, dst_png: Path) -> Tuple[int, int]:
        """
        Download an image from a URL, process it to specification, and save as PNG.

        Args:
            url (str): Image URL.
            tmp_path (Path): Temporary file path for download.
            dst_png (Path): Destination PNG path.

        Returns:
            Tuple[int, int]: The original (width, height) of the image, or (0, 0) if failed.
        """
        async with self.client.semaphore:
            for attempt in range(1, self.client.retries + 1):
                logging.debug(f"Attempt {attempt}/{self.client.retries} downloading {url}")
                try:
                    async with self.client.session.get(url, timeout=aiohttp.ClientTimeout(total=self.client.timeout)) as resp:
                        resp.raise_for_status()
                        content = await resp.read()
                        tmp_path.write_bytes(content)
                        await asyncio.sleep(self.client.delay)
                        w0, h0 = process_image_to_spec(tmp_path, dst_png)
                        tmp_path.unlink(missing_ok=True)
                        rel_path = dst_png.relative_to(self.outdir)
                        logging.debug(f"Successfully downloaded and processed {url} to {rel_path}")
                        return (w0, h0)
                except Exception as e:
                    logging.debug(f"[IMG] attempt {attempt}/{self.client.retries} failed for {url}: {e}")
                    await asyncio.sleep(min(2 * attempt, 5))
            logging.error(f"All attempts failed for {url}")
            return (0, 0)

    def _log_error(self, line: str) -> None:
        """
        Log an error message to the error log file and emit an error log.

        Args:
            line (str): The error message to log.
        """
        with open(self.error_log, "a", encoding="utf-8") as f:
            f.write(f"{now_utc_iso()}\t{line}\n")
        logging.error(line)

async def main_async(args):
    """
    Main asynchronous entry point to coordinate downloading and processing for a list of proteins.

    Args:
        args: Parsed command-line arguments.
    """
    setup_logging(args.debug)

    outdir = Path(args.outdir).resolve()
    ensure_dir(outdir)

    error_log = outdir / "error_log.txt"
    stats_path = outdir / "dataset_stats.json"
    stats = dict(proteins_total=0, proteins_done=0, images_ok=0, images_failed=0)

    protein_ids = [p.strip() for p in open(args.protein_list, encoding="utf-8") if p.strip()]
    stats["proteins_total"] = len(protein_ids)

    timeout = aiohttp.ClientTimeout(total=args.timeout)
    connector = aiohttp.TCPConnector(limit_per_host=args.concurrency, force_close=True)
    semaphore = asyncio.Semaphore(args.concurrency)

    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]{task.description}"),
        BarColumn(),
        TaskProgressColumn(),
        MofNCompleteColumn(),
        TimeElapsedColumn(),
        TimeRemainingColumn(),
        console=console,
        transient=False,
        expand=True,
    ) as progress:
        async with aiohttp.ClientSession(timeout=timeout, connector=connector, headers={"User-Agent": "Mozilla/5.0"}) as session:
            client = HPAClient(session, delay=args.delay, semaphore=semaphore, retries=args.retries, timeout=args.timeout)
            dl = Downloader(outdir, client, progress, stats, error_log)

            progress.update(dl.task_overall, total=len(protein_ids))
            # Compute total number of images for all proteins
            total_images_all = 0
            for protein_id in protein_ids:
                try:
                    # Fetch XML and parse to count images
                    xml_str = await client.fetch_protein_xml(protein_id)
                    xml_map = client.image_map_cache.get(protein_id, {})
                    for ab, cl_dict in xml_map.items():
                        for cl, imgid_map in cl_dict.items():
                            for imgid, ch_map in imgid_map.items():
                                for ch, info in ch_map.items():
                                    total_images_all += 1
                except Exception:
                    pass  # If error, just skip for progress bar
            progress.update(dl.task_images, total=total_images_all)

            for idx, protein_id in enumerate(protein_ids, 1):
                logging.info(f"[{idx}/{len(protein_ids)}] Processing protein: {protein_id}")
                try:
                    await dl.process_protein(protein_id)
                    logging.info(f"[{idx}/{len(protein_ids)}] Finished processing protein: {protein_id}")
                except Exception:
                    dl._log_error(f"{protein_id}\tUnhandled failure\t{traceback.format_exc(limit=1)}")
                finally:
                    stats["proteins_done"] += 1
                    progress.update(dl.task_overall, advance=1)

    with open(stats_path, "w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)

    logging.info("All done.")
    logging.info(json.dumps(stats, ensure_ascii=False, indent=2))

def parse_args():
    """
    Parse command-line arguments for the HPA image downloader.

    Returns:
        argparse.Namespace: The parsed arguments.
    """
    p = argparse.ArgumentParser(description="HPA multi-channel localization image download and processing (XML parsing version)")
    p.add_argument("-p", "--protein-list", required=True, help="Text file with one protein ID per line")
    p.add_argument("-o", "--outdir", default="./HPA_download", help="Output root directory")
    p.add_argument("-c", "--concurrency", type=int, default=6, help="Maximum concurrent requests")
    p.add_argument("-d", "--delay", type=float, default=0.5, help="Request interval seconds")
    p.add_argument("-r", "--retries", type=int, default=3, help="Retry times on failure")
    p.add_argument("-t", "--timeout", type=int, default=30, help="Timeout per request (seconds)")
    p.add_argument("-D", "--debug", action="store_true", help="Enable DEBUG log")
    return p.parse_args()

def main():
    """
    Synchronous entry point. Parse arguments and run the main async workflow.
    """
    args = parse_args()
    try:
        asyncio.run(main_async(args))
    except KeyboardInterrupt:
        print("Interrupted by user.")

if __name__ == "__main__":
    main()