"""
DoE case generation from a Investigation config dictionary.

Transforms the structured JSON produced by ``ivst init`` / :func:`investigation.doeargs.args2config.args2config`
into a flat list of experiment cases ready to be committed and submitted.
"""
from typing import List, Dict
import itertools
import numpy as np
import itertools
from investigation.database.dataframe import format_and_Q_eq_in
import pandas as pd
from numpy.random import Generator, PCG64


def config2jobs(config: Dict[str, Dict[str, Dict]]):
    """Expand a DoE config into a deduplicated list of experiment cases.

    The config has three top-level sections:

    * ``meta`` — ``num_seeds``, ``seed``, and optional ``include`` overrides.
    * ``shared`` — parameters common to every model variant.
    * *<ModelName>* — one section per model variant with model-specific params.

    Args:
        config: Parsed DoE JSON as a nested dictionary.  Typically produced by
            :func:`investigation.doeargs.args2config.args2config` and then edited
            by the user to set ``range`` lists.

    Returns:
        A tuple ``(segs, cases)`` where *segs* is the list of per-segment
        parameter dicts (before Cartesian product) and *cases* is the
        deduplicated flat list of individual experiment parameter dicts.
    """
    shared = shrinkFieldRange(config["shared"])
    meta = config["meta"]
    seed: int = meta["seed"]["value"]
    num_seeds: int = meta["num_seeds"]["value"]

    includes: List[List[List[Dict[str, Dict]]]] = meta["include"]["value"]

    # input(f"Generating jobs with seed {seed} and num_seeds {num_seeds}. Press Enter to continue...")

    # Generate num_seeds seeds according seed
    seeds = Generator(PCG64(seed)).integers(low=0, high=2**31 - 1, size=num_seeds).tolist()

    shared["seed"] = seeds

    shared_group = []
    for include in includes:
        for override in itertools.product(*include):
            override_dict = {}
            for item in override:
                override_dict.update(shrinkFieldRange(item))
            shared_group.append(override_dict)

    segs: List[Dict] = []
    for key, val in config.items():
        if key in ["shared", "meta"]:
            continue
        if (len(shared_group) == 0):
            segs.append({
                **shared,
                **shrinkFieldRange(val),
            })
        else:
            for new_shared in shared_group:
                segs.append({
                    **shared,
                    **shrinkFieldRange(val),
                    **new_shared,
                })

    cases: List[Dict] = []
    for val in segs:
        cases.extend(rangeProduct(val))

    # 去重
    # case_df: pd.DataFrame = pd.DataFrame.from_dict([{**cases[0], "unique_index": 0}], orient='columns')
    case_df = pd.DataFrame.from_dict(cases, orient='columns')
    case_df["unique_index"] = range(len(cases))
    case_df = case_df[0:0]  # 清空 DataFrame，只保留列结构
    for i, case in enumerate(cases):
        query_str, local_dict = format_and_Q_eq_in(case)
        flags = case_df.query(query_str, local_dict=local_dict)

        if flags.shape[0] > 0:
            continue
        # 将 pd.DataFrame.from_dict 移到这里，
        # 仅在需要时创建新的 DataFrame
        new_row = pd.DataFrame.from_dict([{**case, "unique_index": i}], orient='columns')

        case_df = pd.concat(
            [case_df, new_row],
            ignore_index=True,
        )
        # input(case_df[[x for x in case_df.columns if len(case_df[x].unique()) > 1]].to_string(index=True))
    cases = [cases[int(idx)] for idx in case_df["unique_index"].tolist()]
    return segs, cases


def rangeProduct(config: Dict[str, List]):
    """Compute the Cartesian product of all parameter ranges.

    Args:
        config: Mapping of parameter name to list of candidate values.

    Returns:
        List of dicts, each representing one combination of parameter values.
    """
    sorted_keys = sorted(config.keys())
    sorted_fields = list([config[x] for x in sorted_keys])

    cases = itertools.product(*sorted_fields)

    case_list = [{key: val for key, val in zip(sorted_keys, case)} for case in cases]

    return case_list


def shrinkFieldRange(fields: Dict[str, Dict]):
    """Collapse the ``range``/``default`` descriptor format into plain value lists.

    Each field in a DoE config is stored as ``{"type": ..., "default": ...,
    "range": [...]}``.  If ``range`` is empty the ``default`` value is used as
    the sole candidate.

    Args:
        fields: Mapping of parameter name to its descriptor dict.

    Returns:
        Mapping of parameter name to a non-empty list of candidate values.
    """
    new_fields: Dict[str, List] = {}

    for key, field in fields.items():
        rge: List = field["range"]
        if (len(rge) == 0):
            rge.append(field["default"])
        new_fields[key] = rge

    return new_fields
