modules/infographics_generator/template_utils.py

from typing import Dict, List, Tuple, Optional, Union
import random
import json
from modules.infographics_generator.color_utils import get_contrast_color, has_indistinguishable_colors, generate_distinct_palette
import os

# 添加全局字典来跟踪模板使用频率
template_usage_counter = {}
field_order = ['x', 'y', 'y2', 'y3', 'size', 'group', 'group2', 'group3']

# ============================================================
# chart_type 白名单
# ------------------------------------------------------------
# 通过外部 JSON 文件指定允许使用的 chart_type 子集（例如只关心
# "bump chart" / "bar chart"）。analyze_templates 和
# check_template_compatibility 都会读这份白名单：
#   * 文件不存在 / 为空 / 内容不是非空列表  ->  不过滤（保留原全量行为）
#   * 文件存在且为非空列表                   ->  只保留列表内的 chart_type
#
# 文件路径默认为工作目录下的 ``allowed_chart_types.json``，可通过
# 环境变量 ``ALLOWED_CHART_TYPES_FILE`` 覆盖。
#
# 支持的文件格式（任选其一）：
#   ["bump chart", "bar chart"]
#   {"chart_types": ["bump chart", "bar chart"]}
# ============================================================
_ALLOWED_CHART_TYPES_DEFAULT_PATH = "allowed_chart_types.json"
_allowed_chart_types_cache: Optional[set] = None
_allowed_chart_types_cache_mtime: Optional[float] = None
_allowed_chart_types_cache_path: Optional[str] = None


def _get_allowed_chart_types_path() -> str:
    return os.environ.get("ALLOWED_CHART_TYPES_FILE", _ALLOWED_CHART_TYPES_DEFAULT_PATH)


def _load_allowed_chart_types() -> Optional[set]:
    """加载允许的 chart_type 白名单；返回 None 表示不过滤。

    带 mtime 缓存：文件被修改时自动重新读取，多线程/多进程环境下也安全
    （ProcessPoolExecutor 的 worker 各自独立 import，缓存彼此隔离）。
    """
    global _allowed_chart_types_cache, _allowed_chart_types_cache_mtime, _allowed_chart_types_cache_path

    path = _get_allowed_chart_types_path()
    if not os.path.exists(path):
        _allowed_chart_types_cache = None
        _allowed_chart_types_cache_mtime = None
        _allowed_chart_types_cache_path = path
        return None

    mtime = os.path.getmtime(path)
    if (
        _allowed_chart_types_cache_path == path
        and _allowed_chart_types_cache_mtime == mtime
    ):
        return _allowed_chart_types_cache

    with open(path, "r", encoding="utf-8") as f:
        raw = json.load(f)
    if isinstance(raw, dict):
        raw = raw.get("chart_types", [])
    if not isinstance(raw, list) or len(raw) == 0:
        _allowed_chart_types_cache = None
    else:
        _allowed_chart_types_cache = set(raw)

    _allowed_chart_types_cache_mtime = mtime
    _allowed_chart_types_cache_path = path
    return _allowed_chart_types_cache


def _is_chart_type_allowed(chart_type: str) -> bool:
    allowed = _load_allowed_chart_types()
    if allowed is None:
        return True
    return chart_type in allowed

def flatten(lst):
    """Flattens a nested list into a single list."""
    result = []
    for item in lst:
        if isinstance(item, list):  # Check if the item is a list
            result.extend(flatten(item))  # Recursively flatten the sublist
        else:
            result.append(item)  # Add the non-list item to the result
    return result

def get_flatten_fields(required_fields) -> List[str]:
    """Flatten a nested list of fields into a single list"""
    lst = flatten(required_fields)
    lst = [field for field in field_order if field in lst]
    return lst

def get_unique_fields_and_types(
        required_fields: Union[List[str], List[List[str]]],
        required_fields_type: Union[List[List[str]], List[List[List[str]]]],
        required_fields_range: Optional[Union[List[List[int]], List[List[List[int]]]]] = None
    ) -> Tuple[List[str], Dict[str, str], List[List[int]]]:
    """Extract unique fields and their corresponding types from nested structure"""
    field_types = {}
    field_ranges = {}

    # Check if required_fields is a list of lists
    if required_fields and isinstance(required_fields[0], list):
        # Handle list of lists case
        for i, (fields_group, types_group) in enumerate(zip(required_fields, required_fields_type)):
            range_group = required_fields_range[i] if required_fields_range != None else [[float('-inf'), float('inf')] for _ in fields_group]
            for field, type_list, range_list in zip(fields_group, types_group, range_group):
                if field not in field_types:
                    field_types[field] = type_list[0]  # Use first type from the list
                    field_ranges[field] = range_list  # Use first range from the list
    else:
        # Handle simple list case
        range_list = required_fields_range if required_fields_range != None else [[float('-inf'), float('inf')] for _ in required_fields]
        for field, type_list, range_val in zip(required_fields, required_fields_type, range_list):
            if field not in field_types:
                field_types[field] = type_list[0]  # Use first type from the list
                field_ranges[field] = range_val  # Use first range from the list

    # Order fields according to field_order, keeping only those that exist
    ordered_fields = [field for field in field_order if field in field_types]
    for field in field_ranges:
        r = field_ranges[field]
        try:
            if r[0] == "-inf":
                r[0] = float('-inf')
            if r[1] == "inf":
                r[1] = float('inf')
        except:
            pass
    ordered_ranges = [field_ranges[field] for field in ordered_fields]

    return ordered_fields, field_types, ordered_ranges

def analyze_templates(templates: Dict) -> Tuple[int, Dict[str, str], int]:
    """Analyze templates and return count, data requirements and unique colors count"""
    template_count = 0
    template_requirements = {}
    template_list = []
    unique_colors = set()
    requirement_dump = {}

    for engine, templates_dict in templates.items():
        for chart_type, chart_names_dict in templates_dict.items():
            if not _is_chart_type_allowed(chart_type):
                continue
            for chart_name, template_info in chart_names_dict.items():
                if 'base' in chart_name:
                    continue
                if engine == 'vegalite_py':
                    continue
                template_list.append(f"{chart_type} / {chart_name}")
                template_count += 1
                if 'requirements' in template_info:
                    req = template_info['requirements']

                    # Count unique required colors
                    if 'required_other_colors' in req:
                        for color in req['required_other_colors']:
                            unique_colors.add(color)

                    if 'required_fields_colors' in req:
                        for color in req['required_fields_colors']:
                            unique_colors.add(color)

                    if 'required_fields' in req and 'required_fields_type' in req:
                        template_requirements[f"{engine}/{chart_type}/{chart_name}"] = template_info['requirements']
                        requirement_dump[chart_name] = template_info['requirements']

    # print("template_count", template_count)
    if not os.path.exists("template_list.txt"):
        f = open("template_list.txt", "w")
        f.write("\n".join(template_list))
        f.close()
    if not os.path.exists("requirement_dump.json"):
        f = open("requirement_dump.json", "w")
        f.write(json.dumps(requirement_dump, indent=4))
        f.close()

    return template_count, template_requirements

# block_list = ["multiple_line_graph_06", "layered_area_chart_02", "multiple_area_chart_01", "stacked_area_chart_01", "stacked_area_chart_03"]
block_list = []

def check_field_color_compatibility(requirements: Dict, data: Dict) -> bool:
    """Check if the field color is compatible with the template"""
    if len(requirements.get('required_fields_colors', [])) > 0 and len(data.get("colors", {}).get("field", {}).keys()) == 0:
        return False
    data_fields = get_flatten_fields(requirements.get('required_fields',[]))
    for color_field in requirements.get('required_fields_colors', []):
        field_column = None
        for i, field in enumerate(data_fields):
            if field == color_field:
                field_column = data.get("data", {}).get("columns", {})[i]
                break
        if field_column is None:
            return False
        field_name = field_column["name"]
        for value in data.get("data", {}).get("data", []):
            if value[field_name] not in data.get("colors", {}).get("field", {}).keys():
                return False
    return True

def check_field_icon_compatibility(requirements: Dict, data: Dict) -> bool:
    """Check if the field icon is compatible with the template"""
    if len(requirements.get('required_fields_icons', [])) > 0 and len(data.get("images", {}).get("field", {}).keys()) == 0:
        return False
    data_fields = get_flatten_fields(requirements.get('required_fields',[]))
    for icon_field in requirements.get('required_fields_icons', []):
        for i, field in enumerate(data_fields):
            if field == icon_field:
                field_column = data.get("data", {}).get("columns", {})[i]
                break
        if field_column is None:
            return False
        field_name = field_column["name"]
        for value in data.get("data", {}).get("data", []):
            if value[field_name] not in data.get("images", {}).get("field", {}).keys():
                return False
    return True

def check_template_compatibility(data: Dict, templates: Dict, specific_chart_name: str = None) -> List[str]:
    """Check which templates are compatible with the given data"""
    compatible_templates = []

    def normalize_range_bound(value):
        if isinstance(value, str):
            normalized = value.strip().lower()
            if normalized in {"inf", "+inf", "infinity", "+infinity"}:
                return float("inf")
            if normalized in {"-inf", "-infinity"}:
                return float("-inf")
            try:
                return float(normalized)
            except ValueError:
                return value
        return value

    # Get the combination type from the data
    combination_type = data.get("data", {}).get("type_combination", "")
    combination_types = [col["data_type"] for col in data["data"]["columns"]]
    if combination_type == "":
        combination_type = " + ".join(combination_types)

    if not combination_type:
        return compatible_templates

    for engine, templates_dict in templates.items():
        for chart_type, chart_names_dict in templates_dict.items():
            if not _is_chart_type_allowed(chart_type):
                continue
            for chart_name, template_info in chart_names_dict.items():
                if 'base' in chart_name:
                    continue
                if engine == 'vegalite_py':
                    continue

                template_key = f"{engine}/{chart_type}/{chart_name}"

                if specific_chart_name and specific_chart_name != chart_name:
                    continue

                try:
                    if 'requirements' in template_info:
                        req = template_info['requirements']
                        hierarchy = req.get('hierarchy', [])
                        if 'required_fields' in req and 'required_fields_type' in req:
                            ordered_fields, field_types, ordered_ranges = get_unique_fields_and_types(
                                req['required_fields'],
                                req['required_fields_type'],
                                req.get('required_fields_range', None)
                            )
                            data_types = [field_types[field] for field in ordered_fields]
                            data_type_str = ' + '.join(data_types)
                            if len(req.get('required_fields_colors', [])) > 0 and len(data.get("colors", {}).get("field", [])) == 0:
                                # print(f"template {template_key} failed color compatibility check")
                                continue

                            # if len(req.get('required_fields_icons', [])) > 0 and len(data.get("images", {}).get("field", [])) == 0:
                            #     print(f"template {template_key} failed icon compatibility check")
                            #     continue

                            if not check_field_color_compatibility(req, data):
                                # print(f"template {template_key} failed color compatibility check")
                                continue

                            if not check_field_icon_compatibility(req, data):
                                # print(f"template {template_key} failed icon compatibility check")
                                continue
                            # print("data_types", data_types)
                            # print("combination_types", combination_types)
                            # 如果data_types和combination_types相同，或者data_types是combination_types的一个子序列
                            if len(data_types) == len(combination_types):# or all(data_type in combination_types for data_type in data_types):
                                check_flag = True
                                for data_type, combination_type in zip(data_types, combination_types[:len(data_types)]):
                                    if data_type == "categorical" and (combination_type == "temporal" or combination_type == "categorical"):
                                        pass
                                    elif data_type == "numerical" and combination_type == "numerical":
                                        pass
                                    elif data_type == "temporal" and combination_type == "temporal":
                                        pass
                                    else:
                                        check_flag = False
                                        break
                                if not check_flag:
                                    # print(f"template {template_key} failed data type compatibility check")
                                    continue
                            else:
                                # print(f"template {template_key} failed data type compatibility check")
                                continue

                            disallow_temporal_fields = set(req.get('disallow_temporal_fields', []))
                            if disallow_temporal_fields:
                                rejected_for_temporal = False
                                for i, field in enumerate(ordered_fields):
                                    if (
                                        field in disallow_temporal_fields
                                        and i < len(data["data"]["columns"])
                                        and data["data"]["columns"][i].get("data_type") == "temporal"
                                    ):
                                        rejected_for_temporal = True
                                        break
                                if rejected_for_temporal:
                                    continue

                            flag = True
                            # print("check compatibility")
                            for i, range_bounds in enumerate(ordered_ranges):
                                if i >= len(data["data"]["columns"]):
                                    flag = False
                                    break
                                min_bound = normalize_range_bound(range_bounds[0])
                                max_bound = normalize_range_bound(range_bounds[1])

                                if data["data"]["columns"][i]["data_type"] in ["temporal", "categorical"]:
                                    key = data["data"]["columns"][i]["name"]
                                    unique_values = list(set(value[key] for value in data["data"]["data"]))
                                    if len(unique_values) > max_bound or len(unique_values) < min_bound:
                                        flag = False
                                        break
                                    else:
                                        pass
                                        #if specific_chart_name and specific_chart_name == chart_name:
                                        #    print(f"template {template_key} matched", data["name"], len(unique_values), range)
                                elif data["data"]["columns"][i]["data_type"] in ["numerical"]:
                                    key = data["data"]["columns"][i]["name"]
                                    min_value = min(value[key] for value in data["data"]["data"])
                                    max_value = max(value[key] for value in data["data"]["data"])
                                    if min_value < min_bound or max_value > max_bound:
                                        flag = False
                                        break
                                    elif "diverging" in chart_name and min_value >= 0 and min_bound < 0:
                                        flag = False
                                        break
                                    elif "scatterplot" in chart_name and min_value >= 0 and min_bound < 0:
                                        flag = False
                                        break
                            for i, field in enumerate(ordered_fields):
                                if field == "group":
                                    x_col = [j for j, field2 in enumerate(ordered_fields) if field2 == "x"][0]
                                    x_name = data["data"]["columns"][x_col]["name"]
                                    field_name = data["data"]["columns"][i]["name"]
                                    num_unique_x  = len(list(set(value[x_name] for value in data["data"]["data"])))
                                    num_unique_comb = len(list(set(str(value[x_name]) + ' ' + str(value[field_name]) for value in data["data"]["data"])))
                                    if field in hierarchy:
                                        if num_unique_comb > num_unique_x:
                                            flag = False
                                            break
                                    else:
                                        if num_unique_comb == num_unique_x:
                                            flag = False
                                            break
                                elif field == "group2":
                                    x_col = [j for j, field2 in enumerate(ordered_fields) if field2 == "x"][0]
                                    group_col = [j for j, field2 in enumerate(ordered_fields) if field2 == "group"][0]
                                    x_name = data["data"]["columns"][x_col]["name"]
                                    group_name = data["data"]["columns"][group_col]["name"]
                                    field_name = data["data"]["columns"][i]["name"]
                                    num_unique_x  = len(list(set(str(value[x_name]) + ' ' + str(value[group_name]) for value in data["data"]["data"])))
                                    num_unique_comb = len(list(set(str(value[x_name]) + ' ' + str(value[group_name]) + ' ' + str(value[field_name]) for value in data["data"]["data"])))
                                    if field in hierarchy:
                                        if num_unique_comb > num_unique_x:
                                            flag = False
                                            break
                                    else:
                                        if num_unique_comb == num_unique_x:
                                            flag = False
                                            break
                            if flag:
                                if specific_chart_name == None or specific_chart_name == chart_name:
                                    compatible_templates.append((template_key, ordered_fields))
                except:
                    pass
    #print("compatible_templates", compatible_templates)
    return compatible_templates


import fcntl  # 用于文件锁
def select_template(compatible_templates: List[str]) -> Tuple[str, str, str]:
    """
    根据variation.json中的使用统计选择模板
    按照使用频率分为4个level，优先选择使用较少的level
    同level内按照具体使用次数加权随机选择
    使用文件锁确保多线程安全
    """
    # 读取variation.json，使用文件锁
    try:
        with open('variation.json', 'r') as f:
            # 获取文件锁
            fcntl.flock(f, fcntl.LOCK_EX)
            try:
                variation_stats = json.load(f)
            finally:
                # 释放文件锁
                fcntl.flock(f, fcntl.LOCK_UN)
    except:
        variation_stats = {}

    # 获取所有模板的使用次数
    template_counts = []
    for template_info in compatible_templates:
        template_key = template_info[0]
        _, chart_type, chart_name = template_key.split('/')

        # 如果variation_stats为空,所有模板使用次数都为0
        if not variation_stats:
            count = 0
        else:
            if chart_type not in variation_stats:
                variation_stats[chart_type] = {"total_count": 0}

            if chart_name not in variation_stats[chart_type]:
                variation_stats[chart_type][chart_name] = 0

            count = variation_stats[chart_type][chart_name]

        template_counts.append((template_info, count))

    # 按使用次数排序并分level
    template_counts.sort(key=lambda x: x[1])
    n = len(template_counts)
    
    # 如果没有可用模板，返回 None
    if n == 0:
        return None, None, None, None
    
    level_size = max(1, n // 4)
    # 找出使用次数最少的模板
    min_count = min(c for _, c in template_counts)
    min_level_templates = [(t, c) for t, c in template_counts if c == min_count]

    # 固定选择第一个最少使用的模板
    selected_index = 0

    selected_template, _ = min_level_templates[selected_index]
    [template_key, ordered_fields] = selected_template
    print("selected_template", selected_template)

    # 更新variation.json，使用文件锁
    engine, chart_type, chart_name = template_key.split('/')
    if os.environ.get("CHARTPIPELINE_SKIP_VARIATION_STATS_UPDATE") == "1":
        return engine, chart_type, chart_name, ordered_fields
    try:
        with open('variation.json', 'r+') as f:
            # 获取文件锁
            fcntl.flock(f, fcntl.LOCK_EX)
            try:
                # 重新读取以确保获取最新数据
                variation_stats = json.load(f)

                # 初始化如果不存在
                if chart_type not in variation_stats:
                    variation_stats[chart_type] = {"total_count": 0}
                if chart_name not in variation_stats[chart_type]:
                    variation_stats[chart_type][chart_name] = 0

                # 更新计数
                variation_stats[chart_type][chart_name] += 1
                variation_stats[chart_type]["total_count"] += 1

                # 写入更新后的数据
                f.seek(0)
                json.dump(variation_stats, f, indent=2)
                f.truncate()
            finally:
                # 释放文件锁
                fcntl.flock(f, fcntl.LOCK_UN)
    except FileNotFoundError:
        # 如果文件不存在,创建新的variation_stats
        variation_stats = {
            chart_type: {
                "total_count": 1,
                chart_name: 1
            }
        }
        with open('variation.json', 'w') as f:
            json.dump(variation_stats, f, indent=2)
    return engine, chart_type, chart_name, ordered_fields


def process_template_requirements(requirements: Dict, data: Dict, engine: str, chart_name: str) -> None:
    """处理模板的颜色要求"""
    default_colors = {
        "text_color": "#333333",
        "background_color": "#ffffff",
        "field": {},
        "other": {"primary": "#4e79a7"},
        "available_colors": ["#4e79a7", "#f28e2b", "#59a14f", "#e15759", "#76b7b2", "#edc948"]
    }
    default_dark_colors = {
        "text_color": "#ffffff",
        "background_color": "#1f2933",
        "field": {},
        "other": {"primary": "#76b7b2"},
        "available_colors": ["#76b7b2", "#f28e2b", "#59a14f", "#e15759", "#9c755f", "#edc948"]
    }

    if not isinstance(data.get("colors"), dict):
        data["colors"] = json.loads(json.dumps(default_colors))
    else:
        data["colors"].setdefault("text_color", default_colors["text_color"])
        data["colors"].setdefault("background_color", default_colors["background_color"])
        data["colors"].setdefault("field", {})
        data["colors"].setdefault("other", {})
        data["colors"]["other"].setdefault("primary", default_colors["other"]["primary"])
        data["colors"].setdefault("available_colors", default_colors["available_colors"])

    if not isinstance(data.get("colors_dark"), dict):
        data["colors_dark"] = json.loads(json.dumps(default_dark_colors))
    else:
        data["colors_dark"].setdefault("text_color", default_dark_colors["text_color"])
        data["colors_dark"].setdefault("background_color", default_dark_colors["background_color"])
        data["colors_dark"].setdefault("field", {})
        data["colors_dark"].setdefault("other", {})
        data["colors_dark"]["other"].setdefault("primary", default_dark_colors["other"]["primary"])
        data["colors_dark"].setdefault("available_colors", default_dark_colors["available_colors"])

    colors = data["colors"]
    colors_dark = data["colors_dark"]

    if len(colors["field"]) > 1:
        # 检查颜色是否可区分
        field_colors = list(colors["field"].values())
        if has_indistinguishable_colors(field_colors):
            # 如果颜色不可区分,使用主色生成新的调色板
            primary_color = colors["other"]["primary"]
            new_colors = generate_distinct_palette(primary_color, len(field_colors))
            # 更新颜色字典
            for i, field in enumerate(colors["field"].keys()):
                colors["field"][field] = new_colors[i]

    if len(colors_dark["field"]) > 1:
        # 检查颜色是否可区分
        field_colors = list(colors_dark["field"].values())
        if has_indistinguishable_colors(field_colors):
            # 如果颜色不可区分,使用主色生成新的调色板
            primary_color = colors_dark["other"]["primary"]
            new_colors = generate_distinct_palette(primary_color, len(field_colors))
            # 更新颜色字典
            for i, field in enumerate(colors_dark["field"].keys()):
                colors_dark["field"][field] = new_colors[i]


    if len(requirements.get("required_other_colors", [])) > 0:
        for key in requirements["required_other_colors"]:
            if key == "positive" and "positive" not in colors["other"]:
                colors["other"]["positive"] = colors["other"]["primary"]
            elif key == "negative" and "negative" not in colors["other"]:
                colors["other"]["negative"] = get_contrast_color(colors["other"]["primary"])

    colors_dark["text_color"] = "#ffffff"
    # if ('donut' in chart_name or 'pie' in chart_name) and engine == 'vegalite_py':
    #     data["variables"]["height"] = 500
    #     data["variables"]["width"] = 500
    # else:
    #     if "min_height" in requirements:
    #         data["variables"]["height"] = max(600, requirements["min_height"])
    #     elif 'height' in requirements:
    #         data["variables"]["height"] = max(600, requirements["height"][0])

    #     if "min_width" in requirements:
    #         data["variables"]["width"] = max(800, requirements["min_width"])
    #     elif 'width' in requirements:
    #         data["variables"]["width"] = max(600, requirements["width"][0])