import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch
from unsloth import FastLanguageModel
import json
import tqdm

# -----------------------------
#  MODEL CACHE
# -----------------------------
_model_cache = {"model": None, "tokenizer": None}

def load_finetuned_model(model_path: str):
    """Load and cache your fine-tuned subclaim extraction model + tokenizer."""
    if _model_cache["model"] is not None:
        return _model_cache["model"], _model_cache["tokenizer"]

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=model_path,
        max_seq_length=8192,
        load_in_4bit=False,
        load_in_8bit=False,
        full_finetuning=False,
    )
    _model_cache["model"], _model_cache["tokenizer"] = model, tokenizer
    return model, tokenizer


# -----------------------------
#  SUBCLAIM EXTRACTION PROMPT
# -----------------------------
def extraction_prompt(medical_text: str) -> str:
    prompt = f"""
You are an expert medical annotator. Your task is to extract granular, factual subclaims from medical text.
A subclaim is the smallest standalone factual unit that can be independently verified.

Instructions:
1. Read the provided medical text.
2. Break it into clear, objective, atomic subclaims.
3. Each subclaim must come directly from the text.
4. Do not add, guess, or infer information.
5. Each subclaim should be short, specific, and verifiable.
6. Return ONLY a Python-style list of strings.

Medical Text:
{medical_text}

Return your output in JSON list format, like:
[
  "subclaim 1",
  "subclaim 2",
  ...
]
"""
    return prompt


# -----------------------------
#  INFERENCE FUNCTION
# -----------------------------
def infer_subclaims(medical_text: str,
                    model_path: str,
                    temperature: float = 0.2) -> str:
    
    model, tokenizer = load_finetuned_model(model_path)

    prompt = extraction_prompt(medical_text)

    messages = [{"role": "user", "content": prompt}]

    chat_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    inputs = tokenizer(chat_text, return_tensors="pt").to("cuda")

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=temperature,
            top_p=0.9,
            top_k=10,
            do_sample=False,
        )

    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()

    # Remove thinking if model inserts `<think>`
    if "</think>" in output_text:
        output_text = output_text.split("</think>")[-1].strip()

    return output_text


# -----------------------------
#  MAIN EXECUTION
# -----------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, required=True,
                        help="Path to the input JSON file containing medical texts.")
    args = parser.parse_args()
    INPUT_FILE = args.input_file
    file_name=os.path.basename(INPUT_FILE).split(".json")[0]
    SAVE_FOLDER = "/home/mshahidul/readctrl/data/extracting_subclaim"
    MODEL_PATH = "/home/mshahidul/readctrl_model/qwen3-32B_subclaims-extraction-8b_ctx"

    os.makedirs(SAVE_FOLDER, exist_ok=True)

    OUTPUT_FILE = os.path.join(SAVE_FOLDER, f"extracted_subclaims_{file_name}_en.json")

    # Load input dataset
    with open(INPUT_FILE, "r") as f:
        data = json.load(f)

    # Load existing results (resume mode)
    result = []
    if os.path.exists(OUTPUT_FILE):
        with open(OUTPUT_FILE, "r") as f:
            result = json.load(f)

    existing_ids = {item["id"] for item in result}

    # --------------------------------------------------------
    # PROCESS EACH MEDICAL TEXT
    # --------------------------------------------------------
    for item in tqdm.tqdm(data):
        if item["id"] in existing_ids:
            continue

        medical_text = item.get("fulltext", "")

        extracted = infer_subclaims(
            medical_text,
            model_path=MODEL_PATH
        )

        result.append({
            "id": item["id"],
            "medical_text": medical_text,
            "subclaims": extracted,
            "summary": item.get("summary", "")
        })

        # Save every 20 entries
        if len(result) % 20 == 0:
            print(f"Saving intermediate results... Total processed: {len(result)}")
            with open(OUTPUT_FILE, "w") as f:
                json.dump(result, f, indent=4, ensure_ascii=False)

    # Final save
    with open(OUTPUT_FILE, "w") as f:
        json.dump(result, f, indent=4, ensure_ascii=False)

    print("Extraction completed.")