#!/usr/bin/env python3
import os
import re
import json
import argparse
import cv2
import numpy as np
import torch
from torchvision import transforms as tr
from transformers import AutoModel
from huggingface_hub import hf_hub_download

# Image normalization parameters from the training pipeline
CHANNEL_MEAN = np.array([0.485, 0.456, 0.406])
CHANNEL_STD = np.array([0.229, 0.224, 0.225])

# Image transformation pipeline targeting 224x224 inputs
image_transform_pipeline = tr.Compose([
    tr.ToPILImage(),
    tr.Resize((256, 256)),
    tr.CenterCrop(224),
    tr.ToTensor(),
    tr.Normalize(mean=CHANNEL_MEAN, std=CHANNEL_STD),
])


def clean_and_format_radiology_text(text):
    """
    Removes leading spaces before punctuation and enforces capital letters
    at the beginning of sentences to match standard radiological formatting.
    """
    text = re.sub(r"\s+([,.!?:;])", r"\1", text)
    if not text:
        return ""
    text = text.upper() + text[1:]
    text = re.sub(r"([.!?]\s+[a-z])", lambda m: m.group().upper(), text)
    return text


def parse_arguments():
    parser = argparse.ArgumentParser(description="Lab225 ResNet+LSTM Radiology Report CLI Inference Tool")
    parser.add_argument("--image", type=str, required=True, help="Path to the input chest radiograph (PNG/JPG)")
    parser.add_argument("--model_id", type=str, default="lab225/resnet-lstm-belarus-screening",
                        help="Hugging Face Model ID")
    parser.add_argument("--top_p", type=float, default=0.1, help="Nucleus (Top-p) sampling parameter (default: 0.1)")
    parser.add_argument("--top_k", type=int, default=None, help="Top-k sampling parameter (optional)")
    parser.add_argument("--max_len", type=int, default=100, help="Maximum number of tokens to generate")
    return parser.parse_args()


def main():
    args = parse_arguments()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[INFO] Target inference device: {device}")

    # 1. Load the custom Lab225 architecture from Hugging Face Hub
    print(f"[INFO] Downloading and initializing model configurations from {args.model_id}...")
    model = AutoModel.from_pretrained(args.model_id, trust_remote_code=True).to(device)
    model.eval()

    # 2. Download and load the vocabulary map tokens
    print("[INFO] Synchronizing vocabulary lookup tables (vocab.json)...")
    vocab_file_path = hf_hub_download(repo_id=args.model_id, filename="vocab.json")

    with open(vocab_file_path, "r", encoding="utf-8") as f:
        vocab_data = json.load(f)

    tok_to_ind = vocab_data["tok_to_ind"]
    ind_to_tok = {int(k): v for k, v in vocab_data["ind_to_tok"].items()}

    # 3. Read and preprocess the target medical image
    if not os.path.exists(args.image):
        print(f"[ERROR] Target image file at path '{args.image}' does not exist.")
        return

    img_raw = cv2.imread(args.image)
    img_rgb = cv2.cvtColor(img_raw, cv2.COLOR_BGR2RGB)
    img_tensor = image_transform_pipeline(img_rgb).unsqueeze(0).to(device)

    # 4. Autoregressive inference block
    generated_indices = [tok_to_ind["<BOS>"]]
    sequence_tensor = torch.tensor(generated_indices).reshape((1, 1)).to(device)

    for _ in range(args.max_len - 1):
        with torch.no_grad():
            logits = model(img_tensor, sequence_tensor)
            last_token_logits = logits[0, -1, :]

        probabilities = torch.softmax(last_token_logits, dim=-1)

        if args.top_k is not None:
            top_probs, top_indices = torch.topk(probabilities, args.top_k)
            probabilities = torch.zeros_like(probabilities).scatter_(0, top_indices, top_probs)
            probabilities = probabilities / probabilities.sum()
        elif args.top_p is not None:
            sorted_probs, sorted_indices = torch.sort(probabilities, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=0)

            probability_mask = cumulative_probs <= args.top_p
            if not probability_mask.any():
                probability_mask = True

            selected_probs = sorted_probs[probability_mask]
            selected_indices = sorted_indices[probability_mask]

            probabilities = torch.zeros_like(probabilities).scatter_(0, selected_indices, selected_probs)
            probabilities = probabilities / probabilities.sum()

        next_token_index = torch.multinomial(probabilities, 1).item()
        generated_indices.append(next_token_index)
        sequence_tensor = torch.tensor(generated_indices).reshape((1, -1)).to(device)

        if next_token_index == tok_to_ind["<EOS>"]:
            break

    # 5. Decode token indices to text strings excluding technical placeholders
    output_words = [ind_to_tok[idx] for idx in generated_indices[1:-1]]
    raw_joined_text = " ".join(output_words)
    final_radiology_report = clean_and_format_radiology_text(raw_joined_text)

    print("\n" + "=" * 60)
    print("GENERATED RADIOLOGY REPORT (LAB225 PREDICTION):")
    print("=" * 60)
    print(final_radiology_report)
    print("=" * 60 + "\n")


if __name__ == "__main__":
    main()