#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import json
import math
import os
import random
from collections import Counter, defaultdict
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple


CRITERIA = [
    "Color Harmony",
    "Visual Style Consistency",
    "Sharpness",
    "Light and Shadow Modeling",
    "Creativity and Originality",
    "Exposure Control",
    "Application of Classical Composition Principles",
    "Depth of Field and Layering",
    "Visual Center Stability",
    "Visual Flow Guidance",
    "Structural Support Stability",
    "Appropriateness of Negative Space",
    "Subject Integrity",
]


PROMPT_SCORE = (
    "You are an aesthetics expert. Evaluate the image on the following 13 criteria:\n"
    "Color Harmony, Visual Style Consistency, Sharpness, Light and Shadow Modeling, "
    "Creativity and Originality, Exposure Control, Application of Classical Composition Principles, "
    "Depth of Field and Layering, Visual Center Stability, Visual Flow Guidance, "
    "Structural Support Stability, Appropriateness of Negative Space, Subject Integrity.\n\n"
    "For each criterion, output a numeric score in [0.0, 10.0].\n"
    "Then output total_score as an integer in [0, 100].\n"
    "Return STRICT JSON only.\n"
    "JSON format:\n"
    "{\n"
    "  \"criteria\": {\n"
    "    \"Color Harmony\": {\"score\": 0.0}\n"
    "  },\n"
    "  \"total_score\": 0\n"
    "}\n\n"
    "<image>"
)


PROMPT_MULTITASK = (
    "You are an aesthetics expert. Evaluate the image on the following 13 criteria:\n"
    "Color Harmony, Visual Style Consistency, Sharpness, Light and Shadow Modeling, "
    "Creativity and Originality, Exposure Control, Application of Classical Composition Principles, "
    "Depth of Field and Layering, Visual Center Stability, Visual Flow Guidance, "
    "Structural Support Stability, Appropriateness of Negative Space, Subject Integrity.\n\n"
    "For each criterion, output:\n"
    "1) score in [0.0, 10.0]\n"
    "2) level in {A,B,C} where A:<5, B:[5,7), C:>=7\n"
    "Then output total_score as an integer in [0, 100].\n"
    "Return STRICT JSON only.\n"
    "JSON format:\n"
    "{\n"
    "  \"criteria\": {\n"
    "    \"Color Harmony\": {\"score\": 0.0, \"level\": \"A|B|C\"}\n"
    "  },\n"
    "  \"total_score\": 0\n"
    "}\n\n"
    "<image>"
)


@dataclass
class Sample:
    image_path: str
    scores: Dict[str, float]
    levels: Dict[str, str]
    total_score: int
    boundary_count: int


def clamp_score(x: float) -> float:
    return max(0.0, min(10.0, x))


def clamp_int(x: int, lo: int, hi: int) -> int:
    return max(lo, min(hi, x))


def score_to_level(score: float) -> str:
    if score < 5.0:
        return "A"
    if score < 7.0:
        return "B"
    return "C"


def safe_float(x: Any, default: float = 0.0) -> float:
    try:
        return float(x)
    except Exception:
        return default


def safe_int(x: Any, default: int = 0) -> int:
    try:
        return int(round(float(x)))
    except Exception:
        return default


def parse_assistant_payload(payload: str) -> Dict[str, Any]:
    obj = json.loads(payload)
    if not isinstance(obj, dict):
        raise ValueError("assistant payload is not dict")
    return obj


def normalize_record(rec: Dict[str, Any], boundary_margin: float) -> Sample:
    assistant = parse_assistant_payload(rec["messages"][1]["content"])
    criteria = assistant.get("criteria", {})

    scores: Dict[str, float] = {}
    levels: Dict[str, str] = {}
    boundary_count = 0

    for c in CRITERIA:
        v = (criteria.get(c, {}) or {}).get("score", 0.0)
        s = round(clamp_score(safe_float(v, 0.0)), 1)
        lv = score_to_level(s)
        scores[c] = s
        levels[c] = lv
        if abs(s - 5.0) <= boundary_margin or abs(s - 7.0) <= boundary_margin:
            boundary_count += 1

    t = assistant.get("total_score", 0)
    total_score = clamp_int(safe_int(t, 0), 0, 100)

    img = rec.get("images", [""])
    image_path = str(img[0]) if isinstance(img, list) and img else ""

    return Sample(
        image_path=image_path,
        scores=scores,
        levels=levels,
        total_score=total_score,
        boundary_count=boundary_count,
    )


def build_score_record(s: Sample) -> Dict[str, Any]:
    assistant = {
        "criteria": {k: {"score": round(v, 1)} for k, v in s.scores.items()},
        "total_score": int(s.total_score),
    }
    return {
        "messages": [
            {"role": "user", "content": PROMPT_SCORE},
            {"role": "assistant", "content": json.dumps(assistant, ensure_ascii=False, indent=2)},
        ],
        "images": [s.image_path],
    }


def build_multitask_record(s: Sample) -> Dict[str, Any]:
    assistant = {
        "criteria": {
            k: {"score": round(s.scores[k], 1), "level": s.levels[k]} for k in CRITERIA
        },
        "total_score": int(s.total_score),
    }
    return {
        "messages": [
            {"role": "user", "content": PROMPT_MULTITASK},
            {"role": "assistant", "content": json.dumps(assistant, ensure_ascii=False, indent=2)},
        ],
        "images": [s.image_path],
    }


def compute_level_counts(samples: List[Sample]) -> Dict[str, Counter]:
    counts: Dict[str, Counter] = {c: Counter() for c in CRITERIA}
    for s in samples:
        for c in CRITERIA:
            counts[c][s.levels[c]] += 1
    return counts


def compute_sample_dup(
    s: Sample,
    level_counts: Dict[str, Counter],
    max_class_weight: float,
    boundary_bonus: float,
    max_dup: int,
) -> int:
    weights: List[float] = []
    for c in CRITERIA:
        cnt = level_counts[c]
        max_freq = max(cnt.values()) if cnt else 1
        cur = cnt.get(s.levels[c], 1)
        w = math.sqrt(float(max_freq) / float(max(1, cur)))
        w = min(max_class_weight, max(1.0, w))
        weights.append(w)

    avg_w = sum(weights) / len(weights)
    raw_dup = avg_w + boundary_bonus * float(s.boundary_count)
    dup = int(round(raw_dup))
    return clamp_int(dup, 1, max_dup)


def save_json(path: str, obj: Any) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(obj, f, ensure_ascii=False, indent=2)


def build_stats(
    train_samples: List[Sample],
    val_samples: List[Sample],
    level_counts: Dict[str, Counter],
    dup_counts: Counter,
    balanced_size: int,
) -> Dict[str, Any]:
    agg = Counter()
    for c in CRITERIA:
        agg.update(level_counts[c])

    total_lv = sum(agg.values())
    agg_ratio = {k: round(v / total_lv, 6) for k, v in agg.items()} if total_lv else {}

    per_criterion = {}
    for c in CRITERIA:
        cnt = level_counts[c]
        n = sum(cnt.values())
        per_criterion[c] = {
            "counts": dict(cnt),
            "ratio": {k: round(cnt[k] / n, 6) if n else 0.0 for k in ["A", "B", "C"]},
        }

    train_total = [s.total_score for s in train_samples]
    val_total = [s.total_score for s in val_samples]

    return {
        "train_count": len(train_samples),
        "val_count": len(val_samples),
        "balanced_train_count": balanced_size,
        "level_distribution_aggregate": dict(agg),
        "level_distribution_aggregate_ratio": agg_ratio,
        "level_distribution_per_criterion": per_criterion,
        "duplication_histogram": {str(k): v for k, v in sorted(dup_counts.items())},
        "boundary_stats": {
            "train_mean_boundary_count": round(sum(s.boundary_count for s in train_samples) / max(1, len(train_samples)), 4),
            "train_max_boundary_count": max((s.boundary_count for s in train_samples), default=0),
        },
        "total_score": {
            "train_mean": round(sum(train_total) / max(1, len(train_total)), 4),
            "val_mean": round(sum(val_total) / max(1, len(val_total)), 4),
            "train_min": min(train_total) if train_total else None,
            "train_max": max(train_total) if train_total else None,
            "val_min": min(val_total) if val_total else None,
            "val_max": max(val_total) if val_total else None,
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--train_json",
        type=str,
        default="./codabench_portrait_train.json",
    )
    parser.add_argument(
        "--val_json",
        type=str,
        default="./codabench_portrait_val.json",
    )
    parser.add_argument("--output_dir", type=str, default=".")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--boundary_margin", type=float, default=0.4)
    parser.add_argument("--boundary_bonus", type=float, default=0.3)
    parser.add_argument("--max_class_weight", type=float, default=4.0)
    parser.add_argument("--max_dup", type=int, default=5)
    parser.add_argument(
        "--stats_json",
        type=str,
        default="logs/prepare_dataset_v2_stats.json",
    )
    args = parser.parse_args()

    random.seed(args.seed)

    with open(args.train_json, "r", encoding="utf-8") as f:
        train_raw = json.load(f)
    with open(args.val_json, "r", encoding="utf-8") as f:
        val_raw = json.load(f)

    train_samples = [normalize_record(x, args.boundary_margin) for x in train_raw]
    val_samples = [normalize_record(x, args.boundary_margin) for x in val_raw]

    level_counts = compute_level_counts(train_samples)

    train_score = [build_score_record(s) for s in train_samples]
    train_multitask = [build_multitask_record(s) for s in train_samples]
    val_score = [build_score_record(s) for s in val_samples]
    val_multitask = [build_multitask_record(s) for s in val_samples]

    train_balanced: List[Dict[str, Any]] = []
    dup_hist = Counter()
    for s in train_samples:
        dup = compute_sample_dup(
            s=s,
            level_counts=level_counts,
            max_class_weight=args.max_class_weight,
            boundary_bonus=args.boundary_bonus,
            max_dup=args.max_dup,
        )
        dup_hist[dup] += 1
        rec = build_multitask_record(s)
        for _ in range(dup):
            train_balanced.append(rec)

    random.shuffle(train_balanced)

    os.makedirs(args.output_dir, exist_ok=True)

    f_train_score = os.path.join(args.output_dir, "codabench_portrait_train_score_v2.json")
    f_train_multitask = os.path.join(args.output_dir, "codabench_portrait_train_multitask_v2.json")
    f_train_balanced = os.path.join(args.output_dir, "codabench_portrait_train_balanced_v2.json")
    f_val_score = os.path.join(args.output_dir, "codabench_portrait_val_score_v2.json")
    f_val_multitask = os.path.join(args.output_dir, "codabench_portrait_val_multitask_v2.json")

    save_json(f_train_score, train_score)
    save_json(f_train_multitask, train_multitask)
    save_json(f_train_balanced, train_balanced)
    save_json(f_val_score, val_score)
    save_json(f_val_multitask, val_multitask)

    stats = build_stats(
        train_samples=train_samples,
        val_samples=val_samples,
        level_counts=level_counts,
        dup_counts=dup_hist,
        balanced_size=len(train_balanced),
    )
    save_json(args.stats_json, stats)

    print("=" * 60)
    print("saved:", f_train_score)
    print("saved:", f_train_multitask)
    print("saved:", f_train_balanced)
    print("saved:", f_val_score)
    print("saved:", f_val_multitask)
    print("stats:", args.stats_json)
    print("train_count:", len(train_samples), "balanced_count:", len(train_balanced), "val_count:", len(val_samples))
    print("dup_hist:", dict(sorted(dup_hist.items())))
    print("=" * 60)


if __name__ == "__main__":
    main()