"""
Inference script for the aesthetic scoring model (score-token approach).

For each test image, runs 13 criterion queries + 1 total-score query,
maps score tokens to numeric values, then classifies into levels:
    A — Poor   (0 <= score < 5)
    B — Medium (5 <= score < 7)
    C — Good   (7 <= score <= 10)

Also answers the per-image MCQ by feeding the question+options directly.

Usage:
    CUDA_VISIBLE_DEVICES=0,1,2,3 python evaluation/inference_aesthetic.py
"""

import os
import sys
import json
import math
import re
import glob
import time
import torch
import multiprocessing as mp
from pathlib import Path
from tqdm import tqdm
from multiprocessing import Process
from PIL import Image

Image.MAX_IMAGE_PIXELS = None
mp.set_start_method("spawn", force=True)

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from qwenvl.train.score_tokens import ScoreTokenRegistry
from qwenvl.data.aesthetic_dataset import CRITERIA_QUESTIONS, CRITERIA_ORDER, TOTAL_SCORE_QUESTION

# ── Config ────────────────────────────────────────────────────────────────────
CHECKPOINT = ""

IMAGE_DIR = "zijielou/PortraitCraft/main"

TEST_JSON = "PortraitCraft/track_1_test.json"

OUTPUT_JSON = "PortraitCraft/track_1_test_results_2.json"

WORKERS_PER_GPU = 1
MAX_NEW_TOKENS = 1024

# ── Score token registries (must match training) ─────────────────────────────
CRIT_BINS = 20
TOTAL_BINS = 101

crit_registry = ScoreTokenRegistry(n_bins=CRIT_BINS, score_min=1.0, score_max=10.0, prefix="score")
total_registry = ScoreTokenRegistry(n_bins=TOTAL_BINS, score_min=0.0, score_max=100.0, prefix="total")

# ── Level classification ─────────────────────────────────────────────────────
def score_to_level(score: float) -> str:
    if score < 5.0:
        return "A"
    elif score < 7.0:
        return "B"
    else:
        return "C"


# ── Image path resolver (test json has flat names, images live in subdirs) ───
_image_index = None

def build_image_index(image_dir: str) -> dict:
    index = {}
    for subdir in sorted(os.listdir(image_dir)):
        subdir_path = os.path.join(image_dir, subdir)
        if not os.path.isdir(subdir_path):
            continue
        for fname in os.listdir(subdir_path):
            index[fname] = os.path.join(subdir_path, fname)
    return index


def resolve_image_path(image_name: str, image_dir: str) -> str:
    global _image_index
    full = os.path.join(image_dir, image_name)
    if os.path.exists(full):
        return full
    if _image_index is None:
        _image_index = build_image_index(image_dir)
    return _image_index.get(image_name, full)


# ── Model wrapper ─────────────────────────────────────────────────────────────
class AestheticInferenceServer:
    def __init__(self, gpu_id: int):
        from transformers import Qwen3VLForConditionalGeneration, AutoProcessor

        self.device = torch.device(f"cuda:{gpu_id}")
        print(f"Loading model on GPU {gpu_id} from {CHECKPOINT}...")

        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
            CHECKPOINT,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
        ).to(self.device).eval()

        self.processor = AutoProcessor.from_pretrained(CHECKPOINT)

        self.crit_token_ids = crit_registry.get_token_ids(self.processor.tokenizer)
        self.crit_values = crit_registry.get_value_tensor(device=self.device)
        self.total_token_ids = total_registry.get_token_ids(self.processor.tokenizer)
        self.total_values = total_registry.get_value_tensor(device=self.device)

        print(f"Model loaded on GPU {gpu_id}")

    def _generate(self, image_path: str, text: str, max_new_tokens: int = None) -> tuple:
        """Returns (generated_text, generated_ids, scores_tuple)."""
        if max_new_tokens is None:
            max_new_tokens = MAX_NEW_TOKENS

        img = Image.open(image_path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": text},
            ],
        }]

        inputs = self.processor.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                return_dict_in_generate=True,
                output_scores=True,
            )

        generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
        generated_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)
        scores_tuple = outputs.scores

        return generated_text, generated_ids, scores_tuple

    def predict_criterion(self, image_path: str, criterion: str) -> dict:
        question = CRITERIA_QUESTIONS[criterion]
        gen_text, gen_ids, scores_tuple = self._generate(image_path, question)

        score_token_id_set = set(self.crit_token_ids)
        expected_score = None
        argmax_score = None
        argmax_token = None

        gen_ids_list = gen_ids.tolist()
        for step_idx, tid in enumerate(gen_ids_list):
            if tid in score_token_id_set:
                token_str = self.processor.tokenizer.decode([tid])
                argmax_token = token_str.strip()
                argmax_score = crit_registry.token_to_score(argmax_token)
                logits_at_step = scores_tuple[step_idx]
                expected_score = crit_registry.expected_score(
                    logits_at_step[0], self.crit_token_ids, self.crit_values
                ).item()
                break

        final_score = expected_score if expected_score is not None else argmax_score
        if final_score is None:
            print(f"    WARNING: no score token found for [{criterion}], raw output: {gen_text[:200]}", flush=True)
            final_score = 5.0

        return {
            "score": round(final_score, 2),
            "level": score_to_level(final_score),
            "argmax_token": argmax_token,
            "argmax_score": round(argmax_score, 2) if argmax_score is not None else None,
            "expected_score": round(expected_score, 2) if expected_score is not None else None,
            "raw_text": gen_text,
        }

    def predict_total(self, image_path: str) -> dict:
        gen_text, gen_ids, scores_tuple = self._generate(image_path, TOTAL_SCORE_QUESTION)

        total_token_id_set = set(self.total_token_ids)
        expected_score = None
        argmax_score = None
        argmax_token = None

        gen_ids_list = gen_ids.tolist()
        for step_idx, tid in enumerate(gen_ids_list):
            if tid in total_token_id_set:
                token_str = self.processor.tokenizer.decode([tid])
                argmax_token = token_str.strip()
                argmax_score = total_registry.token_to_score(argmax_token)
                logits_at_step = scores_tuple[step_idx]
                expected_score = total_registry.expected_score(
                    logits_at_step[0], self.total_token_ids, self.total_values
                ).item()
                break

        final_score = expected_score if expected_score is not None else argmax_score
        if final_score is None:
            print(f"    WARNING: no total token found, raw output: {gen_text[:200]}", flush=True)
            final_score = 50.0

        return {
            "total_score": round(final_score, 1),
            "argmax_token": argmax_token,
            "argmax_score": round(argmax_score, 1) if argmax_score is not None else None,
            "expected_score": round(expected_score, 1) if expected_score is not None else None,
            "raw_text": gen_text,
        }

    def answer_mcq(self, image_path: str, question: str, options: dict) -> dict:
        options_text = "\n".join([f"{k}. {v}" for k, v in options.items()])
        prompt = (
            f"{question}\n\n"
            f"Options:\n{options_text}\n\n"
            f"Answer with a single letter (A, B, C, or D)."
        )

        img = Image.open(image_path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt},
            ],
        }]
        inputs = self.processor.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(self.device)

        think_ids = self.processor.tokenizer.encode("<think>\n", add_special_tokens=False)
        think_tensor = torch.tensor([think_ids], device=self.device)
        inputs["input_ids"] = torch.cat([inputs["input_ids"], think_tensor], dim=1)
        if "attention_mask" in inputs:
            inputs["attention_mask"] = torch.cat(
                [inputs["attention_mask"], torch.ones_like(think_tensor)], dim=1
            )

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=2048,
                do_sample=False,
                return_dict_in_generate=True,
                output_scores=True,
            )

        generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
        gen_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)

        text_to_search = gen_text.split("</think>")[-1] if "</think>" in gen_text else gen_text
        match = re.search(r'\b([A-D])\b', text_to_search)
        answer = match.group(1) if match else "A"
        return {"answer": answer, "raw_text": gen_text}


# ── Worker ────────────────────────────────────────────────────────────────────
def worker_run(worker_id: int, gpu_id: int, data_chunk: list):
    server = AestheticInferenceServer(gpu_id)

    out_path = OUTPUT_JSON + f".part{worker_id}.json"

    results = []
    if os.path.exists(out_path):
        try:
            with open(out_path, "r", encoding="utf-8") as f:
                results = json.load(f)
            print(f"Worker {worker_id} resuming from {len(results)} existing results")
        except Exception:
            results = []

    done_images = {r["image_path"] for r in results}

    total_to_process = len([it for it in data_chunk if it["image_path"] not in done_images])
    processed_count = 0
    t_start_all = time.time()

    for item_idx, item in enumerate(data_chunk):
        if item["image_path"] in done_images:
            continue

        image_path = resolve_image_path(item["image_path"], IMAGE_DIR)
        if not os.path.exists(image_path):
            print(f"[W{worker_id}] IMAGE NOT FOUND: {image_path}", flush=True)
            continue

        t_img_start = time.time()

        try:
            criteria_results = {}
            print(f"\n[W{worker_id}] ({'='*60})", flush=True)
            print(f"[W{worker_id}] Image {processed_count+1}/{total_to_process}: {item['image_path']}", flush=True)
            print(f"[W{worker_id}] {'-'*60}", flush=True)

            for crit_name in CRITERIA_ORDER:
                t_crit = time.time()
                pred = server.predict_criterion(image_path, crit_name)
                criteria_results[crit_name] = {
                    "level": pred["level"],
                    "score": pred["score"],
                    "argmax_token": pred.get("argmax_token"),
                    "argmax_score": pred.get("argmax_score"),
                    "expected_score": pred.get("expected_score"),
                    "raw_text": pred.get("raw_text", ""),
                }
                elapsed = time.time() - t_crit
                argmax_info = f"argmax={pred.get('argmax_token','?')}({pred.get('argmax_score','?')})" if pred.get('argmax_token') else "NO_TOKEN"
                print(f"[W{worker_id}]   {crit_name:45s} => E[score]={pred['score']:5.2f}  level={pred['level']:6s}  {argmax_info}  ({elapsed:.1f}s)", flush=True)
                print(f"[W{worker_id}]     raw_text: {pred.get('raw_text', '')[:300]}", flush=True)

            total_pred = server.predict_total(image_path)
            t_argmax = f"argmax={total_pred.get('argmax_token','?')}({total_pred.get('argmax_score','?')})" if total_pred.get('argmax_token') else "NO_TOKEN"
            print(f"[W{worker_id}]   {'Total Score':45s} => E[score]={total_pred['total_score']:5.1f}/100  {t_argmax}", flush=True)
            print(f"[W{worker_id}]     raw_text: {total_pred.get('raw_text', '')[:300]}", flush=True)

            mcq_result = server.answer_mcq(
                image_path,
                item.get("question", ""),
                item.get("options", {}),
            )
            mcq_answer = mcq_result["answer"]
            mcq_raw = mcq_result["raw_text"]
            print(f"[W{worker_id}]   {'MCQ Answer':45s} => {mcq_answer}", flush=True)
            print(f"[W{worker_id}]     raw_text: {mcq_raw[:500]}", flush=True)

            t_img_elapsed = time.time() - t_img_start
            processed_count += 1
            avg_time = (time.time() - t_start_all) / processed_count
            eta = avg_time * (total_to_process - processed_count)
            print(f"[W{worker_id}] {'-'*60}", flush=True)
            print(f"[W{worker_id}] Done in {t_img_elapsed:.1f}s | Progress: {processed_count}/{total_to_process} | ETA: {eta/60:.1f}min", flush=True)

            level_summary = {v["level"]: 0 for v in criteria_results.values()}
            for v in criteria_results.values():
                level_summary[v["level"]] += 1
            print(f"[W{worker_id}] Level distribution: {dict(level_summary)} | total={total_pred['total_score']} | mcq={mcq_answer}", flush=True)

            record = {
                "image_path": item["image_path"],
                "criteria": criteria_results,
                "total_score": total_pred["total_score"],
                "total_score_detail": {
                    "argmax_token": total_pred.get("argmax_token"),
                    "argmax_score": total_pred.get("argmax_score"),
                    "expected_score": total_pred.get("expected_score"),
                    "raw_text": total_pred.get("raw_text", ""),
                },
                "question": item.get("question", ""),
                "options": item.get("options", {}),
                "answer": mcq_answer,
                "mcq_raw_text": mcq_raw,
            }

            results.append(record)

            tmp_path = out_path + ".tmp"
            with open(tmp_path, "w", encoding="utf-8") as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            os.replace(tmp_path, out_path)
            print(f"[W{worker_id}] Saved to {out_path} ({len(results)} total records)", flush=True)

        except Exception as e:
            print(f"[W{worker_id}] ERROR on {item['image_path']}: {e}", flush=True)
            import traceback
            traceback.print_exc()
            continue

    total_elapsed = time.time() - t_start_all
    print(f"\n[W{worker_id}] ALL DONE | {len(results)} records | {total_elapsed/60:.1f}min total", flush=True)


# ── Merge ─────────────────────────────────────────────────────────────────────
def merge_results():
    files = glob.glob(OUTPUT_JSON + ".part*.json")
    unique = {}
    for f in files:
        with open(f, "r", encoding="utf-8") as fp:
            for x in json.load(fp):
                unique[x["image_path"]] = x
    all_results = list(unique.values())
    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"Merged {len(all_results)} results to {OUTPUT_JSON}")


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    with open(TEST_JSON, "r", encoding="utf-8") as f:
        data = json.load(f)

    done_set = set()
    for part_file in glob.glob(OUTPUT_JSON + ".part*.json"):
        try:
            with open(part_file, "r") as fp:
                for x in json.load(fp):
                    done_set.add(x["image_path"])
        except Exception:
            continue

    remaining = [item for item in data if item["image_path"] not in done_set]
    print(f"Total: {len(data)}, already done: {len(done_set)}, remaining: {len(remaining)}")

    if len(remaining) == 0:
        print("All done, merging...")
        merge_results()
        return

    NUM_GPUS = torch.cuda.device_count()
    TOTAL_WORKERS = NUM_GPUS * WORKERS_PER_GPU
    print(f"GPUs: {NUM_GPUS}, workers: {TOTAL_WORKERS}")

    chunk_size = math.ceil(len(remaining) / TOTAL_WORKERS)
    chunks = [remaining[i:i + chunk_size] for i in range(0, len(remaining), chunk_size)]

    processes = []
    for i, chunk in enumerate(chunks):
        gpu_id = i // WORKERS_PER_GPU
        p = Process(target=worker_run, args=(i, gpu_id, chunk))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    merge_results()


if __name__ == "__main__":
    main()