"""
CheXNet Compressed Model — Inference

Loads the compressed model (optionally with the fine-tuned classifier head)
and predicts pathology probabilities for a chest X-ray image.

Usage:
    python inference.py <path/to/image.png>
    python inference.py <path/to/image.png> --classifier classifier_finetuned.pt
"""
import sys
sys.dont_write_bytecode = True  # avoid creating __pycache__ in the release dir

import argparse
import os

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchxrayvision as xrv
from PIL import Image


def build_compressed_model(ckpt_path, device, classifier_ft_path=None):
    """Load baseline CheXNet then reshape to compressed dimensions and load weights."""
    model = xrv.models.DenseNet(weights="densenet121-res224-all").to(device).eval()
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
    alive = ckpt["alive_per_block"]
    for block_idx in [1, 2, 3, 4]:
        block = getattr(model.features, f"denseblock{block_idx}")
        block_alive = alive.get(block_idx, alive.get(f"block{block_idx}", {}))
        for dl_key, n_alive in block_alive.items():
            i = int(dl_key[2:]) if dl_key.startswith("dl") else int(dl_key)
            L = getattr(block, f"denselayer{i}")
            in_ch = L.conv1.in_channels
            L.conv1 = nn.Conv2d(in_ch, n_alive, 1, bias=True).to(device).eval()
            L.norm2 = nn.BatchNorm2d(n_alive, eps=L.norm2.eps).to(device).eval()
            L.conv2 = nn.Conv2d(n_alive, 32, 3, padding=1, bias=False).to(device).eval()
    model.load_state_dict(ckpt["state_dict"])
    # intermediate batch norms are folded; replace with identity
    for block_idx in [1, 2, 3, 4]:
        block = getattr(model.features, f"denseblock{block_idx}")
        n_layers = {1: 6, 2: 12, 3: 24, 4: 16}[block_idx]
        for i in range(1, n_layers + 1):
            getattr(block, f"denselayer{i}").norm2 = nn.Identity()

    if classifier_ft_path and os.path.exists(classifier_ft_path):
        cls_ft = nn.Linear(1024, 18).to(device)
        cls_ft.load_state_dict(torch.load(classifier_ft_path,
                                          map_location=device, weights_only=True))
        model.classifier = cls_ft

    model.eval()
    return model


def preprocess(img_np):
    """xrv normalization: scale to [-1024, 1024] and resize to 224x224."""
    arr = img_np.astype(np.float32)
    arr = (arr - arr.min()) / max(arr.max() - arr.min(), 1) * 2048 - 1024
    pil = Image.fromarray(arr, mode="F").resize((224, 224), Image.BILINEAR)
    return np.array(pil)


def predict(model, image_path, device, top_n=5):
    img = Image.open(image_path)
    if img.mode != "L":
        img = img.convert("L")
    img_proc = preprocess(np.array(img))
    t = torch.tensor(img_proc, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)
    with torch.no_grad():
        feat = model.features(t)
        feat = F.relu(feat, inplace=False)
        pooled = F.adaptive_avg_pool2d(feat, (1, 1)).flatten(1)
        logits = model.classifier(pooled)[0]
        probs = torch.sigmoid(logits).cpu().numpy()
    pathologies = list(model.pathologies)
    return sorted(zip(pathologies, probs.tolist()), key=lambda x: -x[1])[:top_n]


def main():
    parser = argparse.ArgumentParser(description="CheXNet compressed model inference")
    parser.add_argument("image_path", help="Path to chest X-ray image (PNG)")
    parser.add_argument("--ckpt", default="compressed_model.pt",
                        help="Compressed model checkpoint")
    parser.add_argument("--classifier", default="classifier_finetuned.pt",
                        help="Optional fine-tuned classifier head")
    parser.add_argument("--top", type=int, default=5, help="Number of predictions to show")
    args = parser.parse_args()

    here = os.path.dirname(os.path.abspath(__file__))
    ckpt_path = args.ckpt if os.path.isabs(args.ckpt) else os.path.join(here, args.ckpt)
    cls_path = args.classifier if os.path.isabs(args.classifier) else os.path.join(here, args.classifier)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    model = build_compressed_model(ckpt_path, device, cls_path)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Model parameters: {n_params:,}")

    ranked = predict(model, args.image_path, device, top_n=args.top)
    print(f"\nTop-{args.top} predictions for {args.image_path}:")
    for pathology, prob in ranked:
        print(f"  {pathology:<28}  {prob:.4f}")


if __name__ == "__main__":
    main()